Technische Samenvatting: Interpretatie en Sturing van Ruimte-toestandmodellen via Activeringsruimte-Verstoppingen

1. Probleemstelling

Ruimte-toestandmodellen (SSM's), met name de Mamba-architectuur, zijn naar voren gekomen als efficiënte alternatieven voor transformers voor taalmodellering, en bieden training met lineaire tijd en inferentie met constante tijd door het gebruik van vaste verborgen toestanden in plaats van groeiende sleutel-waarde caches. Ondanks hun efficiëntie staan SSM's echter voor aanzienlijke uitdagingen op het gebied van interpreteerbaarheid en stuurbaarheid. In tegenstelling tot transformers, die expliciete aandachtskoppen en neuronrepresentaties bezitten die uitvoerig zijn bestudeerd voor interpreteerbaarheid, vertrouwen SSM's op recurrente structuren die deze expliciete componenten missen. Deze ondoorzichtigheid belemmert het begrijpen van faalmodi, zoals slechte prestaties bij leren uit context en bij het ophalen van informatie uit lange contexten. Bovendien heeft het bestaande onderzoek naar interpreteerbaarheid weinig praktische verbeteringen opgeleverd voor SSM's, en zijn huidige sturingsmethoden grotendeels toegespitst op transformer-architecturen.

2. Methodologie

De auteurs stellen een mechanistisch onderzoek voor dat zich richt op het identificeren van Activeringsruimte-Verstoppingen binnen Mamba-modellen. De methodologie bestaat uit drie hoofdbestanddelen:

A. Identificatie van Activeringsruimte-Verstoppingen

De auteurs maken gebruik van een meerstapsproces om functionele verstoppingen te lokaliseren:

• Sparse Auto-encoders (SAE's): SAE's worden getraind om verborgen activeringen te reconstrueren met behulp van schaarse latente representaties, dienend als een comprimerende verstopping om dominante activeringsniveau-kenmerken te identificeren.
• Benadering van Aandachtsmapping: Aangezien Mamba geen expliciete aandachtskoppen heeft, construeren de auteurs een matrix in de stijl van aandacht ($A$) uit de SSM-laag van de Mamba-mixer met behulp van procedures voor impliciete aandachtsmapping. Dit maakt het mogelijk om vectors op token-niveau voor belang af te leiden en Activeringsruimten te definiëren als gewogen sommen van verborgen toestanden.
• Stochastische Parameterontleding (SPD): SPD wordt gebruikt om modelparameters te ontleden in schaars gebruikte vectoren. Metrieken zoals entropie, variantie, effectieve rang en KL-divergentie na ablatie worden laag-voor-laag berekend om verstoppingen te identificeren.
• Delta-gevoelige Ruimten: De auteurs identificeren ruimten die een hoge variantie vertonen als reactie op de $\Delta$-parameter (die de gediscretiseerde tijdstap en input-afhankelijke updates controleert). Deze ruimten worden als kritiek beschouwd voor recurrente dynamiek en worden geselecteerd als sturingsdoelen.

B. Post-hoc Sturing

Zodra verstoppingen zijn geïdentificeerd, introduceren de auteurs een interventie op het moment van testen die geen fijnafstemming vereist:

• Doelselectie: Het proces identificeert specifieke lagen (bijvoorbeeld Laag 20 in Mamba-130M) die pieken in entropie en hoge KL-divergentie na ablatie vertonen. Binnen deze lagen worden de meest kritieke Delta-gevoelige ruimten geselecteerd op basis van hun impact op de nauwkeurigheid van de voorspelling van het volgende token wanneer ze worden geablateerd.
• Interventie: De activeringen van deze geïdentificeerde ruimten worden vermenigvuldigd met een scalair sturingsfactor (bijvoorbeeld 5x voor ruimten met hoge impact, 2x voor gematigde) tijdens inferentie. Dit versterkt de signaalstroom door de geïdentificeerde verstoppingen.

C. Stable-Mamba Architectuur

Om te valideren dat de geïdentificeerde verstoppingen structurele beperkingen zijn in plaats van trainingsartefacten, stellen de auteurs Stable-Mamba voor, een gewijzigde architectuur die inzichten uit interpreteerbaarheid integreert:

• Dynamiek van Toestanden met Meerdere Tijdschalen: Parallelle SSM-takken (snel, gemiddeld, traag) maken gelijktijdige verwerking op korte, middellange en lange afstand mogelijk.
• Geënssembleerde Output: Een gewogen ensemble vervangt de lineaire outputprojectie om adaptieve temporele resolutie mogelijk te maken.
• Schaarse Injectie van Globale Context: Schaarse aandacht injecteert globale samenvattingen in lokale toestanden om de informatiestroom te stabiliseren.
• Gestuurde Gating en Adaptieve Compressie: Ensemble-gates verminderen extreme schaarsheid, en adaptieve compressie moduleert de capaciteit dynamisch.
• Gradientenvergroting en Geschaalde Residuen: Deze modificaties stabiliseren diepe SSM-ketens en verbeteren de gradientenstroom.

3. Belangrijkste Resultaten

De studie evalueert de voorgestelde methoden over 7 SSM's en 6 diverse benchmarks (waaronder TriviaQA, SQuAD, MuSiQue, IFEval, RULER en DROP).

• Sturingsprestaties: De post-hoc sturingsinterventie verbetert de prestaties over 7 verschillende SSM's met gemiddeld 8,27% zonder enige taakspecifieke fijnafstemming. Opmerkelijk is dat de interventie werd geïdentificeerd met behulp van samples uit "The Pile", maar zich generaliseerde naar diverse benchmarks.
• Winsten op Lange Context: Op benchmarks voor lange context (RULER, Long Range Arena, LongBench v2) tonen de gestuurde modellen aanzienlijke verbeteringen. Zo bereikt Stable-Mamba een winst van 35 procentpunten op de PathFinder-taak binnen Long Range Arena in vergelijking met de standaard Mamba.
• Validatie van Verstoppingen:
  • Entropie-analyse: Standaard Mamba vertoont een scherpe piek in entropie bij Laag 20, wat wijst op een routeringsverstopping. Zowel sturing als de Stable-Mamba-architectuur gladde deze entropie en herstellen de informatiestroom.
  • Ablatiestudies: Het verwijderen van Laag 20 in standaard Mamba verbetert de prestaties eigenlijk (77,0% nauwkeurigheid versus 63,5% baseline), wat bevestigt dat het fungeert als een beperkende verstopping. Omgekeerd verslechtert het ablaten van dezelfde laag in Gestuurde of Stable-modellen de prestaties, wat aangeeft dat de verstopping is verzacht.
  • Parametergevoeligheid: De $dt\_proj.bias$-parameter in Laag 20 fungeert als een "hoofdgat" die effectief bevroren raakt tijdens training, 45% van de voorspellingen controleert en catastrofale informatieverlies veroorzaakt (KL-divergentie van 813) wanneer geablateerd.
• Efficiëntie: Stable-Mamba voegt slechts 256 parameters toe. Hoewel het de FLOPs per token met ongeveer 2,8x verhoogt, resulteert parallelisatie in slechts een ~9% toename in inferentielatentie en verwaarloosbare geheugennadelen in vergelijking met standaard Mamba.

4. Betekenis en Beweringen

Het artikel beweert aan te tonen dat mechanistische interpreteerbaarheid kan worden vertaald naar praktische prestatiewinsten voor SSM's. Door modelfalen te lokaliseren tot specifieke activeringsruimten en architectonische componenten, tonen de auteurs aan dat:

1. Interpreteerbaarheid actiegericht is: Het identificeren van verstoppingen maakt gerichte interventies (sturing) mogelijk die de prestaties verbeteren zonder opnieuw te trainen.
2. Structurele beperkingen bestaan: SSM's lijden onder specifieke architectonische beperkingen (bijvoorbeeld verwerking op één tijdschaal, bevroren gating-parameters) die redeneren op lange context en informatiestroom belemmeren.
3. Minimale wijzigingen leveren grote winsten op: Kleine architectonische veranderingen geleid door interpreteerbaarheid (Stable-Mamba) kunnen deze beperkingen oplossen, en prestaties bereiken die vergelijkbaar zijn met of groter zijn dan die van grotere modellen, terwijl de efficiëntie behouden blijft.

De auteurs houden bescheidenheid in het algemeeniseren, en merken op dat hun proces momenteel is gericht op Mamba en het onduidelijk blijft of dezelfde aanpak zich generaliseert naar andere architecturen of zeer grote parametergroottes. Ze erkennen ook dat de set van geëvalueerde taken standaard is en dat de generaliseerbaarheid naar complexe redeneertaken verder onderzoek vereist.