All martial arts in the world are invincible only with speed: speculative decoding

According to the scaling law, the more model parameters and data, the better the performance. However, the side effect of more parameters is large GPU memory usage and slow inference speed.

Generally speaking, current LLM acceleration techniques fall into three categories: model-level, computation-level, and decoding process.

Model-Level Acceleration:

1. Model pruning: Remove unimportant weights
2. Model quantization: High-precision parameters are unnecessary; 8-bit is sufficient
3. Model distillation: In specific domains, only part of the LLM's intelligence is needed, not all of it—so teach that part to a small model

Computation-Level Acceleration:

1. Model parallelism: Too many parameters, split them across multiple GPUs
2. Data parallelism: Too much data, split it across multiple GPUs
3. Mixed precision computation: Use high precision where necessary, low precision where not, wrapped in PyTorch's amp decorator

Decoding Process Acceleration:

1. Speculative decoding: Use a small, fast model to generate multiple candidate sequences, have the LLM evaluate the most reasonable one, then continue generating
2. KV cache: Leverage the computational principles of self-attention to cache previously computed KV products

Speculative Decoding

Speculative decoding seems to draw inspiration from computer architecture—when designing a five-stage pipeline CPU, there's an acceleration technique called branch prediction.

The Speculative Process

1. Draft phase: A small model quickly generates a prefix sequence $x_1, x_2, x_3, ..., x_t$
2. Verification phase: The LLM performs inference on this prefix, computes the probability of the draft, and if it meets expectations, samples the next token $x_{t + 1}$, yielding the sequence $x_1, x_2, ... x_{t}, x_{t + 1}$. Specifically:
   1. For the small model's generated sequence $x_{1 ... t}$, the LLM computes the probability of each token as $p(x_t)$, while the small model computes the probability of each token as $q(x_t)$
   2. If $p(x_t) < q(x_t)$ (the LLM considers this token unreasonable), reject this token with probability $1 - \frac{p(x_t)}{q(x_t)}$, and resample the token from the new probability distribution $p^`(x) = norm(max(0, p(x_t) - q(x_t)))$

Equivalence Proof

The speculative process is listed above—it sounds reasonable, but is it actually valid? Won't it reduce the model's intelligence? Through mathematical proof, we can show that speculative sampling is equivalent to LLM autoregressive sampling.

• $p(x)$: Probability distribution generated by the LLM

• $q(x)$: Probability distribution generated by the small model

• The probability of a token at time t ($x_t$) is $\beta(x)$

$\beta(x_t)$ has only two cases:

1. The LLM verifies and finds it reasonable, i.e., $p(x) >= q(x)$, and directly accepts the token. Then $\beta(x_t) = q(x_t) = min\{p(x_t), q(x_t)\}$
2. The LLM verifies and finds it unreasonable, i.e., $p(x) < q(x)$, rejects it and resamples to obtain the token.
   1. The probability of rejecting all other tokens is $1 - \sum_{x^`_t != x_t} min\{p(x_{t^`} ), q(x_{t^`})\}$
   2. The probability that the LLM's resampling happens to yield this token: $p(x_t)$
   3. Then $\beta(x_t) = p(x_t) \times (1 - \sum_{x^`_t != x_t} min\{p(x_{t^`} ), q(x_{t^`})\})$

Therefore: $\beta(x_t) = min\{p(x_t), q(x_t)\} + p(x_t) \times (1 - \sum_{x^`_t != x_t} min\{p(x_{t^`} ), q(x_{t^`})\})$. I can't continue the proof—my math skills are too weak. Anyway, it roughly belongs to the same distribution as $p(x_t)$.