🤖 machine learning

Task Structure Reverses Layerwise State Encoding in Sequence Models

This paper demonstrates that the layerwise distribution of state encoding in sequence models is not a fixed architectural trait but reverses based on task computational structure, revealing that algebraic properties like commutativity are less predictive of mechanistic signatures than the underlying computational requirements such as prefix updates versus stack operations.

Original authors: Yuhang Jiang

Published 2026-06-02

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Yuhang Jiang

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are trying to understand how different types of "thinking machines" (AI models) keep track of information as they read a story. For a long time, researchers thought that a machine's "thinking style" was fixed, like a person's personality. They believed:

Transformers (the current standard AI) spread their thoughts out evenly across their whole brain.
Recurrent models (older, loop-based AIs) kept their main thoughts tucked away in the very last part of their brain.

This paper says: "Not so fast."

The authors discovered that these machines don't have a fixed personality. Instead, they change how they organize their thoughts depending on what job they are doing. It's like a chef who uses a different knife for chopping vegetables versus filleting a fish. The tool changes based on the task, not just the chef's identity.

Here is the breakdown of their discovery using simple analogies:

1. The Two Types of Jobs

The researchers gave the machines three different puzzles to solve:

The "Running Total" Job (Parity & S3): Imagine counting how many times you've flipped a coin. You just need to remember the current count (0 or 1) and update it as you go. This is a simple, linear update.
The "Stack" Job (Dyck): Imagine checking if a sentence has balanced parentheses, like (( )). You can't just count; you need to remember which opening bracket matches which closing one. You have to build a mental "stack" (like a pile of plates) to keep track of nesting.

2. The Great Reversal

The paper found that the machines flip their strategy depending on the job:

On the "Running Total" Job:
- Recurrent models (like Mamba) act like a secret keeper. They do a lot of work early on, but the final answer only becomes clear at the very last layer of the brain. It's like a magician who does all the sleight of hand in the shadows and reveals the trick only at the very end.
- Transformers act like a team huddle. They build the answer gradually, layer by layer, from the start to the finish. Everyone contributes a little bit along the way.
On the "Stack" Job:
- The roles flip!
- Transformers suddenly become the secret keepers. They figure out the answer almost immediately (in the first few layers) and then just hold onto it.
- Recurrent models become the team huddle. They have to build the solution slowly, layer by layer, to get the answer right.

The Takeaway: You cannot say "Transformer always spreads information" or "Mamba always hides it." The strategy depends entirely on whether the task is a simple update or a complex stack.

3. The "Non-Commutative" Test (The S3 Puzzle)

To prove this wasn't just about math rules (like whether order matters in addition), they added a third, trickier puzzle called S3. This was a "Running Total" job where the order did matter (like putting on socks before shoes is different from shoes before socks).

The Prediction: If the difference was just about math rules, this tricky job should look like the "Stack" job.
The Reality: The machines treated this tricky job exactly like the simple "Running Total" job. They used the same "secret keeper" or "team huddle" strategies as before.
Conclusion: The machines aren't reacting to math rules; they are reacting to the computational structure (is it a simple update or a complex stack?).

4. The "Readable" vs. "Important" Trap

The researchers also looked at where the information is stored. They found a surprising disconnect, especially in larger, pre-trained models:

The "Readable" Layer: This is where you can easily "read" the answer out of the machine's brain (like finding a sticky note with the answer written on it).
The "Important" Layer: This is the part of the brain that, if you broke it, would cause the machine to fail.

The Discovery:

In small, simple models, the "Readable" layer and the "Important" layer are usually the same.
In big, pre-trained models, they often don't match.
- Example: On the "Stack" job, a big model might have the answer clearly written on a sticky note in the middle of the brain (Layer 7), but if you break the end of the brain (Layer 11), the model still works fine. However, if you break the middle, it crashes.
- Example: On the "Running Total" job, the answer might be clearly visible at the very end, but breaking that specific spot doesn't stop the model from working because the information is spread out everywhere else.

The Lesson: Just because you can see the answer in a specific part of the brain doesn't mean that part is the only thing keeping the machine alive. The machine might be holding the answer in many different places at once.

Summary

This paper teaches us that AI architectures are not rigid. They are flexible.

Task matters: A machine's internal strategy changes based on whether it's doing a simple update or a complex stack.
Structure matters: The "shape" of the problem (update vs. stack) dictates the strategy more than the math rules do.
Visibility $\neq$ Necessity: In big models, finding where the answer is "written" doesn't always tell you where the machine is most fragile.

The authors conclude that we can't just ask, "How does this AI work?" We must ask, "How does this AI work for this specific task?"

Technical Summary: Task Structure Reverses Layerwise State Encoding in Sequence Models

Problem Statement

Mechanistic interpretability of sequence models often treats layerwise state encodings as fixed architectural traits. The prevailing assumption is that recurrent models (e.g., LSTMs, GRUs) and State-Space Models (SSMs like Mamba) concentrate readable state in late layers, while attention-based models (Transformers) distribute state accumulation across depth. This paper challenges the view that these encoding strategies are stable architectural invariants. It posits that the same architecture reverses its layerwise state-encoding profile when the computational structure of the task changes. Furthermore, the literature often conflates two potential drivers for these patterns: algebraic structure (specifically commutativity) versus computational structure (prefix updates vs. stack tracking). The authors aim to disentangle these factors and determine whether state encoding is a property of the architecture alone or an interaction between architecture and task.

Methodology

The study employs a controlled experimental suite across five architecture families: Transformers, Mamba, Mamba-2, LSTMs, and GRUs.

Tasks:
The authors utilize three formal state-tracking tasks and one semi-real benchmark:

Parity: A commutative, single-bit prefix-update task.
$S_3$ Permutation Composition: A non-commutative, 6-class prefix-update task. This serves as the critical discriminant: if algebraic commutativity drives the encoding profile, $S_3$ should behave like Dyck; if computational structure drives it, $S_3$ should behave like Parity.
Dyck- $k$ : A bounded-depth stack-tracking task (bracket matching) with a maximum depth of 10.
Code-Depth: A token-level nesting-depth task derived from Python code in CodeSearchNet, serving as a semi-real validation.

Experimental Setup:

Formal Tasks: Experiments use 4-layer models with $d_{model}=128$ . Architectures are parameter-matched (approx. 830K parameters) where possible. Training occurs at sequence length $L=40$ , with evaluation at $L \in \{40, 60, 80, 100, 120\}$ .
Pretrained Scale: Fine-tuning is performed on Mamba-130M (24 layers) and Pythia-160M/410M (12/24 layers).
Measurements:
- Linear Probing: Training logistic regression classifiers on frozen hidden states to measure state readability at each layer.
- Causal Interventions: Zero-ablation (projecting out probe directions) and subspace ablation (removing top PCA or singular directions) to test functional necessity.
- OOD Generalization: Applying ablations trained at $L=40$ to unseen lengths ( $L=80, 100, 120$ ).
- Component Attribution: Ablating Mamba's Conv1D component under a matched pure-PyTorch backend to isolate its contribution across tasks.
- Activation Patching: In pretrained Mamba, patching mid-position activations from clean inputs into corrupted inputs to measure logit recovery.

Key Contributions

Task-Dependent Reversal: The paper demonstrates that the "concentrated late" vs. "distributed buildup" pattern is not an architectural invariant. It flips depending on the task:
- Prefix-Update Tasks (Parity, $S_3$ ): Recurrent models and SSMs concentrate readable state in the final layer, while Transformers build it gradually across layers.
- Bounded Stack-Tracking (Dyck): The pattern reverses. Transformers achieve near-perfect readability immediately (early layers), while recurrent models and SSMs build the representation progressively.
Disentangling Algebra and Computation: By introducing the non-commutative $S_3$ task, the authors show that the encoding profile tracks computational structure (prefix update vs. stack) rather than algebraic structure (commutativity). $S_3$ groups with Parity across all architectures, refuting the hypothesis that commutativity is the primary driver of the encoding profile.
Readability–Causality Decoupling: The study reveals that the layer where state is most linearly readable (best probe layer) does not always coincide with the layer that is causally necessary for the task. This decoupling is particularly pronounced in pretrained models, where representations are more redundant and distributed.
Component-Level Task Dependence: Even within a single architecture (Mamba), the importance of specific components varies by task. For instance, the Conv1D component is critical for $S_3$ and Parity but has no significant effect on Dyck performance, challenging the notion of a single "mechanism" for an architecture.

Key Results

Behavioral Level: Recurrent baselines generalize perfectly on Parity and $S_3$ and outperform SSMs and Transformers on Dyck. S3 performance ordering (Recurrent > SSM > Transformer) mirrors Parity, not Dyck.
Probing Profiles:
- Parity/ $S_3$ : Mamba and LSTMs show a sharp jump in probe accuracy at the final layer ( $L_3$ in 4-layer models). Transformers show a gradual increase from $L_1$ to $L_3$ .
- Dyck: Transformers reach high accuracy at $L_0$ or $L_1$ . Mamba and LSTMs show a gradual increase, reaching saturation only at later layers.
- $S_3$ Confirmation: $S_3$ follows the Parity profile (late concentration for RNNs/SSMs, gradual for Transformers), confirming the computational structure hypothesis.
Causal Interventions:
- In 4-layer formal models, ablation of the best probe direction causes significant accuracy drops, and these directions remain necessary at longer sequence lengths (OOD).
- In Pretrained Models, a split occurs:
  - Pythia (Dyck): A strong middle-layer bottleneck exists (ablation at $L_6$ – $L_7$ drops accuracy by ~81% at 160M). The best probe layer ( $L_{11}$ ) is not the most causally critical.
  - Mamba (Parity/Dyck/ $S_3$ ): The final layer is highly readable, but single-direction ablation causes negligible drops (<0.2%). However, mid-position activation patching at the final layer recovers ~97–98% of the logit gap, indicating the state is sufficient but distributed across many directions, not captured by a single linear probe.
Conv1D Attribution: Removing Conv1D from Mamba causes a large drop on $S_3$ (-12.64 pp) and Parity (-6.83 pp), but no significant drop on Dyck (+4.04 pp, not significant). This indicates that component importance is task-dependent.

Significance and Claims

The paper argues that asking "how does architecture X encode state?" is incomplete without specifying the computational structure of the task. The authors claim that:

Mechanistic signatures are joint properties: State encoding strategies are not fixed architectural traits but emerge from the interaction between the model's inductive biases and the task's computational requirements (prefix update vs. stack).
Methodological caution: Extrapolating mechanistic findings from one task regime (e.g., prefix updates) to another (e.g., bounded stack tracking) risks mistaking task-specific signatures for architectural invariants.
Readability vs. Causality: Probing identifies where information is linearly available, but this does not always map to where the computation is bottlenecked, especially in large pretrained models where redundancy is high.
Algebra vs. Computation: In the context of these formal tasks, computational structure (sequential prefix update vs. hierarchical stack) is a better predictor of layerwise encoding profiles than algebraic properties like commutativity.

The authors conclude that their findings hold across formal tasks, pretrained scales (up to 410M parameters), and semi-real code benchmarks, suggesting a robust interaction between architecture and task structure in sequence modeling.