Causal Evidence of Stack Representations in Modeling Counter Languages Using Transformers
Provides causal evidence for the role of learned stack structures in transformers, advancing mechanistic interpretability for formal language tasks.
The paper shows that transformers trained on counter languages learn stack representations that are causally necessary for performance: ablating the stack direction causes accuracy to drop to near 0%.
Formal languages have proven to be effective conduits to understand the inner mechanisms of transformers. Past work has shown that transformers trained on next token prediction over counter languages learn representations consistent with an underlying stack structure. Beyond representational analysis, this paper investigates the causal role of these representations. Linear probes are trained to predict the stack depth at each token from the model's hidden states, and a principal representation direction is extracted from the probe. Ablation of this direction from the model causes sequential accuracy to collapse to near 0%, providing strong empirical evidence that the stack representation is not just learned, but is causally necessary for model performance.