From Weights to Features: SAE-Guided Activation Regularization for LLM Continual Learning

Technische Samenvatting: Van Gewichten naar Kenmerken: SAE-gestuurde Activatieregulering voor Continual Learning in LLM's

Probleemstelling
Continual learning (CL) in Large Language Models (LLM's) kampt met de uitdaging van catastrofale vergetelheid (catastrophic forgetting), waarbij sequentiële fine-tuning de representaties overschrijft die nodig zijn voor eerdere taken. Hoewel regularisatiemethoden zoals Elastic Weight Consolidation (EWC) standaard zijn voor kleinere modellen, presteren ze aanzienlijk minder goed bij LLM's. De auteurs stellen dat dit falen structureel is, voortkomend uit het "polysemantische" karakter van LLM's. Vanwege de superpositie-hypothese coderen individuele neuronen en hun bijbehorende gewichten gelijktijdig meerdere ongerelateerde concepten. Bijgevolg zijn schattingen van de belangrijkheid in de gewichtsruimte (bijv. Fisher-informatie) te grof; het beschermen van een gewicht om Concept A te behouden, beperkt onbedoeld Concept B, wat een niet-selectieve straf creëert die plasticiteit verstikt. Bovendien falen bestaande methoden vaak in het expliciet afdwingen van de vereiste plasticiteit om nieuwe taken te leren, wat leidt tot "intransigentie" waarbij het model oude kennis behoudt maar niet kan adapteren.

Methodologie
Het artikel stelt SAE-gestuurde activatieregulering voor, een methode die het doel van de regulering verschuift van de verstrengelde parameterruimte naar een meer selectieve activatieruimte met behulp van vooraf getrainde Sparse Autoencoders (SAEs).

1. Verschuiving van Coördinatensysteem: In plaats van gewichtsdrift te bestraffen, gebruikt de methode een bevroren, vooraf getrainde SAE als een vaste, monosemantische kenmerkendictionary. SAE's deconstrueren dichte modelactivaties in ijle (sparse) kenmerkrichtingen die vaak interpreteerbaarder en minder verstrengeld zijn dan individuele neuronen.
2. Geformuleerde Beperkte Optimalisatie: Het trainingsdoel is afgeleid van een geoptimaliseerd probleem met beperkingen (constrained optimization) in plaats van ad-hoc strafontwerp. Het model moet de huidige taakverlies minimaliseren, onderworpen aan twee beperkingen:
   • Stabiliteit: Taakonrelevante kenmerken (lage relevantie) mogen niet verder afwijken dan een specifiek budget ($\varepsilon_s$).
   • Plasticiteit: Taakrelevante kenmerken (hoge relevantie) moeten voldoende afwijken ($\ge \varepsilon_p$) om degeneratieve oplossingen te voorkomen waarbij het model niets leert.
3. Verliesfunctie: Gebruikmakend van Lagrangian relaxatie en een squared-hinge penalty, leidt de auteur een trainingsverlies af dat bestaat uit een protect loss (straft excessieve drift in beschermde kenmerken af) en een guide loss (straft onvoldoende drift in relevante kenmerken af). EWC wordt getoond als een speciaal geval van dit framework, waarbij de representatieruimte gewichten zijn en de plasticiteitsbeperking ontbreekt.
4. Maskerconstructie: Voor elke nieuwe taak wordt een continu relevantiemasker ($M_t$) geconstrueerd met behulp van huidige-taakdata die door het bevroren basismodel en de SAE worden gestuurd. Dit masker identificeert relevante kenmerken via activatiemagnitude en propageert relevantie naar naburige kenmerken via k-NN op decodervectoren. Cruciaal is dat zodra het masker is geconstrueerd, alle taakdata worden verwijderd; alleen het compacte masker wordt bewaard.
5. Training: Tijdens de training worden het huidige model (met LoRA) en het bevroren basismodel beide gecodeerd door de SAE. De kenmerkdrift wordt gewogen door het masker om de protect- en guide-losses te berekenen.

Belangrijkste Bijdragen

• SAE-gestuurde Regulering: Een praktische CL-methode die reguleert in de activatieruimte met behulp van SAE-kenmerken. Het vereist geen gegevens van vorige taken, geen opgeslagen activaties, geen taakspecifieke parameters en geen inference-time routing na de constructie van het masker.
• Geprincipeerde Afleiding: De stabiliteit-plasticiteit trade-off wordt geformuleerd als een geoptimaliseerd probleem met beperkingen in de SAE-kenmerkruimte, wat een squared-hinge trainingsverlies oplevert via Lagrangian relaxatie.
• Empirisch Bewijs voor Polysemantiek: Het artikel levert mechanistisch bewijs dat gewichtsruimte-bescherming niet-selectief is. Experimenten tonen aan dat taakrelevante representaties lineair scheidbaar zijn in de SAE-kenmerkruimte (AUC ~0,88), maar ononderscheidbaar zijn van toeval in de gewichts-/neuronenruimte (AUC ~0,50). Daarnaast beperkt het beschermen van gewichten voor één taak de kenmerken voor een volgende taak met een snelheid van 91–96% van de oorspronkelijke taak, terwijl kenmerkruimte-bescherming deze bijkomende beperking vermindert tot 43–41%.
• Schaalbaarheid: Door de lagere dimensionaliteit van de kenmerkruimte vergeleken met de parameterruimte, wordt de opslag per taak verminderd van gigabytes (gewicht-ankers) naar sub-megabytes (kenmerkmaskers).

Resultaten
De methode werd geëvalueerd op TRACE-5000 (8 cross-domein taken) en MedCL (10 biomedische taken) met Gemma-2 9B-it met LoRA.

• Prestaties: Op TRACE behaalde SAE-gestuurde regulering de sterkste prestaties onder de niet-architecturale benaderingen (OP = 0,545), waarmee het gewichtsruimte-regulariseerders (EWC: 0,447, SI: 0,513), gradiënt-projectiemethoden en replay-baselines (zelfs met 10% data) overtrof. Op MedCL leidde het alle methoden die geen voorgaande taakvoorbeelden behouden (OP = 0,510).
• Stabiliteit-Plasticiteit Trade-off: In tegen tegenstelling tot EWC, dat een hoge retentie (ARA = 0,987) bereikte door een zware opoffering van plasticiteit (0,453), behield de voorgestelde methode een hoge plasticiteit (0,567, vergelijkbaar met onbeschermde fine-tuning) terwijl 95% van de geleerde prestaties werd behouden.
• Efficiëntie: De methode voegt minimale computationele overhead toe (~1,88s/stap vs 1,70s baseline) en vermindert de opslagvereisten drastisch (412 KB per taak vs 6,5 GB voor EWC-ankers).

Betekenis en Claims
Het artikel beweert dat de onderprestatie van gewichtsruimte-regulering in LLM's geen kwestie is van hyperparameter-afstemming, maar een structureel gevolg van polysemantiek. Door de reguleringsbeperking te verplaatsen naar een monosemantische kenmerkbasis, bereikt de methode een meer selectieve bescherming van kennis. De auteurs stellen dat de effectiviteit van regulering kritisch afhangt van de representatieruimte waarin beperkingen worden opgelegd. De methode demonstreert dat het mogelijk is om sterke continual learning-prestaties te bereiken zonder voorgaande data op te slaan of de modelarchitectuur tijdens inference te wijzigen, mits de regulering opereert in een coördinatensysteem dat is uitgelijnd met de interne conceptstructuur van het model. Het artikel concludeert dat SAE-gebaseerde regulering een meer selectief en schaalbaar alternatief biedt voor traditionele gewichtsruimte-methoden voor LLM's.