Target probabilities improve speed and accuracy in speculative decoding
Understanding Target Probability Reuse in Speculative Decoding
Computation and Language
Summary
When computers try to guess the next word or token quickly, they sometimes check many possibilities all at once, called speculative decoding. This paper looks at how the probabilities of those guesses can help pick better paths to try next. The authors found that using a separate calculation of these probabilities early on can make the guessing faster and more accurate. This helps computers spend less time checking bad guesses and more time on good ones.
What this means in practice
- •For language model engineers: Improve decoding efficiency by selecting draft token trees using early probability estimates to reduce computation during text generation.
- •For machine translation developers: Speed up token generation by using target probability reuse to prioritize better translation hypotheses with limited computational budgets.
Authors
Jungseob Lee, Giyoon Jang, Chailin Nam, Jiwoo Kum, Sihyeon Lee, Kyojun Choo, Jihun Song, Eunseo Song, Chanjun Park
Abstract
Speculative decoding uses a target model to verify a tree of draft tokens in parallel. We study whether the resulting next-token distributions can help select the next tree. For unrestricted target models, we show that guaranteed exact reuse requires a distribution computed at the next round's complete context. We characterize when retaining this distribution requires sacrificing accepted tokens and establish conditions under which this loss persists across decoding rounds. Across six targets, we find that reranking first-token branches with a separately computed distribution at the current context improves expected acceptance over draft ranking, especially with small tree budgets. Drawing the first token from this distribution before allocating the tree further improves expected acceptance over pruning a preselected tree, even at equal mean tree size. In our experiments, both selection rules also decode faster than draft ranking when the separate target pass overlaps standalone drafting. Our findings show that making target probabilities available before tree allocation can improve the efficiency of speculative decoding. Our code is available at https://github.com/js-lee-AI/TargetRows.