Resumen Técnico: Atención Lineal de Kalman (KLA)

Planteamiento del Problema

Los modelos actuales de espacio de estados (SSM), como Mamba y la atención lineal con compuertas (GLA), ofrecen alternativas de complejidad lineal y paralelizables frente a la atención softmax, abordando el cuello de botella del escalado cuadrático de los Transformers. Sin embargo, estos modelos dependen de actualizaciones de estado oculto lineales o afines, lo que limita su expresividad en comparación con las interacciones no lineales inducidas por la normalización softmax. Además, los SSM y las RNN lineales existentes carecen de mecanismos explícitos para representar la incertidumbre del estado, tratando los estados ocultos como deterministas. Este artículo aborda la brecha entre la eficiencia de las actualizaciones lineales y la expresividad/robustez del filtrado no lineal, preguntándose si las actualizaciones exactas del filtro de Kalman pueden implementarse de manera eficiente como un mezclador de secuencias.

Metodología: Atención Lineal de Kalman (KLA)

Los autores proponen la Atención Lineal de Kalman (KLA), un mezclador de secuencias probabilístico que formula el modelado de secuencias como inferencia de filtrado bayesiano exacto. La innovación central reside en la reparametrización del filtro de Kalman hacia la forma de información (utilizando precisiones y parámetros naturales) para permitir la computación paralela.

1. Formulación Probabilística

KLA modela la mezcla de secuencias como inferencia en un modelo de espacio de estados lineal-gaussiano:

• Prior de Dinámica: La evolución del estado latente se modela como un proceso de Ornstein-Uhlenbeck (OU), discretizado exactamente. Esto introduce un término de ruido de proceso ($\bar{p}_t$) que captura la incertidumbre en las transiciones de estado, diferenciándose de la dinámica determinista de los SSM estándar.
• Modelo de Observación: Los tokens de entrada se tratan como observaciones ruidosas del estado latente, parametrizadas por un operador de observación ($k_t$) y una precisión de valor ($\Lambda^v_t$) que representa la confianza en la evidencia del token.
• Lectura (Readout): La salida se deriva de la media posterior del estado latente, proyectada mediante un operador de consulta ($q_t$).

2. Paralelización mediante Transformaciones de Möbius

El filtrado de Kalman clásico es inherentemente secuencial. KLA supera esto explotando la forma de información del filtro:

• Recursión de Precisión: Se demuestra que la actualización de la precisión posterior ($\lambda_t$) es una transformación de Möbius (fraccional-lineal) de la precisión anterior ($\lambda_{t-1}$).
  $$\lambda_t = \frac{\alpha_t \lambda_{t-1} + \beta_t}{\gamma_t \lambda_{t-1} + \delta_t}$$
  Crucialmente, las transformaciones de Möbius se componen asociativamente. Esto permite que las actualizaciones de precisión se computen mediante escaneos de prefijo paralelos (escaneos asociativos) con un trabajo de $O(T)$ y una profundidad de $O(\log T)$, igualando la eficiencia paralela de los SSM modernos como Mamba.
• Recursión de la Media: Una vez que la trayectoria de precisión se computa en paralelo, las actualizaciones de la media posterior (media de información) siguen una transformación afín que también es paralelizable.

3. Compuerta No Lineal desde la Incertidumbre

A diferencia de los SSM lineales donde el gating (compuerta) es a menudo lineal o dependiente del token pero determinista, el gating de KLA surge naturalmente de la relación de precisiones. La "compuerta de olvido" (forget gate) y la "compuerta de escritura" (write gate) dependen de la historia y son no lineales, impulsadas por la incertidumbre relativa entre el estado previo y la nueva observación. Esto permite al modelo incorporar información de manera selectiva basada en la confianza sin requerir parámetros de compuerta aprendidos explícitamente para cada token.

4. Arquitectura

KLA está diseñado como un reemplazo directo (drop-in replacement) para las capas estándar de SSM o atención. Sigue una arquitectura de MLP fusionado similar a Mamba, incorporando convoluciones causales 1D y activaciones SiLU. Los autores implementan núcleos CUDA y Triton eficientes para el escaneo de Möbius para evitar la materialización de estados expandidos en la memoria de alto ancho de banda (HBM).

Contribuciones Clave

1. Reparametrización Asociativa: Los autores demuestran que el filtrado de Kalman gaussiano-lineal diagonal en forma de información admite una estructura de Möbius, permitiendo escaneos de prefijo paralelos exactos para el filtrado bayesiano.
2. No Linealidad Impulsada por la Incertidumbre: KLA introduce un gating no lineal derivado de las relaciones de precisión, ofreciendo una expresividad estrictamente mayor que las actualizaciones de SSM lineales/afines, manteniendo al mismo tiempo una estructura de escaneo de tiempo lineal.
3. Capa KLA: Un mezclador de secuencias práctico y directo que produce la incertidumbre del estado de creencia explícita y puede apilarse en redes profundas.
4. Escalabilidad y Validación: El artículo presenta los primeros primitivos de filtrado bayesiano apilados entrenados a la escala de miles de millones de tokens, validando su estabilidad y rendimiento.

Resultados Experimentales

Los autores evalúan KLA a través de tareas sintéticas, pruebas de seguimiento de estado y modelado de lenguaje a gran escala.

• Tareas Sintéticas (MAD-Lab): KLA supera a GDN, GLA, Mamba y mLSTM en tareas que requieren selección bajo incertidumbre, específicamente Fuzzy Recall y Selective Copy. Logra resultados de vanguardia en Compresión, lo que sugiere que su filtrado en forma de información actúa como un compresor de contexto óptimo.
• Seguimiento de Estado (A5 Permutation): En la tarea de composición de permutación A5 (un benchmark para el seguimiento de estado NC1-completo), KLA resuelve la tarea con profundidad constante (1–2 bloques). En contraste, los SSM lineales y los modelos basados en atención requieren una profundidad que crece con la longitud de la secuencia. Esto demuestra que las actualizaciones fraccionales-lineales de KLA rompen los límites de expresividad de los modelos recurrentes lineales manteniendo la capacidad de entrenamiento paralelo.
• Recuerdo Asociativo de Contexto Largo (MQAR): KLA logra una precisión casi perfecta (>95%) en la tarea de Recuerdo Asociativo de Multi-Consulta con una longitud de secuencia de 2048, superando a Mamba y excediendo significativamente a GLA, que falla al aprender la tarea en este entorno.
• Preentrenamiento a Gran Escala:
  • Independiente (Standalone): Con 45M y 180M de parámetros, KLA iguala o mejora a GPT, Mamba y GDN en benchmarks de sentido común de cero disparos (zero-shot) (ej. HellaSwag, PIQA).
  • Híbrido: Reemplazar la última capa de atención de un GPT con KLA (GPT+KLA) produce el rendimiento más fuerte entre los híbridos, superando a GPT+Mamba y GPT+GDN.
• Ablación: Los autores muestran que el prior de Ornstein-Uhlenbeck (OU) es crítico para el apilamiento profundo estable. Sin la dinámica estocástica de reversión a la media, el entrenamiento de pilas profundas de KLA se vuelve inestable.

Significado y Reivindicaciones

El artículo afirma que KLA logra cerrar la brecha entre la eficiencia de los modelos de espacio de estados lineales y la expresividad de los mecanismos de atención no lineales. Al plantear la mezcla de secuencias como un filtrado bayesiano exacto, KLA proporciona:

• Incertidumbre Explícita: A diferencia de los mezcladores deterministas, KLA porta una incertidumbre de estado de creencia explícita que controla directamente el gating de la información.
• Expresividad Estrictamente Superior: Las actualizaciones de Möbius permiten que KLA resuelva tareas de seguimiento de estado (como A5) que los SSM lineales y la atención no pueden, sin sacrificar las capacidades de entrenamiento paralelo.
• Escalabilidad: El trabajo demuestra que los primitivos de filtrado bayesiano exacto pueden apilarse y entrenarse a la escala de miles de millones de tokens, desafiando la noción de que el filtrado probabilístico es inherentemente secuencial o computacionalmente prohibitivo para el modelado de lenguaje a gran escala.

Los autores posicionan a KLA no solo como una curiosidad teórica, sino como un primitivo práctico y directo que mejora las capacidades de robustez y seguimiento de estado de los modelos de lenguaje modernos, manteniendo la eficiencia computacional necesaria para el procesamiento de contextos largos.