Technisches Resümee: Sparse Delta Memory (SDM)

Problemstellung

Lineare rekurrenten neuronale Netze (RNNs), wie etwa Gated DeltaNet (GDN) und Mamba2, bieten den Vorteil eines festen Speicher- und Rechenaufwands pro Token, was sie theoretisch für beliebig lange Kontexte geeignet macht. Ihre praktische Leistung bei der Langzeit-Rekallierung (Long-Context Recall) und dem In-Context Learning ist jedoch durch die geringe Größe des verborgenen Zustands (Hidden State) begrenzt. Die Vergrößerung der Zustandsgröße in Standard-linearen RNNs führt zu einem linearen Anstieg der FLOPs aufgrund dichter Zustandsaktualisierungen (Matrix-Vektor-Multiplikationen), was einen prohibitiven Rechenengpass erzeugt. Im Gegensatz dazu behalten Standard-Transformer-Architekturen mit Softmax-Attention eine hohe Rekallierungsrate bei, leiden aber unter den linear wachsenden Key-Value-Cache-Speicher- und Rechenkosten mit steigender Sequenzlänge.

Die zentrale Herausforderung, die hier adressiert wird, besteht darin, wie die Kapazität des verborgenen Zustands linearer RNNs um Größenordnungen gesteigert werden kann, um die Langzeit-Kontextleistung zu verbessern, ohne die Rechenkosten (FLOPs) pro Token zu erhöhen.

Methodik: Sparse Delta Memory (SDM)

Die Autoren führen Sparse Delta Memory (SDM) ein, eine Architektur, die das Gated DeltaNet erweitert, indem sie das dichte Key-Value-Outer-Product durch ein spärliches Adressierungsschema basierend auf Product-Key Memory (PKM) ersetzt.

Kernmechanismus

1. Expliziter spärlicher Speicher: Anstatt eine dichte Matrix $M_t \in \mathbb{R}^{d_{qk} \times d_v}$ zu führen, pflegt SDM eine explizite Speichertabelle $M_t \in \mathbb{R}^{N \times d_v}$ mit $N$ Slots, wobei $N$ Größenordnungen größer als die verborgene Dimension sein kann (z. B. $N \approx 10^6$).
2. Spärliche Schlüsselwahl (Sparse Key Selection): In jedem Zeitschritt werden die projizierten Keys ($k'$) und Queries ($q'$) in zwei Hälften geteilt. Deren äußere Summe erzeugt $N$ Scores für die Speicher-Slots. Mittels einer Top-$W$ (Write) und Top-$R$ (Read) Selektionsstrategie identifiziert das Modell spezifische Speicher-Slots, die aktualisiert oder aus denen gelesen werden soll. Diese Selektion erfolgt effizient, ohne die vollständige $N$-große Score-Matrix zu materialisieren, indem die Eigenschaft $\text{topk}(s_1 \oplus s_2) = \text{topk}(\text{topk}(s_1) \oplus \text{topk}(s_2))$ genutzt wird.
3. Gated Delta Updates: Für die ausgewählten $W$ Write-Slots wendet das Modell die Gated DeltaNet-Aktualisierungsregel an:
   • 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])$
   • Nicht ausgewählte Slots bleiben unverändert.
4. Spärliches Lesen (Sparse Read): Die Ausgabe ist eine gewichtete Summe der $R$ ausgewählten Read-Slots.

IsoFLOP-Design

Um einen fairen Vergleich zu gewährleisten, entwerfen die Autoren SDM so, dass es die Parameteranzahl und die FLOPs einer dichten GDN-Baseline angleicht:

• FLOPs: Die Kosten pro Token betragen $O((W+R) \times d_v)$ und sind unabhängig von der Speichergröße $N$. Durch das Setzen von $W$ und $R$ gleich der Query-Dimension der GDN-Baseline bleiben die Rechenkosten identisch.
• Parameter: Die linearen Projektionen sind in der Größe identisch mit der Baseline dimensioniert. Die Speichergröße $N$ wird über die Anzahl der Köpfe $H$ und die Projektionsdimensionen gesteuert, sodass die Zustandsgröße als $O(1/H^2)$ skalieren kann, ohne die FLOPs zu beeinflussen.

Gelernter Initialzustand

Im Gegensatz zu GDN, das seinen kleinen Zustand typischerweise auf Null initialisiert, behandelt SDM den initialen Speicherzustand $M_0$ als lernbaren Parameter. Dies ermöglicht es dem Modell, vortrainiertes Wissen direkt in der Speicherstruktur zu speichern, was als parametrischer Speicher fungiert, der über die Inferenz hinweg persistiert.

Wichtigste Beiträge

1. Architektur: Der Vorschlag von SDM, das die Kapazität des Zustands linearer RNNs um drei Größenordnungen skaliert (z. B. von Kilobyte auf Gigabyte), während die FLOPs pro Token konstant bleiben.
2. IsoFLOP-Validierung: Der Nachweis, dass die Steigerung der Zustandskapazität allein (oh-ne Erhöhung der Rechenleistung) die Leistung bei Langzeit-Kontext-Aufgaben signifikant verbessert.
3. Parametrischer Speicher: Die Einführung eines lernbaren Initialzustands $M_0$ für lineare RNNs, der zeigt, dass dieser Mechanismus effektiv verbreitetes Wissen speichert und Denkaufgaben verbessert.
4. Skalierungsgesetze (Scaling Laws): Die empirische Validierung, dass SDM vorhersagbare Skalierungsgesetze folgt und bei größeren Skalen (8B Parameter) Full-Attention-Modelle im Trainingsverlust übertrifft.

Experimentelle Ergebnisse

Die Autoren trainierten Modelle von 1,4B bis 8B Parametern auf über 1 Billion Tokens.

• Skalierungsgesetze: Auf einer Skalierungsleiter erreichte SDM konsistent einen niedrigeren Trainingsverlust als GDN auf allen Rechenstufen. Bei der 8B-Skala erreichte SDM einen niedrigeren Verlust als sowohl GDN als auch ein Full-Attention-Modell.
• Langzeit-Rekallierung (RULER Benchmark): SDM übertraf GDN bei Langzeit-Retrieval-Aufgaben deutlich. Bei 1,4B und 8B Skalen erreichte SDM Werte von 31,2 bzw. 50,2, verglichen mit GDNs 20,0 und 34,2. Bemerkenswerterweise erreichte SDM bei mehreren RULER-Subtasks die Leistung von Full Attention oder übertraf diese, trotz der Verwendung einer fixen Speicherrepräsentation.
• Kurzzeit-Kontext und Reasoning: SDM erzielte eine niedrigere Perplexität bei Code-Daten und eine höhere durchschnittliche Genauigkeit bei Reasoning-Benchmarks (z. B. MMLU, GSM8K, HumanEval) im Vergleich zu GDN. Bei der 8B-Skala übertraf SDM die durchschnittliche Genauigkeit von Full Attention.
• Ablationsstudien:
  • Zustandsgröße: Die Leistung sank monoton, wenn die Speichergröße reduziert wurde, was bestätigt, dass der große Speicher der primäre Treiber der Gewinne ist.
  • Gelernte Initialisierung: Das Lernen von $M_0$ verbesserte die Leistung von SDM, bot jedoch für GDN vernachlässigbare Vorteile, was darauf hindeutet, dass die große Kapazität notwendig ist, um die gelernte Initialisierung effektiv zu nutzen.
  • Effizienz: Während der Trainingsdurchsatz von SDM bei der 8B-Skala 1,49-mal langsamer war als der von GDN (aufgrund von HBM-Speicherzugriffen vs. SRAM), war er während der Inferenz 6-mal schneller als Full Attention.

Bedeutung und Ansprüche

Das Paper behauptet, dass SDM erfolgreich die Lücke zwischen der Effizienz linearer RNNs und den Langzeit-Kontext-Fähigkeiten von Transformern schließt. Durch die Entkopplung von Speicherkapazität und Rechenkosten mittels Sparsity ermöglicht SDM:

• Konstanter Platz/Rechenaufwand: Verarbeitung beliebig langer Sequenzen (bis zu 1 Million Tokens demonstriert) ohne das ungebundene Wachstum der KV-Caches.
• Überlegene Langzeit-Kontext-Leistung: Erreichen von State-of-the-Art-Ergebnissen unter den linearen RNNs und kompetitive Ergebnisse gegenüber Full Attention bei Langzeit-Retrieval und In-Context Learning.
• Parametrische Wissensspeicherung: Die Fähigkeit, vortrainiertes Wissen innerhalb des Speicherzustands zu lernen und zu speichern, was allgemeines Denken und Aufgaben des Allgemeinwissens verbessert.

Die Autoren räumen Einschränkungen ein, insbesondere dass der große Speicherbedarf (vergleichbar mit den Modellparametern) erhebliche HBM-Ressourcen erfordert und aktuelle Kernel-Implementierungen weniger effizient sind als optimierte dichte GDN-Kernel. Sie postulieren jedoch, dass SDM neue Möglichkeiten für die Entwicklung von Agenten mit verbessertem Langzeitgedächtnis und für die Verarbeitung von Langzeitvideos oder erweiterten Reasoning-Traces eröffnet.