An Exact Instrument for State Usage in Selective State-Space Models, and the Input-Driven Migration It Reveals· 状态空间模型中的精确测量工具及其输入驱动迁移
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.
提出状态空间模型中模式使用的精确测量工具,揭示输入驱动的模式迁移现象。
- 核心方法
- 通过构建一个精确的测量工具,该工具利用对角状态矩阵的特点,将每个通道的输出精确分解为每个模式的贡献。并使用每(层,通道,窗口)的Gram张量来预测在任何预算下移除模式集合后的输出误差。
- 适合谁读
- 研究者 / 工程师
- 要解决的问题
- 状态空间模型中,如何精确测量不同模式的使用情况,以及如何理解这些模式随输入变化的迁移现象。
- 关键实验
- 在Mamba-1系列模型(130M--2.8B参数)及7B参数的Falcon-Mamba和Mamba-2模型上验证了工具的准确性和剪枝方法的有效性。
- 主要贡献
- 1. 提供了一种精确测量模式使用的工具,验证其高精度;2. 揭示训练模型在输入变化时重新分配状态空间的现象;3. 提出基于输入调度的模式剪枝方法,优于静态和层适应性剪枝方法。
- 意义与局限
- 这项研究为理解状态空间模型的内部机制提供了新视角,尤其在模型随输入变化的动态行为方面。同时,输入调度的剪枝方法为模型优化提供了新的方向。局限性在于该工具目前没有实际部署的计算或内存节省效果。