Linear Attention Architectures: Mechanisms, Trade-offs, and Cross-Layer Routing

Resumen Técnico: Arquitecturas de Atención Lineal: Mecanismos, Compromisos y Enrutamiento entre Capas

Declaración del Problema
La atención de softmax estándar permite que cada token recupere información de todo el contexto, un mecanismo central para la expresividad de los modelos de lenguaje Transformer. Sin embargo, este intercambio de token a token incurre en un costo cuadrático respecto a la longitud de la secuencia ($O(T^2)$) tanto en entrenamiento como en inferencia, convirtiéndose en un cuello de botella dominante a medida que las ventanas de contexto se expanden. Mientras que las arquitecturas de atención lineal ofrecen una solución al reformular la atención como una actualización recurrente sobre una matriz de memoria de tamaño constante (generando un tiempo de entrenamiento lineal y un tiempo de inferencia constante por token), las variantes tempranas sufrieron pérdidas significativas de precisión debido a la interferencia entre las asociaciones almacenadas y la falta de mecanismos para el olvido o la edición selectiva. Los desarrollos recientes, incluyendo DeltaNet y sus sucesores, han cerrado esta brecha introduciendo reglas delta de corrección de errores y compuertas de control de grano fino, aunque los compromisos específicos entre estos mecanismos, sus comportamientos de escalado y el potencial para el intercambio de información ligero entre capas aún deben ser mapeados sistemáticamente.

Metodología y Marco de Trabajo
Los autores presentan un estudio comparativo de la atención softmax y cuatro arquitectas de atención lineal recurrente recientes: DeltaNet, Gated DeltaNet, Kimi Delta Attention y Gated DeltaNet-2.

1. Notación Unificada: El artículo establece una notación común de memoria recurrente para expresar estos mecanismos. Este marco aísla los roles de la matriz de memoria recurrente $W$, el residuo de la regla delta $r$ (el error entre el valor actual y la predicción de la memoria) y los diversos mecanismos de control de decaimiento, borrado y escritura.

   • DeltaNet: Reemplaza la actualización aditiva ingenua con una regla delta de corrección de errores, escribiendo solo el residuo $r = v - \bar{v}$ en la memoria.
   • Gated DeltaNet: Añade una compuerta de decaimiento escalar aprendida $\alpha$ a la matriz de estado de la memoria antes de calcular el residuo, introduciendo un mecanismo de olvido grueso.
   • Kimi Delta Attention: Refina el olvido mediante el uso de un vector de decaimiento por canal $\alpha$, permitiendo que diferentes dimensiones de características decaigan a diferentes ritmos.
   • Gated DeltaNet-2: Desacopla aún más la actualización mediante la introducción de compuertas de borrado y escritura por canal separadas, permitiendo el control independiente sobre qué canales de clave se eliminan y qué canales de valor se almacenan.

2. Configuración Experimental: El estudio utiliza modelos de solo decodificador de 350 millones de parámetros entrenados con 15 mil millones de tokens (FineWeb-Edu). Los experimentos comparan la pérdida de validación, el rendimiento (throughput) de entrenamiento, la sensibilidad del optimizador (AdamW vs. Muon) y el escalado de la longitud de la secuencia. Los autores también realizan ejecuciones a mayor escala (1.3B y 3B de parámetros) y evalúan tareas de seguimiento (HellaSwag, PIQA, WinoGrande).

3. Enrutamiento entre Capas (Cross-Layer Routing): Para abordar la dilución de información en pilas profundas, el artículo introduce y evalúa mecanismos de enrutamiento ligeros. Explora dos ejes: qué enrutar (error de escritura vs. valor de escritura) y dónde inyectarlo (el objetivo de valor de la siguiente capa vs. el flujo oculto compartido).

   • CLER (Cross-Layer Error Residuals): Reenvía el error de escritura de la regla delta hacia el objetivo de valor de la siguiente capa.
   • CLER-H: Enruta el error de escritura hacia el flujo oculto compartido.
   • CLVR (Cross-Layer Value Routing): Enruta el valor de escritura de la capa (el objetivo antes de la sustracción) hacia el flujo oculto compartido mediante una proyección inicializada en cero.

Resultos Clave

• Compromisos entre Arquitectura y Optimizador: Dentro del barrido de 350M de parámetros y 15B de tokens, Kimi Delta Attention entrenada con el optimizador Muon en una pila híbrida (entrelazada con capas softmax) logra la pérdida de validación final más baja (2.273). Sin embargo, las pilas de Gated DeltaNet puro entrenadas con AdamW demuestran el mayor rendimiento de entrenamiento normalizado (establecido en el 100% de la línea base). Las pilas híbridas generalmente mejoran la pérdida a costa de algo de rendimiento, mientras que las pilas puras preservan el escalado de tiempo de iteración más fuerte para contextos largos.
• Sensibilidad del Optimizador: La elección del optimizador es inseparable del ajuste de la tasa de aprendizaje. Muon prefiere tasas de aprendizaje más bajas ($3 \times 10^{-4}$) para las variantes de atención lineal, mientras que las variantes de AdamW prefieren tasas más altas ($10^{-3}$). Muon reduce consistentemente la pérdida de validación final en relación con AdamW en configuraciones de arquitectura comparables.
• Comportamiento de Escalado: Las pilas recurrentes puras exhiben un escalado superior con la longitud de la secuencia. A 32k tokens, el Gated DeltaNet puro muestra un aumento mínimo en el tiempo de iteración (~1.1x desde 4k), mientras que la atención softmax aumenta aproximadamente ~2.9x. Las pilas híbridas se sitúan en un punto intermedio, con aumentos en el tiempo de iteración aproximadamente proporcionales a su densidad de capas softmax.
• Hallazgos de Enrutamiento entre Capas:
  • CLER (Error al Objetivo de Valor): Reenviar el error de escritura de la regla delta al objetivo de valor de la siguiente capa no mejora el rendimiento sobre las líneas base comparables. Los autores lo atribuyen a un "desajuste de espacio", donde el residuo enrutado compite con la propia geometría de valor aprendida del receptor.
  • CLVR (Valor al Flujo Oculto): Enrutar el valor de escritura hacia el flujo oculto compartido produce una mejora modesta pero consistente en la pérdida de validación final tanto para DeltaNet como para Gated DeltaNet. Esto sugiere que la señal útil entre capas es el objetivo de escritura, no la corrección de la regla delta.
  • Rendimientos Decrecientes: Las ganancias de CLVR son pequeñas (reduciendo la pérdida en ~0.01 en 350M/1B de tokens) y disminuyen a medida que la escala del modelo o el conteo de tokens de entrenamiento aumentan, lo que sugiere que el beneficio es más pronunciado cuando la memoria base aún no ha capturado toda la información necesaria.

Significancia y Reivindicaciones
El artículo no declara una única arquitectura "mejor", sino que busca hacer legible el espacio de diseño. Sus contribuciones principales son:

1. Un Marco Unificado: Proporcionar una notación común que clarifica cómo las variantes recientes de atención lineal difieren en expresividad, decaimiento de memoria y granularidad de control.
2. Mapeo de Compromisos Empíricos: Demostrar que la elección entre pilas lineales y híbridas, así como entre optimizadores, implica un compromiso multiobjetivo entre la pérdida de validación, el rendimiento y el escalado de la longitud de la secuencia.
3. Perspectivas de Enrutamiento: Identificar que, mientras que el enrutamiento de error ingenuo (CLER) falla debido al desalineamiento de la base, enrutar el valor de escritura hacia el flujo oculto alineado (CLVR) ofrece una vía ligera de tiempo lineal para el intercambio de información entre capas que mejora el rendimiento sin romper la estructura recurrente.

Los autores enfatizan que sus conclusiones se basan en comparaciones de ejecución única y presupuestos de tokens específicos, instando a la cautela al generalizar pequeñas diferencias en la pérdida de validación sin una verificación de semillas repetidas. El trabajo sirve como base para comprender los mecanismos específicos que impulsan el rendimiento en la atención lineal y destaca preguntas abiertas respecto a la generalización de las estrategias de enrutamiento a arquitecturas más complejas como Gated DeltaNet-2.