Riepilogo Tecnico: Gated DeltaNet-2

Enunciato del Problema

I meccanismi di attenzione lineare offrono un'alternativa convincente all'attenzione softmax standard sostituendo la cache illimitata con uno stato ricorrente a dimensione fissa, riducendo così la mescolanza delle sequenze a tempo lineare e la decodifica a memoria costante. Tuttavia, un collo di bottiglia fondamentale rimane nei modelli basati sulla regola delta (come DeltaNet, Gated DeltaNet e Kimi Delta Attention/KDA). Sebbene questi modelli migliorino la gestione della memoria consentendo sovrascritture mirate delle associazioni, fanno affidamento su una porta scalare ($\beta_t$) per controllare simultaneamente due operazioni distinte:

1. Cancellazione: Decidere quanto del contenuto vecchio (letto dal lato chiave) deve essere rimosso.
2. Scrittura: Decidere quanto del nuovo contenuto (dal lato valore) deve essere impegnato.

Il documento sostiene che trattare queste operazioni come un'unica decisione scalare costituisca una restrizione di modellazione. La cancellazione agisce sull'asse della chiave (determinando quali coordinate della vecchia lettura rimuovere), mentre la scrittura agisce sull'asse del valore (determinando quali coordinate del valore in arrivo impegnare). L'accoppiamento di queste decisioni limita la capacità del modello di modificare selettivamente la memoria compressa senza disturbare le associazioni esistenti, in particolare negli scenari di contesto lungo dove molte associazioni competono per lo stesso spazio di stato finito.

Metodologia: Gated DeltaNet-2

Gli autori propongono Gated DeltaNet-2, un livello di attenzione ricorrente che disaccoppia le operazioni di cancellazione e scrittura all'interno della regola delta.

1. Regola Delta-2 Portata

L'innovazione centrale è la separazione della porta scalare in due porte indipendenti, canale per canale:

• Porta di Cancellazione ($b_t \in [0, 1]^{d_k}$): Un vettore applicato canale per canale alla dimensione della chiave. Pesa le coordinate della lettura corrente ($e_t = b_t \odot k_t$) per determinare quali parti dello stato vecchio vengono cancellate.
• Porta di Scrittura ($w_t \in [0, 1]^{d_v}$): Un vettore applicato canale per canale alla dimensione del valore. Pesa le coordinate del valore in arrivo ($z_t = w_t \odot v_t$) per determinare quali parti del nuovo valore vengono impegnate.

L'equazione di aggiornamento dello stato diventa:
$$S_t = (I - k_t(b_t \odot k_t)^\top) D_t S_{t-1} + k_t(w_t \odot v_t)^\top$$
Dove $D_t = \text{Diag}(\alpha_t)$ rappresenta il decadimento canale per canale ereditato da KDA.

Questa formulazione generalizza i modelli precedenti:

• Se $b_t$ e $w_t$ sono vincolati allo stesso scalare $\beta_t$, si recupera KDA.
• Se anche il decadimento $\alpha_t$ è vincolato a uno scalare, si recupera Gated DeltaNet.

2. Prospettiva di Aggiornamento dei Pesi Veloci

Gli autori interpretano l'aggiornamento come un passo di ottimizzazione online su una perdita di regressione locale. Il modello minimizza la distanza tra il nuovo stato e una memoria decaduta, soggetto a una modifica associativa che corregge il valore letto ($\bar{S}_{t-1}^\top e_t$) verso il target portato ($z_t$). Questa visione chiarisce che l'aggiornamento sottrae la vecchia lettura portata prima di aggiungere il nuovo valore portato.

3. Addestramento Parallelo Efficiente a Blocchi

Per mantenere l'efficienza dell'addestramento, gli autori derivano un algoritmo WY (simile a KDA) che assorbe il decadimento cumulativo canale per canale nei fattori di cancellazione di rango uno.

• Assorbimento del Decadimento: Il decadimento canale per canale è normalizzato nelle matrici di cancellazione ($\bar{E}$) e scrittura ($Z$), permettendo alla ricorrenza di essere espressa come un puro aggiornamento delta asimmetrico.
• Risoluzione Triangolare: Il calcolo a blocchi comporta la risoluzione di un sistema triangolare inferiore $(I+T)^{-1}$ per calcolare le matrici ausiliarie $Y$ (lato cancellazione) e $U$ (lato scrittura).
• Passo Inverso Consapevole delle Porte: Un contributo tecnico critico è la derivazione del passo inverso. A differenza dei modelli con porte scalari, dove il fattore porta può essere estratto dai prodotti scalari, Gated DeltaNet-2 richiede che le porte canale per canale ($b_t$ e $w_t$) siano incorporate direttamente nell'accumulo dei gradienti per le matrici ausiliarie. Ciò garantisce un corretto flusso dei gradienti senza rompere la struttura efficiente a blocchi.

4. Architettura

Il modello è implementato come un mixer di token in un blocco standard in stile Transformer. Supporta configurazioni sia solo ricorrenti che ibride. La variante ibrida intercala il mixer Gated DeltaNet-2 con l'Attenzione a Finestra Scorrevole (SWA) per gestire le interazioni locali esatte, mentre il livello ricorrente gestisce le dipendenze a lungo raggio.

Risultati Chiave

Il modello è stato valutato con 1,3 miliardi di parametri addestrati su 100 miliardi di token FineWeb-Edu, confrontato con Mamba-2, Gated DeltaNet, KDA, Mamba-3 (SISO/MIMO) e Transformer.

• Modellazione del Linguaggio e Ragionamento: Gated DeltaNet-2 ha ottenuto i risultati complessivi più forti sia nelle impostazioni ricorrenti che ibride, sia per la perplessità (WikiText, LAMBADA) sia per le valutazioni di ragionamento di senso comune a zero-shot (PIQA, HellaSwag, ecc.).
• Recupero in Contesto Lungo (RULER): I guadagni più significativi sono stati osservati nei compiti Multi-Key Needle-In-A-Haystack (MK-NIAH). In questi scenari, dove uno stato a dimensione fissa deve separare associazioni in competizione, Gated DeltaNet-2 ha superato tutti i baseline, in particolare all'aumentare della lunghezza del contesto (fino a 8K token).
• Recupero Reale: Su compiti che coinvolgono estrazione e risposta a domande (SQuAD, TriviaQA, DROP), il modello ibrido Gated DeltaNet-2 ha raggiunto la massima accuratezza media, dimostrando robustezza in contesti reali rumorosi.
• Efficienza: L'analisi del throughput su GPU H100 ha mostrato che Gated DeltaNet-2 mantiene una scalabilità quasi piatta all'aumentare della lunghezza della sequenza, con un sovraccarico costante modesto rispetto a KDA dovuto alle porte aggiuntive canale per canale.

Significato e Affermazioni

Il documento afferma che disaccoppiare la modifica attiva della memoria colpisce direttamente il principale punto critico della ricorrenza a stato fisso: l'interferenza tra molte associazioni compresse.

• Flessibilità di Modellazione: Separando cancellazione e scrittura, il modello può cancellare un ampio contesto tramite decadimento, rimuovere associazioni obsolete specifiche tramite la porta di cancellazione e inserire solo i canali di valore necessari tramite la porta di scrittura. Questa flessibilità permette al modello di recuperare KDA e Gated DeltaNet come casi speciali, imparando al contempo strutture canale più complesse quando necessario.
• Validazione Empirica: Gli studi di ablazione confermano che entrambe le porte contribuiscono alle prestazioni, con la porta di cancellazione canale per canale ($b_t$) che rappresenta la maggior parte dei guadagni nella modellazione del linguaggio e nel recupero.
• Praticità: Il lavoro dimostra che meccanismi avanzati di controllo della memoria possono essere integrati nell'attenzione lineare senza sacrificare l'efficienza dell'addestramento parallelo richiesta per l'addestramento su larga scala.

Gli autori concludono che Gated DeltaNet-2 rappresenta un passo avanti nella "frontiera dell'attenzione ricorrente", offrendo un meccanismo più espressivo ed efficiente per la gestione della memoria compressa negli scenari di contesto lungo.