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:

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:

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 (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

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:

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:

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.