SpecOPD speeds up language model decoding with parallel on-policy distillation
SpecOPD: Branch-Parallel On-Policy Distillation for Speculative Decoding
Artificial Intelligence
Summary
Generating text one word at a time can be slow for large language models because each word depends on the previous one. The authors found a way to speed this up by training the model to predict multiple words at once without needing to wait for sequential steps. They do this by creating many parallel prediction branches that each learn from past examples independently. This method, called SpecOPD, can make text generation up to three and a half times faster while keeping quality high.
What this means in practice
- •For natural language processing engineers: Accelerate inference in large transformer models by enabling faster multi-token predictions through parallel on-policy distillation.
- •For cloud ai service operators: Increase throughput and reduce latency of text generation APIs by deploying models trained with branch-parallel speculative decoding.
Authors
Tianyu Liu, Jun Zhang, Wenjie Huang, Yuhao Shen, Tianqi Xu, Xinyue Hao, Hongyaoxing Gu, Xuyang Liu, Baolin Zhang, Yating Zhu, Shuang Ge, Jun Dai, Lei Chen, Yue Li, Mingchen Wan
Abstract
Speculative decoding (SD) accelerates large language model inference by reducing the sequential cost of token-by-token generation. Advanced SD methods reduce proposal cost with parallel drafters that predict entire blocks in one forward pass, but overall speedup still depends on acceptance length. On-policy distillation (OPD) can push well-trained drafters beyond acceptance-length plateaus, yet conventional pipelines collect supervision through sequential draft-and-verify rollouts. Our key insight is that on-policy supervision need not inherit these dependencies: a draft block can be supervised using only its reference prefix and draft-sampled history. We propose SpecOPD, a branch-parallel OPD framework that parallelizes draft sampling and target supervision across independent branches drawn from the evolving drafter policy at reference anchors. We pack these branches into proposal trees to share target-prefix computation, using tree masks for full attention, branch-consistent indexer groups for sparse attention, and indexed state replay for linear attention to preserve each branch's history. We train the drafter with an E2E-TV surrogate that directly targets acceptance length. Our study spans 8B full-attention (GQA), 27B full-linear (GQA--GDN), and 125B sparse-linear (QSA--GDN) targets, three parallel draft architectures (DFlash, DSpark, and DFlash V2), and four task domains (dialogue, coding, mathematics, and agentic). SpecOPD achieves up to 41.45% higher mean acceptance length, $16.09\times$ supervised-token throughput over sequential-rollout OPD under native configurations, and up to $3.56\times$ generation speedup over autoregressive decoding.