大模型推理中的投机采样用小模型加速大模型生成的原理与实现一、自回归解码的串行瓶颈Transformer的自回归解码是严格串行的生成第t个token依赖于前t-1个token的完整上下文。每次解码步骤需要执行一次完整的前向传播加载全部参数、计算全部层的注意力但仅生成一个token。这种计算产出比极低——LLaMA-70B的每次前向传播需要遍历70B参数、执行数千亿次浮点运算而产出仅是4字节的一个token ID。KV Cache已经大幅减少了重复计算前t-1个token的KV表示被缓存不需重复计算但核心瓶颈依然存在每次前向传播仍需要加载全部参数并在全部层上执行计算。投机采样Speculative Decoding进一步挑战这一瓶颈它使用一个小型草稿模型Draft Model快速生成k个候选token然后用大模型Target Model一次性验证这k个token的正确性。二、投机采样的数学原理投机采样的核心思想来自拒绝采样Rejection Sampling。设小模型draft的分布为q(x)大模型target的分布为p(x)。基本流程小模型自回归生成k个候选tokenx₁, x₂, ..., xₖ每个xᵢ从q(x|prefix, x₁, ..., x_{i-1})采样。大模型对原始prefix k个候选token进行一次前向传播得到p(x|prefix)、p(x|prefix, x₁)、...、p(x|prefix, x₁, ..., xₖ)。从i1到k依次验证以概率min(1, p(xᵢ)/q(xᵢ))接受xᵢ若某步拒绝则从p的残差分布中重新采样。关键数学性质最终生成的token序列的分布恰好等于p大模型的原始分布。这意味着投机采样是一种无损加速技术——它在数学上等价于直接从大模型采样只是减少了计算量。import torch import torch.nn.functional as F from typing import List, Tuple, Optional def speculative_decode( target_model: torch.nn.Module, draft_model: torch.nn.Module, prefix: torch.Tensor, # (1, prefix_len) k: int 5, # 每次生成的候选token数 max_new_tokens: int 100, temperature: float 1.0 ) - Tuple[torch.Tensor, dict]: 投机采样的完整实现。 算法流程 1. Draft Model自回归生成k个候选token 2. Target Model一次前向验证所有候选 3. 逐个接受/拒绝确保输出分布与target一致 Args: target_model: 大模型目标分布p draft_model: 小模型草稿分布q prefix: 输入前缀 (1, L) k: 每次生成的候选token数 max_new_tokens: 最大生成token数 temperature: 采样温度 Returns: (生成的token序列, 统计信息) generated [] stats { total_target_calls: 0, # 大模型前向传播次数 total_draft_tokens: 0, # 小模型生成的候选总数 total_accepted: 0, # 被接受的候选数 total_generated: 0 # 最终生成的总token数 } current_prefix prefix while len(generated) max_new_tokens: # Phase 1: Draft Model 快速生成k个候选 draft_tokens [] draft_prefix current_prefix with torch.no_grad(): for _ in range(k): draft_logits draft_model(draft_prefix) draft_probs F.softmax( draft_logits[0, -1] / temperature, dim-1 ) next_token torch.multinomial(draft_probs, 1) draft_tokens.append(next_token.item()) draft_prefix torch.cat([draft_prefix, next_token.unsqueeze(0)], dim1) if next_token.item() eos_token_id: break stats[total_draft_tokens] len(draft_tokens) if not draft_tokens: break # Phase 2: Target Model 一次性验证 with torch.no_grad(): # 将前缀所有候选一起输入大模型 # 大模型计算 p(x|prefix), p(x|prefix,x₁), ..., p(x|prefix,x₁,...,xₖ) target_prefix torch.cat([ current_prefix, torch.tensor([draft_tokens]).to(current_prefix.device) ], dim1) # (1, Lk) target_logits target_model(target_prefix) # target_logits[:, L-1:Lk-1, :] # 对应每个位置的预测分布 target_logits target_logits[ 0, current_prefix.size(1)-1:, : ] # (1k, vocab_size) target_probs F.softmax( target_logits / temperature, dim-1 ) stats[total_target_calls] 1 # Phase 3: 逐个验证和接受/拒绝 all_accepted True num_accepted 0 for i, draft_token in enumerate(draft_tokens): # 小模型的采样概率 if i 0: draft_logits_for_i draft_model(current_prefix) else: draft_prefix_i torch.cat([ current_prefix, torch.tensor([draft_tokens[:i]]).to(current_prefix.device) ], dim1) draft_logits_for_i draft_model(draft_prefix_i) draft_prob F.softmax( draft_logits_for_i[0, -1] / temperature, dim-1 )[draft_token].item() target_prob target_probs[i, draft_token].item() # 接受概率: min(1, p(x)/q(x)) accept_prob min(1.0, target_prob / (draft_prob 1e-8)) if torch.rand(1).item() accept_prob: # 接受该token generated.append(draft_token) stats[total_accepted] 1 num_accepted 1 else: # 拒绝从残差分布重新采样 # 残差分布: max(0, p(x) - q(x)) / sum(max(0, p - q)) residual torch.clamp( target_probs[i] - draft_prob, min0 ) residual residual / (residual.sum() 1e-8) new_token torch.multinomial(residual, 1).item() generated.append(new_token) all_accepted False break # 更新前缀 current_prefix torch.cat([ current_prefix, torch.tensor([generated[-num_accepted-1:]]).to(current_prefix.device) ], dim1) if not all_accepted else draft_prefix stats[total_generated] len(generated) stats[acceptance_rate] ( stats[total_accepted] / stats[total_draft_tokens] if stats[total_draft_tokens] 0 else 0.0 ) return torch.tensor([generated]), stats三、接受率与加速比的关系投机采样的实际加速比由两个因素决定(1) 接受率α——大模型接受小模型预测token的概率(2) 相对速度比r T_target / T_draft——大模型单次前向时间与小模型k次前向时间的比值。加速比公式$$Speedup \frac{1-\alpha^{k1}}{(1-\alpha)(1 k/r)}$$当小模型与大模型的一致性很高α→1时加速比趋近于r·k/(kr)。当α很低时投机采样的开销小模型生成了候选但被大量拒绝可能反而使整体变慢。def compute_expected_speedup( acceptance_rate: float, k: int, speed_ratio: float # T_target / (k * T_draft_per_step) ) - float: 计算期望加速比。 Args: acceptance_rate: 大模型接受小模型token的概率α k: 候选token数 speed_ratio: r 大模型时间 / (k × 小模型单步时间) Returns: 期望加速比1表示加速1表示减速 # 期望生成的token数 # 几何级数求和: E[accepted] (1 - α^(k1)) / (1 - α) expected_accepted (1 - acceptance_rate ** (k 1)) / (1 - acceptance_rate) # 期望时间: T_target k * T_draft # 归一化: 以标准解码的单步时间为1 standard_time_per_token 1.0 # 投机采样的每步期望时间 spec_time_per_step 1.0 k / speed_ratio # 加速比 标准吞吐 / 投机吞吐 # expected_accepted / spec_time_per_step return expected_accepted / spec_time_per_step实测中以LLaMA-7B为target、LLaMA-68M为draft接受率通常在0.7-0.85之间。取α0.8、k5、r≈20大模型约为小模型的20倍计算量理论加速比约为3.2x。实际部署中包含KV Cache管理等开销通常能达到2.0-2.5x的有效加速。四、树形投机与多头草稿标准投机采样一次生成一条线性候选链。树形投机Tree-structured Speculative Decoding扩展了这一思路小模型在每一步生成多个候选分支而非单个token大模型一次性验证整个候选树。这种方式在保持验证成本不变的情况下增加了候选多样性提高了接受率。多头草稿Multi-Head Drafting则是另一种优化使用大模型自身不同层的输出作为草稿——例如使用第N/2层的输出配合小型预测头来预测第N层的输出。这避免了额外的小模型加载但对模型架构有侵入性修改。五、总结投机采样是当前大模型推理加速中投入产出比最高的技术之一(1) 它是无损的——最终输出的token分布与原始大模型严格一致(2) 它需要的额外资源仅是一个小模型的显存和计算——通常为大模型的5-10%(3) 加速比直接取决于大小模型的一致性——在代码生成和翻译等结构化输出任务上接受率高0.85在创意写作等开放任务上接受率较低0.6-0.7。部署时建议(1) 使用同系列的小模型作为draft如LLaMA-7B配LLaMA-68M它们的tokenizer和训练数据一致天然具有更高的接受率(2) k的值不宜过大——k3-5通常是加速比的最优点更大的k会增加小模型错误积累的风险。
)
)
