Research map / Large language models / Reasoning and chain of thought
OmniDraft: A cross-vocabulary, online adaptive drafter for on-device speculative decoding
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Proposes OmniDraft, a single on-device Llama-68M drafter that works with multiple cross-vocabulary targets via an online n-gram cache, hybrid online distillation, and adaptive drafting.
Problem
Speculative decoding requires draft and target models to share a vocabulary and be aligned offline, which breaks when target models change or use different tokenizers; edge deployment needs one universal, efficient drafter that adapts online.
Approach
An online n-gram cache maps multi-token drafter sequences to single target tokens (with adjusted probabilities) to enable cross-vocabulary verification, combined with a hybrid loss (reverse KL on direct-mapped tokens, NLL on n-gram tokens) for online distillation and an acceptance-prediction head for adaptive draft length.
Key results
- x
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Add the n-gram distillation loss to direct-mapping distillation (ablation) | OmniDraft cross-vocabulary distillation · GSM8K, Llama3-8B target | Speedup: 1.58x -> 1.70x | Helped |
| Add the n-gram cache postprocessing (ablation) | SpD_DM · GSM8K subset, k=3 | Speedup: 1.04x -> 1.16x | Helped |
| Add hybrid n-gram distillation after direct-mapping distillation (ablation) | OmniDraft · GSM8K subset, k=3 | Speedup: 1.59x -> 1.66x | Helped |
| Train with only the n-gram loss instead of combined direct-mapping and n-gram losses (ablation) | OmniDraft · GSM8K subset, temperature 0.01 | Speedup: 1.57x vs 0.86x | Hurt |
| Use approximate KL for the n-gram component rather than NLL (ablation) | Hybrid distillation · GSM8K subset, temperature 1 | Speedup: 1.27x -> 1.31x | Helped |
| Jointly train the acceptance prediction head and drafter rather than distill only (ablation) | OmniDraft online adaptive drafting · GSM8K, Vicuna-7B | Speedup: 2.20x -> 2.08x | Hurt |
| Use interleaved rather than joint adaptive training (ablation) | OmniDraft online adaptive drafting · MBPP+HumanEval, Vicuna-7B | Speedup: 1.91x -> 1.94x | Helped |
| Increase the LoRA rank from 8 to 32 (ablation) | Cross-vocabulary distillation with LoRA · GSM8K, Llama3-8B | Speedup: 1.478x -> 1.595x | Helped |
| Increase the LoRA rank from 32 to 128 (ablation) | Cross-vocabulary distillation with LoRA · GSM8K, Llama3-8B | Speedup: 1.595x -> 1.632x | Helped |
| Use a quarter-size LFU cache instead of no cache (ablation) | Cross-vocabulary speculative decoding · MBPP+HumanEval | Speedup: 1.29x -> 1.36x | Helped |
| Pretrain the acceptance prediction head on Alpaca before online training (ablation) | Online adaptive drafting · GSM8K and XSum | Speedup | Mixed |
| Use the joint adaptive-drafting variant rather than distill-only (ablation) | OmniDraft online adaptive drafting · GSM8K, Vicuna-7B | Acc Rate: 0.42 -> 0.61 | Helped |
Surprising findings
- Training only on the n-gram NLL loss is very unstable, yielding speedup below baseline (0.86x).
- On GSM8K, distillation-only outperforms both adaptive drafting variants in speedup despite adaptive drafting improving acceptance rate.
Limitations
- Online training is limited to a single pass over the data stream, so stability on unseen data is uncertain.
- Cache eviction for memory-constrained edge devices is not fully solved; full cache is used per task and target.
- Special tokens without direct mapping need extra handling, which complicates multimodal use.
Experimental setup
Datasets: GSM8K Alpaca XSum MBPP HumanEval
Benchmarks: GSM8K MBPP+HumanEval Alpaca XSum
Metrics: speedup acceptance rate n-gram hit
Baselines: SpD (direct mapping) UAG Distill Only Adapt Only OSD
Builds on: speculative decoding GKD SpecDec++ UAG Medusa OSD
Closest papers by meaning
- Not-a-Bandit: Provably No-Regret Drafter Selection in Speculative Decoding for LLMs (ICLR 2026 · similarity 0.89)
- PEARL: Parallel Speculative Decoding with Adaptive Draft Length (ICLR 2025 · similarity 0.89)
- Online Speculative Decoding (ICML 2024 · similarity 0.89)
- SPEED-Bench: A Unified and Diverse Benchmark for Speculative Decoding (ICML 2026 · similarity 0.89)
- SWIFT: On-the-Fly Self-Speculative Decoding for LLM Inference Acceleration (ICLR 2025 · similarity 0.89)
- Scaling Up, Speeding Up: A Benchmark of Speculative Decoding for Efficient LLM Test-Time Scaling (ICLR 2026 · similarity 0.89)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.