Balanced State Subspaces for Mamba-2 State Compression
Abstract
Mamba-2 compresses all context into a fixed-size recurrent state that dominates decoding memory and bandwidth. Existing compressors quantize state values or prune state coordinates, scoring each channel independently in the native basis. We examine trained Mamba-2 states and find that energy is almost rank-one, but the output does not read the energetic directions: energy concentrates in a bias-like direction accumulated by the recurrence, while the content the output reads lives in a low-dimensional subspace that is not aligned with the state's coordinates. Even output-aware scores are computed coordinate by coordinate, so they miss the directions that determine the input--output map. Compression should therefore target the write--read subspace rather than individual coordinates. We propose \method{}, a post-training transform that instantiates balanced truncation with full controllability and observability Gramians and applies the rotation exactly via Mamba-2's scalar-identity transition. Because the useful rank varies widely across layers, we further introduce HARP, which allocates ranks per layer from the Hankel spectrum and protects retrieval-sensitive layers, improving both perplexity and retrieval over uniform budgets. Calibration takes a single forward pass, an order of magnitude faster than GHOST's released protocol, with lower peak memory. With HARP, \method{} halves the state dimension at an average WikiText-2 perplexity increase of $0.3$ points (1.8%) across 130M--2.7B, versus $1.6$ (9.3%) for GHOST, and the gap widens at 25\% retention.
Related
- StarBOA: Real-Time Mamba State-Space Unrolling for Sparse Radar Micro-Doppler in ISAC Networks
- Bi-Mamba-Load: A bidirectional structured state space model with sequence relationship adaptive perception for electricity load forecasting
- Balanced Active Subspaces for Optimal Control via Trajectory Optimization
- Delta Redistribution: a variation of State Redistribution for steady-state embedded boundary computations
- Continuous-state branching processes with Lévy-Khintchine drift-interaction: Laplace duality and Fellerian extensions
BibTeX
@misc{tingyu2026balanced,
title = {{Balanced State Subspaces for Mamba-2 State Compression}},
author = {Shang Tingyu and Cao Wenfei},
howpublished = {arXiv},
year = {2026},
month = oct,
doi = {10.48550/arxiv.2610.36617},
eprint = {2610.36617},
archivePrefix = {arXiv},
url = {https://arxiv.org/abs/2610.36617},
}
RIS
TY - UNPB
TI - Balanced State Subspaces for Mamba-2 State Compression
AU - Shang Tingyu
AU - Cao Wenfei
PY - 2026
DA - 2026/10/27
PB - arXiv
DO - 10.48550/arxiv.2610.36617
UR - https://arxiv.org/abs/2610.36617
AB - Mamba-2 compresses all context into a fixed-size recurrent state that dominates decoding memory and bandwidth. Existing compressors quantize state values or prune state coordinates, scoring each channel independently in the native basis. We examine trained Mamba-2 states and find that energy is almost rank-one, but the output does not read the energetic directions: energy concentrates in a bias-like direction accumulated by the recurrence, while the content the output reads lives in a low-dimensional subspace that is not aligned with the state's coordinates. Even output-aware scores are computed coordinate by coordinate, so they miss the directions that determine the input--output map. Compression should therefore target the write--read subspace rather than individual coordinates. We propose \method{}, a post-training transform that instantiates balanced truncation with full controllability and observability Gramians and applies the rotation exactly via Mamba-2's scalar-identity transition. Because the useful rank varies widely across layers, we further introduce HARP, which allocates ranks per layer from the Hankel spectrum and protects retrieval-sensitive layers, improving both perplexity and retrieval over uniform budgets. Calibration takes a single forward pass, an order of magnitude faster than GHOST's released protocol, with lower peak memory. With HARP, \method{} halves the state dimension at an average WikiText-2 perplexity increase of $0.3$ points (1.8%) across 130M--2.7B, versus $1.6$ (9.3%) for GHOST, and the gap widens at 25\% retention.
ER -