Questions on Speculative Decoding in gpt-fast generate.py

[hxer7963]

I'm new to speculative decoding. When I was reading the speculative_decode code (https://github.com/pytorch-labs/gpt-fast/blob/main/generate.py#L88), I have a few questions. Could you please help answer them?

1. When obtaining target_logits, the input_pos passed to model_forward is [input_pos, input_pos + speculate_k + 1), rather than[input_pos, input_pos + speculate_k)? +1 is to generate next token logits?

2. When all draft tokens have been accepted, then fill the last token into the draft model, but the last token from the target model doesn't pass to the draft model. Instead, the draft model independently executes one model_forward? How does this ensure that the next token generated by the draft model does not diverge?

3. For autoregressive generation in large language models, passing a sequence to the model generates logits for each token, which are used to check if the model would produce the same sequence. Is the validation process considered a single forward pass, obtaining logits for all input tokens and generating logits for the next token? Is there no need for autoregressive calculation of each token in this context?

[msaroufim]

The core idea of speculative decoding is that verification is faster than autoregressive generation because the base model can process in batch the entire draft model sample tokens so your intuition in your point 3 is correct

Conceptually this is similar to why prefill is more parallelizable than autoregressive decoding, during the prefill stage you can process the entire context length at once and it would be nice if we could do the same for autoregressive decoding and speculative decoding kinda does that but only for the base model

Back to speculative decoding, how the draft tokens are generated doesn't matter too much it's nowadays popular to generate those draft tokens using a smaller base model but you could guess or have n-gram model, you can train a model etc.. but let's say now you've generated a bunch of sample tokens - how do you verify them?

Well you just base_model(draft_tokens) and get out the probability that the base_model assigns to those tokens. Typically once you have those probabilities you can softmax them to get an actual token out but you don't have to do that.

And then if that probability is above some threshold you accept otherwise you reject the new draft tokens and you make the base model generate them.

target_logits = model_forward(
        model,
        torch.cat([cur_token.view(1), draft_tokens]).view(1, -1),
        torch.arange(input_pos, input_pos + speculate_k + 1, device=cur_token.device)
    )
    target_probs = logits_to_probs(target_logits[0], **sampling_kwargs)

Sidenote: On your question 1 why speculate_k+1 let's say

input_pos = 2
speculate_k = 3

# Generates 2, 3, 4
torch.arange(input_post, input_pos + speculate_k) 

# But we actually want 2,3,4,5 hence the +1

I lied a bit above, the way a token gets accepted in speculative decoding isn't a pure threshold though

1. You get the probability that base model assigns a token q and you compare it to probability of the draft model q
2. If q > p it means the base model finds the draft prediction at least as good if not better than probability that the draft model gives it
3. If q < p then you accept the draft token with a a probability of q / p- you don't have to do it this way you could outright reject but this is the way they do it in the speculative decoding paper

In the happy path all draft tokens get accepted so now the next model_forward can skip k tokens. And if any token is rejected, you stop at the first rejected location

[JacksonCakes]

Hi, I have a question. Since we are accepting the tokens whenever q > p, what happens when there is actually other token that has higher probability than q?