投机解码的核心思路并不复杂:让一个小参数量的草稿模型连续生成若干个候选token,再让目标大模型一次前向并行验证这些token,接受与目标模型分布一致的草案,拒绝并纠正偏差的草案。由于大模型验证k个token的计算量与生成1个token几乎相同(自回归逐个生成才是逐次计算),当草稿模型质量足够好时,整体推理延迟能显著下降。但实践里最常见的翻车场景是:草稿模型与目标模型的输出分布差异太大,草案接受率低到可怜,每轮投机都要浪费大量草稿生成时间,最后比不用投机解码还慢。解决这个问题的正道,就是通过蒸馏和针对性训练,把草稿模型的输出分布拉近到目标模型。

为什么草稿模型接受率总是上不去
先要理解接受率的数学定义。假设草稿模型在某个上下文c下给出token x的概率是q(x|c),目标模型给出的是p(x|c),那么单步接受概率近似为E[min(1, q/p)](在原始的拒绝采样框架中是1 - max(0, 1-q/p)的期望形式)。从这个式子能直接读出两个结论:第一,q与p越接近,接受率越高;第二,如果草稿模型把概率质量押在了目标模型认为概率很低的token上,接受率会断崖式下跌。
接受率低的典型原因有三个。一是领域错配:草稿模型是在通用语料上训练的,而目标模型可能在代码、数学或某个垂直领域被深度微调过,两者的分布差异不是参数量差异造成的,而是训练数据造成的。二是分词器不一致:草稿模型与目标模型的tokenizer不同,即使语义相同,token边界对不齐,草案验证时天然错位。三是上下文对齐问题:投机解码运行时,草稿模型看到的是目标模型已经修正过的上下文,如果训练时草稿模型从未见过这种上下文分布,推理时表现就会偏离。
诊断方法很简单:拿一批目标模型实际服务的请求,统计草稿模型在每个位置的top-1预测与目标模型的一致率,以及草案平均接受长度(通常期望达到3以上才有明显加速收益)。如果一致率长期低于60%,单纯靠投机解码框架的改进是救不回来的,必须从训练入手。
用目标模型做蒸馏:让分布对齐成为训练目标
蒸馏是提升草稿模型质量最直接的方案,因为它的优化目标就是让q逼近p。最常见的做法是软标签蒸馏:对语料中每个位置,让目标模型输出完整的logits分布,草稿模型在这个分布上做KL散度最小化。损失函数通常写成前向KL的形式:
import torch
import torch.nn.functional as F
def distill_loss(student_logits, teacher_logits, temperature=2.0):
# 软标签蒸馏:KL(softmax(t) || softmax(s))
# temperature 越高,分布越平滑,学生学到的是相对结构
t_prob = F.log_softmax(teacher_logits / temperature, dim=-1)
s_log_prob = F.log_softmax(student_logits / temperature, dim=-1)
kl = F.kl_div(s_log_prob, t_prob, reduction="batchmean")
# 蒸馏损失通常配合硬标签的交叉熵一起用,防止过度平滑
hard_loss = F.cross_entropy(
student_logits.view(-1, student_logits.size(-1)),
labels.view(-1), ignore_index=-100
)
return kl * (temperature ** 2) + 0.5 * hard_loss几个实操细节值得注意。第一,温度系数的选择:温度在1到3之间比较常见,温度越高教师分布越平滑,学生模型学到的是token之间的相对排序关系;如果目标模型分布本身很尖锐(比如代码场景中token几乎确定),温度取1更合适。第二,top-k截断:完整词表的KL计算开销不小,实践中常只取教师top-50的logits做蒸馏,剩余概率质量归并到一个虚拟类别,这在几乎不损失效果的前提下能把训练提速数倍。第三,教师logits务必用float32保存,或者直接离线把教师的分布存成数据集,避免训练时反复跑大模型推理。
另一个蒸馏变体是针对投机解码本身的序列级蒸馏:不是逐token对齐,而是让草稿模型去模仿目标模型在真实推理中的采样轨迹,包括目标模型自己生成的完整回复。用目标模型批量生成几十万条对话或代码样本,再用这些样本做普通下一词预测训练,效果往往比逐token蒸馏更稳,因为草稿模型学到的上下文分布与推理时见到的完全一致。
从零训练还是基于现有小模型微调
两条路线各有适用场景。从零训练一个草稿模型的最大好处是分词器可以与目标模型完全一致,这消除了token边界错位这个隐患。做法通常是随机初始化一个小架构(比如目标模型的4到8层浅层版本),先用通用语料训练打底,再进入蒸馏阶段对齐目标模型的分布。缺点是训练成本高,且浅层小模型的表达能力有上限,接受率天花板受参数量限制。
# 基于现有小模型做蒸馏微调的典型流程
from transformers import AutoModelForCausalLM, AutoTokenizer
teacher = AutoModelForCausalLM.from_pretrained("target-model-14b", torch_dtype=torch.float16)
student = AutoModelForCausalLM.from_pretrained("draft-model-1.5b", torch_dtype=torch.bfloat16)
# 关键前提:检查两者tokenizer是否一致
assert teacher.config.vocab_size == student.config.vocab_size, \
"词表大小不一致,需要做词表对齐或扩展"基于现有小模型微调是更常见的路线。选一个与目标模型同源的小尺寸模型(同一模型家族通常共享tokenizer和部分训练配方),先做词表对齐检查,然后进入两阶段微调:第一阶段用目标模型的采样输出做监督微调,让草稿模型适应目标模型的风格和领域;第二阶段做logits级蒸馏,精细对齐分布。如果小模型与目标模型不同源,词表不一致时可以用词表映射矩阵做_embedding初始化,或者直接对小模型的embedding层和输出层做扩容重训。
训练数据的构造比大多数人想象的重要。最有效的数据来源是目标模型的真实服务日志:把线上实际收到的prompt重新采样一遍目标模型的输出,这些(prompt, response)对就是草稿模型最理想的训练语料,因为它精确覆盖了推理时会出现的上下文分布。其次是任务匹配的合成数据,比如目标模型主要服务代码补全,那就用代码仓库构造下一词预测语料。纯通用语料在蒸馏阶段的优先级最低。
训练后的评估与持续迭代
训练完成后不要只看困惑度,要直接测投机解码指标:用一批真实请求跑完整的投机解码流程,统计草案平均接受长度、端到端延迟加速比、每秒生成token数三项。平均接受长度至少要到3以上,加速比要稳定在1.5倍以上才值得上线。同时要观察长尾:如果部分请求接受率极低,多半是训练数据没覆盖到该类输入,需要定向补充数据再迭代一轮。
还有一点容易被忽视:目标模型更新(比如安全对齐微调、版本升级)之后,草稿模型的分布对齐会退化,接受率会悄悄下降。建议把接受率监控做成常驻指标,一旦下降超过阈值就触发一轮轻量蒸馏增量训练。草稿模型不是一次性工程,而是需要跟着目标模型同步演进的配套组件,把它当作推理系统的一部分来运营,投机解码的加速收益才能长期稳定保持。