Speculative Decoding

Speculative decoding is a technique in AI inference optimization where we combine a fast draft model usually smaller in parameter size with the actual large language model to help speed up token generation time. It’s a very useful and widely adopted technique in inference engineering and optimization. Let’s look at how this works and tradeoff one needs to make in order to use speculative decoding in their inference pipeline.

How Speculative Decoding Works

Speculative decoding introduces a small parameter model that can work and generate next token faster than the large LLM model with reasoning capability. The small model is termed as “Draft Model” in the diagram below. The draft model generates sequence of tokens and passes it to the actual model to evalulate and verify. LLM evaluates the tokens and accepts the token which is correct and either replaces or rejects the tokens which are wrong. All of this happens in one forward pass which saves times and improves performance. Post verification stage LLM updates the KV cache with the accepted tokens and proceeds to generate more tokens required to process the prompt. This approach increases TPS (token per second) aggresively.

Rejection Mechanism

How the model accepts or rejects a token generated by the draft model is the core algorithm in speculative decoding. Let me explain this a simple approach rather than tons of math. Here is a role of each models:

• Draft Model: Small, fast, bit sloppy. It’s only job is to guess the next word very quickly.
• Main Model: Big, slow but briilliant. It guesses most words with absolute good probability score.

Let’s assume for token x, draft model has a confident score and let’s call it q(x). Main model also has a probabliy score (after checking parallely the token generated by the Draft model) for the same token if it looked at the text itself, let’s call it p(x). if p(x) >=q(x) then the token is immediately accepted. No further question asked. However, if p(x) < q(x) then the main model performs the mathematical dice roll to give the draft model a partial chance. First a probability score of both is divided to figure out a chance to accept the Draft token. p(x)/q(x) Suppose q(x) is 80% and p(x) is 40% then 40/80=1/2=0.5 Model finds a random number between 0 and 1 and if that’s less than 0.5 as acceptance cutoff then the token is accepted otherwise rejected. Here is a very simple code to understand this.

import random

# Main model confidence: 40%, Draft model confidence: 80%
acceptance_cutoff = 0.40 / 0.80  # This equals 0.5

# The computer rolls the "dice" by picking a random number between 0 and 1
dice_roll = random.random() 

if dice_roll < acceptance_cutoff:
    print("Word Accepted!")  # Only happens 50% of the time
else:
    print("Word Rejected!")

If the token is not accepted in this phase then the chain breaks and main model reject this token (i.e., x) and any other token that is coming after it. Main model then continues from the rejected token and generates the remaining tokens required and samples the response. The chain breaking rule is a key point to consider in production. This can either improve the performance or decreases it really bad. Goal is to select or train a draft model to generate more number of words right hence increasing the acceptance rate. The more the acceptance rate higher your performance.

Production Tradeoffs

Speculative decoding is a good technique to improve the performance but it comes with it’s own complexity. Here are some points to consider in production:

• TPS Improvement: If draft model token acceptance rate by main model is more than or near 50% then you will observe performance improvement but if it falls below this threshold then TPS reduces drastically and is slower than using just one main model.
• vRAM Footprint: We need to host two models with two KV cache. This increases the demand for vRAM which adds up cost. if vRAM is saturated then context window is affected to make room for two model KV cache.
• High Throughput Systems: If your system is handling hundreds or thousands of requests parallely then your GPU is already fully saturated. Speculative decoding will actually lower the TPS.
• Tokenizer Alignment: The draft model and the main model must share the exact same tokenizer. You cannot pair a model that uses OpenAI’s tiktoken with a model that uses Llama’s tokenizer without destroying your system’s performance with messy translation layers.

Regardless of this, speculative decoding a sort of must have technique for low latency applications like code auto-complete, voice assistance, and chat interfaces.