Riassunto Tecnico: Kalman Linear Attention (KLA)

Definizione del Problema

Gli attuali modelli a spazio di stato (SSM) come Mamba e la Gated Linear Attention (GLA) offrono alternative a complessità lineare e parallelizzabili rispetto alla softmax attention, affrontando il collo di bottiglia della scalabilità quadratica dei Transformer. Tuttavia, questi modelli si affidano ad aggiornamenti dello stato nascosto lineari o affini, il che limita la loro espressività rispetto alle interazioni non lineari indotte dalla normalizzazione softmax. Inoltre, gli esistenti RNN lineari e SSM mancano di meccanismi espliciti per rappresentare l'incertezza dello stato, trattando gli stati nascosti come deterministici. Questo articolo affronta il divario tra l'efficienza degli aggiornamenti lineari e l'espressività/robustezza del filtraggio non lineare, chiedendosi se gli aggiornamenti esatti del filtro di Kalman possano essere implementati efficientemente come un mixer di sequenze.

Metodologia: Kalman Linear Attention (KLA)

Gli autori propongono la Kalman Linear Attention (KLA), un mixer di sequenze probabilistico che formula la modellazione delle sequenze come inferenza bayesiana esatta. L'innovazione principale risiede nella riparametrizzazione del filtro di Kalman in forma di informazione (utilizzando le precisioni e i parametri naturali) per consentire il calcolo parallelo.

1. Formulazione Probabilistica

KLA modella il mixing di sequenze come inferenza in un modello spazio-stato lineare-gaussiano:

• Prior della Dinamica: L'evoluzione dello stato latente è modellata come un processo di Ornstein-Uhlenbeck (OU), discretizzato esattamente. Ciò introduce un termine di rumore di processo ($\bar{p}_t$) che cattura l'incertezza nelle transizioni di stato, distinguendosi dalle dinamiche deterministiche degli standard SSM.
• Modello di Osservazione: I token di input sono trattati come osservazioni rumorose dello stato latente, parametrizzate da un operatore di osservazione ($k_t$) e una precisione del valore ($\Lambda^v_t$) che rappresenta la fiducia nell'evidenza del token.
• Readout: L'output è derivato dalla media a posteriori dello stato latente, proiettata tramite un operatore di query ($q_t$).

2. Parallelizzazione tramite Trasformazioni di Möbius

Il filtraggio di Kalman classico è intrinsecamente sequenziale. KLA supera questo limite sfruttando la forma di informazione del filtro:

• Ricorsione della Precisione: Si dimostra che l'aggiornamento della precisione a posteriori ($\lambda_t$) è una trasformazione di Möbius (frazionaria-lineare) della precisione precedente ($\lambda_{t-1}$).
  $$\lambda_t = \frac{\alpha_t \lambda_{t-1} + \beta_t}{\gamma_t \lambda_{t-1} + \delta_t}$$
  Fondamentalmente, le trasformazioni di Möbius si compongono in modo associativo. Ciò consente agli aggiornamenti della precisione di essere calcolati tramite scansioni prefisso parallele (scansioni associative) con un lavoro di $O(T)$ e una profondità di $O(\log T)$, eguagliando l'efficienza parallela dei moderni SSM come Mamba.
• Ricorsione della Media: Una volta calcolato in parallelo il percorso della precisione, gli aggiornamenti della media a posteriori (media dell'informazione) seguono una trasformazione affine che è anch'essa parallelizzabile.

3. Gating Non Lineare dalla Incertezza

A differenza degli SSM lineari dove il gating è spesso lineare o dipendente dal token ma deterministico, il gating di KLA emerge naturalmente dal rapporto tra le precisioni. Il "gate di oblio" (forget gate) e il "gate di scrittura" (write gate) sono dipendenti dalla storia e non lineari, guidati dall'incertezza relativa tra il prior dello stato e la nuova osservazione. Ciò consente al modello di incorporare selettivamente l'informazione in base alla fiducia senza richiedere parametri di gating espliciti e appresi per ogni token.

4. Architettura

KLA è progettato come un sostituto "drop-in" per gli standard strati SSM o di attenzione. Segue un'architettura a MLP fuso simile a Mamba, incorporando convoluzioni causali 1D e attivazioni SiLU. Gli autori implementano kernel CUDA e Triton efficienti per lo scan di Möbius per evitare la materializzazione di stati espansi in memoria ad alta larghezza di banda (HBM).

Contributi Chiave

1. Riparametrizzazione Associativa: Gli autori dimostrano che il filtraggio Kalman lineare-gaussiano diagonale in forma di informazione ammette una struttura di Möbius, consentendo scansioni prefisso parallele esatte per il filtraggio bayesiano.
2. Non Linearità Guidata dall'Incertezza: KLA introduce un gating non lineare derivato dai rapporti di precisione, offrendo un'espressività strettamente superiore rispetto agli aggiornamenti SSM lineari/affini pur mantenendo una struttura di scan a tempo lineare.
3. Layer KLA: Un mixer di sequenze pratico e "drop-in" che produce un'incertezza esplicita dello stato di credenza e può essere impilato in reti profonde.
4. Scalabilità e Validazione: Il documento presenta i primi primitivi di filtraggio bayesiano impilati, addestrati alla scala di miliardi di token, validando la loro stabilità e prestazioni.

Risultati Sperimentali

Gli autori valutano KLA attraverso compiti sintetici, benchmark di tracciamento dello stato e modellazione del linguaggio su larga scala.

• Compiti Sintetici (MAD-Lab): KLA supera GDN, GLA, Mamba e mLSTM in compiti che richiedono la selezione sotto incertezza, specificamente Fuzzy Recall e Selective Copy. Raggiunge risultati state-of-the-art in Compression, suggerendo che il suo filtraggio in forma di informazione agisce come un compressore di contesto ottimale.
• Tracciamento dello Stato (A5 Permutation): Nel compito di composizione di permutazioni A5 (un benchmark per il tracciamento dello stato NC1-completo), KLA risolve il compito con profondità costante (1–2 blocchi). Al contrario, i modelli lineari SSM e quelli basati su attenzione richiedono una profondità che cresce con la lunghezza della sequenza. Ciò dimosta che gli aggiornamenti frazionari-lineari di KLA rompono i limiti di espressività dei modelli ricorrenti lineari pur rimanendo parallelizzabili in fase di addestramento.
• Associazione di Ricordo a Lungo Contesto (MQAR): KLA raggiunge un'accuratezza quasi perfetta (>95%) sul compito di Multi-Query Associative Recall con lunghezza della sequenza 2048, superando Mamba e superando significativamente GLA, che fallisce nell'apprendere il compito in questo setting.
• Preaddestramento su Larga Scala:
  • Standalone: A 45M e 180M di parametri, KLA eguaglia o migliora GPT, Mamba e GDN su benchmark di buon senso zero-shot (es. HellaSwag, PIQA).
  • Ibrido: Sostituire l'ultimo strato di attenzione di un GPT con KLA (GPT+KLA) produce le prestazioni più forti tra gli ibridi, superando GPT+Mamba e GPT+GDN.
• Ablation: Gli autori mostrano che il prior di Ornstein-Uhlenbeck (OU) è critico per la stabilità dell'impilamento profondo. Senza la dinamica stocastica di ritorno alla media, l'addestramento di stack profondi di KLA diventa instabile.

Significato e Rivendicazioni

Il documento afferma che KLA riesce a colmare il divario tra l'efficienza dei modelli a spazio di stato lineari e l'espressività dei meccanismi di attenzione non lineari. Inquadrando il mixing di sequenze come filtraggio bayesiano esatto, KLA fornisce:

• Incertezza Esplicita: A differenza dei mixer deterministici, KLA trasporta un'incertezza esplicita dello stato di credenza che controlla direttamente il gating dell'informazione.
• Espressività Strettamente Superiore: Gli aggiornamenti di Möbius permettono a KLA di risolvere compiti di tracciamento dello stato (come A5) che gli SSM lineari e l'attenzione non possono risolvere, senza sacrificare le capacità di addestramento parallelo.
• Scalabilità: Il lavoro dimostra che i primitivi di filtraggio bayesiano esatto possono essere impilati e addestrati alla scala di miliardi di token, sfidando l'idea che il filtraggio probabilistico sia intrinsecamente sequenziale o computazionalmente proibitivo per la modellazione del linguaggio su larga scala.

Gli autori posizionano KLA non solo come una curiosità teorica, ma come un primitivo pratico e "drop-in" che migliora la robustezza e le capacità di tracciamento dello stato dei moderni modelli di linguaggio, mantenendo l'efficienza computazionale necessaria per l'elaborazione di contesti lunghi.