Resumen Técnico: Transformer Recurrente Latente (LRT)

Declaración del Problema

Los transformers autoregresivos son el estándar para el modelado del lenguaje, pero dependen de cálculos feedforward de profundidad fija para cada token generado. Si bien la recurrencia puede teóricamente aumentar la capacidad computacional para el refinamiento iterativo, los enfoques existentes introducen costos significativos de inferencia. Los métodos que añaden recurrencia en profundidad (aplicando repetidamente bloques al mismo token) o en tiempo (insertando tokens de pausa o pensamiento) aumentan la latencia de inferencia mediante pasos de decodificación adicionales o aplicaciones de bloques.

Los autores observan que la decodificación autoregresiva ya calcula señales recurrentes de alto nivel "gratis": los estados ocultos del token anterior. Específicamente, los estados de las capas superiores entrenados para la predicción del siguiente token contienen un contexto latente útil. La pregunta central de investigación es si estas representaciones de alto nivel ya calculadas del token anterior pueden reutilizarse como memoria recurrente para el siguiente token sin añadir pasos de decodificación extra.

Metodología

Arquitectura: Transformer Recurrente Latente (LRT)

LRT es una ampliación ligera de los transformers autoregresivos estándar que introduce una vía latente transversal a capas y tokens.

• Mecanismo: En la posición del token $t$, LRT reutiliza un estado oculto de una capa fuente de la posición anterior, $m_{t-1} = h^{\ell_{src}}_{t-1}$, como memoria recurrente. Aquí, $\ell_{src}$ es una capa fuente seleccionada (típicamente una capa superior-media).
• Inyección: La memoria $m_{t-1}$ se inyecta en las capas del transformer del token actual $t$ mediante dos mecanismos complementarios:
  1. Proyección KV: Proyecta la memoria en los espacios de clave-valor de la atención ($\tilde{k}_{rec}, \tilde{v}_{rec}$) y los combina con las claves/valores locales utilizando puertas dependientes de la entrada. Esto proporciona una ruta directa para la representación del token anterior dentro del mecanismo de atención.
  2. Inyección Residual: Añade la memoria directamente al flujo residual ($\bar{x}^{(\ell)}_t = \alpha_\ell x^{(\ell)}_t + \gamma_\ell m_{t-1}$), exponiendo todo el bloque a la señal recurrente.
• Eficiencia de Inferencia: Crucialmente, LRT preserva la columna vertebral estándar solo de decodificador, el mecanismo de atención y la interfaz de caché KV. No inserta tokens de memoria, tokens de pausa ni bucles de profundidad. Cada token generado utiliza exactamente una pasada forward normal del transformer, con el estado recurrente simplemente actualizándose y pasándose al siguiente paso. La variante predeterminada añade solo un 0.3% de parámetros.

Estrategia de Entrenamiento: Entrenamiento Paralelo Entrelazado

Un desafío mayor es el preentrenamiento de esta recurrencia a escala. El entrenamiento exacto requeriría desenrollar secuencialmente el transformer sobre toda la secuencia (Retropropagación a través del Tiempo), destruyendo el paralelismo esencial para el preentrenamiento eficiente de transformers. El entrenamiento por bloques (procesando bloques contiguos secuencialmente) preserva cierto paralelismo pero solo propaga la memoria a través de los límites de los bloques, ofreciendo una aproximación gruesa.

Los autores introducen el Entrenamiento Paralelo Entrelazado para aproximar la recurrencia a nivel de token manteniendo el cómputo paralelo:

1. Paso de Inicialización: Una única pasada forward de secuencia completa construye un búfer compartido que contiene claves, valores y estados recurrentes de la capa fuente para todas las posiciones.
2. Pasas de Refinamiento: Las posiciones de la secuencia se dividen en $S$ subconjuntos disjuntos (por ejemplo, particiones con saltos). Estos subconjuntos se refinan en paralelo.
3. Escritura de Vuelta: Los estados actualizados de un subconjunto se escriben de nuevo en el búfer compartido. Los subconjuntos subsiguientes pueden consumir la memoria refinada de los subconjuntos anteriores.
4. Objetivo: La pérdida es el promedio de la pérdida de inicialización y las pérdidas de refinamiento para cada subconjunto.

Esta estrategia asegura que cada token reciba un paso de supervisión consciente de la memoria recurrente. En términos de cómputo, cuesta aproximadamente 2× una actualización estándar de transformer (una inicialización + un refinamiento efectivo), pero imita la recurrencia a nivel de token más de cerca que el entrenamiento por bloques.

Contribuciones Clave

1. Arquitectura LRT: Una extensión recurrente ligera que reutiliza estados ocultos de alto nivel de una capa fuente del token anterior como memoria, habilitando la recurrencia transversal a capas sin pasos de decodificación extra ni cambios arquitectónicos en la interfaz de la caché KV.
2. Entrenamiento Paralelo Entrelazado: Una aproximación de entrenamiento paralelo que refina subconjuntos disjuntos de tokens para imitar la recurrencia a nivel de token bajo un presupuesto de entrenamiento fijo, evitando el desenrollado secuencial requerido por el entrenamiento estilo RNN estándar.
3. Validación Empírica: Demostración de que LRT mejora el equilibrio entre cómputo y calidad en varios tamaños de modelo y presupuestos de tokens por parámetro, con la variante predeterminada de proyección compartida añadiendo una sobrecarga de parámetros mínima (0.3%).

Resultados

Se realizaron experimentos en columnas vertebrales GPT estilo nanochat (20L y 24L capas) preentrenadas en FineWeb-Edu.

• Comportamiento de Escala: Bajo un cómputo de entrenamiento equivalente a la línea base (contabilizando el costo 2× del entrenamiento entrelazado), LRT desplaza las curvas de escala hacia Bits por Byte (BPB) más bajos y CORE más altos (una métrica para el aprendizaje en contexto a través de 22 tareas).
• Mejoras de Rendimiento:
  • En la columna vertebral de 24L con cómputo equivalente a la línea base 80, LRT-compartido logró 0.695 BPB en comparación con los 0.699 de la línea base.
  • El aprendizaje en contexto (CORE) también mejoró; en la columna vertebral de 20L con cómputo 80, LRT-compartido alcanzó 0.274 frente a 0.271 de la línea base.
• Hallazgos de Ablación:
  • Capa Fuente: Una capa superior-media (por ejemplo, capa 12 en un modelo de 20L) funciona mejor que la capa final, sugiriendo que la memoria útil debe ser de alto nivel pero no demasiado especializada para los logits.
  • Fuente Temporal: El estado inmediatamente precedente ($m_{t-1}$) es la única fuente de memoria más fuerte que preserva la restricción de una sola pasada forward.
  • Inyección: La Proyección KV y la Inyección Residual son complementarias; usar ambas produce los mejores resultados.
  • Entrenamiento: El entrenamiento paralelo entrelazado supera significativamente al entrenamiento por bloques, que proporciona una aproximación más gruesa de la recurrencia.

Significado y Afirmaciones

El artículo afirma que LRT desplaza exitosamente el equilibrio entre cómputo y calidad para los transformers autoregresivos. Al reutilizar señales de alto nivel "gratis" del token anterior y entrenarlas mediante una aproximación paralela, LRT habilita el refinamiento iterativo y una mejor utilización del contexto sin las penalizaciones de latencia de inferencia de los métodos de recurrencia en profundidad o de tokens de pausa.

Los autores enfatizan que LRT es una ampliación ligera que preserva la interfaz estándar del transformer. Las ganancias se atribuyen al mecanismo de memoria recurrente en sí mismo y no simplemente a añadir parámetros o puertas extra, ya que adiciones arquitectónicas genéricas (como atención con puertas o escalado de capas) sobre la misma línea base fuerte produjeron mejoras insignificantes. El trabajo sugiere que la memoria recurrente útil debe ser de alto nivel pero no la capa final especializada en logits, y que la recurrencia a nivel de token puede aproximarse efectivamente en paralelo durante el preentrenamiento.