大语言模型 (LLMs) 彻底改变了我们与信息交互的方式,但它们长期遭受一个瓶颈的困扰: 延迟。如果你曾看着 ChatGPT 逐字逐句地敲出答案,你就亲身体验过自回归解码的局限性。由于每一个新词元 (token) 都依赖于前一个词元,模型必须按顺序生成输出。这个过程缓慢且计算效率低下,导致昂贵的 GPU 在等待内存访问时处于空闲状态。

为了解决这个问题,研究人员开发了投机解码 (Speculative Decoding) , 这是一种利用较小的模型来“起草”文本,然后由较大的模型进行验证的技术。然而,即便是投机解码也有其上限。如果草稿太短,节省的时间有限;如果草稿太长,验证往往会失败,导致前功尽弃。

Ouroboros 应运而生,这是清华大学和面壁智能 (ModelBest Inc.) 的研究人员提出的一个新颖框架。Ouroboros 通过将起草过程从*词元 (token) 转变为短语 (phrase) *,重新构想了起草过程。它引入了一种递归机制,模型可以从自身的验证输入中学习短语 (这也是其名称“Ouroboros”——吞食自己尾巴的神话衔尾蛇的由来) 。

该框架通过回收被丢弃的验证结果并并行化起草阶段,实现了高达 3.9倍 于标准解码的速度提升——而且无需任何额外的模型训练。

在这篇文章中,我们将解构 Ouroboros 的工作原理、其效率背后的数学原理,以及为什么它代表了大语言模型推理的一个重大飞跃。

瓶颈: 为什么起草很难

在深入了解 Ouroboros 之前,我们必须理解投机解码 (SD) 面临的“金发姑娘”难题 (Goldilocks problem,即难以找到恰到好处的平衡点) 。

在标准的 SD 设置中,你有两个模型:

  1. 草稿模型 (Draft Model, \(S\)) : 小而快,但准确率较低 (例如 Llama-7B) 。
  2. 目标模型 (Target Model, \(T\)) : 大而准,但速度慢 (例如 Llama-70B) 。

草稿模型预测接下来的几个词元 (比如 5 个) 。然后,目标模型通过一次并行的前向传递处理所有 5 个词元来进行验证。如果草稿是正确的,你就用生成一个词元的时间生成了 5 个词元。

然而,效率在很大程度上取决于草稿长度 (\(\gamma\)) 和草稿模型的大小

图1: 起草效率与效果之间的权衡。

图 1 所示,存在一个艰难的权衡:

  • 左图: 较大的草稿模型更准确 (接受率上升) ,但运行速度更慢。
  • 右图: 较长的草稿提供了更高的潜在加速比,但生成时间更长,且更容易包含错误,从而导致验证失败。

作者指出了当前 SD 方法的两个关键局限性:

  1. 起草不足 (Insufficient Drafting) : 逐个词元地生成草稿仍然太慢。
  2. 计算利用率不足 (Underutilized Computation) : 当草稿失败时 (例如,目标模型拒绝了最后 3 个词元) ,这部分工作通常被完全丢弃,即使其中包含语义上有用的短语。

Ouroboros 解决方案

Ouroboros 通过将起草的原子单位从单个词元转变为短语来解决这些问题。它引入了一种递归机制,模型可以从自身的验证输入中学习短语 (因此得名“Ouroboros”,即吞食自己尾巴的神话之蛇) 。

该框架引入了四个关键创新:

  1. 加速起草: 并行生成候选短语。
  2. 延长草稿: 将短语拼接在一起,以近乎零的成本创建更长的候选序列。
  3. 回收验证: 将被丢弃的词元转化为未来可用的有用短语。
  4. 上下文复用: 利用来自提示词 (prompt) 历史的短语。

让我们看看标准投机解码与 Ouroboros 之间的高层对比:

图2: Ouroboros 框架与投机解码的对比。

图 2 中,上半部分显示了标准 SD 逐个生成词元 (“I”, “like”, “to”…) 。下半部分显示 Ouroboros 使用预先计算好的短语 (“I like to do”, “waste my time on”) 。Ouroboros 不像标准 SD 那样为六个词元运行六次草稿模型,而是可能只需要两次前向传递,通过链接这些短语来构建相同的句子。

加速的数学原理

为了理解理论上的收益,让我们看看标准投机解码的加速公式:

公式5: 标准投机解码的加速计算。

这里,\(A(\gamma)\) 是平均接受的词元数,\(t_T\) 是目标模型的时间,\(t_S\) 是草稿模型的时间。

Ouroboros 通过两种方式改进了这个公式。首先,它将起草成本降低了 \(c\) 倍 (通过短语加速) 。其次,它将草稿长度延长了 \(\beta\) (通过短语拼接) ,且无需额外计算。新的加速公式变为:

公式6: Ouroboros 的加速计算。

Ouroboros 的目标是在保持低开销的同时最大化 \(c\) 和 \(\beta\)。

核心方法: 基于短语的起草与验证

Ouroboros 的核心在于它如何构建和验证这些更长的草稿。

1. 通过短语延长草稿

在现代 GPU 中,LLM 的生成是受限于内存 (memory-bound) , 而非受限于计算 (compute-bound) 。这意味着验证 10 个词元所需的时间与验证 1 个词元大致相同,因为瓶颈在于将模型权重移动到内存中,而不是数学运算本身。

Ouroboros 利用了这一点。它不只是向目标模型提交一个草稿,而是通过将不同的候选短语附加到当前生成内容的末尾,构建多个更长的草稿。

假设草稿模型生成的序列以 \(d_\gamma\) 结尾。Ouroboros 会查找 \(K\) 个以 \(d_\gamma\) 开头的短语,并创建 \(K\) 个不同的候选序列:

公式7: 延长的草稿集合。

这对目标模型来说似乎是很多工作,但作者使用了一种巧妙的注意力掩码 (Attention Masking) 技巧,在一次前向传递中验证所有这些候选序列。

图3: 定制的注意力掩码机制。

图 3 可视化所示,注意力掩码允许目标模型正常验证“前缀 (Prefix) ”和“草稿 (Draft) ”。然后,它为“短语 1 (Phrase 1) ”和“短语 2 (Phrase 2) ”创建并行分支。模型独立计算短语 1 中的词元概率,不受短语 2 的影响。这在 GPU 上物理上属于同一个批次 (batch) ,与验证单个草稿相比几乎不增加额外的时间成本。

2. 从验证中生成短语

这些短语从何而来?令人惊讶的是,许多短语来自前几次迭代的“失败”。

在标准投机解码中,如果目标模型拒绝了草稿的后缀,这些词元就会被删除。然而,作者观察到,即使是被拒绝的词元通常也包含正确的序列——它们只是位置错乱或略有偏移。

公式10: 匹配函数的定义。

使用匹配函数 (公式 10) ,Ouroboros 识别出被丢弃草稿中实际上与目标模型输出分布相匹配的片段。

图4: 从目标模型的验证结果中生成短语的示意图。

图 4 完美地说明了这一点。

  • 左侧 (标准 SD) : 草稿输出 [1, 2, 4, 6, 8]。目标期望 [2, 4, 6, 8]。因为草稿以 1 开头 (错误) ,标准 SD 拒绝了其后的所有内容。模型必须重新开始,浪费了正确的序列 4, 6, 8
  • 右侧 (Ouroboros) : 系统检测到 6, 8 是一个有效的短语,只是位置不对。它将 6, 8 保存为一个短语。在下一步中,它立即应用这个短语,跳过了逐个生成这些词元的需要。

实验与结果

研究人员在标准基准测试上评估了 Ouroboros,包括 HumanEval (代码) 、GSM8K (数学) 和 CNN/DM (摘要) 。他们将其与朴素解码 (Vanilla decoding) 、前瞻解码 (Lookahead decoding) 和标准投机解码进行了比较。

速度对比

结果非常显著。Ouroboros 始终优于所有基线。

图5: HumanEval 和 MBPP 上的贪婪解码速度。

图 5 中,请看蓝色柱状图 (Ouroboros) 。在使用 Yi-34B 模型的 HumanEval 数据集上,Ouroboros 达到了 61.2 词元/秒 , 相比之下,朴素解码仅为 15.6,标准投机解码为 21.5。这是吞吐量的巨大飞跃。

这种性能在不同的任务类型中都保持一致,如表 2 所示:

表2: GSM8K、CNN/DM 和 WMT16 上的贪婪解码速度。

对于 GSM8K (数学推理) 上的 Llama-2-70B 模型,Ouroboros 实现了相对于朴素解码 2.68倍的加速 , 显著高于标准投机解码实现的 1.88倍。

块效率 (Block Efficiency)

论文中引入的一个关键指标是块效率 (\(\eta\)) , 它通过计算生成的词元数与目标模型调用次数的比率来衡量加速的理论上限。

公式29: 块效率的定义。

表13: 块效率对比。

表 13 显示 Ouroboros 实现了高得多的块效率 (每次调用生成更多词元) 。在 HumanEval 上,它的效率达到 13.12 , 而投机解码为 11.16,前瞻解码仅为 3.08。这证实了基于短语的方法成功地让目标模型一次性接受了大块的文本。

消融实验

系统的每个部分都有作用吗?作者分解了每个组件的贡献:

表4: 消融实验。

表 4 所示:

  • 基线速度: 21.46 token/s
  • 添加短语起草: 49.90 token/s (最大的提升)
  • 添加短语延长: 55.92 token/s
  • 添加验证回收: 58.18 token/s
  • 添加历史复用: 61.20 token/s

虽然短语起草提供了巨大的初始提升,但随后的优化 (延长和回收) 挖掘出了显著的额外性能。

结论

Ouroboros 证明了我们尚未达到大语言模型推理速度的物理极限。通过摆脱僵化的逐词元生成范式并拥抱基于短语的生成 , 我们可以更好地利用现代 GPU 的大规模并行计算能力。

Ouroboros 最引人注目的一点是它是无需训练的 。 它不需要蒸馏新的草稿模型或微调目标模型。这是一种算法改进,可以立即应用于现有的起草-验证流程中。

随着大语言模型在规模和能力上不断增长,像 Ouroboros 这样优化推理“软件”侧——使解码过程本身更智能——的技术,对于使这些模型在现实应用中可部署且响应迅速至关重要。