投机采样(Speculative Decoding) 是一种加速大型语言模型(LLM)推理的技术,其核心思想是通过预生成候选token序列并异步校验,从而减少主模型的计算量,同时保持生成结果的准确性。
核心思想是通过并行预测和验证候选Token序列,减少自回归解码的延迟,同时保持生成质量无损。
核心思想
- 双模型协作:
使用一大一小两个模型:- 草案模型(Draft Model):轻量级模型(如TinyLLaMA),快速生成候选Token序列(通常3-5个Token)。
- 目标模型(Target Model):原始大模型(如LLaMA-7B),并行验证草案模型的输出,纠正错误并保留正确部分。
- 并行验证机制:目标模型通过单次前向传播计算所有候选Token的概率分布,对比草案模型的预测概率,按接受准则(如概率一致性)决定是否采纳。
-
分阶段生成:
- 第一阶段(投机阶段):使用一个轻量级模型(如小型LLM、缓存历史结果或规则模型)快速生成若干候选token序列。
- 第二阶段(校正阶段):主LLM(如GPT-3等大模型)仅对候选序列中不确定或低置信度的部分进行校验和修正,而非逐token重新计算。
-
核心假设:
- 轻量模型生成的候选token中,大部分是正确的(尤其在简单或重复性任务中),主模型只需修正少数错误即可保证输出质量。
工作流程
- 草案生成:草案模型以自回归方式生成候选序列(如“sunny, and, warm”)。
- 并行验证:目标模型同时计算这些候选的概率,若草案模型的概率≤目标模型的概率,则接受该Token;否则拒绝并回退至首个错误位置,由目标模型重新生成。
- 循环迭代:更新上下文后重复上述步骤,直至生成完整文本。
-
预生成候选:
- 轻量模型根据当前上下文生成多个候选token(例如预测接下来的5个token)。
- 例如:主模型需要生成
["The", "sky", "is", "blue"]
,轻量模型可能先生成["The", "sky", "is", "red"]
。
-
异步校验:
- 主模型将候选序列作为输入,计算每个token的概率分布,并与轻量模型的结果对比。
- 若轻量模型生成的token在主模型中概率较高(置信度达标),则直接采纳;否则替换为正确token。
-
动态调整:
- 根据任务复杂度调整预生成深度(如简单问题多生成候选,复杂问题少生成)。
技术优势
- 无损加速:通过概率校准确保输出与目标模型自回归生成的结果一致,实验显示加速比达2-3倍(如T5-XXL模型)。
- 无需微调:直接复用现有模型,无需修改架构或权重。
- 硬件效率:减少内存带宽压力,适合边缘设备部署。
-
加速效果:
- 减少主模型的逐token计算次数。例如,若预生成5个候选token,主模型只需校验1次即可完成5步生成。
- 实验表明,端到端推理速度可提升2-5倍(取决于任务类型和模型规模)。
-
无损生成:
- 通过校正机制保证最终输出与主模型直接生成的结果完全一致(Lossless),避免精度损失。
-
通用性:
- 兼容不同架构(如Decoder-only、Encoder-Decoder)和生成任务(如文本、代码)。
与其他方法对比
方法 | 原理 | 加速比 | 准确性损失 | 适用场景 |
---|---|---|---|---|
投机采样 | 预生成+校正 | 高 | 无 | 复杂生成任务 |
模型蒸馏 | 用小型模型替代大模型 | 中 | 有 | 简单任务 |
并行采样 | 多头生成候选token(如Medusa) | 中 | 有(需微调) | 短序列任务 |
静态批处理 | 合并多个请求并行处理 | 低 | 无 | 固定长度生成 |
局限性
- 依赖轻量模型质量:
- 若预生成模型性能差,可能导致主模型校验次数增加,反而降低效率。
- 复杂任务效果受限:
- 对需要全局推理的任务(如数学证明、逻辑推理),轻量模型可能难以生成有效候选。
- 系统复杂度:
- 需实现异步调度、缓存管理等机制,工程实现难度较高。
典型应用
- 实时交互:如聊天机器人、搜索引擎即时问答。
- 长文本生成:如新闻摘要、故事创作。
- 资源受限场景:边缘设备或低成本云服务部署。
总结
投机采样通过“预测-校正”范式,在保证生成质量的前提下显著降低主模型的计算负载,是当前LLM推理加速领域的重要突破。其核心价值在于挖掘推理过程中的冗余性,并通过系统级优化(如缓存复用、动态批处理)进一步释放性能潜力。
改进与变体
- Medusa模型:目标模型集成多个预测头(Medusa Head),直接生成多分支候选序列,替代外部草案模型,进一步提升效率。
- 动态调整:根据目标模型置信度分阶段生成候选,优化生成长度与修正成本的平衡。