Reproduce “SpecDec++: Boosting Speculative Decoding via Adaptive Candidate Lengths” on OpenResearch

About this paper

Kaixuan Huang, Xudong Guo, Mengdi Wang

SpecDec++ boosts speculative decoding performance by introducing an adaptive mechanism for determining candidate lengths based on a theoretically grounded Markov Decision Process formulation. The method uses a trained prediction head to estimate acceptance probabilities, enabling dynamic stopping and achieving speedups of 7.2% to 11.1% over fixed-length speculative decoding on Llama-2-chat models.

Key insights

• Framing speculative decoding's adaptive candidate length selection as a Markov Decision Process (MDP) reveals an optimal threshold policy, where speculation should stop when the probability of any candidate token being rejected exceeds a specific threshold.
• The alignment between draft and target models, which dynamically varies with the input prefix, is a crucial factor influencing optimal candidate length that fixed or heuristic approaches fail to capture adequately.
• A small, lightweight acceptance prediction head can be effectively trained to estimate the conditional acceptance probability of draft tokens, providing the necessary signal for implementing the adaptive threshold policy.

Summary by alphaXiv.