💬 NLP

Mechanistic evaluation of Transformers and state space models

This paper employs causal interventions to reveal that while Transformers and SSMs may achieve similar accuracy on associative recall tasks, they rely on fundamentally different mechanisms—specifically, Transformers and Based models use in-context induction whereas most SSMs rely on single-layer state computations, with Mamba uniquely achieving induction through short convolutions rather than its core SSM component.

Original authors: Aryaman Arora, Neil Rathi, Nikil Roashan Selvam, Róbert Csordás, Dan Jurafsky, Christopher Potts

Published 2026-02-02

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Aryaman Arora, Neil Rathi, Nikil Roashan Selvam, Róbert Csordás, Dan Jurafsky, Christopher Potts

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 teach a robot how to play a memory game. The game is simple: You show the robot a list of pairs, like "Apple = Red," "Banana = Yellow," and "Grape = Purple." Then, you ask, "What is the color of the Banana?" The robot needs to look back at the list, find "Banana," and remember that it is linked to "Yellow."

In the world of AI, this is called Associative Recall. The paper you shared is a deep dive into how different types of AI brains (architectures) solve this game. The researchers didn't just look at who got the highest score; they put the robots under a microscope to see exactly what was happening inside their "heads" while they played.

Here is the story of their findings, broken down into simple concepts.

The Two Main Contenders

The paper compares two main types of AI architectures:

Transformers: The current champions of AI (like the ones powering most chatbots). They use a mechanism called "Attention," which is like having a spotlight that can instantly scan the whole list of pairs to find the answer.
State Space Models (SSMs): The new, faster, and more efficient challengers (like Mamba and DeltaNet). They are designed to be lighter and quicker than Transformers but have been suspected of being worse at memory tasks.

The Big Discovery: It's Not Just About the Score

Usually, if an AI gets 99% on a test, we assume it learned the lesson. But this paper argues that getting the right answer doesn't mean you learned the right way.

The researchers used a special "surgery" technique called Causal Intervention. Imagine you have a robot that solves the memory game perfectly. Then, you secretly swap the word "Banana" with "Car" in the robot's memory after it has read the list but before it gives the answer.

If the robot suddenly says "Yellow is the color of a Car," it means it was actually remembering the link (Induction).
If the robot just guesses or fails, it means it wasn't truly storing the link in a flexible way; it was just looking at the immediate surroundings.

What They Found

1. The "Induction" vs. "Direct Retrieval" Difference

Transformers and "Based" models: These models learned to play the game the "smart" way. They created a mental map where they stored the connection between the item and its value (e.g., "Banana" is stuck to "Yellow") as they read the list. When asked a question, they looked up this stored map. The researchers call this Induction.
Most SSMs (like H3 and Hyena): These models failed the game. They couldn't figure out how to store the links.
Mamba and DeltaNet: These models got high scores, but they cheated! They didn't build a mental map. Instead, they used a "shortcut." When the question "Banana?" came up, they looked immediately at the token right before it or used a very specific, rigid trick to grab the answer. They didn't learn the general rule of "associating" things; they just learned to grab the answer from a specific spot.

2. The Secret Weapon: The "Short Convolution"

The researchers discovered that the models that did succeed (Mamba and Based) had a secret weapon: a tiny, short-term memory tool called a short convolution.

Analogy: Imagine reading a book. A "long convolution" is like reading the whole chapter to understand a sentence. A "short convolution" is like glancing at the word you just read and the one before it.
The Finding: Mamba uses this "glance" (short convolution) to link "Banana" to "Yellow" as it reads. Without this specific tool, Mamba completely forgets how to play the game. It turns out, Mamba isn't using its main "brain" (the SSM part) to do the memory work; it's using this tiny "glance" tool to do the heavy lifting.

3. The Twist: The "Tree" Game (ATR)

To see if these findings were just a fluke of the simple game, the researchers invented a harder version called Associative Treecall (ATR).

The Change: In the simple game, "Banana" and "Yellow" were always right next to each other. In the "Tree" game, they could be far apart, separated by other words, like a family tree. "John had a dog. The dog chased a cat. The cat saw a bird. What did John have?"
The Result:
- Transformers and Based: Still used their "Induction" map. They handled the distance easily.
- Mamba (with the shortcut): Still used its "glance" tool and did well.
- Mamba (without the shortcut): Here is the surprise! When the researchers removed Mamba's "short convolution" tool, Mamba failed the simple game but succeeded at the hard "Tree" game.
- Why? Without the shortcut, Mamba was forced to learn the "Induction" map (the same way Transformers do) to solve the harder problem. It proved that Mamba is capable of learning the smart way, but it only does so when it absolutely has to.

The Takeaway

The paper concludes that accuracy is a liar.

Two AI models might get the same score on a test, but one might be a genius who understands the rules (Induction), while the other is a trickster who just memorized the answer key for that specific test (Direct Retrieval).

Transformers are like librarians who organize books on shelves so they can find any connection later.
Mamba (in its standard form) is like a person who only remembers the last thing they saw, unless they use a specific "sticky note" (short convolution) to help them.
The Lesson: To truly understand AI, we can't just look at the test scores. We have to look under the hood to see how the machine is thinking. This "mechanistic evaluation" is the new tool the authors propose to stop us from being fooled by high scores that hide weak understanding.

Technical Summary: Mechanistic Evaluation of Transformers and State Space Models

Problem Statement

While State Space Models (SSMs) offer an efficient, sub-quadratic alternative to Transformers for language modeling, they exhibit deficiencies in retrieval tasks—specifically, copying information from the context—that standard benchmarks often fail to capture. Previous work has utilized synthetic tasks like Associative Recall (AR) to identify these performance gaps, but reliance on behavioral metrics (e.g., task accuracy) provides little insight into why certain architectures fail or succeed. The core problem is the lack of mechanistic understanding regarding how different architectures implement in-context retrieval and whether they learn robust, generalizable solutions (such as induction) or rely on brittle, position-dependent heuristics.

Methodology

The authors propose a framework for mechanistic evaluation using causal interventions to analyze how models solve retrieval tasks.

Tasks:
- Associative Recall (AR): A synthetic task where a model is presented with a sequence of key-value pairs followed by a query key, requiring the model to output the corresponding value.
- Associative Treecall (ATR): A novel, hierarchical retrieval task generated via a Probabilistic Context-Free Grammar (PCFG). Unlike AR, where keys and values are adjacent, ATR requires associating non-adjacent tokens based on latent tree structures, mimicking real-world natural language retrieval more closely.
Mechanistic Metrics (Interchange Interventions):
- The authors employ interchange interventions (Geiger et al., 2021; 2024) on model internals.
- Procedure: They corrupt a specific token in the input (e.g., the key) to degrade performance. They then intervene at specific hidden states (at the key, value, or query positions) within specific layers, replacing the corrupted representation with a "clean" representation from a counterfactual input.
- Metric: They measure the extent to which this intervention restores the likelihood of the correct answer. High restoration at the value token indicates an induction mechanism (storing associations at the value for later retrieval). High restoration at the query or key indicates direct retrieval (computing the association immediately at the query or a specific layer).
Architectures Evaluated:
- Transformers (Attention)
- SSMs: Mamba, DeltaNet, H3, Hyena
- Hybrid Models: Based, BaseConv, DeltaNet
Ablation Studies:
- Systematic removal or modification of short convolutions (varying kernel sizes $d_{conv}$ ) in Mamba and Based to determine their necessity for association.

Key Contributions and Results

1. Divergent Mechanisms for Similar Accuracy

Behavioral metrics show that Transformers, Based, and Mamba achieve high accuracy on AR, while other SSMs (H3, Hyena) fail. However, mechanistic analysis reveals fundamental differences:

Transformers and Based: Learn an induction mechanism. They store key-value associations at the value token during the forward pass, allowing a subsequent layer to retrieve the answer. This is a robust, two-layer solution.
Mamba and DeltaNet: Primarily use direct retrieval. They compute the association at the query token (often in layer 0 or 1) without explicitly storing the association at the value token.
H3 and Hyena: Fail to learn the task effectively under these conditions.

2. The Critical Role of Short Convolutions

The authors identify short convolutions as the critical component enabling association in Mamba and Based:

In Mamba: Short convolutions are necessary for AR. When the convolution kernel size is reduced to 1 or removed, Mamba's performance on AR collapses to near zero.
Mechanism in Mamba: Contrary to the assumption that Mamba relies solely on the SSM state, the authors find that the short convolution within Mamba performs the induction-like association step (moving key information to the value or next key), while the SSM component handles retrieval.
In Based: Replacing short convolutions with implicit long convolutions significantly degrades performance, confirming the necessity of the short receptive field for this specific task.

3. Generalization to Hierarchical Retrieval (ATR)

The authors introduce ATR to test if these mechanisms generalize beyond adjacent tokens:

Behavioral Results: Mamba outperforms or matches all other architectures on ATR, even without short convolutions.
Mechanistic Shift: On ATR, Mamba without short convolutions successfully learns a two-layer induction mechanism (storing at the value), similar to Transformers.
Implication: This suggests Mamba possesses the capacity for induction but defaults to direct retrieval via short convolutions on simpler, adjacent tasks (AR). The reliance on short convolutions is task-dependent; they are essential for AR but not strictly for ATR, where the model can fallback to a more general induction mechanism.

Significance and Claims

The paper argues that mechanistic evaluations are essential for understanding architectural capabilities, as behavioral accuracy alone can be misleading.

Substantive Differences: Architectures with similar accuracy (e.g., Transformers vs. Mamba) may employ fundamentally different internal mechanisms. Transformers and Based learn a generalizable induction mechanism, while Mamba's success on AR is heavily dependent on a specific architectural component (short convolution) that implements induction locally.
Task Dependency: The mechanisms learned are not static; Mamba adapts its strategy based on the task structure (adjacent vs. hierarchical).
Design Implications: The findings suggest that short convolutions are a key enabler for associative memory in SSMs, but their absence does not preclude the learning of induction mechanisms if the task structure (like ATR) necessitates it.
Future Direction: The authors posit that mechanistic metrics should be adopted as a standard tool for architecture analysis to "open the blackbox" of AI progress, moving beyond intuition and empirical benchmarking to explain how models solve tasks.

The paper concludes that while synthetic tasks like AR have driven architectural innovation, treating them solely as behavioral benchmarks discards valuable signal. By applying causal interventions, the authors reveal that the "success" of SSMs on retrieval tasks often relies on specific, task-sensitive components rather than a universal, robust induction capability inherent to the SSM formulation itself.