Learning to Remember, Learn, and Forget in Attention-Based Models

Technische Samenvatting: Leren om te Herinneren, te Leren en te Vergeten in Attention-gebaseerde Modellen

Probleemstelling
In-Context Learning (ICL) in transformers functioneert als een online associatief geheugen, wat hoge prestaties mogelijk maakt op complexe sequentietaken. Echter, gated linear attention-modellen en state space models (SSMs) vertrouwen doorgaans op een geheugentoestand van vaste grootte om de computationele en geheugentechnische knelpunten van standaard Transformers (die lineaire en kwadratische schaling vereisen voor KV-caches) aan te pakken. Deze vaste capaciteit creëert een "stabiliteit-plasticiteitsdilemma": het model moet nieuwe informatie leren zonder kritieke eerdere kennis te overschrijven (catastrofale vergetelheid) of, omgekeerd, onbekwaam worden in het integreren van nieuwe informatie door een excessief behoud van verouderde data (catastrofale herinnering). Bestaande metaplasticiteitsmethoden uit de neurowetenschap zijn vaak ongeschikt voor transformers omdat ze vertrouwen op afzonderlijke taakgrenzen, groeiend geheugen of een gebrek aan differentieerbaarheid. Bov�elijk lijden huidige lineaire attention-modellen vaak onder beperkte geheugencapaciteit en interferentie naarmate de sequentielengtes toenemen.

Methodologie: Palimpsa
De auteurs stellen Palimpsa voor, een self-attention model dat ICL herformuleert als een continual learning probleem dat oplosbaar is door middel van Bayesiaanse metaplasticiteit. Het kerninzicht is dat de plasticiteit van elke attention-state dynamisch aangepast moet worden op basis van een "belangrijkheidstoestand" die geworteld is in een prior-distributie die geaccumuleerde kennis vastlegt.

1. Bayesiaanse Formulering: De attention head wordt gemodelleerd als een Bayesiaans inferentieprobleem waarbij het doel is om een variational free energy objectief te minimaliseren. De posterior-distributie over de geheugentoestand $S$ wordt recursief bijgewerkt. Om catastrofale herinnering te voorkomen (waarbij de prior te geconcentreerd wordt), introduceren de auteurs een gewogen posterior die effectief het geheugenvenster afkapt. Dit wordt bereikt door een fractie van het geaccumuleerde prior-gewicht bij elke stap weg te gooien, gecontroleerd door een forgetting gate $\alpha_t = \exp(-A d_t)$.
2. Metaplastic Update Regel: Door Variational Inference (VI) toe te passen met een diagonale benadering van de covariantie-matrix, leiden de auteurs een closed-form update regel af. De geheugentoestand $\mu_t$ en een belangrijkheidsmatrix $I_t$ (die de inverse variantie of precisie van elke staat vertegenwoordigt) worden gelijktijdig bijgewerkt:
   • $I_t = \alpha_t I_{t-1} + (1 - \alpha_t)I_{prior} + \beta_t \otimes k_t^2$
   • $\mu_t = \alpha_t \frac{I_{t-1}}{I_t} \odot \mu_{t-1} + \frac{1}{I_t} \odot [(\beta_t \odot v_t) \otimes k_t]$
     Hier fungeert $\beta_t$ als een input-afhankelijke belangrijkheidsfactor, en $I_t$ schaalt de learning rate voor elke staat dynamisch. Hoge belangrijkheid (lage variantie) leidt tot stabiliteit (lage plasticiteit), terwijl lage belangrijkheid sneller leren mogelijk maakt.
3. Architecturale Varianten: De auteurs implementeren twee varianten:
   • Palimpsa-D: Een drop-in vervanging voor Gated Delta Rule modellen (Deltanets), die een projectie voor $\beta_t$ en een residuele term toevoegt.
   • Palimpsa-M: Gebouwd op Mamba2, waarbij de input gate $d_t$ een dubbelrol vervult als zowel de forgetting gate als de input gate.
4. Theoretische Unificatie: Het artikel demonstreert dat diverse gated linear attention-modellen (Longhorn, Deltanets, Mesanet) verschijnen als specifieke architecturale keuzes of posterior-benaderingen binnen dit gemeenschappelijke Bayesiaanse framework. Opvallend genoeg wordt Mamba2 getoond als een speciaal geval van Palimpsa waarbij de forgetting rate zo hoog is dat de belangrijkheidsmatrix constant blijft ($I_t \approx I_{prior}$), wat effectief de metaplasticiteit uitschakelt.

Belangrijkste Bijdragen

• Metaplasticiteit in ICL: De introductie van Bayesiaanse metaplasticiteit om in-context catastrofale vergetelheid en herinnering te mitigeren door de plasticiteit van de staat dynamisch aan te passen op basis van onzekerheid.
• Palimpsa Model: Een schaalbare attention block die catastrofale herinnering voorkomt door verouderde kennis geleidelijk vrij te geven, geïmplementeerd met custom kernels voor efficiënte GPU-training.
• Verenigd Framework: Een wiskundige afleiding die aantoont dat bestaande gated linear attention en state space modellen speciale gevallen zijn van een gemeenschappelijk Bayesiaans framework, waarbij Mamba2 als een limiet van Palimpsa wordt geïdentificeerd.
• Transformatie Pipeline: Een methode om niet-metaplastische modellen (zoals pre-trained Mamba2) te transformeren naar metaplastische modellen via fine-tuning, waarbij gebruik wordt gemaakt van de snelheid van niet-metaplastische kernels voor pre-training en de geheugencapaciteit van metaplasticiteit voor fine-tuning.

Experimentele Resultaten
De auteurs evalueren Palimpsa op synthetische, commonsense reasoning en long-context benchmarks met 170M, 760M en 2.7B parameter modellen.

• Synthetisch (MQAR): Op de Multi-Query Associative Recall benchmark presteert Palimpsa consistent beter dan niet-metaplastische baselines (inclusief Gated Deltanet en Mamba2). De prestatiekloof wordt groter naarmate de sequentielengte toeneemt, wat aantoont dat metaplasticiteit effectief het overschrijven in complexe memorisatie-taken mitigeert.
• Taalmodellering & Commonsense Reasoning: Op de FineWeb-Edu dataset bereiken de metaplastische varianten (Palimpsa-D en Palimpsa-M) een superieure perplexiteit en nauwkeurigheid op taken zoals LAMBADA, PIQA en WinoGrande vergeleken met hun niet-metaplastische tegenhangers en standaard Transformers. Specifiek, op de 760M schaal presteert Palimpsa-D ongeveer 0.8 punten beter dan de Gated Deltanet baseline in gemiddelde nauwkeurigheid.
• Long-Context Prestaties: In experimenten met een 2.7B parameter model, resulteerde het chirurgisch vervangen van de 8 Mamba2 lagen met de zwakste forgetting (grootste effectieve geheugenvenster) door Palimpsa blocks in een verbetering van +0.84 punt in commonsense reasoning en een +1.1 punt verbetering op de LongBench benchmark. Hoewel het metaplastische model iets ondermaats presteerde op short-context retrieval taken, vertoonde het een significant tragere degradatie naarmate de contextlengte toenam, waarmee het de baseline overtrof bij 32k tokens.
• Leerdynamiek: Analyse laat zien dat de metaplasticiteitsratio (differentiatie tussen plastische en geconsolideerde synapsen) toeneemt tijdens training, waardoor het model belangrijke informatie kan stabiliseren terwijl het plastisch blijft voor nieuwe data.

Significantie en Claims
Het artikel claimt dat Palimpsa een fundamentele oplossing biedt voor het stabiliteit-plasticiteitsdilemma in fixed-size geheugenmodellen door concepten uit de neurowetenschap (metaplasticiteit) en Bayesiaanse inferentie te integreren. De significantie ligt in:

1. Theoretisch Inzicht: Het verenigen van uiteenlopende lineaire attention-architecturen onder één Bayesiaans framework en het identificeren van Mamba2 als een specifieke, forgetting-dominante limiet.
2. Praktische Schaalbaarheid: Het mogelijk maken van het gebruik van metaplastische modellen op schaal door pre-training (met gebruik van snelle, niet-metaplastische kernels) te ontkoppelen van fine-tuning (met gebruik van metaplastische kernels om de geheugencapaciteit uit te breiden).
3. Prestaties: Het aantonen dat metaplasticiteit consequent de geheugenretentie en redeneervermogens in long-context scenario's verbetert zonder de kwadratische geheugenkosten van standaard Transformers te vereisen.

De auteurs concluderen dat deze aanpak de inzet van performante sequentiemodellering op hulpbron-beperkte systemen (edge devices) ondersteunt door te optimaliseren hoe informatie wordt geleerd, herinnerd en vergeten binnen strikte geheugenlimieten.