Osprey:通用预训练让推测解码草稿模型更高效

推测解码的草稿模型为何如此脆弱
大语言模型(LLM)的推理加速一直是业界关注的焦点,而推测解码(Speculative Decoding)作为关键技术之一,却面临一个致命问题:现有的草稿模型(drafter)过于脆弱。
推测解码是一种利用"先猜测后验证"思想加速大语言模型推理的技术。其核心思路是:用一个体积更小、推理速度更快的草稿模型一次性生成多个候选token,然后由目标大模型(verifier)并行验证这些token是否可接受。由于大模型验证多个token的并行计算成本与生成单个token相当,当草稿模型的猜测命中率足够高时,整体推理吞吐量可以获得数倍提升。这一技术由Google DeepMind在2023年正式提出(Leviathan et al.),其数学基础建立在拒绝采样(rejection sampling)之上,保证了最终输出的分布与直接使用目标模型自回归解码完全一致——即无损加速。推测解码的核心指标是平均接受长度(Mean Acceptance Length, MAL),即每次验证中被接受的连续token数量,该值越高说明草稿模型与目标模型的分布越匹配,加速效果越好。
传统方法中,草稿模型针对单一目标模型在狭窄分布上训练,一旦工作负载发生变化,其接受率就会急剧下降。这与现代LLM通过大规模预训练获得广泛泛化能力的发展路径形成了鲜明对比。

问题的根源在于:现有的草稿模型训练方案高度依赖特定目标模型。它们需要消费目标模型的隐藏状态,并在目标模型的logits上进行蒸馏。这里的logits蒸馏属于知识蒸馏(Knowledge Distillation)技术——其核心思想是让小模型(学生)模仿大模型(教师)的输出概率分布,而不仅仅是学习硬标签。在推测解码场景下,logits蒸馏要求草稿模型在每个位置的输出概率分布尽可能接近目标模型的softmax分布,通常使用KL散度作为损失函数。这种方法能让草稿模型捕捉到目标模型在不同token上的置信度细节(即"暗知识"),但其致命缺陷在于训练过程必须持续调用目标模型生成logits,计算成本高昂且产生了对目标模型的强绑定依赖。这意味着每更换一个目标模型,就必须重新完成整个预训练过程。这种target-specific的设计严重限制了草稿模型的可复用性和适应性。
Osprey的核心创新:与目标模型无关的通用预训练
Osprey的核心理念是将预训练作为一种可复用的、与目标模型无关的资产。研究团队提出从现成的预训练小型语言模型出发,通过轻量级适配步骤来支持不同的目标模型,而不是为每个目标重复整个预训练过程。
这一方案需要克服两大技术挑战:
深度压缩:从小模型到浅层草稿模型
小型语言模型通常比延迟受限的草稿模型深得多。Osprey采用剪枝策略构建浅层主干网络,然后通过与目标无关的next-token预训练恢复其语言建模能力。
具体而言,这里采用的是结构化剪枝(Structured Pruning)方法。与非结构化剪枝(移除单个权重参数)不同,结构化剪枝直接移除整个Transformer层、注意力头或前馈网络通道,从而在不依赖稀疏硬件的前提下获得实际的推理加速。在Osprey的语境中,"从小模型到浅层草稿模型"意味着将一个已经预训练好的小型语言模型(例如具有24-32层的模型)通过移除若干Transformer层压缩为仅保留少量层的浅层网络。这一过程的关键挑战在于:深层网络中不同层承担不同的语义功能——浅层通常负责局部语法和词法特征提取,中间层处理语义组合,深层负责高层抽象推理。因此需要精心选择保留哪些层,并通过后续的next-token预训练(即在标准语言建模目标上继续训练)来修复因剪枝导致的性能损失。这一步骤确保草稿模型在压缩后依然具备扎实的语言理解基础。
计算保持与目标适配的平衡
如何在保持预训练计算完整性的同时,让草稿模型学会接收目标模型的隐藏状态并在目标词汇表中生成token?Osprey通过三个关键技术实现这一目标:
-
词汇对齐:不同的大语言模型通常使用不同的分词器(Tokenizer)和词汇表(Vocabulary)。例如Llama系列使用SentencePiece BPE分词,Qwen系列使用基于字节级BPE的自定义分词器,两者的词汇表大小、token粒度和编码方式都存在差异。对于推测解码而言,草稿模型必须与目标模型使用相同的token空间——否则草稿模型生成的token序列在目标模型看来毫无意义,验证过程无法进行。Osprey的词汇对齐机制需要将草稿模型的嵌入层和输出层映射到目标模型的词汇表空间,处理不同分词粒度带来的对齐偏差,并在多个目标模型之间共享尽可能多的预训练知识。这是实现"一个主干支持多个目标模型"的基础性技术障碍。
-
零初始化QKV扩展:在Transformer的自注意力机制中,每个输入向量会通过三个投影矩阵分别映射为Query(Q)、Key(K)和Value(V)向量。当草稿模型需要接收来自目标模型的隐藏状态作为额外输入时,Q、K、V的投影矩阵需要扩展维度以适配更宽的输入。直接随机初始化扩展部分会破坏已有预训练权重的内部表示平衡,导致灾难性的性能下降。零初始化策略的巧妙之处在于:将新增维度对应的权重初始化为零,这样在适配训练开始时,扩展部分的贡献为零,模型行为与扩展前完全一致,然后在后续微调中逐步学习如何利用目标模型提供的额外信息。这一思想与残差学习(Residual Learning)和LoRA中的零初始化策略一脉相承,体现了"先保持不变,再渐进学习"的设计哲学。
-
输出分布蒸馏:从目标模型的输出分布中学习,实现高效的知识迁移。与前述的logits蒸馏原理相同,但Osprey的关键区别在于这一蒸馏步骤仅作为轻量级适配阶段存在,而非贯穿整个预训练过程。由于草稿模型已经通过通用预训练获得了扎实的语言建模能力,适配阶段的蒸馏只需在较少的数据和较短的训练时间内完成分布对齐,大幅降低了计算开销。
实验验证:单个主干跨模型迁移的显著性能提升
实验结果充分验证了Osprey方案的有效性。单个预训练的Osprey主干可以跨多个目标模型迁移,并带来显著的推测解码性能提升:
- Qwen3-8B:平均接受长度(Mean Acceptance Length)提升16.1%
- Llama-3.3-70B-Instruct:提升21.2%
- MiniMax-M2.5(229B):提升22.7%,同时每秒生成token数提高17.5%
特别说个细节,Osprey在域外数据(out-of-domain)和多语言数据上表现出最大的性能增益。这正是传统target-specific方法最薄弱的环节。传统方案通常在目标模型的典型使用场景(如英文问答、代码生成等)上进行蒸馏训练,其训练数据的分布与评估时的分布高度重合。然而,当输入切换到训练时未覆盖的领域(如医学文献、法律文档)或非英语语言时,蒸馏得到的分布匹配关系迅速退化,接受率大幅下降。Osprey的通用预训练策略之所以在这些场景中表现优异,根本原因在于其预训练阶段使用的是广泛的语言建模目标,而非特定模型的输出分布,因此草稿模型保留了对自然语言普遍模式的理解能力。这一发现也暗示了一个更广泛的规律:在LLM系统的各个组件设计中,过度专门化(over-specialization)可能以牺牲鲁棒性为代价。这充分证明了通用预训练策略在泛化能力上的优越性。
对LLM推理加速的启示与展望
Osprey的成功揭示了一个重要趋势:推测解码技术正在从"为每个模型定制草稿模型"转向"一次训练,广泛适用"的范式。这种转变不仅降低了部署成本,更重要的是提高了系统对不同应用场景的适应能力。
从工程实践角度看,Osprey的设计思路为LLM推理优化提供了新的方向。通过将预训练和适配解耦,开发者可以构建更加灵活和可维护的推理系统。这一思路与软件工程中的"关注点分离"原则高度一致——将通用的语言建模能力(预训练阶段)与特定目标模型的对齐需求(适配阶段)清晰解耦,使得每个阶段可以独立优化和迭代。随着大模型规模持续增长,这种通用化的推测解码加速方案将变得越来越重要。可以预见,未来可能出现专门的"草稿模型即服务"(Drafter-as-a-Service)生态,开发者只需选择合适的预训练主干并完成轻量级适配,即可为任意目标模型部署高效的推测解码加速。
项目代码已在GitHub开源(https://github.com/LeanModels/Osprey),为研究者和工程师提供了实践和改进的基础。
核心要点
相关推荐

多个Claude Code如何相互对话?AI Agent团队协作底层机制详解
深入解析多个Claude Code实例之间的通信机制,包括Pub/Sub发布订阅模式、Signal文件信号监听、收信箱文件系统等底层架构,帮助开发者搭建可持续运作的AI Agent协作团队。

ACN智能体上下文网络:AI Agent安全共享上下文的开源方案解析
深入解析Agent Context Network(ACN)开源方案,探讨AI智能体如何通过细粒度权限控制实现安全的上下文共享,而非暴露全部记忆。涵盖技术架构、MCP协议集成及多智能体协作的核心权衡。

Busabase深度解析:为AI Agent而生的应用数据库
深入分析Busabase for DeepSeek Harness这款Agent数据库的产品定位、技术架构与行业前景。探讨Agent数据库如何将数据存储与技能执行融为一体,解决Agent基础设施碎片化痛点。