Maglev: Sliding Recurrent Memory

Sintesi Tecnica: Maglev – Sliding Recurrent Memory

Problematica
Le attuali architetture Transformer affrontano un compromesso fondamentale tra capacità di memoria ed efficienza computazionale. La piena auto-attenzione causale abilita una potente ritenzione non lineare e computazione in-context, ma comporta costi che crescono linearmente con la lunghezza del contesto. Al contrario, l'attenzione a finestra scorrevole (sliding-window attention) limita i costi di attenzione e della cache KV, ma scarta le informazioni distanti, fallendo nel fornire una memoria persistente a lungo termine. I modelli ricorrenti classici (es. LSTM) comprimono la storia ma propagano stati ridotti in modo sequenziale, limitandone la capacità e il parallelismo di sequenza. Sebbene i Memory Transformer e i modelli di attenzione lineare tentino di colmare questo divario, essi spesso impongono confini di segmento artificiali, si affidano ad aggiornamenti lineari strutturati o richiedono lo srotolamento (unrolling) sequenziale durante l'addestramento. La sfida centrale è sviluppare un meccanismo di memoria ricorrente non lineare, basato sui token, per i Transformer che permetta un addestramento in parallelo mantenendo costi di inferenza limitati.

Metodologia: Architettura Maglev
Maglev affronta questo problema separando la costruzione della memoria durante l'addestramento dalla propagazione della memoria durante l'inferenza. L'architettura consiste in due modelli accoppiati:

1. Prefiller ($Q$): Un Transformer causale che elabora la sequenza osservata in parallelo utilizzando un campo recettivo più ampio (ad esempio, attenzione intercalata tra full e sliding-window). Genera una traiettoria di vettori di memoria target, $m'_t$, per ogni posizione di token.
2. Decoder ($P$): Un Transformer a finestra scorrevole che predice il token successivo e genera i propri vettori di memoria, $m_t$. Fondamentalmente, $P$ non ha accesso alla storia completa; vede solo una finestra limitata di token recenti e una finestra traslata di vettori di memoria.

Meccanismo di Addestramento: Lifted Parallel Training
Invece di srotolare lo stato ricorrente sequenzialmente (il che impedisce il parallelismo), Maglev impiega una strategia di addestramento a due passaggi:

• Passaggio 1: Il prefiller $Q$ genera l'intera traiettoria di memoria target $m'_{1:T}$ in parallelo.
• Passaggio 2: Il decoder $P$ consuma la traiettoria target traslata $m'_{0:T-1}$ insieme ai token di input per predire il token successivo e produrre la propria memoria $m_{1:T}$.
• Obiettivo: Il modello è ottimizzato utilizzando una funzione di perdita composita:
  • Perdita di Predizione del Prossimo Token ($L_{CE}$): Standard cross-entropy loss per il linguaggio.
  • Perdita di Consistenza della Memoria ($L_{cons}$): Un termine che allinea la memoria in uscita del decoder $m_t$ con il target del prefiller $m'_t$. Questo insegna al decoder a riprodurre lo stato di memoria necessario per il passo successivo senza lo srotolamento sequenziale.

Meccanismo di Inferenza
All'inferenza, il prefiller $Q$ viene scartato. Il decoder $P$ opera in modo ricorrente, rialimentando nel sistema le proprie memorie precedentemente generate $m_{t-1}$.

• Iniezione Ricorrente di K/V: Invece di utilizzare token di memoria dedicati, Maglev inietta la memoria traslata nelle caratteristiche Key/Value (K/V) del meccanismo di attenzione tramite proiezioni apprese e meccanismi di gating.
• Complessità Limitata: Il decoder mantiene una finestra scorrevole di dimensione fissa di token e memorie. Lo stato ricorrente è trasportato in avanti attraverso la cache K/V, garantendo che i costi di inferenza (calcolo e dimensione della cache) rimangano limitati e indipendenti dalla lunghezza totale della sequenza, corrispondendo al profilo della standard sliding-window attention.

Contributi Chiave

1. Ricorrenza Non Lineare per Token: Maglev introduce un meccanismo in cui lo stato del decoder di ogni token funge da memoria, attraversando l'intera profondità non lineare del Transformer senza richiedere token di memoria extra o confini di segmento artificiali.
2. Addestramento in Parallelo per Modelli Ricorrenti: Utilizzando un prefiller per generare i target in parallelo, Maglev consente l'addestramento in parallelo di modelli ricorrenti non lineari, permettendo il throughput standard delle sequenze.
3. Condivisione dei Parametri: Gli autori dimostrano che il prefiller $Q$ e il decoder $P$ possono condividere la maggior parte dei loro parametri (specificamente i blocchi Transformer) senza degradare significativamente le prestazioni, riducendo l'overhead di memoria dei parametri del setup di addestramento.
4. Inferenza Limitata: Il modello distribuito mantiene l'efficienza della sliding-window attention acquisendo al contempo capacità di memoria persistente attraverso l'iniezione ricorrente di K/V.

Risultati Sperimentali
Valutato sull'architettura nanochat d20 (435M di parametri, addestrato su 43.52B di token), Maglev è stato confrontato con Transformer a finestra intercalata (full/sliding-window), modelli puramente a finestra scorrevole e baseline di Latent Recurrent Transformer (LRT).

• Prestazioni: La migliore variante di Maglev (parametri separati, $\lambda=1$) ha migliorato i bits-per-byte (BPB) di validazione FineWeb-Edu da 0.7413 (baseline sliding-window) a 0.7251.
• Task a Valle: Ha raggiunto un'accuratezza media nei task a valle del 56.4%, superando sia il Transformer a finestra scorrevole corrispondente (54.1%) che la baseline LRT corrispondente.
• Peso di Consistenza: Lo studio ha notato che il peso ottimale della consistenza ($\lambda$) dipende dalla strategia di condivisione dei parametri; un $\lambda$ più alto ha giovato al modello a parametri separati ma ha potuto vincolare il modello a parametri condivisi.

Significatività e Rivendicazioni
Il paper sostiene che Maglev fornisce una via pratica per equipaggiare i Transformer non lineari con una memoria persistente per token che può essere riscritta dopo ogni token. Esso riesce a colmare il divario tra l'espressività della full attention e l'efficienza dei modelli a contesto limitato. Separando la generazione della memoria dall'uso della memoria durante l'addestramento, Maglev ottiene l'addestramento in parallelo per architetture ricorrenti, un traguardo precedentemente difficile per la ricorrenza non lineare. Gli autori pongono Maglev come un'indagine preliminare in questo spazio di progettazione, notando che, sebbene migliori rispetto alle baseline di sliding-window e latent recurrent testate, sono necessari ulteriori lavori per scalare l'approccio, ottimizzare l'efficienza dei kernel per il deployment ed esplorare meccanismi alternativi di iniezione della memoria (ad esempio, iniezione nel residual-stream o cross-attention).