引言

如果你曾盯着闪烁的光标等待大型语言模型 (LLM) 完成一个段落，你就体验过所谓的“内存墙 (memory wall) ”。LLM 是自回归的，这意味着它们一次只能生成一个令牌 (Token) 。要生成第 100 个词，模型必须先生成第 99 个词，然后将其反馈给自己。这种顺序依赖性阻碍了我们利用现代 GPU 巨大的并行处理能力。

投机解码 (Speculative Decoding) 已成为解决这一瓶颈的热门方案。其思路很简单: 使用一个更小、更快的“草稿”模型来猜测接下来的几个令牌，然后使用大型“目标”模型并行验证所有这些猜测。

然而，目前验证这些草稿的方法远非最佳。标准方法，如递归拒绝采样 (Recursive Rejection Sampling, RRS) ，往往无法接受草稿序列后部的令牌，浪费了生成它们所消耗的计算资源。而在另一个极端，像最优传输 (Optimal Transport, OTM) 这样的数学框架理论上可以最大化接受率，但计算量过大，无法实时运行。

在这篇文章中，我们将深入探讨 SpecHub , 这是一种新的研究成果，旨在填补这一空白。SpecHub 将复杂的最优传输问题简化为易于管理的线性规划模型。通过引入“枢纽令牌 (Hub Token) ”策略——创建一个专注于最可能令牌的稀疏联合分布——SpecHub 既实现了最优传输的理论严谨性，又具备了实时推理所需的速度。

背景: 草稿与验证范式

要理解 SpecHub，我们首先需要了解多草稿投机解码 (Multi-Draft Speculative Decoding, MDSD) 的机制。

令牌树 (The Token Tree)

在标准的投机解码中，草稿模型生成单链令牌。MDSD 更进一步，生成一个 令牌树 。 草稿模型不再只是猜测“The cat sat”，而是可能同时探索多种可能性，分支出“The cat sat”、“The cat lay”和“The cat is”。

图 2: MDSD 深度为 d=4 的令牌树示例。

如上图 2 所示，草稿模型生成了一棵潜在序列树 (从“I”发出的箭头) 。然后，目标模型并行处理所有这些令牌。目的是在这棵树中找到一条符合目标模型分布的有效路径。

拒绝采样的挑战

我们如何决定保留哪些草稿令牌？我们使用 拒绝采样 (Rejection Sampling) 。 如果草稿令牌与目标模型的预测一致，我们就接受它。

图 3: 拒绝采样图解。

图 3 展示了这个概念。蓝色曲线 (\(q\)) 是草稿分布，红色曲线 (\(p\)) 是目标分布。

  1. 重叠 (Overlap) : 如果草稿采样的令牌位于 \(q\) 和 \(p\) 重叠的区域 (紫色区域) ，我们通常接受它。
  2. 残差 (Residual) : 如果我们拒绝了草稿，我们必须从“残差”分布中采样——即目标分布 (\(p\)) 中草稿 (\(q\)) 未覆盖的部分。

RRS (递归拒绝采样) 的问题: RRS 是一种标准的贪婪方法。它检查第一个草稿令牌；如果接受，则移动到第二个。问题出现在移动到第二个令牌时。第二个令牌的目标分布取决于第一个令牌。如果我们被迫从残差分布中采样 (因为第一个草稿并不完美) ，第二个草稿的分布往往会与原始草稿不对齐。

这种错位会导致后续令牌的接受率骤降。你可能会“免费”得到第一个令牌，但第二个和第三个草稿很可能会被拒绝，从而使“多草稿”变得毫无用处。

核心方法: SpecHub

研究人员提出了 SpecHub，旨在解决 RRS 的效率衰减问题，且不产生全量最优传输的巨大成本。

1. 从最优传输到线性规划

研究人员意识到，最大化多个草稿的接受率实际上是一个 最优传输 (Optimal Transport, OT) 问题。我们试图以最小的“成本” (拒绝) 将“概率质量”从草稿分布 (\(Q\)) 移动到目标分布 (\(p\)) 。

耦合 \(\pi\) 的成本函数的数学定义为:

成本函数 C(pi) 的方程。

最小化这个成本可以确保我们接受尽可能多的令牌。完整的 OTM 解决方案需要解决整个词汇表空间上的传输计划 \(\pi\)。对于 32,000 个令牌的词汇表和 2 个草稿，这涉及解决 \(|\mathcal{V}|^3\) 量级的变量，对于实时文本生成来说，这在计算上是不可能的。

简化: SpecHub 通过将其公式化为 线性规划 (Linear Programming, LP) 问题来简化这一点，该问题 仅 关注我们实际接受令牌的场景。我们不需要显式地映射拒绝的情况；我们只需要最大化接受的概率。

简化的最小化目标如下所示:

LP 的简化最小化目标。

受限于确保有效概率分布的约束条件:

简化 LP 公式的约束条件。

这大大降低了复杂度，但对于大型词汇表来说，它仍然太慢。这引出了论文的第二个、更关键的创新。

2. “枢纽 (Hub) ”令牌策略

为了使数学计算足够快以进行推理，SpecHub 改变了我们构建草稿分布 \(Q\) 的方式。SpecHub 不使用任何令牌都可以跟随任何其他令牌的“稠密”分布，而是构建一个 稀疏 分布。

算法识别出一个单一的 枢纽令牌 (Hub Token) (记为 \(a\)) 。这仅仅是草稿模型中概率最高的令牌。

规则很简单: 我们只采样包含枢纽的对。 如果我们使用 2 个草稿，我们只考虑 \((x, a)\) 或 \((a, x)\) 这样的对。我们从不采样像 \((b, c)\) 这样两个都不是最高概率令牌的对。

联合草稿分布 \(Q\) 定义为:

稀疏联合草稿分布 Q(x1, x2) 的定义。

这种结构非常精妙，因为它将复杂的概率网络变成了“星型”拓扑。枢纽令牌充当了概率质量的中央转运站。

3. 可视化差异

让我们在视觉上比较标准最优传输 (OTM) 方法和 SpecHub。

图 4: OTM (稠密) 和 SpecHub (稀疏) 分布的比较。

  • 左图 (a) OTM: 你看到的是一个稠密的 3D 矩阵。算法试图将每一对可能的草稿令牌映射到目标。它很准确，但很慢。
  • 右图 (b) SpecHub: 分布是稀疏的。网格的大部分是空的。我们只关心“枢纽” (黄色条形) 。这种稀疏性允许线性规划在 \(O(|\mathcal{V}|)\) 线性时间内求解。

4. “枢纽”如何传输概率

通过利用这种稀疏分布求解简化的 LP，我们得到了传输计划的闭式解。这确切地告诉了我们如何验证令牌。

传输计划 pi 的闭式解。

这里有一个具体的例子来说明 SpecHub 如何像 RRS 一样移动概率质量。

想象我们有一个草稿模型和一个目标模型。

  • 草稿: 对令牌 ‘a’ 信心高 (0.5) ，对 ‘b’ 中等 (0.3) 。
  • 目标: 对 ‘a’ 信心低 (0.1) ，对 ‘b’ 高 (0.6) 。

在 RRS (标准方法) 中，因为草稿过度预测了 ‘a’，它被拒绝了。残差分布基本上将 ‘a’ 清零。当我们试图验证第二个草稿时，我们不得不处理一个混乱的残差。

在 SpecHub 中，我们查看这一对令牌。 由于我们采样了涉及 ‘a’ 的对，并且 ‘a’ 被草稿过度采样，SpecHub 使用 ‘a’ 作为 枢纽 来“支付” ‘b’ 的接受，而 ‘b’ 是草稿预测不足但目标需要的。

图 6: SpecHub 表格显示概率通过枢纽令牌 ‘a’ 传输。

在上图 6 中，注意概率质量是如何流动的。草稿过度预测了 ‘a’ (\(q=0.5, p=0.1\)) 。SpecHub 获取这部分多余的草稿概率，并有效地将其重新分配以帮助接受 ‘b’，因为 ‘b’ 被草稿低估了但被目标需要。RRS 无法进行这种“交易”——它独立对待每一步。SpecHub 将序列视为一个联合整体。

5. 算法

SpecHub 的最终实现非常优雅。它不需要外部求解器；从 LP 导出的闭式解允许直接的算法实现。

算法 3: 使用 SpecHub 进行采样和验证。

该算法主要执行以下步骤:

  1. 识别枢纽令牌 \(a\)。
  2. 构建稀疏概率。
  3. 采样对 \((x^{(1)}, x^{(2)})\)。
  4. 使用从上述数学推导出的优化接受概率来验证它们 (论文中的定理 2 证明了这保持了正确的目标分布) 。

实验与结果

研究人员使用 Llama-2 模型在标准数据集 (CNN DailyMail, OpenWebText) 上测试了 SpecHub 与 RRS 和 RRSw (无放回 RRS) 的对比。

批处理效率

主要指标是 批处理效率 (Batch Efficiency) ——即每个模型步骤生成的平均令牌数。数字越高，意味着投机解码的效果越好。

图 1: 批处理效率比较。SpecHub (绿色) 优于 RRS (蓝色) 和 RRSw (红色) 。

如图 1 所示:

  • 图表 (a): 在 Llama2-7B 上，随着令牌树中节点数量的增加，SpecHub (绿线) 的扩展性明显更好。当 RRS 趋于平稳时，SpecHub 继续利用额外的草稿生成更多的令牌。
  • 图表 (b): 在 Vicuna-7B (一个聊天模型) 上，差距更大。SpecHub 每步生成近 3.5 个令牌，而 RRS 很难突破 2.0。

跨温度的鲁棒性

LLM 通常在不同的“温度”下运行。低温使模型具有确定性；高温使模型具有创造性 (且分布更平坦/更难预测) 。

图 7: 不同温度下的批处理效率。

图 7 显示，无论温度如何，SpecHub 的优势都保持不变。

  • 在 低温 (左侧) 下，分布尖锐。SpecHub 模仿贪婪解码，表现与 RRSw 相似。
  • 在 高温 (右侧) 下，分布更平坦且更难匹配，SpecHub 通过枢纽“传输”概率质量的能力使其能够保持比基线更高的效率。

为什么不增加更多草稿？

你可能会想，为什么要停在 2 个草稿 (\(k=2\)) ？为什么不是 5 个或 10 个？

图 9: 不同草稿数的接受率衰减。

图 9 说明了 MDSD 的“收益递减”。在最初的几个草稿之后，接受率 (y 轴) 急剧下降。到了第 4 或第 5 个草稿，所有方法的接受率都接近于零。 由于最优传输的计算成本随着 \(k\) (草稿数量) 呈指数增长，SpecHub 专注于优化 \(k=2\) 的情况。它从最初几个最重要的草稿中榨取了最大可能的效率。

结论

SpecHub 代表了使大型语言模型更快、更高效的重要一步。它成功地识别出多草稿投机解码的瓶颈不仅仅在于生成更多的草稿——而在于 我们如何验证它们 。

通过将验证步骤重新构想为线性规划问题并引入“枢纽”令牌约束，作者成功地:

  1. 降低了复杂度: 将成本从多项式级降低到线性级。
  2. 提高了接受率: 可证明地优于像 RRS 这样的启发式方法。
  3. 保持了准确性: 确保最终输出分布与目标模型完全匹配。

对于 NLP 领域的学生和研究人员来说，SpecHub 是一个完美的例子，展示了如何通过巧妙的算法设计 (稀疏性和线性规划) 来克服理论约束 (如最优传输的繁重计算) ，从而实现现实世界的加速。