Technische Samenvatting: Sparse Delta Memory (SDM)

Probleemstelling

Lineaire Recurrente Neurale Netwerken (RNN's), zoals Gated DeltaNet (GDN) en Mamba2, bieden het voordeel van een vaste geheugenomvang en rekenkracht per token, wat hen theoretisch geschikt maakt voor willekeurig lange contexten. Echter, hun praktische prestaties op het gebied van lang-context herinnering (recall) en in-context leren worden beperkt door hun kleine verborgen toestandsgrootte (hidden state size). Het vergroten van de toestandsgrootte in standaard lineaire RNN's leidt tot een lineaire toename van de FLOPs vanwege dichte state-updates (matrix-vector vermenigvuldigingen), wat een verboden computationele bottleneck creëert. Daartegenover staan standaard Transformer-architecturen met softmax attention, die een hoge recall behouden maar lijden onder lineair groeiende Key-Value (KV) cache geheugen- en rekenkosten naarmate de sequentielengte toeneemt.

De kernuitdaging die wordt aangepakt, is hoe de capaciteit van de verborgen toestand van lineaire RNN's met ordes van grootte kan worden vergroot om de lang-context prestaties te verbeteren, zonder de computationele kosten (FLOPs) per token te verhogen.

Methodologie: Sparse Delta Memory (SDM)

De auteurs introduceren Sparse Delta Memory (SDM), een architectuur die de Gated DeltaNet uitbreidt door de dichte key-value outer product te vervangen door een sparse adresseringsschema gebaseerd op Product-Key Memory (PKM).

Kernmechanisme

1. Expliciet Sparse Geheugen: In plaats van een dichte matrix $M_t \in \mathbb{R}^{d_{qk} \times d_v}$ bij te houden, houdt SDM een expliciete geheugentabel $M_t \in \mathbb{R}^{N \times d_v}$ bij met $N$ slots, waarbij $N$ ordes van grootte groter kan zijn dan de verborgen dimensie (bijv. $N \approx 10^6$).
2. Sparse Key Selectie: Op elk tijdstip worden geprojecteerde keys ($k'$) en queries ($q'$) in twee helften gesplitst. Hun outer sum genereert $N$ scores voor geheugenslots. Met behulp van een top-$W$ (write) en top-$R$ (read) selectiestrategie identificeert het model specifieke geheugenslots om bij te werken of uit te lezen. Deze selectie wordt efficiënt uitgevoerd zonder de volledige $N$-grote scorematrix te materialiseren, door gebruik te maken van de eigenschap dat $\text{topk}(s_1 \oplus s_2) = \text{topk}(\text{topk}(s_1) \oplus \text{topk}(s_2))$.
3. Gated Delta Updates: Voor de $W$ geselecteerde write slots past het model de Gated DeltaNet update regel toe:
   • Forget Gate: $\tilde{M}_t[i] = \alpha_t \cdot M_{t-1}[i]$
   • Delta Update: $M_t[i] = \tilde{M}_t[i] + \beta_t \cdot k^{(i)}_t \cdot (v_t - \tilde{M}_t[i])$
   • Niet-geselecteerde slots blijven ongewijzigd.
4. Sparse Read: De output is een gewogen som van de $R$ geselecteerde read slots.

IsoFLOP Ontwerp

Om een eerlijke vergelijking te garanderen, ontwerpen de auteurs SDM zodanig dat het overeenkomt met het aantal parameters en de FLOPs van een dichte GDN baseline:

• FLOPs: De kosten per token zijn $O((W+R) \times d_v)$, onafhankelijk van de geheugengrootte $N$. Door $W$ en $R$ gelijk te stellen aan de query-dimensie van de GDN baseline, blijft de computationele kosten identiek.
• Parameters: Lineaire projecties zijn qua omvang gelijk aan de baseline. De geheugengrootte $N$ wordt gecontroleerd via het aantal koppen $H$ en de projectiedimensies, waardoor de state-grootte kan schalen als $O(1/H^2)$ zonder de FLOPs te beïnvloeden.

Geleerde Initiële Toestand

In tegen tegenstelling tot GDN, dat zijn kleine state meestal naar nul initialiseert, behandelt SDM de initiële geheugentoestand $M_0$ als een leerbare parameter. Dit stelt het model in staat om pretraining kennis direct in de geheugenstructuur op te slaan, wat fungeert als een parametrisch geheugen dat tijdens inferentie standhoudt.

Belangrijkste Bijdragen

1. Architectuur: De voorstel van SDM, die de capaciteit van de lineaire RNN state met drie ordes van grootte schaalt (bijv. van kilobytes naar gigabytes) terwijl de FLOXs per token constant blijven.
2. IsoFLOP Validatie: Demonstratie dat het vergroten van de state capaciteit alleen (zonder de rekenkracht te verhogen) de prestaties op lang-context taken aanzienlijk verbetert.
3. Parametrisch Geheugen: De introductie van een geleerde initiële toestand $M_0$ voor lineaire RNN's, waarbij wordt aangetoond dat dit mechanisme effectief algemene kennis opslaat en redeneertaken verbetert.
4. Scaling Laws: Empirische validatie dat SDM voorspelbare schaalwetten volgt, waarbij het de dichte GDN verslaat en bij grotere schalen (8B parameters) zelfs de Full Attention modellen overtreft in trainingsverlies.

Experimentele Resultaten

De auteurs trainden modellen variërend van 1.4B tot 8B parameters op meer dan 1 biljoen tokens.

• Scaling Laws: Op een schaalladder behaalde SDM consistent een lager trainingsverlies dan GDN op alle rekenniveaus. Op de 8B schaal bereikte SDM een lager verlies dan zowel GDN als een Full Attention model.
• Lang-Context Recall (RULER Benchmark): SDM presteerde aanzienlijk beter dan GDN op lang-context retrieval taken. Op 1.4B en 8B schaal behaalde SDM respectievelijk scores van 31.2 en 50.2, vergeleken met GDN's 20.0 en 34.2. Opvallend genoeg matched of overtrof SDM de Full Attention prestaties op verschillende RULER sub-taken, ondanks het gebruik van een vaste geheugenrepresentatie.
• Korte-Context en Redeneren: SDM behaalde een lagere perplexiteit op code data en een hogere gemiddelde nauwkeurigheid op redeneerbenchmarks (bijv. MMLU, GSM8K, HumanEval) vergeleken met GDN. Op de 8B schaal presteerde SDM beter dan Full Attention op gemiddelde nauwkeurigheid.
• Ablatie Studies:
  • State Size: Prestaties degradeerden monotoon naarmate de geheugengrootte werd verkleind, wat bevestigt dat de grote state de primaire drijfveer van de winst is.
  • Geleerde Initialisatie: Het leren van $M_0$ verbeterde de prestaties voor SDM, maar leverde verwaarloosbare winst op voor GDN, wat suggereert dat de grote capaciteit noodzakelijk is om de geleerde initialisatie effectief te gebruiken.
  • Efficiëntie: Hoewel de trainingsdoorvoer van SDM 1.49x trager was dan die van GDN op de 8B schaal (door HBM geheugentoegang vs. SRAM), was het 6x sneller dan Full Attention tijdens inferentie.

Betekenis en Claims

Het artikel claimt dat SDM erin slaagt de kloof te overbruggen tussen de efficiëntie van lineaire RNN's en de lang-context capaciteiten van Transformers. Door de geheugencapaciteit te ontkoppelen van de computationele kosten via sparsity, maakt SDM het volgende mogelijk:

• Constante Ruimte/Compute: Het verwerken van willekeurig lange sequenties (tot 1 miljoen tokens gedemonstreerd) zonder de onbegrensde geheugengroei van KV caches.
• Superieure Lang-Context Prestaties: Het behalen van state-of-the-art resultaten onder lineaire RNN's en competitieve resultaten tegenover Full Attention op lang-context retrieval en in-context leren.
• Parametrische Kennisopslag: Het vermogen om pretraining kennis te leren en op te slaan binnen de geheugentoestand zelf, wat algemeen redeneren en taken met algemene kennis verbetert.

De auteurs erkennen beperkingen, specifiek dat de grote geheugenvoetafdruk (vergelijkbaar met modelparameters) aanzienlijke HBM-bronnen vereist, en dat de huidige kernel-implementaties minder efficiënt zijn dan geoptimaliseerde dichte GDN kernels. Ze stellen echter dat SDM nieuwe mogelijkheden opent voor het ontwikkelen van agenten met verbeterd langetermijngeheugen en voor het verwerken van langdurige video's of uitgebreide redeneersporen.