Speculative decoding

Remarque

1. features here refer to the hidden states of the decoder layers second-to-top-layer of the LLM, before the LM head. Not to be confused with features ↩

2. Vanilla autoregressive at token-level is described by $T_{1:j} \rightarrow E_{1:j} \rightarrow f_j \rightarrow p_{j+1} \rightarrow t_{j+1}$:

   • input $T_{1:j}$ is then transformed into embeddings $E_{1:j}$
   • then into features $F_{1:j}$,
   • LM Head maps $f_j$ to a distribution $p_{j+1} = \text{LM\_Head}(f_j)$
   • sampling next token $t_{j+1}$
   ↩

3. See vllm-project/vllm#20078 ↩

4. Aligns with DistillSpec and Medusa ↩

5. This seems like an premature optimization. For use-cases where the batch sizes fluctuates, the calculation for an optimal speculative length would probably too overkill when the improvement could be minimal. ↩

6. Note that we refer to standard sampling to methods such as argmax, top-k, nucleus, temperatures, et al., albeit each have a different ways to process logits. We will consider these as standard sampling from an adjusted distribution ↩

7. This work from DeepMind was performed concurrently and independently from Leviathan et al. (2023). The work at DeepMind focuses more on distributed settings of speculative decoding ↩

8. also referred in practice as num_speculative_tokens ↩

9. or natural measure of the acceptance rate $\beta$ ↩

10. Rejection sampling follows a iterative sampling procedure that might looks superficially similar to speculative sampling:

    1. Sample $x \sim q(x)$ and returns $r \sim U(0,1)$
    2. If $r < \frac{p(x)}{M q(x)}$ return $x$
    3. then go to 1

    Where $M = \operatorname{max}_{x} \frac{p(x)}{q(x)}$

    We could employ non-iterative version of rejection sampling instead of speculative sampling here (go through step 1 and 2, and otherwise sample an unmodified $p(x)$ directly)

    Specifically, the expected accept probability:

    $$E_{x\sim q(x)} \frac{p(x)}{M q(x)} = \sum_{x} p(x) \min_{x^{'}}{\frac{q(x^{'})}{p(x^{'})}} \le \sum_{x} p(x) \min{(1, \frac{q(x)}{p(x)})} = \sum_{x} \min{(p(x), q(x))}$$ ↩

11. A lenience parameter $l \in [0,1]$ to introduce further trade-off. This is useful when the distributions of draft models does not match the target model exactly.

    Specifically we have:

    $$\begin{aligned} \alpha &= \mathbb{E}_{x\sim q(x)} \!\left[ \begin{cases} 1, & \text{if } l\,q(x) \le p(x),\\[6pt] \displaystyle\frac{p(x)}{l\,q(x)}, & \text{if } l\,q(x) > p(x) \end{cases} \right] \\[10pt] &= \mathbb{E}_{x\sim q(x)}\! \frac{p(x)}{\max\!\bigl(p(x),\,l\,q(x)\bigr)} \\[8pt] &= \sum_{x} \frac{p(x)\,q(x)}{\max\!\bigl(p(x),\,l\,q(x)\bigr)} \\[8pt] &= \frac{1}{l}\sum_{x} \min\!\bigl(p(x),\,l\,q(x)\bigr) \\[8pt] &= \sum_{x} \min\!\Bigl(\tfrac{p(x)}{l},\,q(x)\Bigr). \end{aligned}$$

    Important

    this relies on q is sampled from this given distributions, and $l$ increases $\alpha$

    In the case of greedy decoding (temperature=0), the draft essentially outputs $x^{'}_q = \argmax{q(x)}$, so scaling $l q(x)$ becomes a no-op, given that the argmax will be unchanged in this case. ↩

12. Note that we can’t use temperature=0 (i.e argmax sampling):

    • Instead we allow some lenience before standardizing the distribution (accept token $x$ sampled from $M_q$ in case of $p(x) \le l \dot \max{p}$)
    • In this case, then similar empirical increases to $\alpha$ to those of temperature=1
    ↩

13. On Correctness of Speculative Sampling (SpS)

    We will show that $\forall p(x) \text{ and } q(x)$, tokens sampled via speculative sampling from $p(x)$ and $q(x)$ are distributed identically to those sampled from $p(x)$ alone.

    Let $\beta$ be the acceptance probability

    Note that

    $$p'(x) = \operatorname{norm}\!\bigl(\max(0,\;p(x)-q(x))\bigr) = \frac{p(x)-\min\!\bigl(q(x),\,p(x)\bigr)} {\displaystyle \sum_{x'}\!\bigl(p(x')-\min\!\bigl(q(x'),\,p(x')\bigr)\bigr)} = \frac{p(x)-\min\!\bigl(q(x),\,p(x)\bigr)}{1-\beta},$$

    so the normalising constant for the adjusted distribution $p'(x)$ is $1-\beta$; the last equality follows immediately from Lemma 3.3 and Theorem 3.5.

    Now

    $$P(x = x') \;=\; P(\text{guess accepted},\,x = x') \;+\; P(\text{guess rejected},\,x = x').$$

    Where

    $$P(\text{guess accepted},\,x = x') \;=\; q(x')\,\min\!\bigl(1,\tfrac{p(x')}{q(x')}\bigr) \;=\; \min\!\bigl(q(x'),\,p(x')\bigr),$$

    and

    $$P(\text{guess rejected},\,x = x') \;=\; (1-\beta)\,p'(x') \;=\; p(x') - \min\!\bigl(q(x'),\,p(x')\bigr).$$

    Overall

    $$\begin{aligned} P(x = x') &= \min\!\bigl(p(x'),\,q(x')\bigr) \;+\; p(x') - \min\!\bigl(p(x'),\,q(x')\bigr) \\ &= p(x'). \end{aligned}$$

    $\boxed{}$ ↩

14. $$\begin{aligned} D_{LK}(p,q) &= \sum_{x} |p(x) - M(x)| = \sum_{x} \frac{|p-q|}{2} \\ &= 1- \sum_{x} \frac{p+q - |p-q|}{2} \\ &= 1 - \sum_{x} \min{p(x), q(x)} \end{aligned}$$

    $\boxed{}$ ↩

15. $$\begin{aligned} \beta &= \mathbb{E}_{x \sim q(x)} \Biggl[ \begin{cases} 1 & \text{if } q(x) \le p(x), \\[6pt] \displaystyle\frac{p(x)}{q(x)} & \text{if } q(x) > p(x) \end{cases} \Biggr] \\[8pt] &= \sum_{x} \min\!\bigl(p(x),\,q(x)\bigr). \end{aligned} \qquad\square$$ ↩

16. also known as wikipedia/en/Elapsed_real_time. This is different from CPU time, given that it measure the actual time taken from the start of the computer program, where as CPU time only measures time during which processor is actively working on a certain task or process ↩

17. Denote the cost of running single steps of $M_p$ by $T$.

    Each run will then costs $T c \gamma + T = T(c \gamma +1)$ (running $M_q$ $\gamma$ times and running $M_p$ once)

    Given (1) procduces $\frac{1-\alpha^{\gamma +1}}{1-\alpha}$ tokens

    The cost to produces a token with speculative sampling would be $\frac{(c \gamma +1)(1-\alpha )}{1-\alpha^{\gamma +1}} T$

    $\boxed{}$ ↩

18. Denote by $\hat{T}$ the number of arithmetic operations done by standard decoding per tokens, therefore speculative sampling costs $\hat{T} \hat{c} \gamma + \hat{T}(\gamma +1)$ operations. Then divided by the expected tokens we got the desired results $\boxed{}$ ↩