# NCM: modelo, tokenizador y fragmentos ## 1. Qué es una NCM Una **NCM** (Neural Codebook Model) es una arquitectura jerárquica en la que la señal se representa mediante **códigos discretos** aprendidos y se reconstruye con una torre de bloques transformer. Frente a un LLM, donde la unidad es el token de texto y el objetivo es predecir el siguiente, en una NCM la unidad es una **trama** y el objetivo es modelar y reconstruir esa trama con la máxima fidelidad posible. Esa diferencia tiene consecuencias directas para el reparto distribuido: | Propiedad | LLM denso | NCM | | --- | --- | --- | | Unidad de entrada | token de texto | trama continua cuantizada | | Uniformidad entre capas | alta | baja (jerárquica) | | Coste de las primeras capas | medio | alto (más información) | | Sensibilidad a la latencia en la cola | alta | menor (refinamiento) | | Replicar la cabeza | caro | barato y rentable | ## 2. Tokenizador VQ-VAE residual ### 2.1 Frontend de texto → tramas Cada byte UTF-8 se normaliza al intervalo $[-1, 1]$: $$ f_i = \frac{b_i - 127.5}{127.5}, \qquad b_i \in [0, 255] $$ y se agrupa en tramas de `frame_dim` componentes (la última se rellena con ceros). La inversa es directa: $$ \hat{b}_i = \operatorname{round}(127.5\,f_i + 127.5) $$ ### 2.2 Cuantización vectorial Cada etapa $s$ mantiene un libro de códigos $\mathcal{C}^{(s)} = \{c_1, \dots, c_K\}$. Codificar una trama $x$ consiste en una búsqueda exhaustiva del centroide más cercano y el cálculo del residuo: $$ k_s = \arg\min_{k} \lVert r_{s-1} - c_k^{(s)} \rVert_2^2, \qquad r_s = r_{s-1} - c_{k_s}^{(s)}, \qquad r_0 = x $$ Decodificar es sumar los centroides elegidos: $$ \hat{x} = \sum_{s=1}^{S} c_{k_s}^{(s)} $$ Con $S$ etapas y $K$ entradas por etapa, el espacio efectivo de códigos es $K^S$ sin aumentar la memoria de cada libro. Es exactamente el truco que hace que 3 etapas de 32 códigos (32³ = 32 768 combinaciones) reconstruyan texto byte a byte con precisión suficiente. ### 2.3 Entrenamiento `NcmTokenizer.train()` aplica una actualización tipo EMA sobre el centroide ganador: $$ c_{k}^{(s)} \leftarrow c_{k}^{(s)} + \eta \,(r - c_{k}^{(s)}) $$ con $\eta$ = `learning_rate` (0.35 por defecto). Cada etapa se entrena sobre el residuo que deja la anterior, igual que en inferencia; así el libro aprendido es el que realmente se usa. `reconstruction_error()` devuelve el error cuadrático medio por componente y `train()` lo devuelve por época, de modo que se puede comprobar que decrece. ### 2.4 Persistencia `to_json()` guarda configuración y libros; `from_json()` reconstruye el tokenizador exactamente igual (mismo `seed`, mismos centroides). Los ficheros se publican como `*.vqvae` en las *releases* del repositorio. ## 3. Fragmento transformer ### 3.1 Composición de un bloque Cada bloque del fragmento aplica pre-norm y dos subcapas con residual: $$ \begin{aligned} h &= x + \operatorname{Attn}(\operatorname{RMSNorm}(x)) \\ y &= h + \operatorname{SwiGLU}(\operatorname{RMSNorm}(h)) \end{aligned} $$ RMSNorm: $$ \operatorname{RMSNorm}(x) = \frac{x}{\sqrt{\frac{1}{d}\sum_i x_i^2 + \epsilon}} \odot g $$ Atención causal multi-cabeza con $H$ cabezas de dimensión $d_h = d/H$: $$ \operatorname{Attn}(Q,K,V) = \operatorname{softmax}\!\left(\frac{QK^\top}{\sqrt{d_h}} + M\right)V $$ donde $M_{ij} = -\infty$ si $j > i$ (máscara causal). SwiGLU: $$ \operatorname{SwiGLU}(x) = W_2\big(\operatorname{SiLU}(W_1 x) \odot W_3 x\big) $$ ### 3.2 Inicialización determinista Los pesos se inicializan con el generador global de PyTorch sembrado con `FragmentConfig.seed`, y cada `nn.Linear` usa la inicialización Kaiming de PyTorch ($\mathcal{U}(-1/\sqrt{d}, 1/\sqrt{d})$): $$ W_{ij} \sim \mathcal{U}\!\left(-\frac{1}{\sqrt{d}}, \frac{1}{\sqrt{d}}\right) $$ Dos procesos con la misma `FragmentConfig` obtienen **exactamente** los mismos pesos, lo que permite comparar la ejecución distribuida con la local en los tests (`assert_tensors_close(atol=1e-5)`) sin transferir checkpoints. `checksum()` (suma de cuadrados de todos los parámetros) es la huella barata que se usa para detectar si dos nodos cargaron fragmentos distintos. ### 3.3 Rangos de capas `FragmentConfig.layer_start`/`layer_end` definen qué capas globales ejecuta el fragmento. Un nodo puede tener varios fragmentos (por ejemplo, su tramo primario y una réplica de la cabeza). El contrato gRPC transmite `layer_start`, y el plano de datos elige el fragmento que **empieza** exactamente en esa capa: ```python fragment = service.fragments[layer_start] # p. ej. {0: (0-5), 6: (6-9), 10: (10-11)} salida = fragment.forward(activaciones) ``` Un fragmento se ejecuta **completo**: entrar en mitad de uno recalcularía capas que otro nodo ya ha calculado. Si `layer_start` no es el inicio de ningún fragmento, el servicio responde `FAILED_PRECONDITION` con la lista de inicios disponibles. ### 3.4 Empaquetado binario Las activaciones viajan como `bytes` `float32` **little-endian** en orden de filas, sin cabecera ni padding: ```python # Python (plano de datos) tensor = torch.frombuffer(bytearray(payload), dtype=torch.float32).reshape(rows, cols) payload = tensor.detach().to("cpu", dtype=torch.float32).contiguous().numpy().tobytes(order="C") ``` Se usa little-endian porque es el orden nativo de x86-64 y ARM64, y el mismo que escribe `prost` desde Rust: `f32::to_le_bytes` y `numpy.tobytes()` producen los mismos bytes sin *swap* en la frontera. `decode_activations()` valida que el payload mida exactamente `rows × cols × 4` bytes antes de interpretarlo, y `ActivationBuffer::validate` hace la misma comprobación en Rust antes de enviar. ## 4. Por qué el paralelismo por capas encaja con una NCM 1. **Forma del coste**: la cabeza concentra el trabajo pesado. El reparto ponderado (`VRAM × 2 + RAM × 0.5 + NPU × 1.5`) la coloca en el nodo con GPU y reparte el resto. 2. **Replicación barata de la cabeza**: replicar las primeras capas duplica poco trabajo y elimina el punto único de fallo más crítico. 3. **Cola tolerante a latencia**: un salto de red en las últimas capas añade milisegundos sobre un tramo de trabajo pequeño; el impacto en la calidad es menor que en un LLM. 4. **Unidades discretas**: los códigos del tokenizador permiten cortar y reanudar el flujo entre nodos sin ambigüedad de representación. ## 5. Parámetros recomendados | Parámetro | Valor de desarrollo | Valor de LAN doméstica | Motivo | | --- | --- | --- | --- | | `d_model` | 16-24 | 256-512 | Los tests deben correr en milisegundos. | | `n_heads` | 4 | 8-16 | `d_model % n_heads == 0` es obligatorio. | | `d_ff` | 2 × `d_model` | 4 × `d_model` | Compromiso entre capacidad y coste. | | `frame_dim` | 8-16 | 32-64 | Bytes por trama; más grande = menos tramas, más contexto por código. | | `codebook_size` | 16-32 | 512-1024 | Cada etapa corrige el residuo de la anterior. | | `num_stages` | 2-3 | 3-6 | Cada etapa corrige el residuo de la anterior. | > **Aviso de estado (2026-10-29).** La fila `codebook_size` decía «Vocabulario efectivo = > `codebook_size ** num_stages`», que contradice la propiedad `vocab_size` del tokenizador > (publica el **producto**, `tokenizer_wrapper.py:100-102`). Cuál de las dos es correcta es una > decisión de diseño pendiente, no un error de una línea: si las etapas son independientes y el > código es la tupla, el espacio es $K^S$; si lo que se publica es «cuántos símbolos emite por > trama», el producto es correcto y el nombre `vocab_size` es el equivocado. Registrado en > `docs/audit/divergences.md` §1 fila 15. ## 5.bis Qué recupera de verdad: el techo del cuantizador Medido el 2026-10-29 sobre `README.md` (2 010 tramas de `frame_dim=16`), con `scripts/train_ncm_head.py` y el codificador cerrado del tokenizador. **Códigos exactos** —es decir, con una cabeza perfecta y sin error de predicción: | `codebook_size` | `num_stages` | Bytes recuperados exactos | Error medio (unidades de byte) | | --- | --- | --- | --- | | **32** (omisión) | 3 | **0,83 %** | 52,8 | | 128 | 3 | 3,96 % | 8,9 | | 512 | 4 | **13,3 %** | 2,33 | | 2 048 | 4 | **25,8 %** | 1,23 | Para acertar un byte hace falta que la reconstrucción caiga a menos de **0,5** unidades del valor original. Con la configuración por omisión el error medio es 52,8: dos órdenes de magnitud por encima de la tolerancia. ### Qué se puede y qué no se puede concluir * **Se puede** concluir que el techo es de **capacidad, no de diseño**: multiplicar el libro de códigos por 64 divide el error por 43. La arquitectura responde; la configuración es la que no llega. * **Se puede** concluir que el cuantizador **conserva estructura**: su error es **1,87× menor** que el de códigos aleatorios (55,4 frente a 103,6). No es ruido, es una aproximación gruesa. * **No se puede** concluir que la NCM genere texto hoy. Con `K=32` y tres etapas no lo hace, y ninguna mejora de la cabeza lo arregla, porque el límite está antes. ### Consecuencia para `forward_and_decode` `NcmCodeHead.forward_and_decode` y `NcmConditionedCodeHead.forward_and_decode` cierran la composición activaciones → códigos → tramas → texto, que era la pieza que faltaba en el camino de datos. **Devuelven texto; no devuelven texto legible a la capacidad por omisión.** La constante `TEXT_RECOVERY_AT_DEFAULTS` (`layer_fragment.py`) guarda el 0,83 % y `tests/python/test_ncm_decode.py` falla si alguien cambia la capacidad sin actualizarla, de modo que la cifra no pueda quedarse mintiendo. Para que la NCM genere texto de verdad hace falta, en este orden: `codebook_size` de miles, `num_stages` de 4 o más, entrenar los libros sobre un corpus representativo, y **sólo entonces** la cabeza. Lo que hoy está medido es que el paso 1 mueve la aguja 30×. ## 6. Cómo extenderlo a tensores reales `NcmLayerFragment.forward()` es la única frontera que hay que respetar: ```python def forward(self, activations: torch.Tensor) -> torch.Tensor: ... ``` Acepta `[tokens, d_model]` o `[batch, tokens, d_model]` y devuelve la misma forma. Para sustituirlo por un transformer real (Hugging Face, `torch.nn.TransformerEncoder`…) basta con envolver el módulo y mantener esa firma: el worker en Rust, el orquestador, el scheduler, el router y el contrato Protobuf **no cambian**, porque ninguno de ellos sabe qué hay dentro del fragmento. La frontera ya usa `float32` little-endian, que es el orden nativo de NumPy y de `prost`, así que no hay conversión de extremidad en el camino crítico.