Osprey boosts language model drafting speed and accuracy across tasks
Osprey: Target-agnostic Pre-training Makes Stronger Drafters in Speculative Decoding
Computation and Language
Summary
Generating text quickly with AI language models is often slowed down by a part called the drafter, which is usually trained for just one target model. This means when the task changes, performance drops. The authors introduce Osprey, a method that uses a small pretrained language model as a general drafter that can easily adapt to different target models. This approach makes drafting faster and more reliable across various models and types of text, especially for ones that are new or in different languages. Osprey improves both speed and the quality of text acceptance in experiments.
What this means in practice
- •For machine learning engineers: Accelerate inference of large language models by deploying Osprey to improve drafter generalization and speed across multiple target models.
- •For natural language processing teams: Improve robustness and efficiency in multilingual or out-of-domain text generation by using Osprey for target-agnostic drafter adaptation.
Authors
Fengxiang Bie, Yuqing Jian, Yifan Yu, Zhongzhu Zhou, Zelei Shao, Ben Athiwaratkun, Shuaiwen Leon Song, Chenfeng Xu, Xiaoxia Wu, Tianyi Zhang
Abstract
Speculative decoding is critical for accelerating LLM inference. However, the speedup is fragile: drafters are typically trained against a narrow distribution for a single target model, and their acceptance rate collapses under workload shifts. This is a striking inversion of modern LLM development, where target models are valued precisely for the broad generalization they acquire through large-scale pretraining. We argue that the natural remedy, pretraining, has been hard to apply to drafters because existing recipes are target-specific: the drafter consumes the target's hidden states and is distilled on the target's logits, so pretraining must be repeated for each target. We introduce Osprey, which instead bootstraps drafters from off-the-shelf pretrained small language models, treating broad pretraining as a reusable, target-agnostic asset and reducing per-target work to a lightweight adaptation step. Realizing this requires overcoming two challenges: small LMs are far deeper than a latency-bound drafter can afford, and their pretrained computation must remain intact while the drafter learns to ingest target hidden states and emit tokens in the target's vocabulary. Osprey addresses both by pruning to a shallow backbone, restoring its language-modeling capability with target-agnostic next-token pretraining, and adapting it to each target through vocabulary alignment, zero-initialized QKV expansion, and distillation from the target model's output distribution. Empirically, a single pretrained Osprey backbone transfers across targets and improves mean acceptance length by 16.1% for Qwen3-8B, 21.2% for Llama-3.3-70B-Instruct, and 22.7% for the 229B MiniMax-M2.5 (with 17.5% higher tokens per second), with the largest gains on out-of-domain and multilingual data. Our code is available at https://github.com/LeanModels/Osprey.