以 GPT-4 和 LLaMA 为代表的大语言模型 (LLM) 彻底改变了人工智能领域。然而,如果你曾观察过 LLM 生成回复的过程,你可能会注意到一个根本性的瓶颈: 文本是一个词接一个词出现的,就像一个打字很慢的人。

这种迟缓是由自回归 (Auto-Regressive, AR) 解码范式造成的。要生成第 100 个 Token,模型严格需要前 99 个 Token 作为基础。这种顺序依赖性阻碍了生成过程中的并行处理,让强大的 GPU 在等待下一个 Token 确定时处于空闲状态。

研究人员一直在竞相利用推测解码 (Speculative Decoding) 技术来解决这一延迟问题——这种技术先“猜测”草稿 Token,然后并行验证它们。然而,目前大多数方法都代价高昂: 它们需要训练额外的“草稿模型”,添加复杂的头 (如 Medusa) ,或运行独立的、资源密集型的后训练阶段。

如果我们能教主模型在不增加任何额外参数或不需要单独训练流程的情况下自我加速,那会怎样?

在这篇文章中,我们将深入探讨一篇引人入胜的论文,题为 “Make Some Noise: Unlocking Language Model Parallel Inference Capability through Noisy Training” (制造噪音: 通过噪声训练解锁语言模型并行推理能力) 。 研究人员提出了一个名为 MSN (Make Some Noise) 的框架,该框架可以无缝集成到标准的微调过程中。通过训练模型处理“噪声”输入,他们解锁了强大的并行解码能力,实现了比标准模型快 2.3 倍到 2.7 倍的速度——这实际上是为 LLM 推理加速提供了一顿“免费的午餐”。

问题所在: 加速的代价

要理解 MSN 框架的精妙之处,我们首先需要看看目前是如何尝试为 LLM 提速的。

在标准的推测解码中,我们通常需要一个“起草者 (Drafter) ”和一个“验证者 (Verifier) ”。

  1. 起草者: 快速猜测接下来的几个 Token (例如,“The cat sat on the…”) 。
  2. 验证者 (即 LLM) : 在一次并行前向传播中检查这些猜测。

瓶颈通常在于起草者。以往的解决方案通常涉及:

  • 小型草稿模型: 在运行大 LLaMA 的同时运行一个小 LLaMA。这会消耗额外的显存。
  • 额外的头 (例如 Medusa) : 添加新的架构层来预测未来的 Token。这改变了模型结构并增加了参数。
  • 独立的后训练: 在你为了实际工作 (如写代码或聊天) 微调完模型后,还得再次训练它,仅仅是为了让它学会变快。

MSN背后的研究人员意识到,这些方法将“任务能力” (变聪明) 和“推理速度” (变快) 视为两个独立的问题。

图 1: 本文提出的 MSN 框架与现有基于模型的推测解码方法之间的差异说明。书本图标代表特定任务能力,火箭图标代表并行解码能力。

如上图 1 所示,现有方法 (上半部分) 需要一个复杂的流程。首先你需要一个基础模型 (Base Model) ,然后进行有监督微调 (SFT) 得到一个聪明的模型,然后再进行特定训练或添加结构以得到一个快速的模型。

MSN 框架 (下半部分) 简化了这一点。它提出,通过稍微调整标准的 SFT 阶段——具体来说是添加噪声——我们可以立即生成一个聪明快的模型。不需要额外的步骤,也不需要额外的权重。

背景: 雅可比解码 (Jacobi Decoding)

要理解 MSN 如何工作,我们需要了解它用于加速的机制: 雅可比解码

标准解码一次求解一个 Token (\(y\))。雅可比解码将文本生成视为一个非线性方程组。它试图同时求解多个 Token (\(y_1, y_2, ... y_m\))。

\[ \left\{ \begin{array} { r l } { y _ { 1 } } & { { } = \arg \operatorname* { m a x } P _ { \theta } ( y _ { 1 } | x ) } \\ { y _ { 2 } } & { { } = \arg \operatorname* { m a x } P _ { \theta } ( y _ { 2 } | y _ { 1 } , x ) } \\ { ~ } & { { } \vdots } \\ { y _ { m } } & { { } = \arg \operatorname* { m a x } P _ { \theta } ( y _ { m } | y _ { 1 : m - 1 } , x ) } \end{array} \right. \]

上面的方程展示了目标: 在给定前序内容的情况下,找到每个位置的最佳 Token。在雅可比解码中,你从未来 Token 的随机猜测 (或“噪声”) 开始。你将这个序列输入模型。模型查看序列并更新其预测。你迭代地重复这个过程。

理想情况下,序列会迅速“收敛”到正确的文本。然而,标准的 LLM 在这方面表现糟糕。它们是在完美的文本上训练的 (教师强制,Teacher Forcing) 。当它们看到雅可比解码中用作占位符的随机噪声时,它们会感到困惑并产生幻觉,导致无法收敛。

这就是 Make Some Noise (MSN) 发挥作用的地方。

核心方法: 训练鲁棒性

作者提出,并行解码本质上是一项去噪任务。如果一个模型能够在看到中间夹杂着乱码的句子时,正确预测出应该在那里的内容,它就能成功执行雅可比解码。

MSN 训练框架

LLM 的标准训练方式使用基于“教师强制”的损失函数,即模型总是看到正确的历史记录:

\[ L o s s _ { A R } = \sum _ { i = 0 } ^ { n } - \log P ( X _ { i } | X _ { < i } ; \theta ) \]

MSN 对此稍作改动。在有监督微调 (SFT) 阶段——即模型学习遵循指令或编写代码的阶段——研究人员随机将输入 Token 的一段替换为“噪声”。

\[ L o s s _ { M S N } = \sum _ { i = 0 } ^ { n } - \log P ( X _ { i } | \hat { X } _ { < i } ; \theta ) \]

这里,\(\hat{X}\) 代表包含噪声的输入序列。至关重要的是,模型试图预测的目标 (\(X_i\)) 仍然是正确的 Token。模型被迫忽略噪声,并根据上下文重构正确的下一个 Token。

图 2: Make Some Noisy 训练框架和雅可比解码策略的图示。图中的训练阶段使用长度为 2 的噪声片段,推理阶段则展示了噪声片段长度设置为 3 时的示例。

让我们看看图 2 (上半部分) 的“MSN 训练阶段”。

  • 原始句子: “I am always on the lookout for exciting…”
  • 噪声输入: 短语 “the look” 被噪声 Token “always I” 替换。
  • 任务: 模型读取 “…always on always I…",但仍必须预测 “the”、“look” 和 “out”。

这迫使模型停止过度依赖紧邻的上一个 Token (那可能是错的) ,转而观察更广泛的上下文。这正是并行解码所需的能力。

什么样的噪声?

你不能只使用随机静态噪声。作者发现使用前文噪声 (Ahead Noise) 效果最好:

\[ \hat { x } _ { i } = r a n d o m \_ s a m p l e ( X _ { < i } ) \]

他们采样之前在句子中出现过的 Token (“前缀 Token”) 。这很高明,因为:

  1. 它模仿了语言生成中经常看到的“复制机制”行为。
  2. 它比纯随机噪声更难去除,因为这些 Token 与上下文在语义上是相关的,迫使模型更具辨别力。

推理: TR-Jacobi 解码

一旦模型经过 MSN 训练,它就变得“鲁棒”了。当它看到不正确的草稿 Token 时不会惊慌。现在,我们如何利用这一点来提速?

作者提出了 TR-Jacobi (基于树的检索增强雅可比) 解码。

第一步: 带噪声的雅可比解码

回顾图 2 的下半部分 (“推理阶段”) 。 模型不再是生成一个 Token,而是生成包含多个 Token 的“草稿” (填充了噪声或猜测) 。因为模型经过了去噪训练,它看着草稿 _on _always _I 并将其修正为 _on _the _look。它并行地迭代此过程,不断优化序列。

第二步: 检索增强 (TR-Jacobi 中的 “R”)

雅可比解码有一个痛点: “冷启动”。在一个新句子的最开始,随机噪声是非常糟糕的猜测。修正它需要太多的迭代次数。

为了解决这个问题,作者增加了一条检索路径。他们使用“提示词查找解码 (Prompt Lookup Decoding) ”。基本上,模型回顾自己的提示词或之前生成的文本,寻找与当前上下文匹配的模式 (n-gram) 。

例如,如果提示词谈到了 “Newton’s Laws” (牛顿定律) ,而模型生成了 “Newton’s”,那么下一个词很有可能是 “Laws”。检索机制抓取 “Laws” 并将其放入草稿候选列表中。

第三步: 树验证

模型不再验证单个序列,而是验证一棵Token 树

图 7: Token 树验证的图示。模型通过特制的稀疏注意力矩阵实现对多条候选路径的同时验证。

如图 7 所示,系统将不同的猜测 (来自雅可比迭代和检索) 组织成一棵树。

  • 路径 A: 雅可比去噪的结果。
  • 路径 B: 检索查找的结果。
  • 路径 C: 这些路径的延伸。

模型使用特定的注意力掩码在一次前向传播中检查所有这些分支。保留最长的有效分支。这最大化了一次性找到长且正确序列的机会。

图 3: TR-Jacobi 解码的主要流程图。需要注意的是,候选生成和树验证是在同一步骤中执行的。为清晰起见,我们在图中选择 T 时刻的候选生成和 T+1 时刻的树验证进行分析。

图 3 展示了这一流程。在步骤 \(T\),生成候选 (有些来自噪声,有些来自检索) 。在步骤 \(T+1\),验证树,并选择最佳路径 (例如,接受 “out” 和 “_for”) 继续向前。

实验与结果

研究人员在两个主要的基础模型上测试了 MSN: LLaMA3-8B (用于通用任务) 和 DeepSeekCoder-6.7B (用于代码任务) 。

1. 噪声训练会损害智能吗?

这是修改 SFT 阶段最大的担忧。如果在训练时喂给模型垃圾 (噪声) ,它会变笨吗?

表 1: 通用领域和代码领域的任务性能实验结果。通用领域指标使用 MT-bench 分数和 MMLU 加权分。代码领域使用贪婪解码下的 pass@1。对于代码领域,我们选择 HumanEval 和 MBPP 的平均值作为综合指标。’(+)’: 执行 evalplus 额外测试后的结果。‘↑’: 相对于未使用 MSN 模型的提升百分比。

表 1 给出了明确的答案: 不会。 实际上,观察 MSN (Ours) 这一行,模型在 MT-Bench 和 HumanEval 等基准测试上的表现与标准 SFT 基线相当,有时甚至略好。例如,在代码领域,MSN 在综合指标上略微优于标准 SFT (77.0 对比 76.6) 。这表明学习去噪实际上可能帮助模型更好地理解上下文,而不是混淆它。

2. 它有多快?

主要目标是速度。作者在 Spec-Bench 上测试了模型,这是一个专为推测解码设计的基准测试。

表 2: Spec-Bench 各个领域 (多轮对话、翻译、摘要、问答、数学推理、检索增强生成) 的加速比实验结果。虚线下方表示类雅可比解码方法。‘AS’: 额外结构。’#MAT’: 平均接受 Token 数。‘↑’: 相对于未使用 MSN 模型的提升百分比。

表 2 将 MSN 与其他方法进行了比较,包括那些使用额外结构 (AS) 的方法,如 Medusa 和 Eagle。

  • 标准 AR: 42.13 tokens/s。
  • MSN + TR-Jacobi: 91.63 tokens/s。

总体而言,这是 2.17 倍的加速 , 某些任务 (如摘要) 甚至接近 2.8 倍 。 值得注意的是,MSN 实现了与 Medusa2 和 EAGLE 相当的加速比,但不需要像那些方法那样需要额外的 1.6B 或 0.3B 参数。它有效地将一个标准的 8B 模型转变为高速推理引擎,而不会增加其内存占用。

3. 消融研究: 噪声与检索

作者进行了更深入的分析,以理解为什么它有效。

训练噪声的长度重要吗? 他们训练了噪声片段长度为 1、4 和 8 的模型。

表 3: 训练噪声片段长度对加速和任务能力的影响。‘L’ 代表噪声片段的长度。

表 3 显示了一个权衡。长度为 1 (L=1) 没有提供太多加速。长度为 8 (L=8) 损害了模型的准确性 (HumanEval 分数下降) 。 长度 4 是最佳平衡点——在保持准确性的同时实现了高加速。

检索真的有帮助吗? 他们比较了纯雅可比解码和他们的 TR-Jacobi (带检索) 。

图 5: 关于 TR-Jacobi 解码中检索部分的消融实验结果。 图 6: 不同尺寸 StarCoder2 模型的 MSN 训练加速实验结果。

在图 5 (左图) 中,蓝色柱状图 (带检索的 TR-Jacobi) 在不同任务中始终优于橙色柱状图 (不带检索) 。这证实了添加“LookAhead”检索路径显著帮助模型从冷启动问题中恢复,导致每步接受更多的 Token。

结论

“Make Some Noise” 论文提出了一个令人信服的论点,让我们重新思考如何训练 LLM 进行推理。MSN 不再将训练和加速视为分离的阶段,而是将它们融合在一起。

通过简单地在有监督微调阶段引入噪声,模型学习到了双重能力:

  1. 任务解决: 它学会了写代码、聊天或写摘要。
  2. 自我修正: 它学会了观察含有噪声的草稿序列,并将其“修正”为正确的形式。

这种自我修正能力是解锁 雅可比解码 的关键,允许模型以并行爆发的方式而不是单步方式生成文本。加上利用检索来稳定过程的 TR-Jacobi , MSN 提供了 2-3 倍的加速。

对于学生和从业者来说,结论很明确: 加速并不总是需要复杂的新架构。有时,它只需要以稍微不同的方式训练模型——通过制造一点噪音。