An Exact Instrument for State Usage in Selective State-Space Models, and the Input-Driven Migration It Reveals
2026-07-13 • Machine Learning
Machine Learning
AI summaryⓘ
The authors study a special type of neural network called Mamba that processes information through separate channels called modes. They develop a precise way to measure how each mode contributes to the output and how errors change if some modes are removed. Testing their method on various Mamba models shows that the importance of modes changes depending on the input, and adjusting mode selection based on this input can reduce errors significantly. Their approach outperforms other pruning methods and works efficiently without extra computational cost during deployment.
Selective state-space modelsMambaDiagonal state matrixMode pruningInput couplingGram tensorOutput errorInput-dependent write mapHankel-based methodsNeural network pruning
Authors
Raktim Bhattacharya
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.