Osprey: General-Purpose Pretraining Makes Speculative Decoding Draft Models More Efficient

Why Are Speculative Decoding Draft Models So Fragile?

Accelerating inference for large language models (LLMs) has been a major focus across the industry, and speculative decoding—one of the key techniques—faces a critical problem: existing draft models (drafters) are far too fragile.

Speculative decoding is a technique that speeds up LLM inference using a "guess first, verify later" approach. The core idea is to use a smaller, faster draft model to generate multiple candidate tokens at once, then have the target large model (verifier) validate these tokens in parallel. Since the parallel computation cost for the large model to verify multiple tokens is comparable to generating a single token, overall inference throughput can improve several-fold when the draft model's guess accuracy is high enough. This technique was formally introduced by Google DeepMind in 2023 (Leviathan et al.), with its mathematical foundation built on rejection sampling, guaranteeing that the final output distribution is identical to standard autoregressive decoding from the target model—in other words, lossless acceleration. The key metric for speculative decoding is Mean Acceptance Length (MAL), the number of consecutive tokens accepted per verification step. A higher MAL indicates better distribution matching between the draft and target models, resulting in greater speedup.

In traditional approaches, draft models are trained on narrow distributions for a single target model, and their acceptance rates drop sharply once the workload changes. This stands in stark contrast to the development trajectory of modern LLMs, which achieve broad generalization through large-scale pretraining.

The root cause is that existing draft model training pipelines are highly dependent on a specific target model. They need to consume the target model's hidden states and perform distillation on the target model's logits. This logits distillation is a form of Knowledge Distillation—the core idea being to have the small model (student) mimic the large model's (teacher's) output probability distribution, rather than just learning hard labels. In the speculative decoding context, logits distillation requires the draft model's output probability distribution at each position to match the target model's softmax distribution as closely as possible, typically using KL divergence as the loss function. This approach lets the draft model capture the target model's confidence details across different tokens (i.e., "dark knowledge"), but its fatal flaw is that the training process must continuously invoke the target model to generate logits—computationally expensive and creating a tight binding dependency on the target model. This means that switching to a different target model requires repeating the entire pretraining process from scratch. This target-specific design severely limits draft model reusability and adaptability.

Osprey's Core Innovation: Target-Agnostic General-Purpose Pretraining

Osprey's central idea is to treat pretraining as a reusable, target-agnostic asset. Instead of repeating the entire pretraining process for each target, the research team proposes starting from off-the-shelf pretrained small language models and supporting different target models through lightweight adaptation steps.

This approach needs to overcome two major technical challenges:

Deep Compression: From Small Models to Shallow Draft Models

Small language models are typically much deeper than what latency-constrained draft models require. Osprey uses a pruning strategy to build a shallow backbone network, then restores its language modeling capability through target-agnostic next-token pretraining.

Specifically, this employs Structured Pruning. Unlike unstructured pruning (which removes individual weight parameters), structured pruning removes entire Transformer layers, attention heads, or feed-forward network channels, achieving real inference speedup without relying on sparse hardware. In Osprey's context, "from small models to shallow draft models" means compressing an already-pretrained small language model (e.g., one with 24-32 layers) into a shallow network retaining only a few layers by removing several Transformer layers. The key challenge here is that different layers in deep networks serve different semantic functions—shallow layers typically handle local syntactic and lexical feature extraction, middle layers process semantic composition, and deep layers handle high-level abstract reasoning. Therefore, careful selection of which layers to retain is essential, followed by next-token pretraining (i.e., continued training on the standard language modeling objective) to repair performance loss caused by pruning. This step ensures the draft model retains a solid language understanding foundation after compression.

Balancing Computation Preservation with Target Adaptation

How can you maintain the integrity of pretrained computation while enabling the draft model to receive the target model's hidden states and generate tokens in the target vocabulary? Osprey achieves this through three key techniques:

• Vocabulary Alignment: Different LLMs typically use different tokenizers and vocabularies. For example, the Llama family uses SentencePiece BPE tokenization, while the Qwen family uses a custom byte-level BPE tokenizer—they differ in vocabulary size, token granularity, and encoding methods. For speculative decoding, the draft model must operate in the same token space as the target model—otherwise the token sequences generated by the draft model are meaningless to the target model, and the verification process cannot proceed. Osprey's vocabulary alignment mechanism maps the draft model's embedding and output layers to the target model's vocabulary space, handles alignment discrepancies from different tokenization granularities, and shares as much pretrained knowledge as possible across multiple target models. This is the foundational technical barrier to achieving "one backbone supporting multiple target models."

• Zero-Initialized QKV Extension: In Transformer self-attention, each input vector is projected through three projection matrices into Query (Q), Key (K), and Value (V) vectors. When the draft model needs to receive hidden states from the target model as additional input, the Q, K, V projection matrices need dimension expansion to accommodate the wider input. Randomly initializing the expanded portion would disrupt the internal representation balance of existing pretrained weights, causing catastrophic performance degradation. The elegance of zero initialization lies in initializing the weights corresponding to the new dimensions to zero, so at the start of adaptation training, the contribution from the expanded portion is zero—the model behaves identically to before expansion—and then gradually learns how to leverage the additional information from the target model during subsequent fine-tuning. This idea is consistent with the zero-initialization strategies in Residual Learning and LoRA, embodying the design philosophy of "preserve first, then learn incrementally."

• Output Distribution Distillation: Learning from the target model's output distribution for efficient knowledge transfer. The principle is the same as the logits distillation described earlier, but Osprey's key distinction is that this distillation step exists only as a lightweight adaptation phase, rather than spanning the entire pretraining process. Since the draft model has already acquired solid language modeling capabilities through general-purpose pretraining, the distillation during the adaptation phase only needs a small amount of data and short training time to achieve distribution alignment, dramatically reducing computational overhead.

Experimental Validation: Significant Performance Gains from Single-Backbone Cross-Model Transfer

Experimental results thoroughly validate Osprey's effectiveness. A single pretrained Osprey backbone can transfer across multiple target models, delivering significant speculative decoding performance improvements:

• Qwen3-8B: Mean Acceptance Length improved by 16.1%
• Llama-3.3-70B-Instruct: Improved by 21.2%
• MiniMax-M2.5 (229B): Improved by 22.7%, with tokens per second also increasing by 17.5%

One detail worth highlighting: Osprey shows the largest performance gains on out-of-domain and multilingual data—precisely the weakest points of traditional target-specific approaches. Conventional methods typically perform distillation training on the target model's typical use cases (e.g., English Q&A, code generation), where the training data distribution closely overlaps with evaluation distributions. However, when inputs shift to domains not covered during training (e.g., medical literature, legal documents) or non-English languages, the distilled distribution matching rapidly degrades and acceptance rates plummet. The fundamental reason Osprey's general-purpose pretraining strategy excels in these scenarios is that its pretraining phase uses a broad language modeling objective rather than a specific model's output distribution, so the draft model retains its understanding of universal natural language patterns. This finding also hints at a broader principle: in the design of various components within LLM systems, over-specialization may come at the cost of robustness. This conclusively demonstrates the superiority of general-purpose pretraining strategies in generalization capability.

Implications and Outlook for LLM Inference Acceleration

Osprey's success reveals an important trend: speculative decoding technology is shifting from "customizing a draft model for each model" to a "train once, apply broadly" paradigm. This transformation not only reduces deployment costs but, more importantly, improves system adaptability across different application scenarios.

From an engineering practice perspective, Osprey's design philosophy opens new directions for LLM inference optimization. By decoupling pretraining from adaptation, developers can build more flexible and maintainable inference systems. This approach is highly aligned with the "separation of concerns" principle in software engineering—cleanly decoupling general language modeling capabilities (pretraining phase) from alignment requirements for specific target models (adaptation phase), allowing each phase to be independently optimized and iterated. As model scales continue to grow, this generalized approach to speculative decoding acceleration will become increasingly important. It's foreseeable that a dedicated "Drafter-as-a-Service" ecosystem may emerge in the future, where developers simply select an appropriate pretrained backbone and complete lightweight adaptation to deploy efficient speculative decoding acceleration for any target model.

The project code is open-sourced on GitHub (https://github.com/LeanModels/Osprey), providing a foundation for researchers and engineers to experiment and improve upon.

Key Takeaways