Maglev: Sliding Recurrent Memory

Resumen Técnico: Maglev – Memoria Recurrente Deslizante

Planteamiento del Problema
Las arquitecturas Transformer actuales enfrentan un compromiso fundamental entre la capacidad de memoria y la eficiencia computacional. La atención de auto-atención causal completa permite una recuperación no lineal y un cómputo en contexto potentes, pero incurre en costos que crecen linealmente con la longitud del contexto. Por el contrario, la atención de ventana deslizante limita la atención y los costos del caché KV, pero descarta la información distante, fallando en proporcionar una memoria persistente a largo plazo. Los modelos recurrentes clásicos (por ejemplo, LSTMs) comprimen la historia pero propagan estados pequeños de forma secuencial, lo que limita su capacidad y el paralelismo de secuencia. Aunque los Memory Transformers y los modelos de atención lineal intentan cerrar esta breal, a menudo imponen límites de segmento artificiales, dependen de actualizaciones lineales estructuradas o requieren un despliegue secuencial durante el entrenamiento. El desafío central es desarrollar un mecanismo práctico de memoria recurrente no lineal, a nivel de token, para Transformers que permita el preentrenamiento en paralelo manteniendo costos de inferencia acotados.

Metodología: Arquitectura Maglev
Maglev aborda esto desacoplando la construcción de la memoria durante el entrenamiento de la propagación de la memoria durante la inferencia. La arquitectura consta de dos modelos acoplados:

1. Prefiller ($Q$): Un Transformer causal que procesa la secuencia observada en paralelo utilizando un campo receptivo más amplio (por ejemplo, atención intercalada de ventana completa y deslizante). Genera una trayectoria de vectores de memoria objetivo, $m'_t$, para cada posición de token.
2. Decoder ($P$): Un Transformer de ventana deslizante que predice el siguiente token y genera sus propios vectores de memoria, $m_t$. Crucialmente, $P$ no tiene acceso a la historia completa; solo ve una ventana acotada de tokens recientes y una ventana desplazada de vectores de memoria.

Mecanismo de Entrenamiento: Entrenamiento Paralelo Elevado (Lifted Parallel Training)
En lugar de desplegar el estado recurrente secuencialmente (lo que impide la paralelización), Maglev emplea una estrategia de entrenamiento de dos pasadas:

• Pasada 1: El prefiller $Q$ genera la trayectoria completa de la memoria objetivo $m'_{1:T}$ en paralelo.
• Pasada 2: El decoder $P$ consume la trayectoria objetivo desplazada $m'_{0:T-1}$ junto con los tokens de entrada para predecir el siguiente token y producir su propia memoria $m_{1:T}$.
• Objetivo: El modelo se optimiza utilizando una función de pérdida compuesta:
  • Pérdida de Predicción del Siguiente Token ($L_{CE}$): Pérdida de entropía cruzada estándar para el modelado de lenguaje.
  • Pérdida de Consistencia de Memoria ($L_{cons}$): Un término que alinea la memoria de salida del decoder $m_t$ con el objetivo del prefiller $m'_t$. Esto enseña al decoder a reproducir el estado de memoria necesario para el siguiente paso sin el despliegue secuencial.

Mecanismo de Inferencia
Durante la inferencia, el prefiller $Q$ se descarta. El decoder $P$ se ejecuta de forma recurrente, retroalimentando sus propias memorias generadas anteriormente $m_{t-1}$ al sistema.

• Inyección de K/V Recurrente: En lugar de utilizar tokens de memoria dedicados, Maglev inyecta el estado de memoria desplazado en las características de Clave/Valor (K/V) del mecanismo de atención mediante proyecciones aprendidas y mecanismos de compuerta (gating).
• Complejidad Acotada: El decoder mantiene una ventana deslizante de tamaño fijo de tokens y memorias. El estado recurrente se transporta hacia adelante a través del caché K/V, asegurando que los costos de inferencia (cómputo y tamaño de caché) permanezcan acotados e independientes de la longitud total de la secuencia, coincidiendo con el perfil de la atención de ventana deslizante estándar.

Contribuciones Clave

1. Recurrencia No Lineal a Nivel de Token: Maglev introduce un mecanismo donde el estado del decoder de cada token sirve como memoria, atravesando toda la profundidad no lineal del Transformer sin requerir tokens de memoria adicionales o límites de segmento artificiales.
2. Entrenamiento en Paralelo para Modelos Recurrentes: Al utilizar un prefiller para generar objetivos en paralelo, Maglev permite el preentrenamiento en paralelo de modelos recurrentes no lineales, evitando los cuellos de botella secuenciales del entrenamiento tradicional de RNN.
3. Compartición de Parámetros: Los autores demuestran que el prefiller $Q$ y el decoder $P$ pueden compartir la mayoría de sus parámetros (específicamente los bloques Transformer) sin degradar significativamente el rendimiento, reduciendo la sobrecarga de memoria de parámetros del setup de entrenamiento.
4. Inferencia Acotada: El modelo desplegado retiene la eficiencia de la atención de ventana deslizante mientras gana capacidades de memoria persistente a través de la inyección recurrente de K/V.

Resultados Experimentales
Evaluado en la arquitectura nanochat d20 (435M de parámetros, entrenado en 43.52B de tokens), Maglev fue comparado contra Transformers de ventana completa/deslizante intercalada, modelos puramente de ventana deslizante y líneas base de Transformer Recurrente Latente (LRT).

• Rendimiento: La mejor variante de Maglev (parámetros separados, $\lambda=1$) mejoró los bits por byte (BPB) de validación de FineWeb-Edu de 0.7413 (línea base de ventana deslizante) a 0s.7251.
• Tareas de Abajo (Downstream Tasks): Alcanzó una precisión promedio en tareas de abajo del 56.4%, superando tanto al Transformer de ventana deslizante emparejado (54.1%) como a la línea base de LRT emparejada.
• Peso de Consistencia: El estudio señaló que el peso óptimo de consistencia ($\lambda$) depende de la estrategia de compartición de parámetros; un $\lambda$ más alto benefició al modelo de parámetros separados pero pudo restringir al modelo de parámetros compartidos.

Significado y Reivindicaciones
El artículo afirma que Maglev proporciona una vía práctica para equipar a los Transformers no lineales con una memoria persistente a nivel de token que puede reescribirse después de cada token. Logra con éxito cerrar la brecha entre el poder expresivo de la atención completa y la eficiencia de los modelos de contexto acotado. Al separar la generación de memoria del uso de la memoria durante el entrenamiento, Maglev logra el preentrenamiento en paralelo para arquitecturas recurrentes, una hazaña que anteriormente era difícil para la recurrencia no lineal. Los autores posicionan a Maglev como una investigación preliminar en este espacio de diseño, señalando que, si bien mejora sobre las líneas base de ventana deslizante y de recurrencia latente en el entorno probado, se necesita más trabajo para escalar el enfoque, optimizar la eficiencia de los kernels para el despliegue y explorar mecanismos alternativos de inyección de memoria (por ejemplo, inyección en el flujo residual o atención cruzada).