arXiv cs.CL· Shang Tingyu, Cao Wenfei·· 13 小时前AI 评分26
面向 Mamba-2 状态压缩的平衡状态子空间方法
Balanced State Subspaces for Mamba-2 State Compression
AI 导读
研究者提出一种面向 Mamba-2 的状态压缩方法，通过平衡截断构造写—读子空间，并配合 HARP 按层分配秩。在 130M–2.7B 模型上，该方法将状态维度减半，WikiText-2 困惑度平均仅上升 0.3 点（1.8%），而 GHOST 上升 1.6 点（9.3%）。校准只需一次前向传播，速度比 GHOST 公开协议快一个数量级，峰值内存更低。
正文 · 原文
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.
| Comments: | 15 pages, 4 figures |
| Subjects: | Machine Learning (cs.LG); Artificial Intelligence (cs.AI); Computation and Language (cs.CL); Applications (stat.AP); Methodology (stat.ME) |
| Cite as: | arXiv:2610.36617 [cs.LG] |
| (or arXiv:2610.36617v1 [cs.LG] for this version) | |
| https://doi.org/10.48550/arXiv.2610.36617 arXiv-issued DOI via DataCite (pending registration) |
Submission history
From: Wenfei Cao [view email]
[v1]
Tue, 27 Oct 2026 04:21:27 UTC (163 KB)
来源：arXiv cs.CL · arxiv.org