Osprey:目标无关的预训练让推测解码中的草稿模型更强大

arXiv cs.CL 论文

摘要

Osprey引入了一种针对推测解码中草稿模型的目标无关预训练方法,通过从现成模型中引导并适应最少的特定目标工作来提高效率,在多个大语言模型上实现了显著的接受率提升。

arXiv:2609.09338v1 公告类型:新 摘要:推测解码对于加速大语言模型推理至关重要。然而,加速效果是脆弱的:草稿模型通常针对单个目标模型的窄分布进行训练,在工作负载变化下其接受率会崩溃。这是现代大语言模型开发的一个显著反转,其中目标模型因其通过大规模预训练获得的广泛泛化能力而受到重视。我们认为,自然的补救措施——预训练——之所以难以应用于草稿模型,是因为现有方法都是目标特定的:草稿模型消耗目标的隐藏状态并在目标的logits上进行蒸馏,因此预训练必须针对每个目标重复进行。我们引入Osprey,它从现成的预训练小语言模型中引导草稿模型,将广泛的预训练视为可重用的目标无关资产,并将每个目标的工作减少到轻量级的适应步骤。实现这一点需要克服两个挑战:小语言模型比延迟受限的草稿模型所能承受的要深得多,且在草稿模型学习吸收目标隐藏状态并输出目标词汇表中的令牌时,其预训练计算必须保持完整。Osprey通过修剪到一个浅层骨干网络、通过目标无关的下一个令牌预训练恢复其语言建模能力,并通过词汇表对齐、零初始化的QKV扩展和从目标模型输出分布中蒸馏来适应每个目标。实证表明,单个预训练的Osprey骨干网络可以跨目标转移,并为Qwen3-8B、Llama-3.3-70B-Instruct和229B MiniMax-M2.5(每秒令牌数提高17.5%)分别提高平均接受长度16.1%、21.2%和22.7%,在域外和多语言数据上收益最大。我们的代码可在https://github.com/LeanModels/Osprey获取。
查看原文
查看缓存全文

缓存时间: 2026/09/10 08:09

# Osprey:目标无关的预训练使投机解码的草稿模型更强大  
来源:https://arxiv.org/html/2609.09338  
Fengxiang Bie††thanks:Equal Contribution\. Correspondence to: Tianyi Zhang\.Yuqing Jian, Yifan Yu, Zhongzhu Zhou, Zelei Shao  
隶属机构:Together AI  
隶属机构:澳大利亚悉尼大学  
隶属机构:美国伊利诺伊大学厄巴纳-香槟分校  
已被EMNLP 2026收录\.  
Ben Athiwaratkun  
隶属机构:Together AI  
Shuaiwen Leon Song  
隶属机构:Together AI  
隶属机构:澳大利亚悉尼大学  
Chenfeng Xu  
隶属机构:Together AI  
隶属机构:美国德克萨斯大学奥斯汀分校  
Xiaoxia Wu  
隶属机构:Together AI  
Tianyi Zhang  
隶属机构:Together AI  

###### 摘要  
投机解码对于加速大语言模型推理至关重要。然而,其加速效果往往十分脆弱:草稿模型通常仅针对特定目标模型的狭窄分布进行训练,在工作负载变化时其接受率会急剧下降。这与当今大语言模型的发展趋势形成了鲜明反转——目标模型正是因为通过大规模预训练所获得的广泛泛化能力而备受重视。我们认为,自然的补救方法——预训练——之所以难以应用于草稿模型,是因为现有方法都是目标特定的:草稿模型需要消费目标模型的隐藏状态,并在目标模型的logits上进行蒸馏,因此必须为每个目标重复进行预训练。  

我们提出**Osprey**,它从现成的预训练小型语言模型中引导草稿模型,将广泛的预训练视为可重复使用的目标无关资产,并将针对每个目标的工作量减少为一个轻量级的适配步骤。实现这一目标需要克服两个挑战:小型语言模型的深度远超延迟受限的草稿模型所能承受,且其预训练的计算能力必须在草稿模型学习摄取目标隐藏状态并在目标词表中生成token时保持完整。Osprey通过将模型剪枝为浅层骨架、通过目标无关的下一个token预训练恢复其语言建模能力,并通过词汇表对齐、零初始化的QKV扩展以及从目标模型输出分布中进行蒸馏来适配每个目标,从而同时解决了这两个挑战。  

实验表明,单一的预训练Osprey骨架可以在不同目标之间迁移,并将**Qwen3-8B**的平均接受长度提高**16.1%**,将**Llama-3.3-70B-Instruct**的平均接受长度提高**21.2%**,将**229B MiniMax-M2.5**的平均接受长度提高**22.7%**(每秒token数提高**17.5%**),其中在域外和多语言数据上获得了最大的增益。我们的代码已发布在https://github.com/LeanModels/Osprey。  

## 1 引言  
投机解码(Leviathan et al., 2023 (https://arxiv.org/html/2609.09338#bib.bib36); Chen et al., 2023 (https://arxiv.org/html/2609.09338#bib.bib9); Li et al., 2025 (https://arxiv.org/html/2609.09338#bib.bib39); Cai et al., 2024 (https://arxiv.org/html/2609.09338#bib.bib8))已成为加速大语言模型服务最有效的技术之一。一个轻量级的*草稿模型*提出一个未来的token块,而一个更大的*目标*模型并行验证这些token。由于被接受的token完全按照目标分布提交,而被拒绝的位置则回退到标准的自回归解码,投机解码提供了一种罕见的加速与无损性的结合。这使其在延迟敏感和资源受限的部署环境中尤其具有吸引力,先前的系统通常报告**2-5倍**的解码加速(Xia et al., 2024b (https://arxiv.org/html/2609.09338#bib.bib66))。  

然而,这种成功背后隐藏着一个根本性的脆弱性。实际上,草稿模型通常是在特定的数据分布、模型家族或部署领域上进行训练或微调的(Hong et al., 2025 (https://arxiv.org/html/2609.09338#bib.bib28))。当服务工作负载发生变化时,例如从聊天转向代码,从通用问答转向特定领域推理,或从短回复转向长文本生成,草稿模型的预测可能与目标模型严重失配。结果是接受率急剧下降,从而导致投机解码旨在提供的加速效果大幅丧失。正如我们在图2 (https://arxiv.org/html/2609.09338#S4.F2)中所示,一个在特定数据分布上训练的最先进Eagle3(Li et al., 2025 (https://arxiv.org/html/2609.09338#bib.bib39))草稿模型在域内数据上表现良好,但在域外数据上的表现明显较差(例如,在数学数据上训练的Eagle3在数学数据上平均接受长度为5.06,而在代码数据上平均接受长度仅为1.86)。  

这种脆弱性在当今的大语言模型发展背景下(Kaplan et al., 2020 (https://arxiv.org/html/2609.09338#bib.bib33))显得尤为有趣。我们今天服务的目标模型是在日益广泛的数据混合体上训练的,其价值恰恰在于它们跨领域、任务和交互风格泛化的能力(Radford et al. (2018) (https://arxiv.org/html/2609.09338#bib.bib54))。领域偏移曾是经典机器学习中的核心关注点,但对于前沿大语言模型来说已很少成为主导话题;大语言模型领域更强调涌现能力和广泛迁移(Wei et al. (2022) (https://arxiv.org/html/2609.09338#bib.bib63))。然而,用于加速这些模型的草稿模型仍然是一个更狭隘和脆弱的组件。换句话说,投机解码继承了小模型的延迟优势,但并未继承它所加速的大模型的泛化优势。  

这促使我们重新思考草稿模型的设计。如果投机解码的脆弱性源于分布偏移,那么一个自然的问题是:我们能否以机器学习历史上提高泛化能力的方式(例如通过预训练)来提高草稿模型的泛化能力?事实上,预训练在学习可迁移表示方面的作用已被研究了数十年(Erhan et al., 2010 (https://arxiv.org/html/2609.09338#bib.bib17); Radford et al., 2018 (https://arxiv.org/html/2609.09338#bib.bib54); Devlin et al., 2019 (https://arxiv.org/html/2609.09338#bib.bib14))。因此,人们可能会尝试在更多样化的数据上预训练现有的投机草稿模型,例如EAGLE风格的草稿模型。然而,这些方法本质上是目标特定的:草稿模型消费目标模型的激活值,并在目标的表示空间中进行训练。因此,将该方法适配到新的目标模型需要新的、昂贵的训练流程。更重要的是,这样的训练流程仍然无法直接继承现有预训练语言模型已经学到的广泛泛化能力。  

因此,我们建议采用一个不同的方向。如今,小型语言模型(SLMs)(Belcak et al. (2025) (https://arxiv.org/html/2609.09338#bib.bib4))能力日益增强,经过广泛预训练,并且其规模接近投机解码中使用的草稿模型。与从头训练的目标特定草稿模型不同,这些模型已经通过大规模预训练编码了通用的语言、推理和领域知识。因此问题是:投机解码能否使用预训练的小型语言模型作为通用草稿模型,同时仍能将其与目标模型紧密对齐以保持高接受率?  

然而,基于小型语言模型构建投机模型并非易事。现有的小型语言模型(SLMs)通常有16层或更深(Zhang et al. (2024c) (https://arxiv.org/html/2609.09338#bib.bib74))。草稿模型必须浅层,因为深层的草稿模型会抵消投机带来的延迟优势。此外,骨干网络需要与广泛的目标模型兼容。其嵌入层和语言模型头必须能够适配任何目标分词器,并且必须能够摄取目标隐藏状态而不破坏预训练的能力。  

为了解决这些挑战,我们为投机草稿模型提出了一种四阶段训练方案(图1 (https://arxiv.org/html/2609.09338#S1.F1)),命名为**Osprey**,分为一个通用的、目标无关的阶段和一个针对每个目标的适配阶段。首先,我们*剪枝*一个现成的预训练小型语言模型,仅保留其嵌入层、语言模型头和前几个Transformer层,以构建一个浅层骨干网络。其次,我们通过下一个token预测在通用语料库上对该剪枝模型进行离策略*预训练*,从而恢复其鲁棒且与目标无关的语言建模能力。**关键在于**,此预训练阶段完全是*目标无关的*,生成一个可跨不同目标模型重复使用的草稿模型骨干网络。第三,我们通过替换分词器并将嵌入层和语言模型头与目标词表对齐,将该通用检查点*适配*到特定的目标模型。在此适配步骤中,我们还通过零初始化的列扩展QKV投影,允许草稿模型摄取目标隐藏状态而不干扰预训练的初始化。最后,我们通过将草稿模型置于目标激活值的条件下,并训练其预测目标logits,将目标特定的行为*蒸馏*到草稿模型中。因此,每个部署的草稿模型最终都是针对特定目标的,并继承了可重复使用骨干网络的广泛语言先验知识。  

这种解耦的方法允许单一的目标无关检查点在不同架构(从Qwen3(Team, 2025 (https://arxiv.org/html/2609.09338#bib.bib61))、Llama-3.3(Meta AI, 2024 (https://arxiv.org/html/2609.09338#bib.bib45))到MiniMax-M2.5(MiniMax AI (2026) (https://arxiv.org/html/2609.09338#bib.bib47)))之间高效地重复使用。  

在本工作中,我们的主要贡献是双重的。首先,我们为投机草稿模型提出了一种*目标无关的预训练*方案,将通用的语言建模与目标条件化分离,创建了一个单一的、高度可重复使用的草稿模型骨干网络。其次,我们证明了Osprey比强大的EAGLE-3基线模型提供了更高的接受率和更优的加速效果。实验表明,在Qwen3-8B上,Osprey在所有五个评估领域(聊天、代码、常识、金融和数学)中均提高了相对于EAGLE-3的平均接受长度,并且在更远的前瞻位置获得了最显著的增益。对于更大的目标模型,Osprey在匹配的EAGLE-3(Li et al., 2025 (https://arxiv.org/html/2609.09338#bib.bib39))基线模型基础上,将Llama-3.3-70B-Instruct的平均接受长度和每秒token数分别提高了**21.2%**和**17.9%**,将MiniMax-M2.5的这两项指标分别提高了**22.7%**和**17.5%**。此外,尽管仅使用代码数据进行针对特定目标的训练,Osprey在多语言数据(如42种语言的Global-MMLU和11种语言的MGSM基准测试)上表现出优异的迁移能力。  

参考标题  
图1:Osprey概览。该流程分为四个阶段,左侧为仅执行一次的目标无关阶段,右侧为针对每个新目标模型重复执行的阶段。(1) 剪枝:仅保留现成小型语言模型的嵌入层、语言模型头和前几个Transformer层。(2) 预训练(离策略):在通用语料库上通过下一个token预测训练剪枝模型,以恢复语言建模能力。(3) 适配(针对每个目标):替换分词器并将嵌入层/语言模型头与目标词表对齐,然后通过零初始化的tap扩展QKV投影,以便草稿模型能够在不干扰预训练初始化的情况下摄取目标隐藏状态。(4) 蒸馏(在策略):在目标生成的序列上,通过KL散度从目标激活值和输出分布中进行蒸馏。相同的目标无关检查点在从Qwen3-8B到MiniMax-M2.5-229B的一系列目标模型上重复使用。  

## 2 相关工作  

##### 投机解码与目标条件化的草稿模型。  
投机解码通过使用一个更便宜的草稿模型提议多个token,由目标模型并行验证,从而加速自回归推理(Leviathan et al., 2023 (https://arxiv.org/html/2609.09338#bib.bib36); Chen et al., 2023 (https://arxiv.org/html/2609.09338#bib.bib9))。现有的草稿模型范围从外部提议模型(Xia et al., 2023 (https://arxiv.org/html/2609.09338#bib.bib64); Miao et al., 2024 (https://arxiv.org/html/2609.09338#bib.bib46))到消费目标隐藏状态的目标条件化架构,包括EAGLE(Li et al., 2024b (https://arxiv.org/html/2609.09338#bib.bib38))、EAGLE-2(Li et al., 2024a (https://arxiv.org/html/2609.09338#bib.bib37))和EAGLE-3(Li et al., 2025 (https://arxiv.org/html/2609.09338#bib.bib39)),以及Medusa、Hydra、ReDrafter、Clover、P-EAGLE、Falcon和HASS(Cai et al., 2024 (https://arxiv.org/html/2609.09338#bib.bib8); Ankner et al., 2024 (https://arxiv.org/html/2609.09338#bib.bib3); Bhendawade et al., 2024a (https://arxiv.org/html/2609.09338#bib.bib6); Xiao et al., 2024a (https://arxiv.org/html/2609.09338#bib.bib68); Xiao et al., 2024b (https://arxiv.org/html/2609.09338#bib.bib69); Hui et al., 2026 (https://arxiv.org/html/2609.09338#bib.bib30); Gao et al., 2024 (https://arxiv.org/html/2609.09338#bib.bib19); Zhang et al., 2024b (https://arxiv.org/html/2609.09338#bib.bib73))。另一条并行的研究路线则避免了单独的草稿模型,而是通过早期退出、自适应层跳过、适配器或多流注意力来复用目标模型本身(Elhoushi et al., 2024 (https://arxiv.org/html/2609.09338#bib.bib16); Zhang et al., 2024a (https://arxiv.org/html/2609.09338#bib.bib72); Xia et al., 2024a (https://arxiv.org/html/2609.09338#bib.bib65); Liu et al., 2024a (https://arxiv.org/html/2609.09338#bib.bib40); Bhendawade et al., 2024b (https://arxiv.org/html/2609.09338#bib.bib7))。这些方法降低了推理成本,但要么从头训练草稿模型,要么将草稿行为紧密耦合到特定目标,要么需要目标侧的架构/训练变更。Osprey则保持目标模型不变,构建一个独立的浅层草稿模型,其主体由一个目标无关的预训练下一个token预测器初始化。我们采用EAGLE-3的TTT范式以及三层目标隐藏状态的tap,但通过零初始化的QKV扩展注入目标特征,从而在适配阶段保留预训练的草稿模型功能。  

##### 鲁棒且可重复使用的草稿模型训练。  
最近的研究表明,投机解码的接受率可能因领域和用户群体而急剧变化(Xia et al., 2024b (https://arxiv.org/html/2609.09338#bib.bib66); Sandler et al., 2025 (https://arxiv.org/html/2609.09338#bib.bib56); Hong et al., 2025 (https://arxiv.org/html/2609.09338#bib.bib28)),这促使人们开发能够超越目标特定蒸馏分布的草稿模型。先前的工作通过预训练的SLM初始化、状态空间草稿模型、在线更新、跨词汇表验证或通用验证规则来提高鲁棒性(Berdoz et al., 2025 (https://arxiv.org/html/2609.09338#bib.bib5); Choi et al., 2025 (https://arxiv.org/html/2609.09338#bib.bib12); Ramakrishnan et al., 2025 (https://arxiv.org/html/2609.09338#bib.bib55); Timor et al., 2025 (https://arxiv.org/html/2609.09338#bib.bib62); Liu et al., 2024b (https://arxiv.org/html/2609.09338#bib.bib41))。Osprey的不同之处在于,它在通用网络文本(Penedo et al., 2024 (https://arxiv.org/html/2609.09338#bib.bib53))上显式预训练一个可重复使用的浅层语言模型,然后针对每个目标进行轻量级适配,而不是依赖在线修正或从现有草稿模型中选择。我们的贡献不是一个新的蒸馏目标,而是一个可重复使用的初始化和适配路径,在相同的投机解码接口下产生更强的领域鲁棒性和跨目标迁移能力。关于相关子领域——深度剪枝、词汇表迁移、继续预训练、白盒蒸馏等——更完整的讨论...

相似文章

Draft-OPD:面向推测式草稿模型的在线策略蒸馏

Hugging Face Daily Papers

Draft-OPD 引入在线策略蒸馏,结合目标辅助展开和错误重放,克服了训练用于推测解码的草稿模型时存在的离线到推理不匹配问题,实现了超过5倍的无损加速,相较于EAGLE-3和DFlash分别提升了23%和13%。

跨语言的推测解码

arXiv cs.CL

本文比较了三种策略以提高非英语语言的推测解码效率,发现任务特定蒸馏能提高接受率但泛化性差,而n-gram草稿模型尽管接受率较低,却能提供持续的加速。