An Exact Instrument for State Usage in Selective State-Space Models, and the Input-Driven Migration It Reveals
arXiv:2607.11796
The paper introduces an exact measurement tool that quantifies how selective state‑space models like Mamba use their internal modes, enabling precise prediction of pruning errors and revealing that mode usage shifts with the input. Using this tool, the authors show that input‑dependent scheduling of mode pruning outperforms static methods across model sizes up to a 7B Falcon‑Mamba.
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 on the Mamba-1 family where it is exact, the instrument predicts a layer's deployed pruning error to a median relative deviation of over 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 ; 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.