SpecDec++: Boosting Speculative Decoding via Adaptive Candidate Lengths (2405.19715v2)

Abstract: Speculative decoding reduces the inference latency of a target LLM via utilizing a smaller and faster draft model. Its performance depends on a hyperparameter K -- the candidate length, i.e., the number of candidate tokens for the target model to verify in each round. However, previous methods often use simple heuristics to choose K, which may result in sub-optimal performance. We study the choice of the candidate length K and formulate it as a Markov Decision Process. We theoretically show that the optimal policy of this Markov decision process takes the form of a threshold policy, i.e., the current speculation should stop and be verified when the probability of getting a rejection exceeds a threshold value. Motivated by this theory, we propose SpecDec++, an enhanced version of speculative decoding that adaptively determines the candidate length on the fly. We augment the draft model with a trained acceptance prediction head to predict the conditional acceptance probability of the candidate tokens. SpecDec++ will stop the current speculation when the predicted probability that at least one token gets rejected exceeds a threshold. We implement SpecDec++ and apply it to the llama-2-chat 7B & 70B model pair. Our adaptive method achieves a 2.04x speedup on the Alpaca dataset (an additional 7.2% improvement over the baseline speculative decoding). On the GSM8K and HumanEval datasets, our method achieves a 2.26x speedup (9.4% improvement) and 2.23x speedup (11.1% improvement), respectively.

An Expert Review of "SpecDec++: Boosting Speculative Decoding via Adaptive Candidate Lengths"

The paper presents a novel method called SpecDec++, which aims to enhance speculative decoding by adaptively tuning the candidate token length for LLMs. Speculative decoding exploits a smaller, more efficient draft model to reduce the inference latency of larger target models by predicting and verifying sequences of tokens. The choice of candidate length, represented by the hyperparameter $K$, is crucial as it dictates the trade-off between speed and accuracy—larger values enable quicker token generation at the risk of increased rejection probability. Unlike previous approaches that relied on heuristics for setting $K$, this work formulates the decision-making process as a Markov Decision Process (MDP) and theoretically demonstrates that an optimal stopping threshold exists for when speculation should be paused for verification.

Key Contributions

1. Formulation as a Markov Decision Process: The authors model the candidate length selection as an MDP, providing a structured framework for analyzing and determining the optimal time to stop token speculation based on predicted acceptances and rejections.
2. Adaptive Speculative Decoding (SpecDec++): SpecDec++ extends speculative decoding by incorporating a trained acceptance prediction head into the draft model. This prediction head dynamically evaluates whether to continue or halt decoding based on the likelihood of candidate token rejection, according to a threshold derived from theoretical insights.
3. Practical Implementation and Performance: SpecDec++ was implemented and tested on several datasets (Alpaca, GSM8K, HumanEval) using the llama-2-chat model pair. The empirical results demonstrate a significant improvement over baseline methods, achieving up to a 2.04x speedup on Alpaca and over 2.23x on GSM8K and HumanEval datasets.

Numerical Outcomes

The evaluation of SpecDec++ is grounded in rigorous empirical experiments. On the Alpaca dataset, SpecDec++ achieved a speedup factor of 2.04x, indicating a 7.2% increase over the baseline's 1.90x. For GSM8K and HumanEval, SpecDec++ improved the speedup to 2.26x and 2.23x, representing increases of 9.4% and 11.1%, respectively. Such consistent gains across diverse benchmark datasets underscore the method's robustness and efficiency in reducing inference latency without compromising on acceptance rates.

Implications and Future Directions

The successful integration of an adaptive mechanism within speculative decoding holds substantial implications for real-world applications of LLMs, where inference efficiency is critical. By minimizing inference delays, SpecDec++ enhances the feasibility of deploying ultra-large models in latency-sensitive environments. The approach sets a precedent for further refinement of speculative execution frameworks through adaptive learning methods.

Future avenues might explore the extension of this methodology to other model architectures and beyond token-based processing, such as expanding the application of MDPs in speculative execution. Additionally, due to its reliance on a prediction head, research could assess the effects of different head architectures and training regimes to maximize the efficacy of adaptive speculative decoding.

The theoretical grounding in MDPs paired with empirical validations makes this work a substantive contribution to the ongoing effort to optimize LLM inference, aligning well with broader trends in AI towards faster and more resource-efficient model deployment. SpecDec++ thus stands as an effective strategy for accelerating LLMs, balancing the high computational demands of large models with practical performance considerations.