Causal Evidence of Stack Representations in Modeling Counter Languages Using Transformers

Technical Summary: Causal Evidence of Stack Representations in Modeling Counter Languages Using Transformers

Problem Statement
While mechanistic interpretability has successfully identified latent features in language models (LMs) trained on formal languages, a critical gap remains in establishing the causal role of these features. Previous work, such as Tiwari et al. (2025), demonstrated that transformers trained on counter languages (e.g., Dyck-1, Shuffle-k) learn representations consistent with an underlying stack structure. However, these studies did not determine whether these representations are merely correlated with the task or are causally necessary for the model's computation. This paper addresses the question: Are the learned stack representations causally involved in the transformer's ability to predict the next token?

Methodology
The study employs a three-stage experimental pipeline involving training, probing, and causal intervention:

1. Model Training: An encoder-only transformer is trained on next-token prediction for Shuffle-k languages (where $k \in \{2, 4, 6, 8\}$). Shuffle-k is defined as the interleaving of $k$ independent Dyck-1 languages, requiring the model to track the depth of $k$ distinct bracket stacks simultaneously. The model architecture consists of a single layer with 4 attention heads, an embedding dimension of 32, and a hidden dimension of 64. The model achieves perfect validation accuracy by the 25th epoch, ensuring all subsequent analyses are performed on a fully converged model.
2. Probing for Stack Depth: Linear probes are trained to predict the depth of specific stacks ($c_i(t)$) from the model's hidden states ($h^{(l)}_t$) at layer $l$. To ensure the probes capture genuine structural information rather than spurious correlations, a control probe is trained on perturbed labels. The metric Selectivity is defined as the accuracy difference between the experimental probe and the control probe. High selectivity indicates that stack depth is genuinely encoded in the hidden space.
3. Causal Intervention (Ablation):
   • Direction Extraction: The probe's weight matrix $W$ is decomposed via Singular Value Decomposition (SVD). The top right singular vector, $w$, is identified as the principal direction corresponding to the stack depth representation.
   • Ablation: During inference, the component of the hidden state aligned with $w$ is removed using the operation $h' = h - \alpha(w \cdot h)w$, where $\alpha \in [0, 1]$ represents the ablation strength.
   • Control: Random directions are similarly ablated to serve as a baseline.
   • Evaluation: The ablated states are passed directly to the decoder head (without re-running the transformer's internal layers) to measure the impact on Positional Accuracy ($A_{pos}$) and Sequential Accuracy ($A_{seq}$).

Key Results

• Representation Existence: Linear probes successfully extract stack depth information with high selectivity (e.g., 85.8% for $k=6$, Stack 2), confirming that the model learns representations consistent with a stack structure.
• Causal Necessity: Targeted ablation of the principal stack direction ($w$) causes a dramatic collapse in model performance. As the ablation strength $\alpha$ increases from 0 to 1, Sequential Accuracy drops to near 0%.
• Specificity: Ablating random directions has no significant effect on accuracy, confirming that the performance collapse is specific to the removal of the stack representation and not a general artifact of vector manipulation.
• Metric Sensitivity: The study observes that while positional accuracy dips, sequential accuracy collapses more sharply. This is attributed to the compounding nature of errors in sequence generation, where a single incorrect token invalidates the entire string.

Significance and Claims
The paper claims to provide strong empirical evidence that stack representations in transformers trained on counter languages are not just emergent artifacts but are causally necessary for the model's internal computation and next-token prediction.

By demonstrating that the removal of a specific, learned representation direction leads to a near-total failure of the model, the authors argue that the transformer explicitly utilizes these stack-like structures to solve the logical reasoning tasks required by formal languages. This moves the field beyond correlational analysis (probing) toward causal mechanistic understanding, validating that the model's "black box" behavior relies on specific, identifiable computational primitives.

Future Directions
The authors suggest that future work should explore:

• Intervening mid-network by passing ablated or edited states through the feed-forward network to observe changes in subsequent outputs.
• Investigating the role of different architectures in the emergence, redundancy, and causality of these representations.
• Understanding how transformers handle the manipulation of other emergent computational primitives like counters.