推测解码：实现示例指南

了解什么是推测解码、其工作原理、何时使用它，以及如何使用 Gemma2 模型实现。

已更新 2026年10月6日 · 12分钟阅读

使用 AI 探索

LLM 功能强大，但在需要速度的场景中，它们往往有些慢，这并不理想。推测解码是一种在不牺牲质量的前提下，加快 LLM 生成响应速度的技术。

本质上，它是在文本生成过程中进行“超前猜测”，提前预测可能出现的下一个词，同时仍然保证我们对 LLM 的准确性和深度预期。

在本文中，我将解释什么是推测解码、它如何工作，以及如何在 Gemma 2 模型上实现它。

什么是推测解码？

推测解码通过引入一个更小、更快的模型来生成初步预测，从而加速 LLM。这个较小的模型通常被称为“草稿”模型，会一次性生成一批 token，供主要、更强大的 LLM 确认或修正。草稿模型相当于第一遍生成，产生多个 token，以加速生成过程。

与其让主 LLM 逐个顺序生成 token，不如由草稿模型提供一组可能的候选，主模型并行评估它们。这种方法将初始预测的计算负载转移到主 LLM 之外，使其只需专注于纠正或验证。

可以把它想象成作家与编辑的配合。主 LLM 就像作家，能产出高质量文本，但速度较慢。更小、更快的“草稿”模型像编辑，能迅速给出可能的续写建议。主 LLM 随后评估这些建议，采纳准确的部分，舍弃其余。这样 LLM 就能并行处理多个 token，从而加快文本生成。

我们把推测解码流程拆解为几个简单步骤：

草稿生成：较小模型（如 Gemma2-2B-it）根据输入提示生成多个 token 建议。这些 token 是“推测性”生成的，模型并不完全确定其正确性，只是作为“草稿”提供。
并行验证：更大模型（如 Gemma2-9B-it）并行验证这些 token，将其概率与模型学习到的分布进行比对。若 token 可接受，则用于最终输出；否则由大模型进行修正。
最终输出：一旦完成验证（或修正），这些 token 将作为最终输出返回给用户。整个流程远快于传统的一次一个 token 的解码方式。

传统解码 vs. 推测解码

传统解码逐个处理 token，延迟较高；推测解码则先由小模型批量生成 token，再由大模型验证。这可以将响应时间降低 30–40%，把延迟从 25–30 秒缩短至约 15–18 秒。

此外，推测解码通过将大部分 token 生成转移给小模型来优化内存使用，将内存需求从 26 GB 降至约 14 GB，使端侧推理更易实现。

最后，它可将算力需求降低 50%，因为大模型只需验证而非生成 token，从而在算力受限的移动设备上也能更流畅地运行，并避免过热。

上手示例：使用 Gemma2 模型进行推测解码

下面用 Gemma2 模型实现一个推测解码的动手示例。我们将从延迟和性能两方面比较推测解码与标准推理。

步骤 1：模型与分词器设置

首先，导入依赖并设置随机种子。

接着，检查您所使用的机器是否可用 GPU。对于 Gemma 2-9B-it 或 LLama2-13B 等大模型，这一点尤为重要。

最后，加载大小两个模型及其分词器。此处我们使用 Gemma2-2b-it（instruct）作为草稿模型，Gemma2-9b-it 作为验证模型。

也有其他可替代的组合，例如：

Gemma 7B（主）& Gemma 2B（草稿）
Mixtral-8x7B（主）& Mistral 7B（草稿）
Pythia 12B（主）& Pythia 70M（草稿）
Llama 2 13B（主）& TinyLlama 1.1B（草稿）

import torch
from transformers import AutoTokenizer, AutoModelForCausalLM, set_seed

# Set Seed
set_seed(42)

# Check if GPU is available
device = "cuda" if torch.cuda.is_available() else "cpu"

# Load the smaller Gemma2 model (draft generation)
small_tokenizer = AutoTokenizer.from_pretrained("google/gemma-2-2b-it", device_map="auto")
small_model = AutoModelForCausalLM.from_pretrained("google/gemma-2-2b-it", device_map="auto", torch_dtype=torch.bfloat16)

# Load the larger Gemma2 model (verification)
big_tokenizer = AutoTokenizer.from_pretrained("google/gemma-2-9b-it", device_map="auto")
big_model = AutoModelForCausalLM.from_pretrained("google/gemma-2-9b-it", device_map="auto", torch_dtype=torch.bfloat16)

步骤 2：自回归（常规）推理

首先，仅在大模型（Gemma2-9b-it）上进行推理并生成输出。先分词输入提示，并将 token 移动到正确的设备（若可用则为 GPU）。generate 方法会生成模型输出，最多生成 max_new_tokens。然后将 token ID 解码回可读文本。

def normal_inference(big_model, big_tokenizer, prompt, max_new_tokens=50):
    inputs = big_tokenizer(prompt, return_tensors='pt').to(device)
    outputs = big_model.generate(inputs['input_ids'], max_new_tokens=max_new_tokens)
    return big_tokenizer.decode(outputs[0], skip_special_tokens=True)

步骤 3：推测解码

接下来尝试推测解码方法，步骤如下：

草稿生成：小模型根据给定提示生成文本草稿。
验证：大模型通过计算草稿中每个 token 的对数似然来验证草稿。
对数似然计算：计算平均对数似然，以判断大模型认为小模型草稿正确的可能性。

def speculative_decoding(small_model, big_model, small_tokenizer, big_tokenizer, prompt, max_new_tokens=50):
    # Step 1: Use the small model to generate the draft
    inputs = small_tokenizer(prompt, return_tensors='pt').to(device)
    small_outputs = small_model.generate(inputs['input_ids'], max_new_tokens=max_new_tokens)
    draft = small_tokenizer.decode(small_outputs[0], skip_special_tokens=True)

    # Step 2: Verify the draft with the big model
    big_inputs = big_tokenizer(draft, return_tensors='pt').to(device)

    # Step 3: Calculate log-likelihood of the draft tokens under the large model
    with torch.no_grad():
        outputs = big_model(big_inputs['input_ids'])
        log_probs = torch.log_softmax(outputs.logits, dim=-1)

    draft_token_ids = big_inputs['input_ids']
    log_likelihood = 0
    for i in range(draft_token_ids.size(1) - 1):
        token_id = draft_token_ids[0, i + 1]
        log_likelihood += log_probs[0, i, token_id].item()

    avg_log_likelihood = log_likelihood / (draft_token_ids.size(1) - 1)

    # Return the draft and its log-likelihood score
    return draft, avg_log_likelihood

说明：对数似然是模型为特定 token 序列分配概率的对数。这里它反映了模型对该 token 序列（生成文本）在既有上下文下的可信度或“信心”。

步骤 4：测量延迟

在实现两种技术后，我们即可测量各自的延迟。对于推测解码，我们通过查看对数似然值来评估表现。对数似然接近 0（尤其是负值区间内接近 0）通常说明生成文本较为准确。

def measure_latency(small_model, big_model, small_tokenizer, big_tokenizer, prompt, max_new_tokens=50):
    # Measure latency for normal inference (big model only)
    start_time = time.time()
    normal_output = normal_inference(big_model, big_tokenizer, prompt, max_new_tokens)
    normal_inference_latency = time.time() - start_time
    print(f"Normal Inference Output: {normal_output}")
    print(f"Normal Inference Latency: {normal_inference_latency:.4f} seconds")
    print("\n\n")

    # Measure latency for speculative decoding
    start_time = time.time()
    speculative_output, log_likelihood = speculative_decoding(
        small_model, big_model, small_tokenizer, big_tokenizer, prompt, max_new_tokens
    )
    speculative_decoding_latency = time.time() - start_time
    print(f"Speculative Decoding Output: {speculative_output}")
    print(f"Speculative Decoding Latency: {speculative_decoding_latency:.4f} seconds")
    print(f"Log Likelihood (Verification Score): {log_likelihood:.4f}")

    return normal_inference_latency, speculative_decoding_latency

返回结果如下：

对数似然（验证分数）：-0.5242
常规推理延迟：17.8899 秒
推测解码延迟：10.5291 秒（约快 70%）

更低的延迟得益于小模型生成文本所需时间更少，而大模型仅负责验证所需时间也更短。

在五个提示上测试推测解码

我们用以下五个提示，比较推测解码与自回归推理：

# List of 5 prompts
prompts = [
    "The future of artificial intelligence is ",
    "Machine learning is transforming the world by ",
    "Natural language processing enables computers to understand ",
    "Generative models like GPT-3 can create ",
    "AI ethics and fairness are important considerations for "
]

# Inference settings
max_new_tokens = 200

# Initialize accumulators for latency, memory, and tokens per second
total_tokens_per_sec_normal = 0
total_tokens_per_sec_speculative = 0
total_normal_latency = 0
total_speculative_latency = 0

# Perform inference on each prompt and accumulate the results
for prompt in prompts:
    normal_latency, speculative_latency, _, _, tokens_per_sec_normal, tokens_per_sec_speculative = measure_latency_and_memory(
        small_model, big_model, small_tokenizer, big_tokenizer, prompt, max_new_tokens
    )
    total_tokens_per_sec_normal += tokens_per_sec_normal
    total_tokens_per_sec_speculative += tokens_per_sec_speculative
    total_normal_latency += normal_latency
    total_speculative_latency += speculative_latency

# Calculate averages
average_tokens_per_sec_normal = total_tokens_per_sec_normal / len(prompts)
average_tokens_per_sec_speculative = total_tokens_per_sec_speculative / len(prompts)
average_normal_latency = total_normal_latency / len(prompts)
average_speculative_latency = total_speculative_latency / len(prompts)

# Output the averages
print(f"Average Normal Inference Latency: {average_normal_latency:.4f} seconds")
print(f"Average Speculative Decoding Latency: {average_speculative_latency:.4f} seconds")
print(f"Average Normal Inference Tokens per second: {average_tokens_per_sec_normal:.2f} tokens/second")
print(f"Average Speculative Decoding Tokens per second: {average_tokens_per_sec_speculative:.2f} tokens/second")

Average Normal Inference Latency: 25.0876 seconds
Average Speculative Decoding Latency: 15.7802 seconds
Average Normal Inference Tokens per second: 7.97 tokens/second
Average Speculative Decoding Tokens per second: 12.68 tokens/second

可见推测解码更高效，单位时间生成的 token 数量多于常规推理。提升的原因在于小模型承担了大部分文本生成工作，而大模型的职责主要是验证，从而在延迟与内存方面都降低了整体计算开销。

在这样的内存需求下，我们可以轻松在边缘设备上部署推测解码技术，并在端侧应用（如聊天机器人、翻译器、游戏等）中获得加速效果。

通过量化优化推测解码

上述方法虽高效，但在端侧推理时，延迟与内存优化之间存在取舍。为此，我们对大小模型都应用量化。您也可以只对大模型进行量化试验，因为小模型本身占用空间已经较小。

我们使用 Hugging Face transformers 库中的 BitsAndBytesConfig 配置对大小模型进行量化。量化通过将模型权重转换为更紧凑的形式，显著降低内存占用，并在很多情况下提升推理速度。

将以下代码片段加入上述代码中，即可观察量化带来的效果。

bnb_config = BitsAndBytesConfig(
    load_in_4bit=True,  # Enables 4-bit quantization
    bnb_4bit_quant_type="nf4",  # Specifies the quantization type (nf4)
    bnb_4bit_compute_dtype=torch.bfloat16,  # Uses bfloat16 for computation
    bnb_4bit_use_double_quant=False,  # Disables double quantization to avoid additional overhead
)

# Load the smaller and larger Gemma2 models with quantization applied
small_model = AutoModelForCausalLM.from_pretrained("google/gemma-2-2b-it", device_map="auto", quantization_config=bnb_config)
big_model = AutoModelForCausalLM.from_pretrained("google/gemma-2-9b-it", device_map="auto", quantization_config=bnb_config)

下面对比一下启用与未启用量化时，推测解码的输出差异：