Balanced State Subspaces with Spectral Rank Allocation for Mamba-2 State Compression

원문 초록

Mamba-2 compresses all context into a fixed-size recurrent state that dominates decoding memory and bandwidth. Existing compressors quantize state values or prune coordinates with per-channel scores. Trained states are nearly rank-one and bias-like, aligned with the mean write; but the output's energy spans an order of magnitude more directions. Per-channel scores see only the diagonal of the write--read operator, so they cannot select the subspace that governs the input--output map, which is a rotation of the coordinates. We therefore propose \method{}, which keeps that subspace by applying balanced truncation exactly via Mamba-2's scalar-identity transition. Layers differ just as sharply: most are bias-dominated, while the few that carry long-range associations are the flattest and need most of their state. We therefore introduce \spar{}, which redistributes a fixed total rank budget across layers according to the Hankel spectrum. On Mamba-2 checkpoints from 130M to 2.7B, \method{}+\spar{} halves the state at an average cost of $0.3$ WikiText-2 perplexity points, against $1.6$ for GHOST, while retaining $97\%$ of the dense model's five-digit in-context recall. At a quarter of the state, giving every layer the same rank keeps $5\%$ of that recall and \spar{} $77\%$, for a $7\%$ perplexity cost.