Improving draft model training to speed up large language model decoding
Acceptance-Aware Draft Model Training for Speculative Decoding
Machine Learning
Summary
Large language models can be slowed down because they check each generated word carefully. Speculative decoding uses a quicker model to guess many words at once, which the main model then verifies, making things faster. The authors found that current training methods don’t directly focus on making these guesses accepted more often, which would boost speed. They created new ways to train the quick model to better match what the main model will accept, leading to faster and more efficient language generation.
What this means in practice
- •For ai inference engineers: Train draft models with acceptance-aware losses to accelerate large language model output without sacrificing quality.
- •For machine learning infrastructure teams: Optimize large-scale language model serving pipelines by integrating acceptance-length-aware trained draft models for efficient speculative decoding.
Authors
Tianhua Xia, Mugilan Ganesan, Yifei Feng, Haiyu Wang, Maximilian Egger, Sai Qian Zhang
Abstract
Speculative decoding accelerates large language model (LLM) inference by using a lightweight draft model to generate multiple candidate tokens that are verified by the target model in a single forward pass. Its speedup is largely determined by the acceptance length, yet existing draft-model training methods mainly optimize cross-entropy or Kullback-Leibler (KL) divergence as proxies. These objectives encourage distribution matching but do not directly optimize acceptance length, and the acceptance mechanism also differs between greedy and sampling-based decoding. In this work, we propose acceptance-length-aware training losses that directly optimize the expected number of accepted tokens within a speculative window. For greedy verification, we derive an expected accepted length (EAL) loss that explicitly maximizes expected acceptance length. For sampling-based decoding, we introduce a window total variation (WTV) loss that optimizes the overlap between temperature-scaled draft and target distributions while accounting for sequential acceptance dependencies. Both objectives can be further combined with a group-relative reinforcement learning stage (GRPO) using simulated acceptance length as the reward. Experiments across different target and draft models, tasks, and decoding settings show that our losses consistently improve acceptance length over KL-based training. WTV provides particularly strong gains under sampling-based decoding, while EAL better matches greedy verification. These results show that directly optimizing the acceptance objective, with losses tailored to the decoding mode, is more effective than conventional distribution-matching objectives.