像 GPT-4 和 LLaMA-2 这样的大型语言模型 (LLM) 是语言奇才，能够轻松创作诗歌、代码和文章。然而，一旦问它们一个多步骤的小学数学题，它们往往会跌跌撞撞，表现不佳。

解决这个问题的标准方案是 思维链 (Chain-of-Thought, CoT) 提示——在要求模型解答新问题之前，先给它几个逐步解决类似问题的例子。这被称为少样本学习 (few-shot learning)。直觉上，你展示的例子越多，模型的表现应该越好。但这里存在一个硬性上限: 上下文窗口 (context window) 。

如果你在提示词中塞入太多例子，就会达到 token 限制。如果尝试用标准方法压缩这些例子，往往会丢失关键的数字或运算符，导致这些例子在数学上毫无用处。

在香港科技大学和微软亚洲研究院的论文 “Fewer is More: Boosting Math Reasoning with Reinforced Context Pruning” 中，研究人员介绍了 CoT-Influx 。 这种新颖的方法采用从粗到细的剪枝机制，通过移除无用的例子和冗余的词语，将更多高质量的推理例子塞入上下文窗口。结果呢？数学推理能力显著提升，使得像 LLaMA2-70B 这样的开源模型在没有任何微调的情况下，表现超越了 GPT-3.5。

先导研究: 为什么我们需要更好的剪枝

在构建系统之前，研究人员进行了一项先导研究，以了解例子数量 (shots) 与推理性能之间的关系。他们在 GSM8K 数学数据集上测试了 LLaMA2-7B。

观察 1: 多多益善 (通常情况下)

标准做法是使用 8 个手动设计的 CoT 例子。研究人员发现，增加例子数量通常会提高准确率。

图 1: 随着 TopK 检索到的 CoT 例子数量增加，LLaMA2-7B 的推理准确率也在变化。图 2: 使用 LLMLingua 提示压缩工具压缩后的 CoT 例子。被剪枝的 token 中包含真正冗余的 token (蓝色) 和关键 token (红色) 。

如上图 (图 1) 顶部的图表所示，当我们从 0 个例子增加到 16 个例子时，准确率在攀升。然而，一旦达到 20 个例子，我们就会撞上上下文窗口的限制 (由红框标出) 。模型没有足够的剩余空间来生成答案，导致性能崩溃。

观察 2: 质量至关重要

数量不是一切。与精心挑选的例子相比，简单地将随机例子堆砌到提示词中实际上会损害性能。

表 1: CoT 例子的选择严重影响 LLM 在 GSM8K 上的数学推理性能。

表 1 显示，随机选取的 16-shot 例子表现甚至不如标准的手动 8-shot 例子。这意味着我们需要一种聪明的方法来选择哪些例子应该被包含进去。

观察 3: 冗余普遍存在

最后，研究人员查看了例子中的文本。自然语言是冗余的。我们使用了许多填充词 (如 “Let’s think step by step”, “The answer is”, “There are”) ，这些词对数学逻辑并没有严格的贡献。

然而，现有的压缩工具对数学来说是危险的。请看第一张图片 (图 2) 的下半部分。它展示了一个被 LLMLingua 压缩的例子。该工具剪掉了“不重要”的 token，但意外地删除了数字“15” (标为红色) 。在数学中，如果删除了数字，推理就会变成幻觉。

假设: 如果我们能选择最有帮助的例子 (观察 2) ，并在保持数学逻辑完整的前提下仅剪除真正冗余的词语 (观察 3) ，我们就能在上下文窗口中放入明显更多的例子 (观察 1) ，从而提升性能。

进入 CoT-Influx: 方法论

CoT-Influx 是一个介于用户和 LLM 之间的即插即用模块。它不需要对庞大的 LLM 本身进行微调。相反，它处理输入提示词以最大化信息密度。

图 3: 上图: CoT-Influx 的总体流程；下图: 展示使用 CoT-Influx 先剪枝整个 CoT 例子再剪枝 token 的示例。

如图 3 所示，该系统遵循 从粗到细 (Coarse-to-Fine) 的剪枝策略:

  1. 样本剪枝器 (Shot Pruner): 从大量检索到的例子池 (例如，top-k 相关问题) 中，选择最关键的完整例子进行保留。
  2. 词元剪枝器 (Token Pruner): 针对步骤 1 中保留下来的例子，剪除其中冗余的单个 token (单词、空格) 。

数学公式化表达

目标是选择一个例子子集 (\(\hat{\mathcal{D}}\)) 并对其进行压缩 (\(\hat{x}\))，使得 LLM 的准确率最大化，同时总 token 数量 (\(t\)) 保持在上下文限制 (\(T\)) 之内。

转换过程如下所示:

方程 1: 从数据集到样本剪枝再到 token 剪枝的流程。

优化目标是双重的: 最小化 LLM 的损失 (困惑度) 并最大化推理准确率 (\(R_{Acc}\))，同时受到长度约束:

方程 2: 结合损失、准确率和 token 约束的优化目标函数。

架构

由于 CoT-Influx 是一个外部模块，它需要自己的方式来“阅读”文本。作者使用 BERT-Large 作为特征提取器。

  1. 输入: 一批潜在的 CoT 例子。
  2. 嵌入: BERT 创建这些例子的向量表示 (\(H_{s}\))。
  3. 策略网络: 两个轻量级的多层感知机 (MLP) 作为操作的大脑。

样本剪枝器 (Shot Pruner) 计算保留每个例子的概率 (\(a_{shot}\)):

方程 3: 用于样本剪枝的策略网络。

词元剪枝器 (Token Pruner) 计算在保留下来的例子中保留每个 token 的概率 (\(a_{token}\)):

方程 4: 用于 token 剪枝的策略网络。

输出是简单的二元决策: 1 (保留) 或 0 (剪除) 。

使用强化学习进行训练

这里存在主要的技术挑战。“剪枝”是一个离散操作——你要么保留一个 token，要么不保留。这切断了计算图，意味着我们无法根据 LLM 的输出使用标准的反向传播 (梯度下降) 来训练剪枝器。

为了解决这个问题，作者采用了 强化学习 (RL) , 具体来说是 REINFORCE 算法。

奖励函数

代理 (剪枝器) 需要知道它是否做得好。奖励函数 (\(R\)) 由三个因素组成:

  1. 有效性: LLM 对答案的预测有多好？ (通过 \(\frac{1}{1 + L_{LLM}}\) 衡量) 。
  2. 准确性: LLM 是否得到了正确的最终数学答案？(\(R_{Acc}\))。
  3. 长度约束: 提示词是否在窗口范围内？ (项 \([\frac{t}{T}]^w\)) 。

方程 5: 多目标奖励函数。

如果剪枝器剪掉了关键数字 (导致 LLM 失败) 或留下了太多废话 (超出窗口) ，奖励就会下降。

策略网络的参数 (\(\theta\)) 使用策略梯度进行更新:

方程 6: 策略网络的梯度更新规则。

稳定训练

训练用于文本生成的 RL 是出了名的不稳定。作者使用了两个巧妙的技巧:

  1. 难度感知过滤: 他们使用 GPT-4 创建了一个具有不同难度级别的数据集 (MRD³)。在训练期间，他们过滤出较简单的问题。如果一个问题太难，无论提示词质量如何，LLM 都会做错，这就无法为剪枝器提供有用的信号。
  2. 锚点样本: 在训练的早期阶段，剪枝器可能会随机删除所有内容。为了防止这种情况，他们附加了一些“安全的”、手动设计的例子，以确保基准水平的性能。

实验结果

研究人员在各种 LLaMA2 模型 (7B, 13B, 70B) 以及包括 GSM8K 在内的多个数学数据集上评估了 CoT-Influx。

拓展少样本学习的边界

CoT-Influx 真的允许放入更多例子吗？是的。

图 4: CoT-Influx 输入不同数量 CoT 例子时在 GSM8K 上的 EM(%) 准确率。

图 4 显示了 CoT-Influx (蓝线) 与标准 TopK 检索 (灰线) 和固定少样本 (黑线) 的准确率对比。

  • 性能峰值: LLaMA2-13B 在 48 shots 时达到性能峰值。
  • 效率: 即使有 48 个例子，token 剪枝也能保持输入在可控范围内。
  • 影响: 与基线相比，准确率提升显著。

与最先进技术的比较

表 2 将 CoT-Influx 与其他方法进行了比较，包括像 LLMLingua 和 Selective Context 这样的激进压缩技术。

表 2: 在 GSM8K 上与最先进基线的 EM (%) 准确率比较。

主要结论:

  • 压缩基线失败: 注意像 TopK+Selective Context 这样的方法导致了极低的准确率 (LLaMA2-7B 为 0.45%) 。这证实了通用压缩会破坏数学推理逻辑。
  • CoT-Influx 占据主导地位: 它在 LLaMA2-70B 上达到了 59.59% 的准确率，显著高于标准少样本方法 (55.42%)。

击败巨头

也许最令人印象深刻的结果是，这种方法让开源模型能够越级挑战。

表 4: 在 few-shot-CoT 设置下与更大 LLM 在 GSM8K 上的 EM (%) 准确率比较。

如表 4 所示, 配备 CoT-Influx 的 LLaMA2-70B (59.6%) 优于 GPT-3.5 (57.1%) , 并且非常接近 Google Minerva 540B 模型。这是在没有微调模型权重的情况下实现的——仅仅是通过优化上下文。

长上下文模型怎么样？

有人可能会争辩说: “为什么我们需要剪枝？新模型的上下文窗口有 32k 或 200k tokens。直接把所有例子放进去不就行了！”

作者在 Mistral-7B-32K 和 Yi-6B-200K 上测试了这个假设。

图 5: 提示压缩对长上下文 LLM 的影响。X 轴表示输入 token 的数量。

图 5 揭示了一个反直觉的发现:

  1. 收益递减: 简单地增加上下文 (橙线) 并不能线性提高准确率；它经常会波动或下降。
  2. 效率制胜: CoT-Influx (蓝线) 用一小部分 token 实现了更高的准确率。对于 Yi-6B-200K，它以 少 15 倍的输入 token 实现了高出约 2.5% 的准确率 。 这转化为巨大的推理成本和延迟节省。

剪掉了多少？

该系统激进但有选择性。

图 6: 我们的剪枝器每个阶段后的 token 长度。

图 6 显示了剪枝率。样本剪枝器完成了大部分繁重的工作 (3.87 倍压缩) ，过滤掉了不太相关的例子。词元剪枝器精炼其余部分 (使总压缩率达到 4.28 倍) 。

消融研究 (表 9) 证实了这两个阶段都是必要的。移除 token 剪枝会显著损害性能，证明词级冗余是一个主要瓶颈。

表 9: 不同剪枝策略的比较。

为什么有效？数据集的重要性

这一成功的一个关键组成部分是用于训练剪枝器的数据集: MRD³ (具有不同难度的数学推理数据集) 。

利用 GPT-4，作者进化了标准数学问题，创造了具有不同约束和推理深度的变体。

图 7: 难度分布 (第一行) 和推理步骤数量分布 (第二行) 。

这种多样性至关重要。有趣的是，研究人员发现 不同的 LLM 偏好不同的例子 。

  • 小模型 (7B): 偏好更简单、更容易的例子。
  • 大模型 (70B): 从复杂、多步骤的例子中获益更多。

表 17: 较小、能力较弱的 LLM 偏好更简单的 CoT 例子，而较大的 LLM 偏好更复杂的例子。

CoT-Influx 自动学习这些偏好，为其服务的特定模型选择“恰到好处”的例子。

结论

CoT-Influx 证明，在上下文学习方面, 质量和密度胜过数量 。 通过智能地选择有帮助的例子并精准地移除冗余词语，我们可以为 LLM 提供“超级提示词”，从而大幅提升它们的数学推理技能。

这种方法为部署 LLM 提供了一条切实可行的前进道路。与其依赖昂贵的微调或庞大、缓慢的上下文窗口，我们可以使用轻量级、经过强化学习训练的剪枝模块来释放现有模型的潜在推理能力。

正如论文标题所暗示的，在高效 AI 的世界里，更少的 token——如果是正确的 token——确实意味着更多。