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

2026-07-13

AI Summary

This work addresses the challenge of precisely quantifying how selective state space models, such as Mamba, utilize input-dependent state modes. The authors propose an offline analysis framework based on Gram tensors that enables, for the first time, accurate prediction of output changes resulting from pruning arbitrary subsets of state modes, through per-layer, per-channel, and per-window error computation. By combining analytical decomposition of diagonal state matrices with counterfactual analysis, they demonstrate that state transitions are predominantly driven by the input-dependent write mapping \( B_t \). Experiments on Mamba-1, Falcon-Mamba 7B, and Mamba-2 show that input-scheduled pruning retains full model performance with only 50% of the original state budget, achieving a median relative prediction error as low as \( 5 \times 10^{-7} \).

Abstract

Selective state-space models such as Mamba route information through a bank of first-order modes whose input coupling is set by a learned selection mechanism. We give an exact instrument for measuring how a trained model uses these modes. Because the state matrix is diagonal, each channel's output decomposes exactly into per-mode contributions, and a per-(layer, channel, window) Gram tensor yields the exact output error of dropping any subset of modes, offline, at any budget. Validated against the reference implementation to a relative error of $2.3\times10^{-7}$ on the Mamba-1 family where it is exact, the instrument predicts a layer's deployed pruning error to a median relative deviation of $5\times10^{-7}$ over $4{,}464$ configurations, its floor set by the reconstruction. Applying the instrument across the Mamba-1 family (130M--2.8B), the deployed 7B Falcon-Mamba, and Mamba-2, we find that trained models re-allocate their state space with the input: which modes carry the signal migrates across contexts, and at the most affected layers a per-input oracle roughly halves the output error of a fixed mode set. Frozen-signal counterfactuals attribute the migration primarily to the input-dependent write map $B_t$; the timestep usually identified with selectivity carries almost none of it. Input-scheduled mode pruning on this measurement outperforms static, Hankel-based, and layer-adaptive rankings at every scale from 130M to the deployed 7B Falcon-Mamba, and at half the state budget it matches the unpruned model. Because the scheduler reads each window's mode usage from a first pass, this demonstrates realizable headroom; we claim no deployed compute or memory saving.

Problem

• selective state-space models
• state usage
• input-driven migration
• mode pruning
• Mamba

Innovation

• selective state-space models
• exact mode usage measurement
• input-driven state migration
• input-scheduled pruning
• Gram tensor analysis