Causal Evidence of Stack Representations in Modeling Counter Languages Using Transformers

Technisches Resümee: Kausale Evidenz für Stack-Repräsentationen bei der Modellierung von Counter-Sprachen mittels Transformer-Modellen

Problemstellung
Während die mechanistische Interpretierbarkeit erfolgreich latente Merkmale in Sprachmodellen (LMs) identifiziert hat, die auf formalen Sprachen trainiert wurden, bleibt eine kritische Lücke in der Feststellung der kausalen Rolle dieser Merkmale bestehen. Vorherige Arbeiten, wie etwa von Tiwari et al. (2025), zeigten, dass Transformer, die auf Counter-Sprachen (z. B. Dyck-1, Shuffle-k) trainiert wurden, Repräsentationen erlernen, die konsistent mit einer zugrunde liegenden Stack-Struktur sind. Diese Studien konnten jedoch nicht bestimmen, ob diese Repräsentationen lediglich mit der Aufgabe korreliert oder kausal notwendig für die Berechnung des Modells sind. Diese Arbeit adresset die Frage: Sind die gelernten Stack-Repräsentationen kausal an der Fähigkeit des Transformers zur Vorhersage des nächsten Tokens beteiligt?

Methodik
Die Studie verwendet eine dreistufige experimentelle Pipeline bestehend aus Training, Probing und kausaler Intervention:

1. Modelltraining: Ein Encoder-only-Transformer wird auf Next-Token-Prediction für Shuffle-k-Sprachen trainiert (wobei $k \in \{2, 4, 6, 8\}$). Shuffle-k ist definiert als das Verflechten von $k$ unabhängigen Dyck-1-Sprachen, was das Modell dazu zwingt, die Tiefe von $k$ verschiedenen Stacks gleichzeitig zu verfolgen. Die Modellarchitektur besteht aus einer einzelnen Schicht mit 4 Attention-Heads, einer Embedding-Dimension von 32 und einer Hidden-Dimension von 64. Das Modell erreicht eine perfekte Validierungsgenauigkeit bis zum 25. Epochen, wodurch sichergestellt ist, dass alle nachfolgenden Analysen auf einem vollständig konvergierten Modell durchgeführt werden.
2. Probing der Stack-Tiefe: Lineare Probes werden trainiert, um die Tiefe spezifischer Stacks ($c_i(t)$) aus den Hidden States des Modells ($h^{(l)}_t$) in Schicht $l$ vorherzusagen. Um sicherzustellen, dass die Probes echte strukturelle Informationen erfassen und keine spuriösen Korrelationen, wird eine Kontroll-Probe auf perturbierten Labels trainiert. Die Metrik Selektivität ist definiert als die Differenz zwischen der Genauigkeit der experimentellen Probe und der Kontroll-Probe. Eine hohe Selektivität deutet darauf hin, dass die Stack-Tiefe tatsächlich im Hidden Space kodiert ist.
3. Kausale Intervention (Ablation):
   • Extraktion der Richtung: Die Gewichtsmatrix $W$ der Probe wird mittels Singularwertzerlegung (SVD) dekomponiert. Der oberste rechts stehende Singulärvektor, $w$, wird als die Hauptrichtung identifiziert, die der Stack-Tiefe-Repräsentation entspricht.
   • Ablation: Während der Inferenz wird die Komponente des Hidden State, die mit $w$ ausgerichtet ist, durch die Operation $h' = h - \alpha(w \cdot h)w$ entfernt, wobei $\alpha \in [0, 1]$ die Ablationsstärke darstellt.
   • Kontrolle: Zufällige Richtungen werden ähnlich ablatiert, um eine Baseline zu dienen.
   • Evaluierung: Die ablatierten States werden direkt an den Decoder-Head übergeben (ohne die internen Schichten des Transformers erneut zu durchlaufen), um die Auswirkungen auf die Positionsgenauigkeit ($A_{pos}$) und die Sequenzgenauigkeit ($A_{seq}$) zu messen.

Kernergebnisse

• Existenz der Repräsentation: Lineare Probes extrahieren erfolgreich Informationen über die Stack-Tiefe mit hoher Selektivität (z. B. 85,8 % für $k=6$, Stack 2), was bestätigt, dass das Modell Repräsentationen lernt, die konsistent mit einer Stack-Struktur sind.
• Kausale Notwendigkeit: Die gezielte Ablation der primären Stack-Richtung ($w$) verursacht einen dramatischen Zusammenbruch der Modellleistung. Wenn die Ablationsstärke $\alpha$ von 0 auf 1 steigt, sinkt die Sequenzgenauigkeit auf nahezu 0 %.
• Spezifität: Die Ablation zufälliger Richtungen hat keinen signifikanten Effekt auf die Genauigkeit, was bestätigt, dass der Leistungszusammenbruch spezifisch für die Entfernung der Stack-Repräsentation ist und kein allgemeines Artefakt der Vektormanipulation darstellt.
• Metrik-Sensitivität: Die Studie beobachtet, dass während die Positionsgenauigkeit sinkt, die Sequenzgenauigkeit deutlich stärker einbricht. Dies wird auf die kumulative Natur von Fehlern in der Sequenzgenerierung zurückgeführt, bei der ein einziges falsches Token die gesamte Zeichenkette ungültig macht.

Bedeutung und Ansprüche
Das Paper behauptet, starke empirische Evidenz dafür zu liefern, dass Stack-Repräsentationen in Transformern, die auf Counter-Sprachen trainiert wurden, nicht bloß emergente Artefakte sind, sondern kausal notwendig für die interne Berechnung und die Next-Token-Prediction des Modells.

Durch den Nachweis, dass die Entfernung einer spezifischen, gelernten Repräsentationsrichtung zu einem nahezu totalen Versagen des Modells führt, argumentieren die Autoren, dass der Transformer explizit diese Stack-ähnlichen Strukturen nutzt, um die logischen Aufgaben zu lösen, die von formalen Sprachen gefordert werden. Dies führt das Feld über die korrelative Analyse (Probing) hinaus hin zu einem kausalen mechanistischen Verständnis und validiert, dass die „Black Box“-Verhaltensweise des Modells auf spezifischen, identifizierbaren Rechenprimitiven beruht.

Zukünftige Richtungen
Die Autoren schlagen vor, dass zukünftige Arbeiten untersuchen sollten:

• Die Intervention mitten im Netzwerk, indem ablatierte oder editierte States durch das Feed-Forward-Netzwerk geleitet werden, um Veränderungen in den nachfolgenden Outputs zu beobachten.
• Die Untersuchung der Rolle unterschiedlicher Architekturen bei der Emergenz, Redundanz und Kausalität dieser Repräsentationen.
• Das Verständnis darüber, wie Transformer die Manipulation anderer emergenter Rechenprimitives wie Counter handhaben.