Speculative Decoding

1 minute read

A discussion around tokenizers in slack leads to following comment:
Would be the natural progression after you arrive at the fact taking three steps to get sm art er may be inefficient when there is a high proba it’s going to be smarter

That’s a good explanation of the intuition behind speculative decoding.

1 Speculative Sampling and Speculative Decoding

DeepMind and Google published two paper around the same time regarding this algorithm.

Define notations as follow:
$M_p$ = draft model (llama-7b)
$M_q$ = target model (llama-70b)
$pf$ = prefix, $K$ = 5
Do following two steps in parallel.

Run draft model for K steps to generate K tokens
Run target model once to get distributions of K+1 tokens (No sampling here) The key algorithm is reject sampling, accept or reject tokens based on values of p/q
Even when rejected, still sample the final token from $(q(x)-p(x))_+$. This step is important to make sure we generate at least 1 token.
The worse case is to generate 1 token in each pass, and the best case is to generate K+1 tokens, so speedup is garanteened.