💬 NLP

Detection vs. Execution: Single-Bucket Probes Miss Half the Mamba-2 State Sink

This paper demonstrates that in Mamba-2, representational signatures identified by single-bucket probes often conflate distinct detection and execution circuits for the state sink phenomenon, revealing that only a specific subset of heads (BOS-specialists) causally drives critical long-context retrieval performance while others with similar representations do not.

Original authors: Yuhang Jiang

Published 2026-06-02

📖 6 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

The Big Idea: Seeing the "Shadow" vs. the "Actor"

Imagine you are trying to figure out how a massive, complex orchestra plays a specific song. You have a microphone (a probe) that can detect when a specific instrument is playing loudly.

In the world of AI research, scientists often assume that if their microphone picks up a loud sound from a specific instrument, that instrument is the one actually playing the melody.

This paper argues that in a new type of AI model called Mamba-2, this assumption is wrong. The microphone detects a huge "shadow" of activity, but the actual "actor" doing the heavy lifting is a tiny, hidden group within that shadow. If you only look at the loud sound, you miss the real mechanism.

The Setting: The "State Sink"

To understand the problem, we need to understand what the AI is doing.

The Analogy: Imagine a long line of people passing a bucket of water down a chain. Every time a new person joins the line (a new word in a sentence), the person at the very front (the "Start of Sentence" or BOS token) has to hold the bucket extra tight to keep the water from spilling.
The Science: In AI models, this "holding tight" is called a State Sink. The model focuses a massive amount of energy on the first few tokens (like the start of a sentence or a new line) to keep its memory stable.

The Discovery: Two Groups, One Signal

The researchers looked at Mamba-2 models and found that this "State Sink" isn't just one big group of workers. It's actually two distinct groups that look the same to a simple detector but do very different jobs.

1. The "BOS-Specialists" (The Execution Layer)

Who they are: A tiny group of workers (only about 5% of the total).
What they do: They are the doers. They are the ones actually holding the bucket and keeping the water flowing. If you remove them, the whole line collapses, and the AI forgets everything.
The Paper's Claim: These are the "Execution" circuit. They are the ones actually performing the calculation.

2. The "Dual Heads" (The Detection Layer)

Who they are: A much larger group (about 30% of the total).
What they do: They are the watchers. They are standing next to the doers, looking at the same thing, and getting just as excited. To a simple microphone (a "single-bucket probe"), they look exactly like the doers. Their signals are almost identical.
The Twist: If you remove these watchers, the AI still works fine! They are not doing the heavy lifting. They are just "detecting" that something important is happening.
The Paper's Claim: These are the "Detection" circuit. They have the right signature (they look like they are working) but not the right function (they aren't actually doing the work).

The Mistake: The "Single-Bucket" Trap

The paper calls the mistake of confusing these two groups the "Single-Bucket Probe" error.

The Analogy: Imagine you are trying to find the best chefs in a kitchen by asking, "Who is holding a knife?"
- You find a tiny group of Head Chefs (the BOS-Specialists) who are actually cooking the meal.
- You also find a huge group of Sous Chefs (the Dual Heads) who are holding knives but just chopping vegetables for the garnish.
- If you use a simple test that just says "Anyone holding a knife is a Head Chef," you will mistakenly think the Sous Chefs are the ones cooking the main dish.
- The Result: You think you found the whole kitchen crew, but you missed the fact that the real cooking is done by a tiny, elite team.

In Mamba-2, the "knives" are the mathematical signals. The simple probe sees both groups holding them, but only the tiny group is actually cooking the meal.

Why Does This Happen? (The Architecture)

Why does Mamba-2 have this "Shadow" group while older models don't?

Mamba-1 (The Old Model): Imagine a kitchen where every chef has their own private counter. If a chef wants to focus on the Start of the Sentence, they can do it without bothering anyone else. The "Doers" and the "Watchers" are the same people.
Mamba-2 (The New Model): Imagine a kitchen where chefs are forced to share a single, small counter. Because they have to share space, they have to do multiple things at once.
- The model forces a group of heads to "multiplex" (do two jobs at once).
- One job is to detect the start of a sentence (the Watchers).
- The other job is to execute the reset (the Doers).
- Because they share the same space, they look identical to a simple probe, even though their jobs are different.

The Proof: The "Needle in a Haystack" Test

To prove that the tiny group (Specialists) is the real hero and the big group (Dual Heads) is just a bystander, the researchers did a "Needle in a Haystack" test.

The Test: They hid a specific sentence (the "needle") deep inside a huge book (the "haystack") and asked the AI to find it.
The Result:
- When they removed the Tiny Group (Specialists): The AI failed completely. It couldn't find the needle at all. Accuracy dropped from 100% to 0%.
- When they removed the Big Group (Dual Heads): The AI still found the needle perfectly fine.
- Conclusion: The big group was just watching; the tiny group was the one actually remembering the needle.

Summary of the Paper's Claims

Detection $\neq$ Execution: Just because a part of the AI lights up (is detected) doesn't mean it's doing the work (execution).
The Split: In Mamba-2, the "State Sink" splits into a small Execution Layer (5% of heads) and a large Detection Layer (30% of heads).
The Trap: Simple probes (single-bucket) only see the big Detection Layer and miss the tiny Execution Layer.
The Fix: To understand how Mamba-2 works, you can't just look at who is "loud." You have to test who is actually necessary by removing them and seeing if the AI breaks.
Architecture Matters: This happens because of how Mamba-2 is built (shared projections), not just because of the size of the model.

In short: The paper warns us that in the new Mamba-2 models, the "loud" parts of the brain aren't always the ones doing the thinking. There is a hidden, tiny team doing the real work, and a much larger team just standing around looking busy. If you only listen to the noise, you'll miss the real story.

Technical Summary: Detection vs. Execution in Mamba-2 State Sinks

Problem Statement
Mechanistic interpretability often operates under the assumption that probes identifying a specific representational signature (e.g., high activation on boundary tokens) also identify the specific circuit executing the corresponding computation. This paper challenges that assumption within the context of Selective State Space Models (SSMs), specifically Mamba-2. The authors investigate the "state sink"—analogous to the "attention sink" in Transformers—where disproportionate activation of the token-conditional gate ( $\Delta$ ) occurs on boundary tokens (Begin-of-Sequence/BOS and newlines). The core problem is whether single-bucket probes (classifying units based on activation on a single class) accurately localize the causal circuit responsible for state management, or if they miss a larger, functionally distinct set of units that share the same representational signature.

Methodology
The study employs a rigorous intervention grid involving 12,096 cells (6,060 evaluated) across four model checkpoints: Mamba-1 (130M, 2.8B) and Mamba-2 (1.3B, 2.7B), alongside Pythia (160M–2.8B) as a cross-architecture baseline.

Specialist Classification: Units (channels in Mamba-1, heads in Mamba-2) are classified as "specialists" if their mean pre-softplus $\Delta$ activation on a specific token class (BOS, newline, etc.) exceeds a threshold ( $\tau=0.5$ ) relative to a reference set. This yields "single-bucket" specialists (one class) and "dual" specialists (two classes).
Causal Interventions: The authors perform ablation studies using six intervention types (e.g., gate_zero, delta_zero, gate_one) on specific sets of units. They measure the differential in Negative Log-Likelihood (NLL) on the first eight tokens after a boundary token, comparing the specialist set against size-matched random complements.
Granularity Comparison: The study contrasts Mamba-1 (channel granularity, independent $\Delta$ projections) with Mamba-2 (head granularity, head-shared $\Delta$ projections) to isolate architectural effects.
Validation: Results are validated across two corpora (wikitext-2, Pile-10k), multiple scales, and a long-context retrieval task (RULER Needle-in-a-Haystack).

Key Contributions and Results

Granularity-Conditional Causality:
- In Mamba-1 (channel granularity), single-bucket BOS-specialists carry the causal mass for the state sink. Ablating them causes significant performance degradation, and the probe accurately tracks the execution circuit.
- In Mamba-2 (head granularity), single-bucket probes recover only a small "execution layer" (~5% of heads). However, a much larger set of "dual heads" (27–35% of heads), which share the same representational signature (high activation on both BOS and newlines), is missed by single-bucket classification.
Detection vs. Execution Dichotomy:
- Execution Layer: The single-bucket BOS-specialists are the causal execution layer. Ablating them collapses retrieval accuracy from 1.00 to 0.00 in RULER NIAH tasks and significantly increases NLL on BOS-context and newline targets.
- Detection Layer: The dual heads exhibit high representational similarity to the specialists (cosine similarity $\approx 0.89$ on BOS vs. newline) but show substantially weaker causal effects under ablation. In fact, ablating dual heads often has negligible or even negative effects on performance compared to the baseline, indicating they are a "representational detection layer" rather than the functional execution layer.
- Misleading Probes: The paper demonstrates that representational similarity does not imply functional equivalence. Single-bucket probes systematically undercount the causal substrate by 3–7× in set size at the head granularity of Mamba-2.
Architectural Driver:
- The asymmetry is attributed to Mamba-2's architecture, where the $\Delta$ projection is shared across channels within a head. This forces heads to multiplex across boundary classes, creating the large dual-head population.
- A control experiment involving random channel bucketing of Mamba-1 to match Mamba-2's head-to-class ratio failed to reproduce the dual-head dominance, implicating the specific head-shared $\Delta$ projection as the driver, rather than substrate granularity alone.
Cross-Architecture Baseline:
- Pythia (Transformer) attention heads show single-bucket selectivity but lack the dominant dual-head set found in Mamba-2. This confirms that the detection/execution split is specific to the Mamba-2 architecture's state-space mechanics, not a universal property of head-granular models.

Significance and Claims
The paper claims that the standard "label-to-locus" move in mechanistic interpretability—assuming the highest-activating units identified by a probe are the causal circuit—is systematically unreliable at the head granularity of Mamba-2.

Re-audit Necessity: Interpretability findings localized to Mamba-2 heads using single-bucket probes warrant re-audit under multi-class aggregation rules.
Functional Decomposition: The state sink in Mamba-2 decomposes into two distinct functional layers: a small execution layer (single-bucket specialists) and a large detection layer (dual heads). Separating them requires class-conditional ablation rather than relying on class-conditional cosine similarity.
Active Anchor: Unlike attention sinks in Transformers which may be passive or suppressible, the Mamba-2 state sink is an active anchor where specific heads are necessary for long-context retrieval and boundary state management.

The authors conclude that while probes are effective for representational analysis, they can fail to identify the true causal circuit in architectures with shared projections and multiplexed activation profiles, necessitating intervention-based verification to distinguish between detection and execution.