From Gradients to Riccati Geometry: Kalman World Models for Single-Pass Learning

1. Problema y Motivación

El aprendizaje automático moderno está dominado por la retropropagación (backpropagation) y el descenso de gradiente estocástico (SGD). Aunque eficaz, este paradigma presenta limitaciones estructurales y computacionales críticas:

• Dependencia de la retropropagación: Requiere almacenar todas las activaciones intermedias durante el paso forward para calcular gradientes en el paso backward, lo que genera una sobrecarga de memoria proporcional a la profundidad de la red y la longitud de la secuencia.
• Optimización determinista: Trata los parámetros como cantidades fijas y desconocidas optimizadas sobre una función de pérdida estática, ignorando la incertidumbre y la naturaleza dinámica de los datos en escenarios de aprendizaje continuo (online).
• Sincronización en sistemas distribuidos: La sincronización de gradientes introduce costos de comunicación y posibles inestabilidades.
• Historial limitado: Los intentos previos de usar Filtros de Kalman Extendidos (EKF) en redes neuronales (años 80/90) fracasaron debido a limitaciones de escalabilidad y la falta de conexión con la geometría de información moderna.

El artículo propone un cambio de paradigma: en lugar de minimizar una pérdida mediante gradientes, el aprendizaje se reformula como un problema de estimación de estado secuencial bayesiana en sistemas dinámicos.

2. Metodología: Modelos Mundiales de Kalman (KWM)

El núcleo de la propuesta es Kalman World Models (KWM), un marco que reemplaza la actualización de gradientes por adaptaciones de ganancia de Kalman.

A. Formulación de Espacio de Estados

Las redes neuronales se modelan como sistemas dinámicos estocásticos:

• Estado latente ($x_t$): Representaciones internas de la red.
• Parámetros ($\theta_t$): Tratados no como constantes, sino como estados dinámicos que evolucionan en el tiempo ($\theta_{t+1} = \theta_t + \xi_t$).
• Observaciones ($y_t$): Los datos de entrada/salida.
• Ecuación de actualización: En lugar de $\theta_{t+1} = \theta_t - \eta \nabla L$, se utiliza:
  $$\theta_{t+1} = \theta_t + K_t (y_t - \hat{y}_t)$$
  Donde $K_t$ es la ganancia de Kalman y $(y_t - \hat{y}_t)$ es la innovación (error de predicción).

B. Equivalencia Kalman-Gradiente Natural

El artículo demuestra teóricamente que, bajo modelos de observación gaussianos locales, la actualización de Kalman es equivalente a un paso de gradiente natural con precondicionamiento por la matriz de información de Fisher ($F$):

• La covarianza del estado ($P_t$) actúa como la inversa de la información de Fisher ($P_t \approx F^{-1}$).
• La ganancia de Kalman se deriva de las ecuaciones de Riccati, rastreando implícitamente la curvatura del espacio de parámetros sin necesidad de calcular Hessianos explícitos.

C. Corrección de Innovación a Nivel de Activación

Una innovación clave para los Modelos de Lenguaje (LLMs) es la aplicación de correcciones de Kalman no solo a los parámetros, sino a las activaciones ocultas ($h_t$) de los transformadores.

• Las activaciones se tratan como estados latentes corregidos en tiempo real mediante términos de innovación basados en los tokens observados.
• Esto convierte al transformador en un observador no lineal, permitiendo adaptación en línea y corrección de alucinaciones sin modificar los pesos globales.

D. Elevación Koopman (Koopman Lifting)

Para evitar la linealización local repetida (típica del EKF) que falla en sistemas altamente no lineales, el marco utiliza la teoría del operador Koopman.

• Se mapean las dinámicas no lineales a un espacio de características de mayor dimensión donde la evolución es exactamente lineal ($\phi(x_{t+1}) = K \phi(x_t)$).
• Esto permite aplicar filtrado de Kalman exacto en el espacio elevado, eliminando la necesidad de calcular Jacobianos en cada paso y mejorando la estabilidad.

3. Contribuciones Clave

1. Equivalencia Kalman-Gradiente Natural: Prueba formal de que el filtrado recursivo es un descenso de gradiente natural en una variedad riemanniana definida por la métrica de información de Fisher, unificando el control y la geometría de la información.
2. Principio de Filtrado como Entrenamiento: Reinterpretación del aprendizaje como inferencia bayesiana secuencial, unificando la actualización de parámetros y la corrección de estados latentes.
3. Corrección de Innovación en Activaciones: Introducción de actualizaciones de Kalman directamente en los estados ocultos de los transformadores, permitiendo adaptación en tiempo real y robustez ante cambios de distribución.
4. Modelos Mundiales Elevados por Koopman: Uso de operadores Koopman para convertir dinámicas no lineales en lineales en un espacio elevado, permitiendo filtrado exacto sin linealización local.
5. Escalabilidad mediante Estructuras de Covarianza: Propuesta de aproximaciones de covarianza estructurada (bloque-diagonal, baja rango, Kronecker-factored) para hacer viable el método en redes de miles de millones de parámetros, reduciendo la complejidad de $O(p^2)$ a $O(pr)$.

4. Resultados Empíricos

Los experimentos se realizaron en tres regímenes: modelado de memoria secuencial (MNIST secuencial), modelado de lenguaje (WikiText-2) y aprendizaje continuo bajo cambios de distribución.

• Rendimiento Estacionario: En tareas estándar, KWM logra un rendimiento (precisión/perplejidad) comparable a optimizadores de primer orden (Adam) y de segundo orden estructurados (K-FAC).
• Robustez ante Cambios de Distribución: Bajo perturbaciones sintéticas (dropout de tokens, cambios de vocabulario), KWM con corrección de activaciones muestra una estabilidad superior en la perplejidad en comparación con Adam y K-FAC.
• Aprendizaje Continuo: En benchmarks de "olvido catastrófico" (MNIST permutado), KWM reduce el olvido promedio en un 30% respecto a Adam y un 18% respecto a K-FAC, gracias a que la matriz de covarianza actúa como un regularizador que protege los parámetros con alta certeza.
• Estabilidad de Convergencia: La evolución de la covarianza tiende a una solución estacionaria de Riccati, indicando convergencia estable sin necesidad de ajuste fino de la tasa de aprendizaje.

5. Significado e Implicaciones

El trabajo representa un cambio conceptual fundamental en la arquitectura de los sistemas de aprendizaje:

• De la Optimización a la Estimación: Transforma el entrenamiento de redes neuronales de un problema de minimización determinista a un problema de estimación de estado estocástica.
• Incertidumbre Intrínseca: Proporciona cuantificación de incertidumbre y precondicionamiento adaptativo de forma natural, sin heurísticas externas.
• Robustez y Adaptabilidad: Ofrece un mecanismo principista para la adaptación en línea y la mitigación de alucinaciones en LLMs, corrigiendo representaciones internas en tiempo real.
• Fundamentos Teóricos Sólidos: Establece garantías de estabilidad global bajo convexidad fuerte, estabilidad exponencial para Jacobianos variantes en el tiempo (bajo excitación persistente) y estabilidad robusta bajo aproximaciones de baja rango.
• Más allá del EKF Clásico: A diferencia de los intentos de los años 90, este marco integra la geometría de información, la teoría de operadores (Koopman) y arquitecturas modernas de gran escala, ofreciendo una ruta viable hacia sistemas de aprendizaje online, escalables y conscientes de la incertidumbre.

En conclusión, Kalman World Models proponen que las redes neuronales no deben verse solo como aproximadores de funciones estáticos, sino como estimadores dinámicos estocásticos gobernados por la geometría de Riccati, capaces de aprender de manera recursiva y robusta en entornos no estacionarios.