An Exact Instrument for State Usage in Selective State-Space Models, and the Input-Driven Migration It Reveals

想象一下，Mamba 模型就像一个巨大且超级聪明的管弦乐团。在这个乐团的每一层内部，都有一组由 16 种微小的、单音节乐器（称为“模态”，modes）组成的乐器库。在旧有的思维方式中，我们假设指挥家（模型）会为整首曲子挑选出一套固定的、例如最好的 8 种乐器，而不论音乐正在发生什么变化。

但本论文揭开了帷幕，揭示了一个令人震惊的秘密：指挥家并没有挑选一套固定的组合。

相反，这位指挥家是一位大师级的即兴演奏者。对于模型读到的每一个单词（或“标记”，token），它都会瞬间重新决定哪些 8 种乐器实际上是承载旋律所必需的。有时是长笛；有时是小提琴；有时是鼓。这些“重要”的乐器会根据输入进行迁移。如果你强迫乐团死守一套固定的 8 种乐器（一种“静态”的选择），你就是在要求他们用进行曲乐队的乐谱来演奏爵士独奏。这样做虽然可行，但听起来效果极差，远不如真实演奏。

魔法工具：“精确乐器”

他们是如何得知这一点的？他们构建了一个数学上的“精确乐器”。

由于这个乐团的内部结构是一种特殊的“对角线”设置（即乐器之间互不干扰），作者可以将输出分解为每种乐器贡献的完美总和。他们创建了一个“格拉姆张量”（Gram tensor，可以将其想象为一个极其精确的计分卡），它能精确地告诉我们，如果你丢弃任何特定的乐器组合，你会产生多少误差。

他们用这个工具对真实的模型进行了测试，发现其相对误差仅为 2.3 × 10⁻⁷。这就像是在测量地球到月球的距离时，误差竟然不到一根头发丝的宽度。这不是一个估计值，而是一个精确的测量值。

重大发现：“迁移间隙”

利用这个工具，他们观察了从微型模型（1.3 亿参数）到巨型模型（70 亿参数，如部署的 Falcon-Mamba）的各种模型。

他们发现，在最活跃的层中，固定的一组乐器所产生的误差是随输入变化的乐器组的两倍。

• 数据统计： 在受影响最严重的层中，“迁移间隙”（即固定集合与变化集合之间的误差比率）在 0.44 到 0.57 之间。
• 含义： 如果你允许模型为每一个特定时刻挑选最合适的乐器（即“输入调度预言机”，input-scheduled oracle），那么与仅仅挑选一次固定列表并始终坚持使用的做法相比，你可以将误差减半。

这种情况出现在他们测试的每一个模型中：Mamba-1 系列、70 亿参数的 Falcon-Mamba，甚至包括 Mamba-2。

导致迁移的原因是什么？（“为什么”）

作者问道：模型的哪一部分在进行这一切切换？
Mamba 层中有三个主要信号：

1. 写入映射 ($B_t$)： 决定哪些乐器接收输入信号。
2. 读取输出 ($C_t$)： 决定哪些乐器被听见。
3. 时间步长 ($\Delta$)： 通常被认为是“选择性”的旋钮。

他们进行了一项“冻结信号”实验。他们将每个信号分别冻结在其平均值，以观察迁移是否会停止。

• 结果： 当他们冻结写入映射 ($B_t$) 时，迁移消失了。模型停止了乐器的切换。
• 惊喜： 当他们冻结时间步长 ($\Delta$) 时，迁移依然保持原样。

结论： “时间步长”信号——许多人曾认为它是选择性的关键——对这种迁移信号几乎没有任何贡献。真正的英雄是写入映射 ($B_t$)。它是守门员，逐个标记地决定哪些乐器可以演奏。

后果：我们可以利用这一点吗？

作者尝试利用这些知识来对模型进行剪枝（精简）以节省空间。

• 静态剪枝： 基于平均活动度或测试集挑选出一组最佳的 8 种乐器，并永久保留它们。
• 输入调度剪枝： 查看当前的句子，测量当前活跃的 8 种乐器，并只保留这些乐器。

结果：
在减半状态预算（仅保留 16 种模态中的 8 种）的情况下，输入调度方法在原始准确度方面表现得如同（在某些情况下甚至略优于）完整的、未剪枝的模型。

• 在 1.3 亿参数模型上，调度方法的困惑度（perplexity）为 11.84，而未剪枝的模型为 12.38。
• 在 70 亿参数的 Falcon-Mamba 上，调度方法的表现为 4.44，击败了未剪枝的 4.48。

然而，有一个至关重要的陷阱： 论文明确指出，这种“调度”方法是一个两遍预言机（two-pass oracle）。它先阅读整个窗口一次以决定保留哪些乐器，然后再进行第二次运行以生成输出。这意味着在实际部署中，它并不能节省计算量或内存（你仍然需要读取两次数据）。

作者澄清说，这个结果展示的是可实现的性能上限（realizable headroom），而非已部署的效率提升。它证明了存在这样的潜力：通过构建一个廉价、快速的预测器，在不需要昂贵的第一遍处理的情况下，就能猜出正确的乐器，从而实现一个微小且超高效的模型。目前的这种方法只是向我们展示了天花板在哪里。

他们排除了什么

论文非常明确地说明了哪些方法不起作用：

• 静态排名： 基于平均活动度或“汉克尔能量”（Hankel energy）挑选固定模态集合的方法（如 GHOST 或 LAST）明显较差。它们无法跟上移动的目标。
• 时间步长活跃度： 基于时间步长信号活跃度进行的剪枝是一条死路，因为时间步长并不驱动迁移。
• 简单预测器： 他们尝试仅根据句子的前半部分或“领域”（如代码 vs 散文）来预测掩码。这些简单的技巧仅回收了 2–6% 的潜在增益。因为“重要”的乐器集合变化得非常快（在几百个标记之内），所以你确实需要测量用于评分的具体标记才能获得全部收益。

核心结论

论文证明了经过训练的选择性状态空间模型是动态的、鲜活的生命体，它们会根据输入不断重新分配内部资源。“写入映射”是这场迁移的指挥家。虽然我们目前还无法构建一个能在不产生“两遍处理”成本的情况下实时切换乐器的模型，但这项研究为我们提供了一张精确的地图。它表明，目前的“静态”剪枝方法正让大量的性能白白流失，而未来的路径在于构建能够实时预测这些迁移的调度器。