Speculative Decoding

Speculative Decoding

paper

Algorithm of Speculative Decoding Step

---

Algorithm SpeculativeDecodingStep

---

Inputs: $M_p,M_q,\text{prefix}$.

Sample $\gamma$ guesses $x_{1,\cdots ,\gamma }$ from $M_q$ autoregressively.

for i = 1 to $\gamma$ do
 $q_i(x)\leftarrow M_q(\text{prefix}+[x_1,\cdots ,x_{i-1}])$
 $x_i\sim q_i(x)$
end for

Run $M_p$ in parallel.

$p_1(x),\cdots ,p_{\gamma +1}(x)\leftarrow M_p(\text{prefix}),\cdots ,M_p(\text{prefix}+[x_1,\cdots ,x_{\gamma }])$

Determine the number of accepted guesses $n$.

$r_1\sim U(0,1),\cdots ,r_{\gamma }\sim U(0,1)$

$n\leftarrow \min \left(\left\{i-1|1\le i\le \gamma ,r_i>\frac{p_i(x)}{qi(x)}\right\}\cup \{\gamma \}\right)$

Adjust the distribution from $M_p$ if needed.

$\tilde{p}(x)\leftarrow p_{n+1}(x)$

if $n<\gamma$ then  $\tilde{p}(x)\leftarrow \text{norm}(\max \{0,p_{n+1}(x)-q_{n+1}(x)\})$

end if

Return one token from $M_p$, and $n$ tokens from $M_q$.

$t\sim \tilde{p}(x)$

return $\text{prefix}+[x_1,\cdots ,x_n,t]$

---

Correctness of Speculative Sampling

We will now show that for any distributions $p(x)$ and $q(x)$, the 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 (Definition).

Note that as

$$\begin{array}{rl}\tilde{p}(x)& =\text{norm}(\max \{0,p(x)-q(x)\})\\ & =\frac{p(x)-\min \{q(x),p(x)\}}{{\sum }_{x^{\prime }}p(x^{\prime })-\min \{q(x^{\prime }),p(x^{\prime })\}}\\ & =\frac{p(x)-\min \{q(x),p(x)\}}{1-\beta },\end{array}$$

the normalizing constant for the adjusted distribution $\tilde{p}(x)$ is $1-\beta$, where the last equation follows immediately from Lemma 3.3 and Theorem 3.5.

Now:

$$P(x=x^{\prime })=P(\text{guess accepted},x=x^{\prime })+P(\text{guess rejected},x=x^{\prime })$$

Where:

$$P(\text{guess accepted},x=x^{\prime })=q(x^{\prime })\min \left\{1,\frac{p(x^{\prime })}{q(x^{\prime })}\right\}=\min \{q(x^{\prime }),p(x^{\prime })\}$$

And:

$$P(\text{guess rejected},x=x^{\prime })=(1-\beta )\tilde{p}(x^{\prime })=p(x^{\prime })-\min \{q(x^{\prime }),p(x^{\prime })\}$$

Overall:

$$P(x=x^{\prime })=\min \{p(x^{\prime }),q(x^{\prime })\}+p(x^{\prime })-\min \{p(x^{\prime }),q(x^{\prime })\}=p(x^{\prime }).$$

As desired.

Definition

The acceptance rate ${\beta }_{x_{<t}}$ , given a prefix $x_{<t}$, is the probability of accepting $x_t\sim q(x_t|x_{<t})$ by speculative sampling.

Lemma

Define

$$D_{LK}(p,q)=\sum _x|p(x)-M(x)|=\sum _x|q(x)-M(x)|$$

where $M(x)=\frac{p(x)+q(x)}{2}$. Then

$$D_{LK}(p,q)=1-\sum _x\min \{p(x),q(x)\}$$

Proof.

$$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)\}$$