PaperlayerView exact revision on arXiv

Source-linked AI summary

Osprey: Target-agnostic Pre-training Makes Stronger Drafters in Speculative Decoding

Fengxiang Bie, Yuqing Jian, Yifan Yu, Zhongzhu Zhou, Zelei Shao, Ben Athiwaratkun, Shuaiwen Leon Song, Chenfeng Xu, Xiaoxia Wu, Tianyi Zhang

arXiv:2609.09338v1cs.CL

TL;DR

Speculative decoding drafters are often narrow and target-specific, so their acceptance can degrade under domain or workload shifts. Osprey reuses broadly pretrained small language models through shallow pruning, target-agnostic restoration, and lightweight target adaptation. It improves mean acceptance length across three targets, with the strongest advantages on out-of-domain and multilingual workloads, but its compute benefits depend on sustained serving or reuse.

  • Problem

    Target-specific drafters have limited robustness under distribution shifts, motivating reusable drafters that generalize beyond their distillation data.

  • Method

    Osprey prunes a pretrained small language model, restores it with target-agnostic next-token pretraining, and adapts it using target vocabulary, hidden states, and output distributions.

  • Results

    Osprey achieves 16.1%, 21.2%, and 22.7% higher mean acceptance length than EAGLE-3 on Qwen3-8B, Llama-3.3-70B-Instruct, and MiniMax-M2.5, respectively.

  • Takeaways & Limitations

    A single pretrained Osprey backbone transfers across targets, with strongest advantages on out-of-domain and multilingual workloads.

  • Takeaways & Limitations

    The compute advantages require upfront pretraining cost to be amortized through continuous high-volume serving or multi-target reuse, and evaluations cover limited targets and workloads.

Abstract

from arXiv · show

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.

1 Introduction

Osprey addresses the fragility of target-specific speculative drafters by reusing broadly pretrained small language models as adaptable, shallow drafters. Its target-agnostic backbone and target-specific adaptation improve acceptance across targets and workloads, including multilingual data.

  • Motivation: Target-specific drafters can lose substantial acceptance under workload shifts, with Eagle3 averaging 5.06 accepted tokens on math but only 1.86 on code.This brittleness contrasts with frontier target models’ emphasis on broad transfer across domains and tasks.
  • Motivation: Osprey asks whether pretrained small language models can serve as general-purpose drafters while remaining aligned with target models.The approach treats pretrained linguistic, reasoning, and domain knowledge as reusable drafter initialization.
  • Osprey: Osprey prunes a pretrained small language model into a shallow backbone, restores language modeling through target-agnostic next-token pretraining, and adapts it to each target.The shallow design preserves latency benefits while the reusable pretraining separates general language modeling from target conditioning.
  • Results: Osprey also transfers to multilingual benchmarks despite using only code data for target-specific training.The reported multilingual evaluations include 42-language Global-MMLU and 11-language MGSM.

2 Related Work

Prior work studies robustness through pretrained initialization, alternative drafter architectures, online updates, vocabulary-aware verification, and universal verification rules. Osprey instead pretrains a reusable shallow language model on generic web text and lightly adapts it to each target.

  • Robust and reusable draft-model training: Prior studies show that speculative acceptance varies sharply across domains and user populations, motivating drafters that generalize beyond target-specific distillation data.Existing approaches include pretrained SLM initialization, state-space drafters, online updates, cross-vocabulary verification, and universal verification rules.
  • Osprey’s distinction: Osprey differs by explicitly pretraining a reusable shallow language model on generic web text before adapting it to each target.This pathway contrasts with online correction and selecting among existing drafters.
  • Osprey’s distinction: The contribution is a reusable initialization and adaptation pathway rather than a new distillation objective.The stated goal is stronger domain robustness and cross-target transfer under the same speculative decoding interface.

3 Method

Osprey builds a reusable shallow drafter in a target-agnostic phase, then converts and distills it for each target model. The pipeline preserves pretrained language-modeling capability while adding target vocabulary and hidden-state conditioning.

  • Target-agnostic pretraining: Osprey first prunes a pretrained small language model to a shallow backbone, retaining its embedding, LM head, and first few transformer layers.The first two stages produce a reusable next-token-prediction backbone independent of any target model.
  • Target-agnostic pretraining: The pruned backbone is continued-pretrained with next-token prediction on generic text to restore capability lost through aggressive pruning.The recipe uses a 100B-token FineWeb sample for five epochs, totaling approximately 500B training tokens.
  • Target-specific conversion: For each target, Osprey converts the backbone by aligning its tokenizer and vocabulary, projecting multi-layer target hidden states, and expanding attention projections.The source model and target may belong to different families and use different tokenizers.
  • Target-specific conversion: Zero-initialized QKV expansion initially leaves pretrained attention behavior unchanged, while training later enables the drafter to use target hidden states.The latest available projected target representation is concatenated with the drafter’s token-side representation before the expanded attention projections.
  • Target-specific conversion: Vocabulary resizing preserves shared pretrained token rows and initializes target-only rows from the mean representations of their source-token decompositions.The same initialization procedure is applied to both the input embedding and language-modeling head.
  • Target-specific adaptation: Osprey finally distills the converted drafter from the frozen target using target logits and hidden states during autoregressive training-time unrolling.The discounted distillation loss emphasizes earlier drafting steps.

4 Experiments

Experiments evaluate Osprey across domains and target models, testing transfer, robustness, ablations, and computational justification. Osprey consistently improves accepted length over EAGLE-3, including out-of-domain and multilingual evaluations.

  • Experimental setup: Experiments compare Osprey with EAGLE-3 across Qwen3-8B, Llama-3.3-70B-Instruct, and MiniMax-M2.5.The evaluation covers target models from different families and scales, with domain-shift and transfer tests.
  • Cross-domain evaluation: 16.1% higher mean AL: Osprey reaches 3.686 versus EAGLE-3’s 3.176 across 25 Qwen3-8B adaptation–evaluation domain pairs.Osprey achieves higher AL for every domain pair.
  • Cross-domain evaluation: 17.3% higher OOD AL: Osprey reaches 3.493 versus EAGLE-3’s 2.977, while in-domain AL is 12.4% higher.The OOD-to-ID AL ratio is 0.782 for Osprey versus 0.750 for EAGLE-3.
  • Transfer to larger target models: 21.2% higher mean AL and 17.9% higher throughput: Osprey outperforms the reproduced EAGLE-3 baseline on Llama-3.3-70B-Instruct.Osprey achieves the highest AL in all five evaluation domains.
  • Transfer to larger target models: 22.7% higher mean AL and 17.5% higher throughput: Osprey reaches 3.232 AL and 219.8 TPS versus EAGLE-3’s 2.634 AL and 187.1 TPS on MiniMax-M2.5.The largest relative AL advantages occur on LIVECODEBENCH, HUMANEVAL, and MATH-500.
  • Transfer to larger target models: 23.8% higher AL on GLOBAL-MMLU and 12.4% higher AL on MGSM show transfer to multilingual evaluations after coding adaptation.When answering in the question’s language, Osprey’s AL is 15.4% higher than EAGLE-3’s.
  • Robustness and ablations: Pretrained initialization loses 0.11 AL when the learning rate drops tenfold, compared with 1.07 AL for the from-scratch baseline.The reported values are 3.686 → 3.578 for Osprey and 3.176 → 2.110 for the baseline.
  • Robustness and ablations: Adding Stage 2 pretraining raises overall AL from 3.605 to 3.790, improving both in-domain and OOD performance.The improvement is +0.185 overall, +0.211 in-domain, and +0.176 OOD across 16 adaptation–evaluation pairs.

5 Conclusion

Osprey makes speculative drafters more reusable by combining a pretrained small-language-model backbone with target-specific adaptation. It improves mean acceptance length across three targets and increases MiniMax-M2.5 throughput.

  • Osprey combines a pruned pretrained small language model, target-agnostic pretraining, vocabulary alignment, zero-initialized QKV expansion, and target distillation.
  • 16.1%, 21.2%, and 22.7% higher mean AL are achieved on Qwen3-8B, Llama-3.3-70B-Instruct, and MiniMax-M2.5, respectively.
  • 17.5% higher throughput is reported on MiniMax-M2.5.
  • The strongest advantages appear on out-of-distribution and multilingual workloads, supporting reuse of a broadly pretrained language prior for speculative drafters.

Limitations

Osprey’s evaluation and efficiency claims have bounded scope: results depend on the tested targets, serving configuration, workload, adaptation coverage, and amortization of pretraining compute.

  • Results were evaluated on three targets under a fixed SGLang configuration and may vary with other targets, hardware, batch sizes, or workloads.
  • Stage 2’s upfront compute is repaid only under continuous high-volume serving or multi-target reuse, not low-traffic or one-off settings.
  • The cross-domain study covers five domains, while MiniMax-M2.5 adaptation uses code-only target-generated data and does not cover all languages, tasks, or deployment distributions.
  • Most results come from single runs rather than multiple random seeds, so small differences should be interpreted cautiously.
  • Osprey improves inference efficiency but does not itself mitigate model bias, factuality, or safety issues.

A Detailed Related Work

The related work situates Osprey at the intersection of speculative decoding, drafter reuse, distillation, depth pruning, vocabulary transfer, and continued pretraining. Its distinguishing contribution is target-agnostic pretrained initialization combined with preserved computation during target adaptation.

  • Speculative decoding and target-conditioned drafters: Speculative decoding uses a cheaper drafter to propose future tokens that the target verifies in parallel, while Osprey leaves the target unchanged.
  • Self-speculative and target-internal drafters: Unlike target-internal methods such as LayerSkip and SWIFT, Osprey uses an auxiliary drafter while preserving the target model.
  • Domain generalization and cross-target drafter reuse: Prior work studies domain variation, cross-target reuse, online adaptation, and cross-vocabulary verification, whereas Osprey acquires its broad prior through explicit generic-web-text pretraining.
  • Domain generalization and cross-target drafter reuse: Zero-initialized QKV expansion preserves the pretrained drafter function at adaptation start and supports reuse of one checkpoint across targets.
  • Knowledge distillation for autoregressive drafters: Osprey inherits EAGLE-3’s TTT loss but contributes a pretrained shallow-language-model initialization rather than a new distillation objective.
  • Compact pre-training: Osprey’s stages combine established depth pruning, vocabulary transfer, and continued-pretraining techniques into one speculative-drafter pipeline.

B Implementation Notes

The implementation uses SpecForge on SGLang, preserving pretrained components while supporting target-conditioned training and speculative-serving evaluation. The ablation matrix controls target, budgets, domains, and serving configuration across architectures.

  • SpecForge implements Osprey on SGLang with online target hidden-state generation, FSDP draft training, and optional target tensor parallelism.
  • Conversion scripts preserve pretrained layers, expand Q/K/V projections for the EAGLE-3 dual-stream interface, and optionally transfer tokenizer rows into the target vocabulary.
  • The evaluation path launches the target with EAGLE-3 speculative decoding and records serving results.
  • The ablation architectures share the same target, training budget, training and evaluation domains, and speculative-serving configuration.

C Pretraining Implementation Details

The pretraining pipeline uses large-scale FineWeb training for aggressively pruned small-language-model backbones, with loss and perplexity largely saturating after about 10,000 steps.

  • Training setup: 503B tokens are processed over 60,000 optimizer steps using 4096-token sequences, effective global batching, and approximately 80–110 hours on 32 H100 GPUs.Training uses a 1% warmup followed by cosine decay, with a peak learning rate of 1×10−4.
  • Data pipeline: The corpus is FineWeb sample-100BT, tokenized once with each draft’s tokenizer and served through memory-mapped shards with random fixed-length windows.This avoids re-tokenization during training and supports reuse across draft families.
  • Layer-pruning configurations: Two shallow drafts retain 4 of 16 LayerSkip-Llama blocks or 2 of 36 Qwen3-4B blocks, with approximately 506M and 591M parameters, respectively.Both inherit embeddings, final RMSNorm, and the language-modeling head.
  • Pretraining budget: The 10,000-step checkpoint nearly matches the 55,000-step checkpoint, with only 0.026 additional overall acceptance length from the longer run.Out-of-domain/in-domain retention remains approximately 0.80.
  • Pretraining curves: Held-out perplexity falls from approximately 520 to a plateau near 37 after step 10,000, while cross-entropy declines from about 19 to 3.7 nats.The curves indicate that most pretraining benefit arrives early.

D Adaptation Implementation Details

Adaptation uses target-specific hidden states and logits under target-dependent hardware, while matched Qwen3-8B comparisons keep data, steps, and hardware aligned with the scratch baseline.

  • Qwen3-8B: Qwen3-8B adaptation and evaluation use SGLang with bf16 targets on one or two H100 GPUs, depending on whether the stage is evaluation or training.Evaluation uses one H100 with tensor parallelism 1; adaptation uses two H100 GPUs.
  • MiniMax-M2.5: MiniMax-M2.5 adaptation and evaluation use four B200 GPUs with FP8 serving, tensor parallelism 4, and expert parallelism 4.The 229B-parameter MoE target uses target-generated code data for adaptation.
  • Checkpoint selection: The 10,000-step pretraining checkpoint achieves nearly the same downstream performance as the 55,000-step checkpoint under identical Qwen3-8B adaptation.The additional 45,000 steps improve overall acceptance length by only 0.026.
  • Matched comparisons: For each Qwen3-8B domain, Osprey and the matched scratch EAGLE3 baseline use identical data, steps, hardware, and near-matched drafter bodies.The reported body sizes are 265M and 269M parameters.
  • Serving protocol: The same draft checkpoint is reused during evaluation without further inference-time fine-tuning, using batch size 1, five steps, six draft tokens, and top-k 1.The MiniMax configuration uses tensor parallelism 4 and expert parallelism 4.

D.1 Serving-Configuration Sensitivity

Serving sensitivity experiments vary one axis at a time and show that Osprey’s throughput and accepted-length advantages persist across batch size, draft-chain length, and sampling temperature.

  • Batch size: 24% throughput gain at batch size 1 and 13% at batch size 16 show that Osprey remains faster across tested batch sizes.Accepted length is stable across batch sizes.
  • Draft-chain length: Increasing the chain from 2 to 8 draft tokens widens Osprey’s accepted-length gain from 7% to 53%, while EAGLE-3 saturates near 6 tokens.SGLang couples chain steps and draft-token count when top-k is 1.
  • Sampling: 41% gain at temperature 0 and 35% at temperature 1.0 show that Osprey’s throughput advantage persists under sampling.The experiments sweep temperature while retaining the serving protocol’s other settings.

G.5 Results

The DFlash ablations show that diffusion warm starts provide only small gains under the original target interface, whereas preserving the pretrained attention path improves out-of-domain acceptance substantially.

  • Aligned interface: The EAGLE-style DFlash variant preserves the pretrained attention path and adds target conditioning through zero-initialized projections.This design keeps the initial attention computation identical to a pretrained-only forward.
  • Original DFlash interface: At most 0.07 accepted tokens separate IDLM warm-start and random-body initialization under the original DFlash-style interface.The four-domain OOD mean rises only from 1.64 to 1.68.
  • Interface mismatch: Key/value distribution drift may limit warm-start reuse because fine-tuning supplies target-derived streams instead of the drafter-produced representations seen during diffusion pretraining.The authors frame this as a hypothesis explaining the small gain.
Loading 2609.09338v1…