Balanced State Subspaces for Mamba-2 State Compression

📄 arXiv: 2610.36617v1 📥 PDF

作者: Shang Tingyu, Cao Wenfei

分类: cs.LG, cs.AI, cs.CL, stat.AP, stat.ME

发布日期: 2026-10-27

备注: 15 pages, 4 figures

💡 一句话要点

提出平衡状态子空间方法以优化Mamba-2状态压缩

🎯 匹配领域: 支柱二：RL算法与架构 (RL & Architecture)

关键词: 状态压缩 Mamba-2 后训练变换 平衡截断 HARP 自然语言处理 机器学习

📋 核心要点

  1. 现有的Mamba-2压缩方法在处理状态时未能有效利用状态的能量分布，导致压缩效果不佳。
  2. 论文提出了一种新的后训练变换方法，旨在通过平衡截断优化状态压缩，关注写入-读取子空间而非单个坐标。
  3. 实验结果显示，采用HARP方法后，状态维度减半，困惑度仅增加0.3点，相较于GHOST方法显著降低了性能损失。

📝 摘要（中文）

Mamba-2将所有上下文压缩为固定大小的递归状态，这在解码时占用了大量内存和带宽。现有的压缩方法独立量化状态值或修剪状态坐标，未能有效捕捉输入输出映射的方向。本文提出了一种后训练变换方法，通过平衡截断实现全可控和可观测的Gramian，并通过Mamba-2的标量-恒等转换精确应用旋转。此外，HARP方法根据Hankel谱为每层分配秩，保护检索敏感层，显著提高了困惑度和检索性能。实验表明，使用HARP的方案在WikiText-2上平均仅增加0.3点困惑度，且在25%保留率下提升更为明显。

🔬 方法详解

问题定义：本文旨在解决现有Mamba-2状态压缩方法在状态能量分布利用上的不足，现有方法未能有效捕捉输入输出映射的方向，导致压缩效果不理想。

核心思路：提出了一种后训练变换方法，通过平衡截断实现全可控和可观测的Gramian，关注写入-读取子空间，从而提高压缩效率。

技术框架：整体架构包括状态分析、Gramian计算、平衡截断应用和HARP层级秩分配，确保在不同层级上优化状态压缩。

关键创新：最重要的创新在于引入了HARP方法，根据Hankel谱为每层分配秩，保护检索敏感层，从而显著提升了压缩效果和检索性能。

关键设计：在参数设置上，采用了全可控和可观测的Gramian，损失函数设计关注写入-读取子空间，网络结构上则通过Mamba-2的标量-恒等转换实现状态旋转。

📊 实验亮点

实验结果表明，使用HARP的方案在WikiText-2数据集上，状态维度减半，平均困惑度仅增加0.3点，相比于GHOST方法的1.6点（9.3%）提升显著，尤其在25%保留率下效果更为明显。

🎯 应用场景

该研究的潜在应用领域包括自然语言处理、机器翻译和语音识别等需要高效状态压缩的场景。通过优化状态压缩方法，可以在保持模型性能的同时，显著降低内存和带宽需求，具有重要的实际价值和未来影响。

📄 摘要（原文）

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.