投机性解码：机器学习从业者的实践指南

📢 转载信息

原文链接：https://machinelearningmastery.com/the-machine-learning-practitioners-guide-to-speculative-decoding/

原文作者：Jason Brownlee

---

投机性解码（Speculative Decoding）是现代大型语言模型（LLM）加速推理过程的一种有效方法，特别是在自回归解码阶段。本文将作为一份面向机器学习从业者的实践指南，详细介绍投机性解码的工作原理、实现方法以及如何在实际应用中部署它。

什么是投机性解码？

自回归模型（如GPT模型）在生成文本时，必须一个Token一个Token地进行采样，这是一个序列化的过程，限制了生成速度。投机性解码旨在通过并行化这个过程来克服这一瓶颈。

其核心思想是使用一个小型、快速的“草稿模型”（Draft Model）来预测未来的一系列Token，然后由大型的“目标模型”（Target Model）并行地验证这些预测。目标模型只需一步就能确认或拒绝草稿模型生成的多个Token，从而显著减少了目标模型进行大量计算的次数。

工作流程

投机性解码的工作流程可以分解为以下几个关键步骤：

1. 草稿生成：使用小型草稿模型快速生成 $K$ 个Token的序列。
2. 并行验证：将这 $K$ 个Token的序列作为输入提供给大型目标模型，让它并行地预测下一个Token的概率分布。
3. 接受与拒绝：通过对比草稿模型和目标模型对序列中每个位置的预测，确定哪些Token被接受（与目标模型预测一致）和哪些被拒绝（需要目标模型重新采样）。
4. 更新上下文：将所有被接受的Token添加到已生成的序列中，并更新上下文，然后重复该过程。

如果目标模型接受了所有 $K$ 个Token，那么推理速度理论上可以提高 $K$ 倍。

模型选择

投机性解码的性能在很大程度上取决于草稿模型和目标模型的选择：

• 目标模型（Target Model）：通常是需要加速的大型、高性能LLM（如GPT-3、Llama等）。
• 草稿模型（Draft Model）：需要比目标模型小得多、推理速度快得多，但仍然要足够准确，以确保生成序列的质量。理想情况下，草稿模型的准确性应该足够高，以保证大部分Token能被接受。

投机性解码的实现

实现投机性解码需要细致地协调两个模型的输入和输出。以下是关键的步骤和考虑因素。

1. 准备模型

首先，您需要加载并准备好两个模型：一个大型目标模型和一个小型草稿模型。两个模型通常基于相同的架构（例如，都是Transformer），但草稿模型的参数量更少。

例如，如果您使用Hugging Face Transformers库，这可能涉及加载两个不同大小的预训练模型。

2. 草稿生成

在每一步迭代中，使用草稿模型生成一个短序列的Token。这个过程是标准的自回归生成，但速度很快。

3. 目标模型并行验证

这是投机性解码的核心和难点。目标模型需要同时评估草稿序列的每一步预测。这通常涉及到构造一个输入序列，该序列由已确定的Token前缀加上草稿序列的Token组成。

目标模型为这个序列生成下一个Token的概率分布。例如，如果目标模型接收到上下文 $C$ 和草稿序列 $D = [d_1, d_2, ..., d_K]$，它会预测 $P(x | C, d_1, ..., d_i)$。

4. 接受/拒绝采样

现在，我们需要逐个检查草稿Token是否被目标模型“接受”。

对于草稿Token $d_i$：

• 我们比较目标模型在位置 $i$ 处预测的概率 $P_{target}(d_i)$ 与草稿模型在位置 $i$ 处预测的概率 $P_{draft}(d_i)$。
• 接受一个Token的标准方式是，根据目标模型的概率分布进行采样。如果采样的Token是 $d_i$，则接受 $d_i$。
• 通常，我们使用一种一致性采样策略。我们接受序列中尽可能多的Token，直到出现第一个不一致的Token，或者直到达到草稿序列的长度 $K$。

关键点：如果草稿模型预测的Token $d_i$ 恰好是目标模型最有可能采样的Token，那么它就被接受。如果目标模型采样到了一个不同的Token，那么序列就会中断，我们只保留被接受的部分。

5. 迭代与整合

将所有被接受的Token添加到最终输出中。然后，使用这些新确认的Token作为新的上下文前缀，重复整个过程，直到生成结束标记（EOS）。

示例：

假设上下文为 $C$。草稿模型生成 $D=[d_1, d_2, d_3]$。

1. 目标模型评估 $P(x | C, d_1)$。如果目标模型采样到 $d_1$，则接受 $d_1$。
2. 目标模型评估 $P(x | C, d_1, d_2)$。如果目标模型采样到 $d_2$，则接受 $d_2$。
3. 目标模型评估 $P(x | C, d_1, d_2, d_3)$。如果目标模型采样到 $d_3'$（与 $d_3$ 不同），则拒绝 $d_3$，并使用目标模型在这一步的采样结果 $d_3'$ 作为最终输出。

最终，我们接受了 $[d_1, d_2, d_3']$ 并更新上下文，然后开始下一轮投机。

实践考量与优化

虽然投机性解码提供了显著的加速潜力，但实践中需要注意以下几点：

1. 草稿模型的准确性

草稿模型越准确（即它生成的Token与目标模型采样的结果越一致），加速效果越好。一个非常差的草稿模型甚至可能拖慢速度，因为它产生的许多Token都会被目标模型拒绝。

2. 批量大小（Batch Size）

投机性解码的加速效果在大批量推理时最为显著，因为并行验证步骤可以高效地利用目标模型的计算能力。

3. 引入噪声（Noisy Draft）

一些研究表明，即使草稿模型不是最优的，也可以通过引入一些随机性（例如，不总是选择草稿模型概率最高的Token进行验证）来提高整体性能，因为它可能避免陷入局部最优路径。

4. 硬件与并行化

要获得最佳性能，验证步骤必须在GPU上高效并行执行。这意味着目标模型的推理必须能够有效地处理多个输入序列（前缀+草稿Token的各个前缀）。

总结

投机性解码是实现高吞吐、低延迟LLM推理的关键技术之一。它通过“猜测并验证”的范式，有效弥补了自回归模型的序列化限制。对于希望将大型模型部署到生产环境中的机器学习工程师和研究人员来说，理解并正确实现草稿生成、并行验证和一致性采样，是释放LLM全部性能潜力的重要一步。

---

🚀 想要体验更好更全面的AI调用？

欢迎使用青云聚合API，约为官网价格的十分之一，支持300+全球最新模型，以及全球各种生图生视频模型，无需翻墙高速稳定，文档丰富，小白也可以简单操作。