FG$^2$-GDN: Enhancing Long-Context Gated Delta Networks with Doubly Fine-Grained Control

1. Problemstellung

Lineare Aufmerksamkeitsmechanismen (Linear Attention) haben sich als vielversprechende Alternative zur Softmax-Aufmerksamkeit etabliert, da sie eine lineare Zeitkomplexität $O(L)$ während der Inferenz bieten, im Gegensatz zur quadratischen Komplexität $O(L^2)$ bei Transformer-Modellen. Dies ist entscheidend für Anwendungen mit langen Kontexten (z. B. Agenten, Multimodalität).

Trotz Fortschritten durch Modelle wie Gated DeltaNet (GDN) und Kimi Delta Attention (KDA), die das Delta-Regel-Prinzip (Online-Gradientenabstieg) mit Gating-Mechanismen kombinieren, bestehen weiterhin Einschränkungen:

• GDN führt einen skalaren Vergessensfaktor $\alpha_t$ ein, der für alle Feature-Dimensionen gleich ist.
• KDA verfeinert dies zu einem kanalweisen Vergessensvektor $\alpha_t \in \mathbb{R}^{d_k}$, ermöglicht also unterschiedliche Vergessensraten pro Dimension.
• Das verbleibende Problem: Der Lernraten-Parameter $\beta_t$ im Delta-Update bleibt in beiden Modellen ein skalärer Wert, der über alle Dimensionen hinweg geteilt wird. Dies limitiert die Fähigkeit des Modells, sich dimensionspezifisch an den Kontext anzupassen, da alle Features mit der gleichen „Schreibstärke" (Write Strength) aktualisiert werden, unabhängig von ihrer individuellen Relevanz.

2. Methodik: FG2-GDN und FG2-GDN+

Die Autoren schlagen FG2-GDN (Fine-Grained Gated Delta Network) vor, um diese Asymmetrie zu beheben, indem sie die Lernrate $\beta_t$ ebenfalls kanalweise machen.

Kerninnovation: Kanalweises $\beta_t$

Statt eines skalaren $\beta_t$ wird ein Vektor $\beta_t \in \mathbb{R}^{d_k}$ eingeführt. Dies entspricht dem Übergang von SGD zu adaptiven Optimierern pro Koordinate (wie AdaGrad oder Adam).

• Update-Regel: Die Schlüssel ($k_t$) und Werte ($v_t$) werden vor dem Delta-Update skaliert:
  $$\tilde{k}_t = \sqrt{\beta_t} \odot k_t, \quad \tilde{v}_t = \sqrt{\beta_t} \odot v_t$$
  Der Zustandsupdate lautet dann:
  $$S_t = (I - \tilde{k}_t \tilde{k}_t^\top) \text{Diag}(\alpha_t) S_{t-1} + \tilde{k}_t \tilde{v}_t^\top$$
• Strukturelle Kompatibilität: Durch die symmetrische Skalierung (Wurzel aus $\beta_t$ auf beiden Seiten) bleibt die Rang-1-Struktur (Householder-Transformation) erhalten. Dies ist entscheidend, da es die effiziente Chunkwise-Parallelisierung (WY-Repräsentation) ermöglicht, die auch in KDA verwendet wird. Die Komplexität bleibt $O(L)$.

Erweiterung: FG2-GDN+

Eine Variante namens FG2-GDN+ entkoppelt die Skalierung für Schlüssel und Werte vollständig:

• Es werden zwei separate Vektoren eingeführt: $\beta^k_t$ für die „Löschstärke" (Erasure) und $\beta^v_t$ für die „Schreibstärke" (Write Strength).
• Dies ermöglicht dem Modell, unabhängig zu steuern, wie aggressiv alte Informationen gelöscht werden, im Vergleich dazu, wie stark neue Werte geschrieben werden.

Architektur

Die Modelle werden in einer hybriden Architektur eingesetzt, die FG2-GDN-Schichten mit Multi-head Latent Attention (MLA)-Schichten im Verhältnis 3:1 abwechseln lässt, um die Vorteile linearer Rekursion mit der vollen Paarinteraktion von MLA zu kombinieren.

3. Wichtige Beiträge

1. Einführung von Fine-Grained Learning Rates: Erster Ansatz, der die Lernrate im Delta-Update von einem Skalar auf einen kanalweisen Vektor erweitert, analog zu adaptiven Optimierern.
2. Entkopplung von Lösch- und Schreibmechanismen: Durch FG2-GDN+ wird die Kontrolle über das Vergessen und das Schreiben von Informationen unabhängig gemacht.
3. Effizienz: Die Methode behält die DPLR-Struktur (Diagonal-Plus-Low-Rank) bei, was bedeutet, dass bestehende parallele Implementierungen wiederverwendet werden können. Der Overhead gegenüber KDA beträgt weniger als 5% beim Prefill-Throughput.
4. Umfassende Evaluation: Evaluation auf synthetischen und realen Benchmarks bei zwei Skalierungen (340M und 1,3B Parameter).

4. Ergebnisse

Die Experimente zeigen konsistente Verbesserungen gegenüber GDN und KDA:

• Sprachmodellierung (Language Modeling):

  • FG2-GDN erreicht die beste durchschnittliche Genauigkeit auf Commonsense-Reasoning-Benchmarks (LAMBADA, ARC, etc.) und die niedrigste Perplexität auf LAMBADA bei 1,3B Parametern.
  • Die feinkörnige Lernrate verbessert das Verständnis kurzer Kontexte signifikant.

• Langkontext-Fähigkeiten (Long-Context):

  • RULER Benchmark: FG2-GDN und FG2-GDN+ übertreffen die Baselines deutlich, insbesondere bei Aufgaben wie Single-Key Retrieval (MQ, CWE) und Multi-Key Retrieval (MK).
  • LongBench: Auf Aufgaben wie Dokumenten-QA, Zusammenfassung und Few-Shot-Learning erzielt FG2-GDN+ die besten Ergebnisse (z. B. 18,3% Durchschnitt auf LongBench vs. 16,4% bei KDA).
  • Extrapolation: Die Modelle extrapolieren auch bei Sequenzlängen von 16k (über die Trainingslänge von 8k hinaus) besser als die Baselines.

• Ablationsstudie:

  • Die Einführung des vektoriellen $\beta$ allein verbessert die Sprachmodellierung.
  • Die Entkopplung von $\beta^k$ und $\beta^v$ (auch als Skalare) hilft bereits bei Langkontext-Aufgaben.
  • Die Kombination beider (FG2-GDN+) liefert die beste Gesamtleistung für Langkontext-Szenarien.

• Effizienz:

  • Der Throughput-Verlust gegenüber KDA ist minimal (<5%).
  • Die lineare Skalierung ($O(L)$) bleibt erhalten, während reine MLA-Modelle bei langen Sequenzen stark an Performance verlieren.

5. Bedeutung und Fazit

FG2-GDN demonstriert, dass die Verfeinerung der Lernraten-Granularität in linearen rekurrenten Modellen ein kritischer Hebel für die Leistungsfähigkeit ist. Durch die Einführung von dimensionspezifischen Anpassungen (ähnlich wie bei Adam im Training) wird die assoziative Erinnerung präziser gesteuert.

Die Arbeit zeigt, dass man die Effizienz linearer Attention-Modelle nicht opfern muss, um deren Ausdruckskraft zu erhöhen. Die Fähigkeit, unabhängig zu steuern, welche Informationen wie stark vergessen oder geschrieben werden, führt zu robusteren Modellen für lange Kontexte, die in zukünftigen Agenten- und Multimodal-Anwendungen essenziell sind. FG2-GDN+ stellt dabei einen neuen State-of-the-Art für hybride lineare/rekurrente Architekturen dar.