🤖 machine learning

An Exact Instrument for State Usage in Selective State-Space Models, and the Input-Driven Migration It Reveals

This paper introduces an exact instrument for measuring mode usage in selective state-space models, revealing that input-dependent state reallocation driven by the write map $B_t$ enables input-scheduled mode pruning to significantly outperform static methods and match unpruned performance at half the state budget.

Original authors: Raktim Bhattacharya

Published 2026-07-14

📖 6 min read🧠 Deep dive

CC BY 4.0

Original authors: Raktim Bhattacharya

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 a Mamba model as a giant, super-smart orchestra. Inside every layer of this orchestra, there's a bank of 16 tiny, one-note instruments (called "modes"). In the old way of thinking, we assumed the conductor (the model) would pick a fixed set of, say, the best 8 instruments to play for the entire song, no matter what the music was doing.

But this paper pulls back the curtain and reveals a shocking secret: The conductor doesn't pick a fixed set at all.

Instead, the conductor is a master improviser. For every single word (or "token") the model reads, it instantly re-decides which 8 instruments are actually needed to carry the melody. Sometimes it's the flutes; sometimes it's the violins; sometimes it's the drums. The "important" instruments migrate depending on the input. If you force the orchestra to stick with one fixed set of 8 instruments (a "static" choice), you're asking them to play a jazz solo with a marching band's sheet music. It works, but it sounds terrible compared to the real thing.

The Magic Tool: The "Exact Instrument"

How did the authors know this? They built a mathematical "exact instrument."

Because the orchestra's internal structure is a special kind of "diagonal" setup (where the instruments don't mess with each other), the authors could break down the output into a perfect sum of every single instrument's contribution. They created a "Gram tensor" (think of it as a super-precise scorecard) that tells them, exactly, how much error you would make if you dropped any specific group of instruments.

They tested this tool against the real model and found it was accurate to a relative error of 2.3 × 10⁻⁷. That's like measuring the distance from the Earth to the Moon and being off by less than the width of a human hair. It's not an estimate; it's a precise measurement.

The Big Discovery: The "Migration Gap"

Using this tool, they looked at models ranging from tiny ones (130M parameters) to massive ones (7B parameters, like the deployed Falcon-Mamba).

They found that at the most active layers, a fixed set of instruments makes twice as much error as a set that changes with the input.

The Stat: At the most affected layers, the "migration gap" (the ratio of error between a fixed set and a changing set) was between 0.44 and 0.57.
The Meaning: If you let the model pick the best instruments for each specific moment (an "input-scheduled oracle"), you cut the error in half compared to just picking a static list once and sticking with it.

This happens in every model they tested: the Mamba-1 family, the 7B Falcon-Mamba, and even Mamba-2.

What Causes the Migration? (The "Why")

The authors asked: What part of the model is doing all this switching?
There are three main signals in a Mamba layer:

The Write Map ( $B_t$ ): Decides which instruments get the input signal.
The Readout ( $C_t$ ): Decides which instruments get heard.
The Timestep ( $\Delta$ ): Often thought to be the "selectivity" knob.

They ran a "frozen-signal" experiment. They froze each signal to its average value, one by one, to see if the migration stopped.

The Result: When they froze the Write Map ( $B_t$ ), the migration vanished. The model stopped switching instruments.
The Surprise: When they froze the Timestep ( $\Delta$ ), the migration stayed exactly the same.

The Verdict: The "timestep" signal, which many people thought was the key to selectivity, carries almost none of the migration signal. The real hero is the Write Map ( $B_t$ ). It's the gatekeeper that decides, token by token, which instruments get to play.

The Consequence: Can We Use This?

The authors tried to use this knowledge to prune (cut down) the model to save space.

Static Pruning: Pick the best 8 instruments based on an average test and keep them forever.
Input-Scheduled Pruning: Look at the current sentence, measure which 8 instruments are active right now, and keep only those.

The Outcome:
At half the state budget (keeping only 8 out of 16 modes), the Input-Scheduled method performed as well as (and in some cases slightly better than) the full, unpruned model in terms of raw accuracy.

On the 130M model, the scheduled method had a perplexity of 11.84, while the unpruned model was 12.38.
On the 7B Falcon-Mamba, the scheduled method was 4.44, beating the unpruned 4.48.

However, there is a crucial catch: The paper explicitly states that this "scheduled" method is a two-pass oracle. It reads the whole window once to decide which instruments to keep, then runs the second pass to generate the output. This means it does not save compute or memory in a real-world deployment (you still have to read the data twice).

The authors clarify that this result demonstrates realizable headroom, not a deployed efficiency win. It proves that the potential exists for a tiny, super-efficient model to match the giant one, but only if we can build a cheap, fast predictor that guesses the right instruments without needing that expensive first pass. The current method just shows us the ceiling of what's possible.

What They Ruled Out

The paper is very clear about what doesn't work:

Static Rankings: Methods that pick a fixed set of modes based on average activity or "Hankel energy" (like GHOST or LAST) are significantly worse. They can't keep up with the moving target.
Timestep Activity: Pruning based on how active the timestep signal is is a dead end because the timestep doesn't drive the migration.
Simple Predictors: They tried predicting the mask based on just the first half of a sentence or the "domain" (like code vs. prose). These cheap tricks only recovered 2–6% of the potential gain. The "important" set of instruments changes so fast (within a few hundred tokens) that you really need to measure the specific tokens you are scoring to get the full benefit.

The Bottom Line

The paper proves that trained selective state-space models are dynamic, living things that constantly re-allocate their internal resources based on the input. The "Write Map" is the conductor of this migration. While we can't yet deploy a model that switches its own instruments in real-time without a "two-pass" cost, this research gives us an exact map of the territory. It shows us that the current "static" pruning methods are leaving a massive amount of performance on the table, and that the path forward lies in building schedulers that can predict these migrations on the fly.

Technical Summary: An Exact Instrument for State Usage in Selective State-Space Models

Problem Statement
Selective State-Space Models (SSMs), such as Mamba, utilize a bank of fixed first-order modes where the coupling to the signal is dynamically scheduled by input-dependent matrices. While the state matrix is diagonal (decoupling modes), the specific usage of these modes by a trained model remains opaque. Existing methods for pruning or reducing the state space rely on static, approximate criteria (e.g., output-aware scores from calibration statistics or activity-based rankings). These methods assume a fixed set of important modes, failing to account for the possibility that the "important" set may shift depending on the input context. Without an exact measurement of state usage, it is impossible to determine if static pruning incurs unnecessary error or if the model's internal degrees of freedom are being utilized dynamically.

Methodology: The Exact Instrument
The paper introduces an exact instrument for measuring state usage, leveraging the diagonal structure of the SSM state matrix.

Exact Decomposition: Because the state matrix is diagonal, the output of a channel is an exact sum of per-mode contributions. The authors define a per-mode output $m_{h,t,i}$ for each mode $i$ in channel $h$ at timestep $t$ .
Gram Tensor Construction: By accumulating the outer products of these contributions over a window, the method constructs a small Gram tensor ( $G_h$ ) for each (layer, channel, window).
Closed-Form Error Calculation: The squared output error of pruning any subset of modes $S$ is given exactly by a quadratic form of the dropped modes in the Gram matrix. This allows for the offline calculation of the exact pruning error for any budget without re-running the model.
Validation: The instrument was validated against reference implementations with a relative error of $\approx 2.3 \times 10^{-7}$ on Mamba-1. Furthermore, it predicted the actual deployed pruning error of a pruned layer to a median relative deviation of $10^{-6}$ across 4,464 configurations.

Key Findings: Input-Driven Mode Migration
Using this instrument, the authors discovered a phenomenon termed input-driven mode migration:

Dynamic Allocation: The set of modes carrying a layer's output is not fixed; it migrates across different input contexts. At the most affected layers, an input-scheduled oracle (which selects the best modes for each specific window) halves the output error compared to a static top- $r$ set.
Mechanism: Through frozen-signal counterfactuals, the paper attributes this migration primarily to the input-dependent write map ( $B_t$ ). Freezing $B_t$ to its corpus mean eliminates most of the migration gap. The readout map ( $C_t$ ) contributes partially, while the timestep ( $\Delta$ ), often associated with selectivity, carries almost none of the migration signal.
Slow Modes: The migrating modes are predominantly the "slow" modes (those with long time constants), which integrate input-gated writes over long horizons.
Universality: This behavior is observed across the Mamba-1 family (130M–2.8B), the deployed 7B Falcon-Mamba, and Mamba-2.

Experimental Results
The authors evaluated input-scheduled mode pruning against static baselines (static energy ranking, modal-Hankel singular value ranking, and layer-adaptive LAST) across multiple scales:

Performance: Input-scheduled selection (a two-pass oracle that measures mode energies in a first pass and prunes in a second) outperformed all static rules at every scale and budget.
Budget Efficiency: At half the state budget ( $r=8$ for $N=16$ ), the input-scheduled model matched or slightly exceeded the perplexity of the unpruned model across all tested scales (130M to 7B).
Static Baseline Failure: Static pruning incurred significantly higher error. For example, at the 130M scale with $r=4$ , static pruning reached a perplexity of 19.8, while the scheduled approach achieved 14.2 (compared to 12.4 for unpruned).
Downstream Tasks: On LAMBADA, static pruning collapsed accuracy, while the two-pass oracle recovered 86% of the lost performance.
Headroom vs. Deployment: The authors emphasize that the scheduled approach requires a "first pass" to read mode energies, meaning it demonstrates realizable headroom (the potential for improvement) but does not constitute a deployed compute or memory saving in its current form. Cheap schedulers (e.g., using a prefix or domain regime) captured only a fraction (2–26%) of the oracle's gain.

Significance and Claims
The paper claims three primary contributions:

An Exact Instrument: It provides the first exact, offline measurement of state usage and pruning error for diagonal SSMs, replacing heuristic assumptions with direct measurement.
Discovery of Migration: It establishes that trained selective SSMs dynamically re-allocate their state space based on input, driven by the write map $B_t$ . This challenges the assumption that a fixed set of modes is optimal for all contexts.
Proof of Headroom: It demonstrates that the gap between static and input-scheduled pruning is large and consistent across scales. While the paper does not claim a deployed saving (as the scheduler requires a first pass), it proves that significant performance recovery is possible if the mode selection can be adapted to the input.

The authors conclude that the exact instrument replaces heuristic choices with a direct measurement that varies with the input, and that the next steps involve understanding the realization-theoretic reasons for this migration and developing learned schedulers that can exploit it cheaply.