Resumen Técnico: Transformer de Ancho de Banda Completo (Full-Bandwidth Transformer)

Planteamiento del Problema

Los transformers autorregresivos computan a lo largo de dos ejes: horizontalmente a través de los tokens generados y verticalmente a través de la profundidad del modelo. Mientras que la atención densa proporciona un acceso horizontal amplio a los tokens pasados, el canal de retroalimentación vertical entre los pasos de decodificación permanece estrecho. En la decodificación estándar, el modelo comprime todo su estado oculto de $D$ dimensiones de la capa superior en un único token discreto (que transporta, como máximo, $\log_2 |V|$ bits) para servir como entrada para el siguiente paso. En consecuencia, la computación no verbalizada (activaciones intermedias, incertidumbre, planes parciales) queda "congelada en la profundidad" en el caché KV; permanece accesible para las capas por encima de su punto de producción, pero no puede regresar al fondo de la pila para un mayor procesamiento. Esto obliga al modelo a elegir entre verbalizar estados intermedios token por token o recomputarlos desde cero, limitando la profundidad efectiva de razonamiento por token.

Metodología: Decodificación de Retroalimentación Latente

Los autores proponen el Transformer de Ancho de Banda Completo, que ensancha el canal de retroalimentación inter-paso alimentando el estado oculto de la capa superior anterior en la entrada junto con el embedding del token muestreado.

1. Mecanismo de Retroalimentación Latente

En cada paso de decodificación $t$, la entrada $u_t$ se forma fusionando el embedding del token muestreado $e_t$ con el estado oculto de la capa superior anterior $h^L_{t-1}$:
$$u_t = e_t \otimes h^L_{t-1}$$
El operador de fusión $\otimes$ es una Unidad Lineal Puerta (GLU):
$$e_t \otimes h^L_{t-1} = (W_U h^L_{t-1}) \odot \sigma(W_G e_t)$$

• Asimetría: El estado oculto ocupa la vía de valor, mientras que el token actúa como una puerta multiplicativa. Este diseño evita que el modelo aprenda a ignorar la vía del estado (un posible atajo en la fusión simétrica), obligándolo a utilizar el estado oculto completo.
• Arquitectura: Esto añade solo dos multiplicaciones de matrices $D \times D$ por token. La arquitectura estándar del transformer, el diseño del caché KV y la pila de servicio permanecen intactos.
• Efecto: Esto permite que los estados no verbalizados vuelvan a entrar en la base de la pila con un presupuesto de profundidad renovado. Cada capa, incluso la más superficial, puede ahora acceder a todo el historial procesado por toda la pila, en lugar de solo por las capas debajo de ella.

2. Estrategia de Entrenamiento en Paralelo

Entrenar esta recurrencia directamente perdería el paralelismo de teacher forcing, ya que la entrada en la posición $t$ depende de la salida en $t-1$. Los autores introducen un objetivo de múltiples pasadas para resolver esto:

• Paralelismo Temporal: En lugar de desplegar secuencialmente sobre los tokens, el modelo ejecuta $K$ pasadas hacia adelante paralelas sobre la secuencia.
  • Pasada 1: Pasada hacia adelante estándar con embeddings de tokens simples.
  • Pasada $k$ ($k > 1$): Los estados ocultos de la Pasada $k-1$ se desplazan una posición a la derecha, se fusionan con los embeddings de los tokens y se retroalimentan en la pila en paralelo.
• Función de Pérdida: El objetivo combina la pérdida estándar de Predicción del Siguiente Token (NTP) de la primera pasada con pérdidas NTP auxiliares de las pasadas subsiguientes. Los gradientes de las pasadas posteriores se retropropagan hacia los estados latentes anteriores, fomentando que sean entradas reutilizables.
• Programación (Scheduling): Para asegurar la estabilidad y gestionar los costes de cómputo, los autores utilizan una programación progresiva:
  • El entrenamiento comienza con objetivos de una sola pasada estándar.
  • La retroalimentación latente se introduce a mitad del entrenamiento (por ejemplo, lotes de 2 pasadas).
  • Se mezcla una pequeña fracción (por ejemplo, 3%) de lotes con un mayor número de pasadas (por ejemplo, 3 pasadas). Empíricamente, esta pequeña fracción estabiliza el mapa de retroalimentación aprendido, convirtiéndolo en una contracción hacia un punto fijo que permanece estable incluso cuando se extrapola a horizontes largos (cientos de pasadas) durante la inferencia.
• Mezcla de Prefijo (Prefix Mix-in): Para manejar el desajuste de distribución entre el entrenamiento (entradas totalmente fusionadas) y la inferencia (prompt simple + generación fusionada), el proceso de entrenamiento randomiza la longitud del prefijo donde comienzan las entradas fusionadas, coincidiendo con la estructura de "prompt-luego-generación" de la inferencia.

3. Técnicas de Estabilidad

Para evitar la divergencia durante la inferencia de largo horizonte, los autores emplean:

• Escalamiento de Profundidad: Asegurar que la norma del estado de la capa superior permanezca $O(1)$ en lugar de crecer con la profundidad.
• Vinculación de Pesos (Weight Tying): Vincular los pesos de la capa de embedding y la capa de lectura para asegurar la compatibilidad entre las entradas de token simple y las de estado fusionado.
• Regularización de Ruido: Añadir ruido de fluctuación (jitter noise) al estado oculto transportado durante el entrenamiento para mejorar la robustez.

Resultados Clave

Los autores entrenaron modelos de 1B de parámetros hasta 400B de tokens y los evaluaron contra transformers estándar entrenados con 100B a 1T de tokens.

• Eficiencia de Datos: Los transformers de ancho de banda completo igualan o se acercan al rendimiento de los transformers estándar entrenados con aproximadamente 1.5× a 2× más tokens. Por ejemplo, un modelo de 200B de tokens de ancho de banda completo se acerca al rendimiento de un baseline estándar de 400B de tokens.
• Ganancias de Rendimiento:
  • Pérdida de Validación: Mejorada en todas las escalas.
    im
  • Razonamiento y Codificación: En tareas de generación libre (GSM8K, Math500, HumanEval, MBPP), la decodificación de retroalimentación latente (regímenes SOFT y FUSED) supera consistentemente a la decodificación estándar de los mismos pesos.
    • SOFT (Retroalimentación Latente): Es mejor para tareas matemáticas, sugiriendo que el estado oculto transporta mejor el razonamiento que los tokens verbalizados.
    • FUSED (Prefill de Múltiples Pasadas): Es mejor para tareas de codificación, sugiriendo que las representaciones refinadas del prompt son cruciales.
  • Ajuste de Instrucciones (Instruction Tuning): Las ganancias persisten tras la extensión de contexto largo y el ajuste de instrucciones.
• Concisión: En modelos base (antes del ajuste de instrucciones), la retroalimentación latente produce trazas de razonamiento notablemente más cortas con una precisión igual o mejor. El modelo realiza la computación "silenciosamente" en el estado oculto en lugar de narrarla token por token. Este efecto disminuye tras el ajuste de instrucciones, probablemente porque los datos de ajuste imitan un razonamiento verboso estándar.
• Accesibilidad del Estado: Los experimentos de sondeo confirman que el pre-llenado (prefilling) recurrente expone la información del estado global a las capas superficiales (Capa 0) de manera mucho más efectiva que el pre-llenado estándar, validando la afirmación de "ancho de banda completo".

Significado y Reivindicaciones

El artículo afirma que el Transformer de Ancho de Banda Completo ofrece un nuevo eje de escalado para los Modelos de Lenguaje de Gran Escala (LLMs): asignar más computación por token en lugar de simplemente más tokens o parámetros.

• Revisión del Escalado: A medida que los datos únicos de alta calidad se convierten en una restricción, los autores argumentan a favor de extraer más señal de aprendizaje de cada token. La retroalimentación latente logra esto permitiendo que el modelo itere sobre su estado interno sin el cuello de botella de la verbalización.
• Eficiencia de Inferencia: El enfoque proporciona mejoras de rendimiento con un sobrecoste por token insignificante (solo dos multiplicaciones de matrices). El cómputo adicional se incurre principalmente durante el pre-llenado de múltiples pasadas opcional o el entrenamiento, no durante la generación secuencial de cada token.
• Minimalismo Arquitectónico: El método preserva la arquitectura estándar del transformer y el caché KV, requiriendo solo un cambio en la construcción de la entrada y un esquema de entrenamiento específico.
• Estabilidad: El trabajo demuestra que los bucles de retroalimentación recurrentes pueden entrenarse eficazmente en transformers utilizando objetivos de múltiples pasadas en paralelo, superando la inestabilidad asociada frecuentemente con el entrenamiento de redes recurrentes.

Los autores concluyen que los transformers de ancho de banda completo representan un cambio de depender únicamente de la escala de datos hacia el aprovechamiento de la computación recurrente basada en retroalimentación para mejorar la calidad de la representación y la eficiencia del razonamiento.