引言

序列建模的格局正在发生转变。多年来，Transformer 架构一直占据主导地位，推动了大型语言模型 (LLM) 的革命。然而，一个新的竞争者已经出现: 状态空间模型 (State Space Models, SSMs) , 其中最引人注目的是 Mamba 架构。

Mamba 引发了巨大的轰动，因为它解决了 Transformer 最大的瓶颈: 注意力的二次方计算成本。Mamba 随序列长度线性扩展，使其成为处理海量上下文时潜在的“Transformer 终结者”。然而，将 Mamba 扩展到数十亿参数仍然面临巨大的计算挑战。为了在实际应用中部署这些模型，我们需要提高它们的效率。

在 Transformer 的世界里，一种流行的效率提升技术是 Token 缩减 (Token Reduction) ——移除或合并冗余的 token (如“the”或“um”) ，从而使模型处理更少的数据。这看起来是一个合乎逻辑的步骤: 采用 Transformer 中最先进的 token 缩减方法，并将其应用到 Mamba 上。

但这其中有一个陷阱。正如*“Rethinking Token Reduction for State Space Models”*这篇论文的研究人员所发现的，直接将这些基于 Transformer 的技术应用于 Mamba 会导致性能的灾难性下降。

在这篇文章中，我们将探讨为什么标准的优化技术会破坏 SSM，分析导致这种敏感性的 Mamba 底层机制，并详细介绍研究人员提出的新颖解决方案: 一种统一 Token 缩减 (Unified Token Reduction, UTR) 策略，它成功地使 Mamba 变得更轻、更快，且不牺牲其智能。

背景: 理解挑战

在深入解决方案之前，我们需要了解两个主要组成部分: 状态空间模型架构和 Token 缩减的概念。

状态空间模型 (SSM)

与同时查看所有 token (全局注意力) 的 Transformer 不同，SSM 像循环神经网络 (RNN) 一样按顺序处理数据。它们通过隐藏状态 \(h(t)\) 将输入序列 \(x(t)\) 映射到输出序列 \(y(t)\)。

基本的连续系统由以下方程描述:

Continuous State Space Equation

这里，\(\mathbf{A}\)、\(\mathbf{B}\) 和 \(\mathbf{C}\) 是定义状态如何演变以及如何投影到输出的矩阵。为了使该系统能在数字硬件上计算，Mamba 将其离散化。它引入了一个时间尺度参数 \(\Delta\)，将连续参数转换为离散参数 (\(\bar{\mathbf{A}}, \bar{\mathbf{B}}\)) 。这产生了一个递归公式:

Discretized Recurrence Equation

这种递归允许模型将历史“记忆”在压缩状态 \(h_t\) 中。至关重要的是，Mamba 还可以使用全局卷积进行并行训练，避免了传统 RNN 在训练时的缓慢顺序处理:

Global Convolution Equation

这种双重特性——递归推理和并行训练——是 Mamba 的超能力。然而，递归特性创建了一个依赖链。每个状态都依赖于前一个状态。

Token 缩减的概念

在自然语言处理 (NLP) 和视觉领域，并非所有的 token 都是生而平等的。在一张狗的照片中，代表背景墙的像素不如代表狗眼睛的像素重要。在一个句子中，核心主语比填充词更重要。

Token 缩减方法通常分为两类:

  1. 剪枝 (Pruning) : 识别“无用”的 token 并直接删除它们 (例如 EViT) 。
  2. 合并 (Merging) : 识别相似的 token 并将它们合并为一个单一的表示 (例如 PuMer 或 ToMe) 。

这些方法在 Transformer 上效果奇佳。但当应用于 Mamba 时，它们却失效了。

为什么 Transformer 方法在 Mamba 上会失效

研究人员首先将 Transformer 领域最先进的缩减方法——EViT (剪枝) 和 PuMer (合并) ——应用于 Mamba-2.8B 模型。结果非常明显。

Performance Drop with Standard Methods

如上图 1 所示，应用 EViT (红色柱状图) 导致准确率下降了 20% 。 应用 PuMer (绿色柱状图) 导致下降了 26% 。

原因分析

为什么 Mamba 比 Transformer 如此脆弱？研究人员指出了两个主要罪魁祸首:

1. 不可恢复的信息丢失 (剪枝失败) : 在 Transformer 中，如果你剪掉一个 token，下一层的自注意力机制可能仍然能够从其他全局 token 中获取上下文。而在 Mamba 中，计算是顺序的 (方程 2) 。隐藏状态 \(h_t\) 是历史信息的累积。如果你剪掉一个 token \(x_t\)，你不仅仅是丢失了那个词；你打断了状态演化的链条。这引入了一个信息缺口，并随着序列的推进被放大。

2. 忽视重要性 (合并失败) : 现有的合并策略 (如 ToMe 或 PuMer) 基于相似性进行操作。它们将 token 分成两组，并严格根据它们看起来有多相似来进行合并。它们在合并之前并不检查一个 token 是否重要。在 SSM 中，将一个高度重要的 token 合并到一个不太重要的 token 中，会稀释顺序更新所需的信号，从而破坏隐藏状态。

解决方案: 统一 Token 缩减 (UTR)

为了解决这个问题，研究人员提出了一个专为 SSM 量身定制的新框架。他们的方法建立在两个支柱之上: 识别重要性和混合缩减策略 。

第一步: 重新思考 Token 重要性

我们不能随意剪枝或合并 token。我们需要知道哪些 token 携带了“信号”。研究人员分析了 SSM 层输出的隐藏状态 \(y\)。

SSM Layer Output

他们发现有一个特定的指标最能预测重要性: 截断特征值的平均值 。 与 Transformer 不同，后者的注意力图提供了清晰的重要性评分，Mamba 的隐藏状态是高维的。研究人员发现，对特征维度 \(D'\) 上的正激活值求和 (将负值截断为零) 可以产生最佳的“重要性评分” \(S\):

Importance Metric Equation

通过使用 \(\max(0, \dots)\)，该指标关注的是强特征的存在，而不是负抑制的幅度。

第二步: UTRC 工作流程

计算出重要性评分后, 基于分类的统一 Token 缩减 (UTRC) 方法按照一个结构化的流程进行。

Overview of UTRC Method

如图 2 所示，该过程涉及三个不同的阶段:

  1. 重要性分类: Token 根据其评分 \(S\) 进行排序。它们被分为两个集合:
  • 集合 \(M_A\) (较不重要) : 被移除或合并的候选者。
  • 集合 \(M_B\) (较重要) : 必须保留的“锚点”。
  1. 连接: 系统不仅仅是删除集合 \(M_A\)。它试图保存它们的信息。对于集合 \(M_A\) 中的每个“不重要” token \(a_i\)，算法使用余弦相似度在集合 \(M_B\) 中找到其最相似的“重要”对应项 \(f_i\)。

Similarity Connection Equation

  1. 统一缩减: 这是该方法的亮点所在。它不严格执行剪枝或合并，而是使用由相似度阈值控制的混合方法 。
  • 合并: 如果一个不重要的 token \(a_i\) 与一个重要的 token \(b_j\) 非常相似，它们的特征会被平均 (合并) 。历史信息在融合后的 token 中得以保留。
  • 剪枝: 如果一个不重要的 token 与任何重要 token 都不够相似，它就会被剪枝。

第三步: 细粒度的设计选择

研究人员发现，对网络的所有部分一视同仁是次优的。他们引入了一种针对隐藏状态和残差连接的解耦策略。

  • 隐藏状态: 使用剪枝和合并的混合策略 (例如，剪枝 50% 的候选者，合并另外 50%) 。这平衡了去除噪声与保留基本上下文。
  • 残差连接: 主要使用合并 。 残差连接对于梯度流和携带前几层的信息至关重要。在这里进行剪枝是危险的。通过合并残差路径中的 token，即使序列长度缩短，模型也能保持信号的完整性。

实验结果

这种量身定制的方法有效吗？结果表明答案是肯定的。

准确率恢复

研究人员在六个标准基准测试 (如 LAMBADA、HellaSwag 和 PIQA) 上，使用 Mamba-2 (1.3B 和 2.7B) 和原始 Mamba 模型测试了他们的方法。

将提出的方法与标准的 EViT 和 PuMer 基线进行比较:

  • 在 Mamba-2-1.3B 上，该方法在 20% FLOPs 缩减下达到了 54.6% 的平均准确率，而 EViT 仅为 44.2%。
  • 在 Mamba-2-2.7B 上，差距更大。在 30% 缩减下，该方法保持了 54.7% 的准确率，而现有方法下降到了约 41%。

在许多情况下，该方法成功地大幅减少了计算负载，同时保持了较低的困惑度 (PPL) ，这意味着模型并没有因为缺失 token 而感到“困惑”。

效率提升

Token 缩减的主要目标是节省速度和内存。该方法在这两方面都表现出色。

GPU 峰值显存: 长序列会消耗大量的 VRAM。通过在各层之间分层减少 token，该方法显著降低了峰值显存使用量。

GPU Memory Reduction

如图 3 所示，对于 Mamba-2.8B，该方法最多可减少 40% 的峰值显存 (在 30% FLOPs 缩减时) 。这对于在消费级硬件或边缘设备上部署这些模型来说是一个游戏规则改变者。

吞吐量 (速度) : 降低 token 数量也直接转化为更快的生成速度。

Throughput Comparison

图 4 展示了 Mamba-2.8B 的生成吞吐量实现了 1.29 倍的加速 。 这种加速使模型在聊天机器人或代码生成等实时应用中响应更加迅速。

验证设计选择

研究人员进行了消融实验，以证明他们特定的设计选择是必要的。

重要性指标: 他们比较了“截断求和”指标与标准的 \(L1\) 和 \(L2\) 范数。“截断”版本 (方程 5) 始终产生更低的困惑度和更高的准确率 (论文中的表 3) ，证明了正激活值是 SSM 中衡量重要性的更好代理指标。

混合策略: 他们测试了“仅剪枝”、“仅合并”以及各种混合比例 (\(q\)) 。

  • 仅合并对于残差连接效果最好 (保留信息流) 。
  • 混合 (50/50) 对于隐藏状态效果最好 (平衡噪声去除与上下文保留) 。

Ablation of Design Choices

如表 5 所示，这种组合 (隐藏层 \(q=0.5\)，残差层仅合并) 实现了最高的准确率 (54.7%) ，验证了解耦方法的有效性。

结论

从 Transformer 到像 Mamba 这样的状态空间模型的转变，代表了向更高效的长序列建模的转变。然而，效率技术不能简单地在架构之间复制粘贴。

标准 token 缩减在 Mamba 上的失败凸显了 SSM 的独特敏感性: 你不能在不产生后果的情况下打断顺序链条。

通过引入一种尊重 token 重要性并采用混合剪枝-合并策略的统一 Token 缩减方法，研究人员为高效 SSM 的未来提供了蓝图。他们的方法证明，我们可以两全其美: 既拥有 Mamba 的长上下文能力，又拥有减少 token 数量带来的轻量级效率。

随着 SSM 的不断成熟，像这样的技术将是将这些模型从研究实验室推向实际应用部署的关键。