投机采样（Speculative Decoding）详解：用小模型带路给大模型加速的黑科技

随着大型语言模型（LLM）的尺寸不断增大，推理速度成为了制约其广泛应用的关键瓶颈。标准的自回归（Autoregressive）采样模式要求模型每生成一个Token，就进行一次完整的全模型前向计算，效率低下。

投机采样（Speculative Decoding，或称 Speculative Sampling）是一种颠覆性的加速技术，它利用一个更小、更快的“草稿模型”（Draft Model）来预测后续的多个Token，然后让计算量更大的“目标模型”（Target Model）一次性并行验证这些预测，从而显著提高推理速度。

核心原理：草稿与验证

投机采样将传统的串行推理过程转换为一个高效的并行验证过程。它涉及到两个核心组件：

1. 草稿模型 (Draft Model, D): 一个参数量较小、推理速度极快、但质量可能稍逊的LLM。
2. 目标模型 (Target Model, T): 原始的、参数量大、输出质量最高的LLM（即我们最终希望加速的模型）。

Speculative Decoding 的三步流程

假设我们已经生成了序列 $X$。

步骤 1: 草稿生成（Drafting）

Draft Model (D) 利用 $X$ 快速连续生成 $K$ 个投机性（speculative）的Token序列 $\hat{A} = (\hat{a}_1, \hat{a}_2, …, \hat{a}_K)$。由于D很小，这一步骤极快。

步骤 2: 并行验证（Verification）

Target Model (T) 接收 $X$ 加上整个草稿序列 $\hat{A}$ 作为输入。T 仅需进行一次前向计算，就能并行计算出在每个位置 $x_i$ 之后的正确Token的概率分布。

$$P_T(\hat{a}i | X, \hat{a}_1, …, \hat{a}{i-1})$$

注意：虽然是单次前向传播，但T在计算时内部会并行地计算所有位置的注意力机制和输出 logits。

步骤 3: 接受与拒绝（Acceptance/Rejection Sampling）

系统从 $i=1$ 到 $K$ 依次检查：D 预测的 Token $\hat{a}_i$ 是否与 T 在该位置预测的 Token 相同（或符合一定的概率接受标准）。

• 接受： 如果 $\hat{a}i$ 被接受，我们立即将它添加到最终输出序列中，并检查下一个Token $\hat{a}{i+1}$。
• 拒绝： 如果 $\hat{a}i$ 被拒绝，系统停止接受，丢弃所有后续的草稿Token $(\hat{a}{i}, \hat{a}_{i+1}, …)$。此时，T在该位置预测出的正确Token将取代被拒绝的Token，并作为新的上下文，重新开始步骤 1 的草稿生成。

通过这种方式，如果 D 的预测准确，T 可以用相当于生成一个Token的时间，确认 $K$ 个Token，实现 $K$ 倍的加速。

实操：如何在 Hugging Face Transformers 中启用投机采样

自从 Transformers 库 v4.37 版本及以后，主流模型（如 Llama、Mistral）已经开始原生支持Speculative Decoding。用户只需加载两个模型：目标模型和草稿模型，并在生成时进行配置。

1. 环境准备

pip install transformers torch accelerate

2. Python 代码示例

假设我们使用一个小型模型（例如，一个1B级别的模型）作为草稿模型，并使用一个大型模型（例如，Llama 7B）作为目标模型。

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

# 检查设备
device = "cuda" if torch.cuda.is_available() else "cpu"

# --- 步骤 1: 加载目标模型 (Target Model, T) ---
target_model_id = "meta-llama/Llama-2-7b-hf" # 目标模型，请替换为实际可访问的模型ID
target_model = AutoModelForCausalLM.from_pretrained(target_model_id, torch_dtype=torch.float16).to(device)
tokenizer = AutoTokenizer.from_pretrained(target_model_id)

# --- 步骤 2: 加载草稿模型 (Draft Model, D) ---
# 注意: 实际应用中，草稿模型需要与目标模型有相似的词汇表（Tokenizer）
draft_model_id = "facebook/opt-125m" # 更小的模型作为草稿
draft_model = AutoModelForCausalLM.from_pretrained(draft_model_id, torch_dtype=torch.float16).to(device)

# --- 步骤 3: 配置并运行 Speculative Decoding ---

prompt = "请写一篇关于量子计算的短文，并介绍其应用前景："
inputs = tokenizer(prompt, return_tensors="pt").to(device)

print("\n--- 启用 Speculative Decoding 推理 ---")

# 通过传入 assistant_model 参数，Hugging Face 框架自动启用 Speculative Decoding
# 注意：运行此代码需要足够的VRAM同时加载两个模型
output_speculative = target_model.generate(
    **inputs,
    assistant_model=draft_model, # 传入草稿模型
    max_new_tokens=256,
    do_sample=False, # 通常在贪婪解码或束搜索中效果最佳
    num_beams=1
)

print("加速后的输出:")
print(tokenizer.decode(output_speculative[0], skip_special_tokens=True))

# 传统自回归推理 (作为对比)
print("\n--- 传统自回归推理 (对比) ---")
output_regular = target_model.generate(
    **inputs,
    max_new_tokens=256,
    do_sample=False,
    num_beams=1
)

# 在实际的生产环境中，你会观察到 'assistant_model' 启用的推理速度显著提升。

性能与局限性

性能优势

1. 显著提速: 在最佳情况下，Speculative Decoding 可以实现 2x 到 3x 的加速，因为它减少了目标模型的计算次数。
2. 保证质量: 最终输出的Token仍然来自于目标模型（T）的概率分布，这意味着输出的质量与传统的自回归采样完全相同。它只改变了计算路径，不改变结果。

局限性

1. 硬件要求: 运行 Speculative Decoding 需要同时加载和运行两个模型（D和T），这会增加显存（VRAM）的占用。
2. 草稿模型质量: 加速效果高度依赖于 Draft Model 的准确性。如果 D 的预测能力很差，大量的 Token 会被拒绝，导致加速效果不明显，甚至可能比传统方法更慢（因为引入了额外的验证开销）。因此，D 通常选择经过良好训练的小型模型，或者甚至是 T 模型的一个量化或蒸馏版本。