Causal Evidence of Stack Representations in Modeling Counter Languages Using Transformers

Technische Samenvatting: Causaal bewijs van stack-representaties bij het modelleren van counter-talen met Transformers

Probleemstelling
Hoewel mechanistische interpreteerbaarheid succesvol latente kenmerken heeft geïdentificeerd in taalmodellen (LM's) die getraind zijn op formele talen, blijft er een kritieke kloof bestaan in het vaststellen van de causale rol van deze kenmerken. Vorig werk, zoals Tiwari et al. (2025), toonde aan dat transformers getraind op counter-talen (bijv. Dyck-1, Shuffle-k) representaties leren die consistent zijn met een onderliggende stack-structuur. Deze studies bepaalden echter niet of deze representaties louter gecorreleerd zijn met de taak of dat ze causaal noodzakelijk zijn voor de berekening van het model. Dit artikel adresseert de vraag: Zijn de geleerde stack-representaties causaal betrokken bij het vermogen van de transformer om het volgende token te voorspellen?

Methodologie
De studie maakt gebruik van een experimentele pipeline in drie fasen bestaande uit training, probing en causale interventie:

1. Modeltraining: Een encoder-only transformer wordt getraind op next-token voorspelling voor Shuffle-k talen (waarbij $k \in \{2, 4, 6, 8\}$). Shuffle-k is gedefinieerd als de interleaving van $k$ onafhankelijke Dyck-1 talen, wat vereist dat het model de diepte van $k$ verschillende stack-structuren gelijktijdig bijhoudt. De modelarchitectuur bestaat uit een enkele laag met 4 attention heads, een embedding dimensie van 32 en een verborgen dimensie van 64. Het model bereikt een perfecte validatienauwkeurigheid tegen de 25e epoch, waardoor alle daaropvolgende analyses worden uitgevoerd op een volledig geconvergeerd model.
2. Probing voor Stack-diepte: Lineaire probes worden getraind om de diepte van specifieke stacks ($c_i(t)$) te voorspellen vanuit de verborgen toestanden ($h^{(l)}_t$) van het model op laag $l$. Om ervoor te zorgen dat de probes werkelijke structurele informatie vastleggen in plaats van spuriële correlaties, wordt een control probe getraind op geperturbeerde labels. De metriek Selectiviteit wordt gedefinieerd als het verschil in nauwkeurigheid tussen de experimentele probe en de control probe. Een hoge selectiviteit geeft aan dat stack-diepte werkelijk is gecodeerd in de verborgen ruimte.
3. Causale Interventie (Ablatie):
   • Richtingsextractie: De gewichtsmatrix $W$ van de probe wordt ontleed via Singular Value Decomposition (SVD). De belangrijkste rechter enkelvoudige vector, $w$, wordt geïdentificeerd als de hoofdrichting die overeenkomt met de stack-diepte representatie.
   • Ablatie: Tijdens de inferentie wordt het component van de verborgen toestand dat is uitgelijnd met $w$ verwijderd met behulp van de operatie $h' = h - \alpha(w \cdot h)w$, waarbij $\alpha \in [0, 1]$ de ablatiekracht vertegenwoordigt.
   • Controle: Willekeurige richtingen worden op vergelijkbare wijze geablateerd om als baseline te dienen.
   • Evaluatie: De geablateerde toestanden worden direct doorgegeven aan de decoder head (zonder de interne lagen van de transformer opnieuw uit te voeren) om de impact op de Positionele Nauwkeurigheid ($A_{pos}$) en Sequentiële Nauwkeurigheid ($A_{seq}$) te meten.

Belangrijkste Resultaten

• Bestaan van Representaties: Lineaire probes extraheren succesvol stack-diepte informatie met een hoge selectiviteit (bijv. 85,8% voor $k=6$, Stack 2), wat bevestigt dat het model representaties leert die consistent zijn met een stack-structuur.
• Causale Noodzakelijkheid: Gerichte ablatie van de primaire stack-richting ($w$) veroorzaakt een dramatische instorting van de modelprestaties. Naarmate de ablatiekracht $\alpha$ toeneemt van 0 naar 1, daalt de Sequentiële Nauwkeurigheid naar bijna 0%.
• Specificiteit: Het ablateren van willekeurige richtingen heeft geen significant effect op de nauwkeurigheid, wat bevestigt dat de prestatie-instorting specifiek is voor de verwijdering van de stack-representatie en geen algemeen artefact is van vectormanipulatie.
• Metriekgevoeligheid: De studie observeert dat terwijl de positionele nauwkeurigheid daalt, de sequentiële nauwkeurigheid scherper instort. Dit wordt toegeschreven aan het cumulatieve karakter van fouten in sequentiegeneratie, waarbij één onjuist token de gehele string ongeldig maakt.

Betekenis en Claims
Het artikel claimt sterk empirisch bewijs te leveren dat stack-representaties in transformers getraind op counter-talen niet slechts emergente artefacten zijn, maar causaal noodzakelijk zijn voor de interne berekening en next-token voorspelling van het model.

Door aan te tonen dat het verwijderen van een specifieke, geleerde representatierichting leidt tot een bijna totale mislukking van het model, beargumenteren de auteurs dat de transformer expliciet deze stack-achtige structuren gebruikt om de vereiste logische redeneertaken voor formele talen op te lossen. Dit brengt het veld voorbij correlationele analyse (probing) naar een causale mechanistische verstandhouding, waarbij wordt gevalideerd dat het "black box"-gedrag van het model steunt op specifieke, identificeerbare computationele primitieven.

Toekomstige Richtingen
De auteurs suggereren dat toekomstig werk moet onderzoeken:

• Interveniëren halverwege het netwerk door geablateerde of bewerkte toestanden door het feed-forward netwerk te sturen om veranderingen in de daaropvolgende outputs te observeren.
• Onderzoeken welke rol verschillende architecturen spelen bij de emergentie, redundantie en causaliteit van deze representaties.
• Begrijpen hoe transformers de manipulatie van andere emergente computationele primitieven zoals counters afhandelen.