Technische Samenvatting: SAE-FD voor Continue Lering van Grote Taalmodellen

Probleemstelling

Grote Taalmodellen (LLM's) staan voor de uitdaging van catastrofaal vergeten bij het sequentieel aanpassen aan meerdere taken. Bestaande methoden voor continue lering (CL) werken doorgaans in dichte representatieruimten (gewicht-, gradiënt- of outputruimten). Een fundamentele beperking van deze benaderingen is feature-superpositie, waarbij meerdere concepten worden gecodeerd in overlappende dimensies binnen de dichte activeringsruimte. Deze verstrengeling maakt het moeilijk om selectief eerder geleerde kennis te beschermen zonder onbedoeld het leren van nieuwe taken te beperken, wat leidt tot een suboptimale afweging tussen plasticiteit en stabiliteit.

Methodologie: SAE-FD

De auteurs stellen SAE-FD (Sparse Autoencoder Feature Distillation) voor, een methode die modelrepresentaties verankert in de spare feature-ruimte van een vooraf getrainde Sparse Autoencoder (SAE). Door dichte activeringen te ontleden in een sparsere, overcomplete basis, streeft SAE-FD ernaar superponerende representaties te ontkoppelen, waardoor gerichter regularisatie mogelijk wordt.

De methode werkt als een proces van drie fasen:

1. SAE-training (Fase 1)

Voordat de reeks voor continue lering begint, wordt een Gated SAE getraind op diverse tekstactiveringen van het basismodel.

• Input: MLP-activeringen van de laatste laag ($h \in \mathbb{R}^d$) uit diverse bronnen (classificatie, QA, kennisbenchmarks).
• Architectuur: De Gated SAE splitst de codering op in een gating-pad ($g$) en een magnitude-pad ($m$), wat pre-ReLU-features ($f_{pre}$) en post-ReLU-features ($f$) oplevert.
• Doel: Minimaliseren van reconstructieverlies ($\|\hat{h} - h\|_2^2$) en een L1-sparsiteitsstraf ($\lambda_{\ell1}\|f\|_1$).
• Resultaat: Een bevroren SAE die dichte activeringen ontbindt in sparsere features ($D \gg d$), waardoor concepten effectief worden ontkoppeld.

2. Taaktraining en Verankering (Fase 2)

Voor elke taak $t$ in de reeks:

• Fine-tuning: Het model wordt gefine-tuned met LoRA en de standaard taalmodelleringdoelstelling.
• Verankering: Na training op een taak voert het model een forward pass uit op een kleine set ankerstalen ($N_a$). De per-token MLP-outputs ($H_i$) worden geëxtraheerd en opgeslagen.
• Buffer: Deze ankers accumuleren over taken heen. Tijdens daaropvolgende training wordt een batch uniform uit de buffer gesampled om distillatieverlies te berekenen, waardoor de kosten over alle vorige taken worden afgeschreven.

3. Sparsiteitsbewuste Feature-distillatie (Fase 3)

Tijdens de training van een nieuwe taak ($t+1$) wordt het model geregulariseerd door zijn huidige representaties te vergelijken met de opgeslagen ankers, beide gecodeerd via de bevroren SAE. Het distillatieverlies ($L_{FD}$) bestaat uit twee complementaire termen:

• Cosinus-richtingsverlies ($L_{cos}$): Straft hoekafwijking tussen de huidige pre-ReLU-featurevector en de anker-featurevector. Dit behoudt de globale richting van de representatie.
• Actieve Magnitude-verlies ($L_{mag}$): Past Kwantiteit van het Gemiddelde Kwadratische Fout (MSE) alleen toe op features die actief waren (post-ReLU $> 0$) in de anker. Dit beperkt de magnitude van features die bijdragen aan vorige taken, terwijl inactieve features vrij blijven om voor nieuwe taken te worden gerekruteerd.

Adaptief $\lambda$-mechanisme

Om taakleren en regularisatie in evenwicht te brengen, introduceren de auteurs een adaptief mechanisme voor de wegingscoëfficiënt $\lambda$ in het totale verlies $L = L_{task} + \lambda \cdot L_{FD}$.

• In plaats van een vaste $\lambda$, streeft het systeem naar een vaste bijdrageratio $\rho$ van het distillatieverlies aan het totale verlies.
• $\lambda$ wordt bijgewerkt via exponentiële voortschrijdend gemiddelde (EMA) om te voldoen aan $\rho = \frac{\lambda L_{FD}}{\lambda L_{FD} + L_{task}}$.
• Gedrag: $\lambda$ begint hoog wanneer een nieuwe taak start (wanneer $L_{task}$ groot is en het risico op vergeten hoog is) en ontspant naarmate het model convergeert, waardoor excessieve beperking tijdens fine-tuning wordt voorkomen.

Belangrijkste Bijdragen

1. SAE-FD Framework: De eerste methode voor continue lering die regularisatie uitvoert in de spare, ontkoppelde feature-ruimte van SAE's, waarmee het probleem van feature-superpositie dat inherent is aan dichte ruimten, wordt aangepakt.
2. Sparsiteitsbewuste Distillatieverlies: Een nieuwe verliesfunctie die cosinus-richtingsbehoud combineert met selectieve magnitude-aanpassing op actieve features, waarbij gebruik wordt gemaakt van de sparsiteit van SAE-representaties om interferentie met nieuwe taken te minimaliseren.
3. Adaptieve Regularisatie: Een adaptief $\lambda$-mechanisme dat de regularisatiesterkte dynamisch aanpast op basis van de relatieve magnitude van taak- en distillatieverliezen, waardoor sterke bescherming tijdens risicovolle periodes wordt gewaarborgd zonder convergentie te belemmeren.

Experimentele Resultaten

De auteurs hebben SAE-FD geëvalueerd op twee benchmarks over drie modelarchitecturen (LLaMA-2-7B, Vicuna-7B, Mistral-7B).

• TRACE Benchmark (8 Heterogene Taken):
  • Op LLaMA-2-7B behaalde SAE-FD 52,70% Gemiddelde Nauwkeurigheid (AA) met slechts −0,46% Backward Transfer (BWT).
  • Dit presteert beter dan de sterkste gepubliceerde regularisatie-gebaseerde baseline (GORP) met +2,3% AA en vermindert vergeten met 34% (−0,46 versus −0,70 BWT).
  • Het presteert aanzienlijk beter dan andere LoRA-gebaseerde methoden (EWC, O-LoRA, TreeLoRA), die een BWT vertoonden die slechter was dan −3%.
• Generalisatie:
  • SAE-FD behaalde consistent de beste AA en BWT over alle drie de modelfamilies.
  • Op Vicuna-7B overtrof het O-LoRA met +10,1% AA en behaalde de beste BWT (−0,69%) onder alle methoden, inclusief baselines voor volledige fine-tuning.
• Ablatiestudies:
  • Verliescomponenten: Zowel cosinus-richting als actief magnitude-verlies zijn noodzakelijk; het gebruik van slechts één van beide resulteerde in lagere AA of meer vergeten.
  • Adaptief $\lambda$: Het vervangen van het adaptieve mechanisme door een vaste $\lambda=1,0$ degradeerde de BWT van −0,46% naar −2,38%.
  • SAE versus Raw MSE: Distilleren in de SAE feature-ruimte presteerde beter dan ruwe MSE-distillatie op dichte activeringen, wat het voordeel van de ontkoppelde ruimte bevestigt.
• Efficiëntie: De rekenkundige overhead is ongeveer 1,85× die van SeqLoRA, met een constante kost per stap ongeacht het aantal vorige taken. Opslag vereist ongeveer 0,8 GB per taak voor ankeractiveringen.

Betekenis en Claims

Het artikel claimt dat SAE-FD succesvol mechanistische interpreteerbaarheid (via SAE's) verbindt met continue lering. Door gebruik te maken van de sparsere, ontkoppelde representaties die oorspronkelijk zijn ontwikkeld voor het begrijpen van neurale netwerken, biedt de methode een effectieve inductieve bias voor het beheersen daarvan.

De auteurs stellen dat werken in een sparsere feature-ruimte een meer gerichte bewaring van kennis mogelijk maakt, waarmee het "feature-superpositie"-probleem wordt opgelost dat dichte ruimteregularisatie teistert. De methode bereikt een superieure afweging tussen plasticiteit en stabiliteit, waardoor LLM's zich kunnen aanpassen aan evoluerende taken zonder opnieuw vanaf nul te hoeven trainen, terwijl ze toch hoge prestaties behouden op eerder geleerde taken.

Beperkingen opgemerkt door de auteurs:

• Vereist een eenmalige trainingskosten voor een model-specifieke SAE.
• Opslagvereisten voor ankeractiveringen groeien lineair met het aantal taken (hoewel comprimeerbaar).
• Experimenten zijn uitgevoerd op een vaste taakvolgorde voor de TRACE-benchmark, hoewel multi-model consistentie robuustheid suggereert.
• Vergelijkingen met sommige baselines moeten met voorzichtigheid worden geïnterpreteerd vanwege verschillen in experimentele instellingen (bijv. volledige fine-tuning versus LoRA, aantal epochs).