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

Titolo: Dai Gradienti alla Geometria di Riccati: Modelli del Mondo di Kalman per l'Apprendimento in Singolo Passaggio

1. Il Problema

L'apprendimento automatico moderno è dominato dalla retropropagazione dell'errore (Backpropagation), che utilizza la differenziazione automatica in modalità inversa per ottimizzare i parametri tramite discesa del gradiente stocastica (SGD) e sue varianti. Sebbene efficace, questo paradigma presenta limitazioni strutturali e computazionali critiche:

• Vincoli di Memoria: Richiede la conservazione di tutte le attivazioni intermedie durante il passaggio in avanti per calcolare i gradienti nel passaggio inverso, con un costo di memoria proporzionale alla profondità della rete e alla lunghezza della sequenza.
• Staticità: Tratta i parametri come quantità fisse da ottimizzare su un obiettivo statico, rendendo difficile l'adattamento online e continuo (continual learning) senza meccanismi euristici aggiuntivi.
• Mancanza di Quantificazione dell'Incertezza: I metodi del primo ordine non forniscono una quantificazione intrinseca dell'incertezza sui parametri o sulle rappresentazioni latenti.
• Sincronizzazione: In ambienti distribuiti, la sincronizzazione dei gradienti introduce costi di comunicazione e potenziali instabilità.

Il paper propone un cambio di paradigma: sostituire l'ottimizzazione basata sui gradienti con l'inferenza bayesiana ricorsiva (filtraggio), trattando l'apprendimento come un problema di stima dello stato in sistemi dinamici stocastici.

2. Metodologia: Kalman World Models (KWM)

L'autore introduce i Kalman World Models (KWM), una classe di modelli che sostituisce gli aggiornamenti del gradiente con adattamenti del guadagno di Kalman.

A. Riformulazione come Sistema Dinamico

Invece di vedere le reti neurali come approssimatori di funzioni statiche, KWM le modella come sistemi dinamici stocastici:

• Stato Aumentato: I parametri del modello $\theta_t$ e gli stati latenti interni $x_t$ sono trattati come componenti di uno stato aumentato $z_t = [x_t, \theta_t]^T$.
• Dinamica: L'evoluzione del sistema è governata da equazioni di stato e osservazione:
  $$x_{t+1} = f_\theta(x_t, u_t) + w_t$$
  $$y_t = h_\theta(x_t) + v_t$$
  dove $w_t$ e $v_t$ sono rumore di processo e di osservazione.
• Apprendimento come Filtraggio: L'aggiornamento dei parametri non è una minimizzazione di una funzione di perdita, ma una correzione dello stato basata sull'innovazione (la differenza tra l'osservazione reale $y_t$ e la previsione $\hat{y}_t$):
  $$\theta_{t+1} = \theta_t + K_t(y_t - \hat{y}_t)$$
  Qui, $K_t$ è il guadagno di Kalman, calcolato dinamicamente dalla propagazione della covarianza (equazione di Riccati), non un tasso di apprendimento euristico.

B. Equivalenza con il Gradiente Naturale

Il paper dimostra teoricamente che, sotto modelli di osservazione gaussiani localmente lineari, l'aggiornamento del guadagno di Kalman è equivalente a un passo di gradiente naturale con precondizionamento tramite l'informazione di Fisher:

• La matrice di covarianza $P_t$ agisce come l'inverso della matrice di informazione di Fisher ($F^{-1}$).
• L'aggiornamento diventa: $\theta_{t+1} = \theta_t + F^{-1} \nabla_\theta \ell_t$.
• Questo unifica il filtraggio ricorsivo, l'ottimizzazione di secondo ordine e la geometria dell'informazione in un unico quadro matematico.

C. Correzione dell'Innovazione a Livello di Attivazione

Un contributo innovativo è l'applicazione del filtraggio non solo ai parametri, ma anche alle attivazioni interne (es. negli stati nascosti dei Transformer).

• Le attivazioni $h_t$ sono trattate come stati latenti soggetti a correzione tramite termini di innovazione basati sui token osservati.
• Questo trasforma i LLM in osservatori non lineari, permettendo un adattamento online e una correzione delle rappresentazioni interne senza modificare i pesi globali, mitigando così le allucinazioni e gli shift di distribuzione.

D. Sollevamento Koopman (Koopman Lifting)

Per evitare le approssimazioni di linearizzazione locale (tipiche dei filtri EKF classici che falliscono in regimi fortemente non lineari), il paper introduce il sollevamento Koopman.

• Le dinamiche non lineari vengono mappate in uno spazio di funzioni di dimensione superiore (spazio delle osservabili) dove l'evoluzione è esattamente lineare.
• Il filtraggio di Kalman viene eseguito in questo spazio sollevato, rendendo l'inferenza esatta (a meno dell'errore di proiezione del sollevamento) ed eliminando la necessità di ricalcolare i Jacobiani ad ogni passo.

3. Contributi Chiave

1. Equivalenza Kalman-Gradiente Naturale: Dimostrazione formale che il filtraggio di Kalman implementa un gradiente naturale su una varietà riemanniana definita dalla metrica di Fisher, aggiornata dinamicamente.
2. Principio "Filtraggio come Addestramento": Sostituzione della retropropagazione con l'inferenza bayesiana sequenziale, eliminando la necessità di memorizzare le attivazioni per il passaggio inverso.
3. Correzione delle Attivazioni: Introduzione di aggiornamenti di innovazione direttamente sugli stati nascosti dei Transformer, permettendo adattamento contestuale e robustezza.
4. Modelli del Mondo Sollevati da Koopman: Integrazione della teoria degli operatori Koopman per gestire le non linearità globali senza linearizzazione locale, trasformando le reti neurali in sistemi dinamici lineari in coordinate latenti apprese.
5. Analisi di Stabilità e Scalabilità: Derivazione di condizioni di stabilità (es. osservabilità uniforme, persistenza di eccitazione) e proposte di approssimazioni di covarianza strutturata (bassa rango, fattorizzata a Kronecker) per scalare a modelli con miliardi di parametri.

4. Risultati Sperimentali

Gli esperimenti confrontano KWM con ottimizzatori del primo ordine (SGD, Adam) e di secondo ordine strutturato (K-FAC, EKF classico):

• Sequential MNIST: KWM mostra un tasso di convergenza comparabile ad Adam ma con una sensibilità ridotta all'inizializzazione del tasso di apprendimento e una maggiore stabilità.
• Modellazione Linguistica (WikiText-2): KWM raggiunge una perplessità finale comparabile ad Adam. Tuttavia, sotto shift di distribuzione (es. dropout di token, perturbazioni del vocabolario), KWM con correzione delle attivazioni mostra una stabilità significativamente superiore.
• Apprendimento Continuo (Continual Learning): Su benchmark come Permuted MNIST, KWM riduce l'oblio catastrofico (catastrophic forgetting) del 30% rispetto ad Adam e del 18% rispetto a K-FAC, grazie alla rappresentazione esplicita dell'incertezza nella matrice di covarianza che protegge i parametri consolidati.
• Robustezza: I risultati indicano che KWM è più robusto agli shift di distribuzione rispetto ai metodi basati su gradienti tradizionali.

5. Significato e Implicazioni

Questo lavoro rappresenta un cambiamento concettuale fondamentale nell'architettura dell'apprendimento automatico:

• Oltre la Retropropagazione: Dimostra che l'ottimizzazione basata su gradienti non è l'unico metodo principiato per l'apprendimento; il filtraggio bayesiano offre un'alternativa solida basata sulla teoria del controllo.
• Apprendimento Online e Adattivo: Poiché non richiede un passaggio inverso e mantiene una distribuzione a posteriori sui parametri, KWM è intrinsecamente adatto all'apprendimento online e continuo, adattandosi ai dati man mano che arrivano.
• Quantificazione dell'Incertezza Nativa: L'incertezza non è un'aggiunta post-hoc (come in alcune reti bayesiane), ma è parte integrante della dinamica di aggiornamento attraverso la matrice di covarianza.
• Scalabilità Teorica: Sebbene la covarianza completa sia $O(p^2)$, le approssimazioni a bassa rango e strutturate permettono di scalare a $O(pr)$, rendendo il metodo competitivo con gli ottimizzatori moderni per modelli di grandi dimensioni.

In sintesi, il paper propone di vedere le reti neurali non come semplici approssimatori statici, ma come estimatori dinamici stocastici governati dalla geometria di Riccati, aprendo la strada a sistemi di apprendimento più robusti, adattivi e privi di gradienti.