投机解码（Speculative Decoding）是什么？

前言

近年来，大语言模型（LLM）以其惊人的能力改变了自然语言处理的格局。然而，这些强大的模型也带来了巨大的计算挑战，尤其是在推理（Inference）阶段。模型的规模越大，生成文本所需的时间就越长，这限制了它们在实时交互、大规模部署等场景下的应用。为了解决这个瓶颈，研究人员提出了各种优化技术，其中，“投机解码”（Speculative Decoding）是一种极具潜力且备受关注的方法。

本文将深入浅出地介绍投机解码的原理、优势、挑战以及它为何能显著加速 LLM 推理过程。

LLM 推理的瓶颈：自回归生成

要理解投机解码，首先需要了解传统 LLM 是如何生成文本的。大多数 LLM（如 GPT 系列）采用**自回归（Autoregressive）**的方式生成文本。这意味着模型一次只生成一个词元（Token），并且下一个词元的生成依赖于之前所有已生成的词元。

这个过程可以简化为：

1. 输入一个提示（Prompt）。
2. 模型根据提示预测下一个最可能的词元。
3. 将新生成的词元添加到现有序列中。
4. 将更新后的序列再次输入模型，预测再下一个词元。
5. 重复步骤 3 和 4，直到生成结束符或达到最大长度。

这种串行（Sequential）的生成方式存在一个核心问题：延迟。每生成一个词元，都需要完整地运行一次庞大的模型进行前向传播（Forward Pass）。对于包含数十亿甚至数千亿参数的模型来说，这个过程非常耗时，并且主要受限于内存带宽（Memory Bandwidth）——即从内存中加载模型权重所需的时间。因此，即使硬件计算能力很强，生成速度也常常不尽人意。

投机解码：大胆假设，小心求证

投机解码的核心思想是打破“一次只生成一个词元”的限制，尝试一次性预测和验证多个词元，从而减少调用大模型的次数，提高生成效率。

它引入了一个巧妙的机制，通常需要两个模型：

1. 目标模型（Target Model）/验证模型（Verifier Model）：这就是我们最终想要使用的大型、高质量但速度较慢的 LLM。它的输出是我们信任的“黄金标准”。
2. 草稿模型（Draft Model）/提议模型（Proposer Model）：这是一个规模小得多、速度快得多的模型。它可以是目标模型的蒸馏版本、早期检查点，甚至是完全不同的轻量级模型。它的任务是快速地“猜测”或“提议”接下来可能出现的多个词元。

工作流程详解

投机解码的流程可以概括为以下几个步骤：

1. 草稿生成 (Drafting)：

   • 给定当前已生成的文本序列 $x$。
   • 使用草稿模型快速地、自回归地生成一个包含 $k$ 个候选词元的序列 $\gamma =(t_1,t_2,...,t_k)$。例如，草稿模型可能预测接下来是 " is a large language"。这里 $k=4$。

2. 并行验证 (Parallel Verification)：

   • 将原始序列 $x$ 和草稿模型生成的整个序列 $\gamma$ (或者说 $x$ 后面拼接上 $\gamma$ 的部分序列) 一次性输入到目标模型中。
   • 目标模型进行一次前向传播，并行地计算出对于输入序列 $x$ 加上 $\gamma$ 中每个位置的词元时，它自己会生成的下一个词元的概率分布。
     • 即计算 $P_{t}arget(token|x)$
     • 计算 $P_{t}arget(token|x,t_1)$
     • 计算 $P_{t}arget(token|x,t_1,t_2)$
     • …
     • 计算 $P_{t}arget(token|x,t_1,...,t_{k-1})$

3. 接受或拒绝 (Acceptance/Rejection)：

   • 现在，我们需要比较草稿模型的“猜测” ($t_i$) 和目标模型的“判断”。

   • 从草稿序列的第一个词元 $t_1$ 开始，逐个进行验证：

     • 比较 $t_1$：目标模型在输入 $x$ 后，是否也将 $t_1$ 作为最可能的下一个词元？
       • 如果一致：接受 $t_1$。继续比较 $t_2$。
       • 如果不一致：拒绝 $t_1$ 以及之后的所有草稿词元 ($t_2$ 到 $t_k$)。
     • 比较 $t_2$：(如果 $t_1$ 被接受了) 目标模型在输入 $x,t_1$ 后，是否也将 $t_2$ 作为最可能的下一个词元？
       • 如果一致：接受 $t_2$。继续比较 $t_3$。
       • 如果不一致：拒绝 $t_2$ 以及之后的所有草稿词元 ($t_3$ 到 $t_k$)。
     • 以此类推，直到找到第一个不匹配的词元 $t_i$，或者所有 $k$ 个草稿词元都被成功验证。

   • 更精确的验证方法（概率视角）：实际应用中，不仅仅是看最高概率的词元是否匹配。一种常用的方法是基于目标模型和草稿模型对词元的概率分布进行比较和采样。如果目标模型认为草稿词元 $t_i$ 的概率足够高（相对于草稿模型自己的概率或者一个阈值），则接受它；否则拒绝。这种基于概率的拒绝采样方法可以数学上保证最终生成的文本分布与单独使用目标模型完全一致。

4. 处理结果与迭代：

   • 假设验证过程中，草稿序列的前 $n$ 个词元 ($t_1$ 到 $t_n$，其中 $0<=n<=k$) 被接受了。
   • 将这 $n$ 个接受的词元 $t_1,...,t_n$ 追加到最终生成的序列 $x$ 后面。
   • 如果 $n<k$ (即在第 $n+1$ 个词元处发生不匹配)：目标模型在验证第 $n+1$ 个位置时，已经计算出了它自己认为最可能的词元 $t_n^{\prime }+1$ (可能与草稿模型的 $t_{n+1}$ 不同)。将这个由目标模型修正或确认的词元 $t_n^{\prime }+1$ 追加到序列末尾。
   • 如果 $n==k$ (即所有草稿词元都被接受)：目标模型在验证最后一个草稿词元 $t_k$ 时，也计算出了它认为在 $x,t_1,...,t_k$ 之后最可能的词元 $t_{k+1}^{\prime }$。将这个词元追加到序列末尾。
   • 现在，我们得到了一个新的、更长的序列。以此作为新的起点，重复步骤 1-4，继续生成后续文本。

为什么投机解码能加速？

加速的关键在于：

1. 并行处理：虽然草稿模型是串行生成 $k$ 个词元，但验证这 $k$ 个词元（以及计算修正词元）只需要目标模型进行一次前向传播。
2. 多次命中：如果草稿模型预测得比较准（例如，在生成常见短语、代码片段或遵循简单模式时），目标模型一次调用就能确认多个词元 ($n>1$)。这意味着目标模型的单次调用“摊销”到了多个词元的生成上。
3. 减少大模型调用次数：相比传统方法需要调用 $N$ 次目标模型来生成 $N$ 个词元，投机解码理想情况下可以用远少于 $N$ 次的目标模型调用来生成相同数量的词元。

核心优势

• 显著加速：尤其对于长文本生成、高并发请求场景，可以带来 2-4 倍甚至更高的推理速度提升。
• 无损质量：理论上，通过恰当的接受/拒绝机制（如概率拒绝采样），投机解码生成的文本质量与完全使用目标模型生成的结果在分布上是等价的。它改变的是生成过程，而不是最终结果的统计属性。
• 模型兼容性：可以应用于各种自回归的 Transformer 模型。

挑战与考量

• 草稿模型的选择与开销：
  1. 草稿模型需要足够快，否则其生成时间会抵消部分收益。
  2. 草稿模型的质量也很关键。如果预测太差，大部分词元都会被拒绝，导致 $n$ 值很小，加速效果不明显。
  3. 运行两个模型会增加内存占用。
• 参数 $k$ 的选择：草稿序列的长度 $k$ 需要权衡。太小则加速潜力有限，太大则草稿生成的耗时增加，且后面词元的预测准确率可能下降，导致拒绝率升高。
• 实现复杂度：相比标准解码，投机解码的实现逻辑更复杂，需要仔细处理模型间的交互和验证逻辑。

变种与扩展

投机解码的思想激发了许多后续研究，例如：

• 块状投机解码 (Blockwise Speculative Decoding)：尝试并行验证多个独立的草稿块。
• 树状注意力 (Tree Attention) / Medusa：让模型一次性预测多个“头”，形成一个预测树，然后并行验证这些分支，选择最优路径。这些方法不一定需要独立的草稿模型，而是修改了目标模型的结构或解码策略。

结论

投机解码是一种非常有效的 LLM 推理加速技术。它通过引入一个轻量级的草稿模型来“投机性”地预测多个未来词元，然后利用强大的目标模型进行并行验证，从而在不牺牲生成质量的前提下，大幅减少了对昂贵大模型的调用次数。随着 LLM 规模的持续增长和应用场景的不断扩展，类似投机解码这样的优化技术对于释放大模型潜力、降低应用门槛至关重要。理解其工作原理，有助于我们更好地利用和部署这些强大的 AI 工具。