How to speed up LLMs with speculative decoding

TL;DR

Speculative decoding speeds up inference by running a small draft model in parallel with a large target model and then using a simple acceptance rule to keep quality. It leverages token speculation, parallel verification, and rejection sampling to generate more tokens per round, achieving 2 to 4x faster speeds without sacrificing output quality.

Transcript

So you want your large language model to be fast. Let me show you how. Speculative decoding is an effective technique for speeding up LLM inference times without sacrificing quality of output. This approach follows the slogan draft and verify by using a smaller draft model to speculate about future tokens while a larger target model verifies them i... Read More

Key Insights

• Speculative decoding uses a smaller draft model to guess several next tokens, creating a pool of candidate tokens and their draft probabilities that are then evaluated for usefulness.
• Parallel verification runs the larger target model on the draft tokens to obtain target probabilities and confirm whether the draft tokens would be produced by the target model.
• Rejection sampling compares target probabilities to draft probabilities and accepts tokens where the target probability is greater or equal, discarding unlikely guesses to preserve quality.
• The approach can produce two to four tokens in the same time a normal LLM would produce one, effectively increasing throughput.
• The method relies on the editor writer analogy, where a fast draft model provides candidates and the larger model double checks and corrects as needed.
• By running draft and target models concurrently, resource utilization on GPUs is improved and computation is better distributed.
• Speculative decoding reduces latency and memory pressure while maintaining the same quality of output according to the description.
• The technique is described as a frontier in AI optimization with ongoing research and potential for broader adoption.

Questions & Answers

Q: What is speculative decoding and how does it speed up LLM inference?

Speculative decoding speeds up inference by using a smaller draft model to predict multiple next tokens, while a larger target model verifies those predictions in parallel. This creates a pool of candidate tokens and then a rejection sampling step filters them based on how their target probabilities compare to the draft probabilities. The result is more tokens generated per forward pass with maintained quality.

Q: What are the three main steps of speculative decoding?

The three main steps are token speculation, parallel verification, and rejection sampling. Token speculation uses a smaller model to generate draft tokens; parallel verification has the larger model assess those tokens concurrently; rejection sampling compares probabilities from both models and accepts tokens where the target probability meets or exceeds the draft probability, discarding the rest.

Q: How does token speculation work in this technique?

Token speculation involves a draft model predicting several next tokens instead of just one, along with their draft probabilities. This builds a set of candidate tokens and associated confidence levels that could potentially appear next in the sequence. These candidates are then tested by the larger model to confirm their viability.

Q: What is the role of the larger target model in speculative decoding?

The larger target model verifies the draft tokens by evaluating the next token and its own confidence for the preselected candidates. It provides target probabilities and confirms whether the draft predictions align with what it would generate. This verification ensures the final output preserves the distribution of the larger model.

Q: What is rejection sampling in this context?

Rejection sampling is the final step where each draft token is accepted or rejected by comparing the target probability to the draft probability. If the target probability is greater than or equal to the draft probability, the token is accepted; otherwise, the token and following guesses are discarded and re-sampled.

Q: What performance benefit does speculative decoding claim to provide?

Speculative decoding can generate two to four tokens in the same time a normal LLM would generate just one. In the best case, multiple tokens are produced per round, while in the worst case, at least one token is produced, thanks to the verification step. Overall, it can be two to four times faster on average.

Q: Why is the rejection sampling step considered important?

The rejection sampling step is crucial because it ensures quality is preserved by discarding draft tokens that would not be preferred by the target model. This step prevents the pipeline from degrading output quality and allows the faster draft predictions to be corrected or redirected by the larger model.

Q: What practical benefits does this technique offer beyond speed?

Beyond speed, speculative decoding improves GPU resource utilization and memory efficiency by allowing the draft model to handle heavy lifting while the target model verifies, reducing wasted computation. It enables better overall throughput without sacrificing the distribution or quality of the generated text.

Summary & Key Takeaways

• Speculative decoding speeds up inference by using a smaller draft model to predict future tokens while a larger model verifies them in parallel, enabling more tokens per forward pass.

• The process includes three main steps: token speculation by a draft model, parallel verification by a target model, and rejection sampling to select acceptable tokens, maintaining output quality.

• Overall, this approach improves GPU efficiency and reduces latency by coordinating two models and a verification mechanism to keep results aligned with the larger model’s distribution.