How speculative decoding lets a small model draft and a large one check

The speculative decoding paper (Leviathan, Kalman and Matias, arXiv 2211.17192) notes that decoding K tokens takes K serial runs of the model.

Inference from large autoregressive models like Transformers is slow - decoding K tokens takes K serial runs of the model.

The paper observes that inference is often limited by memory bandwidth and communication rather than arithmetic, so spare compute is often available.

We additionally observe that inference from large models is often not bottlenecked on arithmetic operations, but rather on memory bandwidth and communication, so additional computation resources might be available.

In the paper's algorithm a cheaper approximation model autoregressively proposes gamma guess tokens, and the target model scores the prefix and all guesses in one parallel pass, producing gamma+1 distributions.

Sample γ guesses x1,...,xγ from Mq autoregressively. ... Run Mp in parallel. p1(x), ..., pγ+1(x) <- Mp(prefix), ..., Mp(prefix + [x1, ..., xγ])

Each guess is checked against a random threshold using the ratio of the target's and the approximation model's probabilities, and guesses that pass are kept in order.

r1 ~ U(0, 1), ..., rγ ~ U(0, 1); n <- min({i - 1 : 1 <= i <= γ, ri > pi(x)/qi(x)} U {γ})

At the first rejection the later guesses are discarded and the rejected position is resampled from an adjusted distribution; if all guesses are accepted, one extra token is sampled from the target model.

sampling an additional token from an adjusted distribution to fix the first one that was rejected, or to add an additional one if they are all accepted.

The paper says each target run yields at least one new token and up to gamma+1 tokens.

That way, each parallel run of the target model Mp will produce at least one new token ... but it can potentially generate many new tokens, up to γ + 1, depending on how well Mq approximates Mp.

The paper says the output distribution is the same as sampling from the target model alone, for any approximation model, with no retraining or architecture changes needed.

Speculative sampling, and therefore speculative decoding, guarantee an identical output distribution for any choice of approximation model Mq without restriction ... Our method can accelerate existing off-the-shelf models without retraining or architecture changes.

The paper reports 2X to 3X acceleration with identical outputs on T5-XXL (11B), at batch size 1 on a single TPU-v4, compared with the T5X implementation.

We demonstrate it on T5-XXL and show a 2X-3X acceleration compared to the standard T5X implementation, with identical outputs. ... We measure walltime improvements with a batch size of 1 on a single TPU-v4

The paper says the speedup comes from running more computation concurrently, so total arithmetic operations can rise, and that the method is not helpful for configurations where additional computation resources are not available.

latency is improved through increased concurrency at the cost of an increased number of arithmetic operations. Thus, our method is not helpful for configurations where additional computation resources are not available.

In the paper's experiments, where the draft model was typically a couple of orders of magnitude smaller than the target, its cost coefficient c (the time for one run of the draft model divided by the time for one run of the target) was always below 0.05 and often close to 0.

experiments where Mq is typically a couple of orders of magnitude smaller than Mp, c was always less than 0.05 and often negligibly close to 0.

Mistral's Large 4 model card describes an embedded DSpark draft head that proposes 3-token blocks and reports up to 1.9x higher throughput, with gains varying by hardware and workload.