BlockPilot: Instance-Adaptive Policy Learning for Diffusion-based Speculative Decoding
Hao Zhang, Yiming Hu, Yong Wang, Mingqiao Mo, Xin Xiao, Xiangxiang Chu
arXiv:2606.31315 · 75▲ · cs.CL
View on arXiv →Premium readers get an interactive explainer for this paper — a figure you can poke at, not just read.
What it is
BlockPilot is a lightweight predictor that picks a per-sample block size for diffusion-based speculative decoding, replacing the fixed block size inherited from training. It uses the target model's last-token predictive distribution after prefilling as input to a two-layer MLP classifier that selects the block size from a small local range around the training value.
Why it matters
In speculative decoding with a diffusion draft model, larger blocks add parallelism but accumulate drafting errors, so no single block size is best for all inputs. BlockPilot squeezes out extra speedup on top of an existing state-of-the-art method (DFlash) by choosing the block size per request, adding only millisecond-level latency and about 0.6 GB memory, without changing output quality.
Practical takeaway
If you already run diffusion-based speculative decoding, you can bolt on a small classifier that runs once after prefilling to tune block size per sample, rather than hand-picking one global value. Watch for the offline data construction cost: labeling each training sample requires running full decoding across all candidate block sizes (roughly 25 seconds per sample on a 32B model).
Key result
On Qwen3-4B at temperature 1, BlockPilot reached 4.20x speedup and acceptance length 5.92 averaged across Math, Code, and Chat benchmarks, versus 3.80x for the best fixed-block baseline DFlash(16); gains over the strongest fixed-block variant are consistent but modest (roughly 0.1 to 0.4x speedup) across models and temperatures.
Subscribe
Get the next issue.
Free. One email a week. Unsubscribe any time: no account, no dark patterns.