Riepilogo Tecnico: Sparse Delta Memory (SDM)

Problema

Le Reti Neurali Ricorrenti (RNN) lineari, come Gated DeltaNet (GDN) e Mamba2, offrono il vantaggio di una memoria e di un calcolo fissi per token, il che le rende teoricamente adatte a contesti arbitrariamente lunghi. Tuttavia, le loro prestazioni pratiche sul richiamo a lungo contesto e sull'apprendimento in-context sono limitate dalle loro piccole dimensioni dello stato nascosto. Aumentare la dimensione dello stato nelle standard RNN lineari porta a un aumento lineare di FLOPs a causa degli aggiornamenti densi dello stato (moltiplicazioni matrice-vettore), creando un collo di bottiglia computazionale proibitivo. Al contrario, le architetture Transformer standard con attenzione softmax mantengono un alto richiamo, ma soffrono di costi di memoria e di calcolo della cache Key-Value (KV) che crescono linearmente con l'aumentare della lunghezza della sequenza.

La sfida centrale affrontata è come scalare la capacità dello stato nascosto delle RNN lineari di ordini di grandezza per migliorare le prestazioni su contesti lunghi senza aumentare il costo computazionale (FLOPs) per token.

Metodologia: Sparse Delta Memory (SDM)

Gli autori introducono Sparse Delta Memory (SDM), un'architettura che estende la Gated DeltaNet sostituendo il prodotto esterno denso chiave-valore con uno schema di indirizzamento sparso basato su Product-Key Memory (PKM).

Meccanismo Core

1. Memoria Sparsa Esplicita: Invece di mantenere una matrice densa $M_t \in \mathbb{R}^{d_{qk} \times d_v}$, SDM mantiene una tabella di memoria esplicita $M_t \in \mathbb{R}^{N \times d_v}$ con $N$ slot, dove $N$ può essere ordini di grandezza superiore alla dimensione nascosta (ad esempio, $N \approx 10^6$).
2. Selezione Sparsa delle Chiavi: Ad ogni timestep, le chiavi proiettate ($k'$) e le query ($q'$) vengono divise in due metà. La loro somma esterna genera $N$ punteggi per gli slot di memoria. Utilizzando una strategia di selezione top-$W$ (scrittura) e top-$R$ (lettura), il modello identifica specifici slot di memoria da aggiornare o da leggere. Questa selezione viene eseguita efficientemente senza materializzare l'intera matrice di punteggio di dimensione $N$, sfruttando la proprietà $\text{topk}(s_1 \oplus s_2) = \text{topk}(\text{topk}(s_1) \oplus \text{topk}(s_2))$.
3. Aggiornamenti Gated Delta: Per i $W$ slot di scrittura selezionati, il modello applica la regola di aggiornamento Gated DeltaNet:
   • Gate di Oblio (Forget Gate): $\tilde{M}_t[i] = \alpha_t \cdot M_{t-1}[i]$
   • Aggiornamento Delta: $M_t[i] = \tilde{M}_t[i] + \beta_t \cdot k^{(i)}_t \cdot (v_t - \tilde{M}_t[i])$
   • Gli slot non selezionati rimangono invariati.
4. Lettura Sparsa: L'output è una somma pesata dei $R$ slot di lettura selezionati.

Design IsoFLOP

Per garantire un confronto equo, gli autori progettano SDM per corrispondere al numero di parametri e ai FLOPs di una baseline GDN densa:

• FLOPs: Il costo per token è $O((W+R) \times d_v)$, indipendente dalla dimensione della memoria $N$. Impostando $W$ e $R$ uguali alla dimensione di query della baseline GDN, il costo computazionale rimane identico.
• Parametri: Le proiezioni lineari sono dimensionate identicamente alla baseline. La dimensione della memoria $N$ è controllata tramite il numero di teste $H$ e le dimensioni di proiezione, permettendo alla dimensione dello stato di scalare come $O(1/H^2)$ senza influenzare i FLOPs.

Stato Iniziale Appreso

A differenza di GDN, che tipicamente inizializza il suo piccolo stato a zero, SDM tratta lo stato di memoria iniziale $M_0$ come un parametro apprendibile. Ciò consente al modello di memorizzare la conoscenza pre-addestrata direttamente nella struttura della memoria, agendo come una memoria parametrica che persiste durante l'inferenza.

Contributi Chiave

1. Architettura: La proposta di SDM, che scala la capacità dello stato delle RNN lineari di tre ordini di grandezza (ad esempio, da kilobyte a gigabyte) mantenendo costanti i FLOPs per token.
2. Validazione IsoFLOP: Dimostrazione che aumentare la capacità dello stato da solo (senza aumentare il calcolo) migliora significativamente le prestazioni sui compiti a lungo contesto.
3. Memoria Parametrica: L'introduzione di uno stato iniziale $M_0$ appreso per le RNN lineari, mostrando come questo meccanismo memorizzi efficacemente la conoscenza comune e migliori i compiti di ragionamento.
4. Leggi di Scaling: Validazione empirica che SDM segue leggi di scaling prevedibili, superando GDN denso e, a scale maggiori (8B parametri), superando i modelli Full Attention in termini di perdita di addestramento.

Risultati Sperimentali

Gli autori hanno addestrato modelli da 1.4B a 8B parametri su oltre 1 trilione di token.

• Leggi di Scaling: Su una scala di scaling, SDM ha costantemente ottenuto una perdita di addestramento inferiore rispetto a GDN a tutti i livelli di calcolo. Alla scala 8B, SDM ha raggiunto una perdita inferiore sia a GDN che a un modello Full Attention.
• Richiamo a Lungo Contesto (Benchmark RULER): SDM ha superato significativamente GDN nei compiti di recupero a lungo contesto. Alle scale 1.4B e 8B, SDM ha ottenuto punteggi rispettivamente di 31.2 e 50.2, rispetto ai 20.0 e 34.2 di GDN. Notevolmente, SDM ha eguagliato o superato le prestazioni di Full Attention in diversi sottotesti di RULER nonostante l'uso di una rappresentazione di memoria fissa.
• Contesto Breve e Ragionamento: SDM ha ottenuto una perplessità inferiore sui dati di codice e un'accuratezza media più alta nei benchmark di ragionamento (es. MMLU, GSM8K, HumanEval) rispetto a GDN. Alla scala 8B, SDM ha superato Full Attention in accuratezza media.
• Studi di Ablazione:
  • Dimensione dello Stato: Le prestazioni sono degradate monotonicamente al ridursi della dimensione della memoria, confermando che il grande stato è il driver principale dei guadagni.
  • Inizializzazione Appresa: L'apprendimento di $M_0$ ha migliorato le prestazioni per SDM ma ha fornito guadagni trascurabili per GDN, suggerendo che la grande capacità è necessaria per utilizzare efficacemente l'inizializzazione appresa.
  • Efficienza: Sebbene il throughput di addestramento di SDM sia stato 1.49x più lento di GDN alla scala 8B (a causa dell'accesso alla memoria HBM rispetto a SRAM), è stato 6x più veloce di Full Attention durante l'inferenza.

Significato e Rivendicazioni

L'articolo sostiene che SDM colmi con successo il divario tra l'efficienza delle RNN lineari e le capacità a lungo contesto dei Transformer. Decoppiando la capacità di memoria dal costo computazionale attraverso la sparsità, SDM consente:

• Spazio/Calcolo Costante: Elaborare sequenze arbitrariamente lunghe (fino a 1 milione di token dimostrati) senza la crescita illimitata della memoria della cache KV.
• Prestazioni a Lungo Contesto Superiori: Raggiungere risultati allo stato dell'arte tra le RNN lineari e risultati competitivi contro Full Attention nel recupero a lungo contesto e nell'apprendimento in-context.
• Memorizzazione della Conoscenza Parametrica: La capacità di apprendere e memorizzare la conoscenza pre-addestrata all'interno dello stato di memoria stesso, migliorando il ragionamento generale e i compiti di conoscenza comune.

Gli autori riconoscono i limiti, specificamente che l'ampia impronta di memoria (comparabile ai parametri del modello) richiede risorse HBM significative, e che le attuali implementazioni di kernel sono meno efficienti dei kernel densi ottimizzati di GDN. Tuttavia, essi pongono che SDM apra nuove possibilità per lo sviluppo di agenti con una migliore memoria a lungo termine e per l'elaborazione di video lunghi o tracce di ragionamento estese.