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

Probleemstelling

Hoewel Transformer-modellen revolutionair zijn voor natuurlijke taalverwerking, ondervinden ze beperkingen bij lange contexten door de kwadratische complexiteit ($O(L^2)$) van de softmax-attention. Lineaire attention-mechanismen bieden een oplossing met lineaire complexiteit ($O(L)$), maar eerdere varianten presteerden vaak slechter dan softmax-attention.

Recente ontwikkelingen zoals Gated DeltaNet (GDN) en Kimi Delta Attention (KDA) hebben de kloof verkleind door twee mechanismen te combineren:

1. Gating: Voor adaptieve geheugenverval (forgetting).
2. De Delta-regel: Voor precieze updates via online gradient descent.

Echter, een fundamentele beperking blijft bestaan in de huidige state-of-the-art modellen (zoals KDA):

• KDA introduceerde kanaalgewijs verval (channel-wise decay) via een vector $\alpha_t$, waardoor elke feature-dimension een onafhankelijk vergetenssnelheid kan hebben.
• Maar de leerfactor (learning rate) $\beta_t$ in de delta-update blijft een scalar (één getal) die door alle dimensies wordt gedeeld.
• Dit beperkt het model in zijn vermogen om dimensie-specifieke aanpassingen te doen. Niet alle feature-dimensies zijn even relevant voor een specifieke context, en een uniforme schrijfstrength (write strength) is suboptimaal.

Methodologie: FG2-GDN

De auteurs introduceren FG2-GDN (Fine-Grained Gated Delta Network) om deze asymmetrie op te lossen. De kern van de methode is het vervangen van de scalar $\beta_t$ door een kanaalgewijze vector $\beta_t \in \mathbb{R}^{d_k}$.

1. De Update Regel:
In plaats van de leerfactor alleen op de buitenproduct-term toe te passen, wordt $\beta_t$ symmetrisch geabsorbeerd in de key- ($k_t$) en value- ($v_t$) vectoren:
$$\tilde{k}_t = \sqrt{\beta_t} \odot k_t$$
$$\tilde{v}_t = \sqrt{\beta_t} \odot v_t$$

De update regel wordt dan:
$$S_t = (I - \tilde{k}_t \tilde{k}_t^\top) \text{Diag}(\alpha_t) S_{t-1} + \tilde{k}_t \tilde{v}_t^\top$$

Hierdoor krijgt elke dimensie een eigen effectieve leerfactor, analoog aan de overstap van SGD naar per-coördinaat adaptieve optimalizers zoals AdaGrad en Adam.

2. FG2-GDN+ Variant:
De auteurs stellen ook een variant voor, FG2-GDN+, die de schaling voor keys en values ontkoppelt door twee aparte vectoren in te voeren: $\beta^k_t$ en $\beta^v_t$.

• $\beta^k_t$ controleert de "wissterkte" (erasure strength).
• $\beta^v_t$ controleert de "schrijfstrength" (write strength).
  Dit geeft het model onafhankelijke controle over hoe agressief het oude informatie verwijdert versus hoe sterk het nieuwe waarden schrijft.

3. Efficiëntie en Parallelisatie:
Een cruciaal technisch aspect is dat de update regel de Diagonal-Plus-Low-Rank (DPLR) structuur behoudt.

• De term $(I - \tilde{k}_t \tilde{k}_t^\top)$ blijft een geldige, rang-1 Householder-transformatie.
• Hierdoor kan het bestaande chunkwise parallelle trainingsalgoritme (gebaseerd op de WY-representatie en UT-transformatie) van KDA direct worden hergebruikt.
• De extra rekentijd voor de kanaalgewijze vector $\beta_t$ is verwaarloosbaar (minder dan 5% overhead) omdat deze slechts element-wise vermenigvuldiging vereist, terwijl de zware matrixbewerkingen onveranderd blijven.

Belangrijkste Bijdragen

1. Fine-Grained Learning Rates: De eerste toepassing van per-coördinaat adaptieve leerfactoren binnen de delta-regel voor lineaire attention, wat de expressiviteit van het geheugen vergroot zonder de complexiteit te verhogen.
2. Decoupling van Schrijven en Vergeten: Met FG2-GDN+ wordt de controle over het wissen van oude informatie onafhankelijk gemaakt van het schrijven van nieuwe informatie, wat essentieel is voor complexe retrieval-taken.
3. Architecturale Compatibiliteit: Het bewijzen dat deze verfijning volledig compatibel is met efficiënte chunkwise parallelle implementaties, waardoor het direct toepasbaar is op grote schaal.

Resultaten

De auteurs hebben modellen getraind op schalen van 340M en 1.3B parameters en geëvalueerd op diverse benchmarks:

• Taalmodellering (Language Modeling):

  • FG2-GDN behaalde de beste gemiddelde nauwkeurigheid op commonsense redeneringstaken (LAMBADA, ARC, etc.) en de laagste perplexiteit op LAMBADA bij de 1.3B schaal.
  • Dit suggereert dat fijne-granulariteit in schrijfstrength zeer gunstig is voor korte context taken.

• Lange Context Retrieval (RULER & LongBench):

  • RULER: FG2-GDN en FG2-GDN+ overtroffen zowel GDN als KDA significant, vooral bij taken met unieke zoekopdrachten maar variërende query-formuleringen (bijv. Multi-Query, Common Word Extraction).
  • Multi-Key Retrieval: FG2-GDN+ toonde de grootste verbetering, wat aantoont dat het ontkoppelen van wis- en schrijfkracht helpt bij het beheren van meerdere concurrerende associaties.
  • Extrapolatie: Hoewel getraind op 8k tokens, behielden de modellen hun prestaties tot 16k tokens beter dan de baselines, wat wijst op een gestructureerdere staatrepresentatie.

• Efficiëntie:

  • De prefill-throughput van FG2-GDN is vergelijkbaar met KDA (<5% overhead).
  • Alle lineaire varianten behielden een lineaire schaling ($O(L)$), terwijl volledige attention (MLA) sterk degradeerde bij langere sequenties.

Significantie

Dit paper biedt een belangrijke stap in de evolutie van lineaire recurrente modellen. Het toont aan dat de beperkingen van eerdere modellen niet inherent waren aan de delta-regel, maar aan de te grove parametrisatie van de leerfactor.

Door de overgang van een scalair naar een vectoriële leerfactor (en het ontkoppelen van schrijf/wis-krachten), halen de auteurs het beste van twee werelden:

1. De expressiviteit en geheugenkwaliteit van complexe attention-mechanismen.
2. De rekenefficiëntie van lineaire recurrente netwerken.

Deze bevindingen suggereren dat "fine-grained adaptivity" een essentieel ontwerpprincipe is voor de volgende generatie efficiënte lange-context modellen, met name voor toepassingen zoals agenten en multimodale scenario's waar lange contexten cruciaal zijn.