Paper 2603.27027
TAPS: Task Aware Proposal Distributions for Speculative Sampling
- Published
- Mar 2026
- Research lab
- Independent
- Citations
- 0
- GitHub
- 8 stars
01 In brief
Summary
This paper investigates how the training distribution of draft models affects speculative decoding performance, using HASS and EAGLE-2 drafters trained on MathInstruct, ShareGPT, and mixed data, evaluated on MT-Bench, GSM8K, MATH-500, and SVAMP.
Results show task-specific training yields clear specialization: MathInstruct-trained drafts excel on reasoning benchmarks, while ShareGPT-trained drafts excel on MT-Bench.
Mixed-data training improves robustness but larger mixtures are not uniformly better across temperatures.
For combining specialists, checkpoint averaging performs poorly, while confidence-based routing improves over single-domain drafts, and merged-tree verification achieves the highest acceptance length overall.
Confidence is a more effective routing signal than entropy, and acceptance declines with speculative depth, with task-matched specialists becoming more dominant at deeper levels.
The study concludes that speculative decoding quality depends on both draft architecture and the alignment between draft training data and downstream tasks, and that specialized drafters are best combined at inference time rather than in weight space.
02 From the paper
Abstract
Speculative decoding accelerates autoregressive generation by letting a lightweight draft model propose future tokens that a larger target model then verifies in parallel. In practice, however, draft models are usually trained on broad generic corpora, which leaves it unclear how much speculative decoding quality depends on the draft training distribution. We study this question with lightweight HASS and EAGLE-2 drafters trained on MathInstruct, ShareGPT, and mixed-data variants, evaluated on MT-Bench, GSM8K, MATH-500, and SVAMP. Measured by acceptance length, task-specific training yields clear specialization: MathInstruct-trained drafts are strongest on reasoning benchmarks, while ShareGPT-trained drafts are strongest on MT-Bench. Mixed-data training improves robustness, but larger mixtures do not dominate across decoding temperatures. We also study how to combine specialized drafters at inference time. Naive checkpoint averaging performs poorly, whereas confidence-based routing improves over single-domain drafts and merged-tree verification yields the highest acceptance length overall for both backbones. Finally, confidence is a more useful routing signal than entropy: rejected tokens tend to have higher entropy, but confidence produces much clearer benchmark-level routing decisions. These results show that speculative decoding quality depends not only on draft architecture, but also on the match between draft training data and downstream workload, and that specialized drafters are better combined at inference time than in weight space.