LLM: Speculative Decoding Triples Generation Speed

The Standard Text Generation Process by LLMs

In the context of text generation by a large language model (LLM), two essential steps are generally followed. First, there is the forward pass. During this phase, the input text is converted into tokens and fed into the model. Each layer of the model transforms these tokens, and at the end, a probability distribution is produced for the possible next tokens. This means that each potential token is associated with a probability.

Next comes the decoding phase. In this step, the next token is selected from the probability distribution. This choice can be made either by opting for the token with the highest probability, a process known as greedy decoding, or by sampling from the most probable tokens, a method called top-p sampling or nucleus sampling.

After selecting a token, it is added to the input sequence, and a new forward pass is performed to generate the next token. For a large model, for example with 70 billion parameters, each token requires a complete pass through the model, making the generation process lengthy and costly in terms of time.

Autoregressive models operate sequentially, like dominos. Thus, token 100 can only be generated once all previous tokens have been processed. Each token requires a full pass through the network, and generating 100 tokens at a speed of 20 ms per token results in a delay of about 2 seconds, with each token needing to wait for the processing of its predecessors. This leads to significant latency.

The Innovation of Speculative Decoding

To address these limitations, speculative decoding uses two models: a large LLM, referred to as the target model, and a smaller model, often a distilled version, designated as the draft model. The main idea is that the draft model quickly generates tokens that are simpler and more predictable, while the target model checks for accuracy, especially for more complex or domain-specific tokens.

The draft model mimics the behavior of the target model for the majority of tokens, but the target model acts as a quality controller to ensure the overall accuracy of the output.

The speculative decoding strategy relies on three key steps:

• Draft: Quickly generate K tokens with the draft model.
• Verification: Perform a single forward pass of the target model on all K tokens in parallel.
• Acceptance/Rejection: Accept the correct tokens and replace the incorrect ones through rejection sampling.

This method was proposed by Google Research and Google DeepMind in their paper "Accelerating LLM Decoding with Speculative Decoding."

Generally, a model generates one token per forward pass. However, it is possible to feed multiple tokens into an LLM and evaluate them in parallel during a single forward pass. Checking a sequence of tokens is roughly as costly as generating a single token while producing a probability distribution for each token in the sequence.

The Acceptance and Rejection Principle

The next step is rejection sampling, where it is decided which tokens will be kept and which will be rejected. Each token is examined individually, comparing the probabilities p(x) and q(x) assigned respectively by the draft and target models.

Acceptance or rejection is based on a simple rule:

• Case 1: if q(x) >= p(x), the token is accepted.
• Case 2: otherwise, the token is rejected.

Pessimistic Scenario

In the worst case, the first token is rejected, meaning that all subsequent tokens, although potentially relevant, are also rejected. Each token being dependent on its predecessor, the target model must correct token x1, after which the draft model generates a new set of 5 tokens, which will again be checked by the target model.

Optimistic Scenario

In the best case, all tokens generated by the draft model are accepted by the target model. Additionally, the target model generates an extra token after x5. Thus, in this scenario, we obtain K+1 tokens at once, representing a significant speed gain.

Speculative decoding allows for a speedup of about 2 to 3 times by generating and verifying tokens in parallel. Rejection sampling is crucial to ensure that the quality of the output remains consistent with the target model, despite the use of draft tokens.

Efficiency of the Pass

• Worst case: The first token is rejected, meaning that only one token from the target model is accepted.
• Best case: All draft tokens are accepted, resulting in (draft tokens) + (token from the target model) [K+1].

According to the DeepMind paper, it is advisable to maintain K = 3 and 4, which has often led to a speedup of 2 to 2.5 times compared to autoregressive decoding. In the Google paper, the figure 3 was recommended, allowing for a speedup of 2 to 3.4 times.