Fast Inference from Transformers via Speculative Decoding

论文地址：https://arxiv.org/pdf/2211.17192

speculative sampling

为了从分布 $p(x)$ 中采样，我们实际上是从分布 $q(x)$ 中采样 $x$，如果 $q(x)\le p(x)$，则保留该样本；如果 $q(x)>p(x)$，则以概率 $1-\frac{p(x)}{q(x)}$ 拒绝该样本，并重新从调整后的分布 $p^{\prime }(x)=\text{norm}(\max (0,p(x)-q(x)))$ 中采样。对于任何分布 $p(x)$ 和 $q(x)$，以及以此方式采样的 $x$，确实有 $x\sim p(x)$。

给定通过在条件前缀上运行 $M_q$ 获得的分布 $q(x)$，我们可以采样一个标记 $x_1\sim q(x)$。然后，我们通过在前缀上运行 $M_p$ 来计算分布 $p(x)$，同时并行地推测性地计算下一个标记 $x_2$ 的分布，即在前缀上追加 $x_1$ 后运行 $M_p$。一旦两项计算都完成，我们就按上述方式处理：如果 $x_1$ 被拒绝，我们丢弃 $x_2$ 的计算，并从调整后的分布中重新采样 $x_1$；如果$x_1$ 被接受，我们就保留两个标记。算法 1 将这一想法推广为一次采样 1 到 $\gamma +1$ 个标记。

分析

有几个证明需要注意一下：

单次算法期望能生成的token

1. 单次算法期望能生成的token数量服从几何分布，但是求和项是有限制的，这里推导下​

2. ​​接受率β的定义​​
   设目标模型分布为 p(x)，草稿模型分布为 q(x)。草稿模型生成的单个token被目标模型接受的概率为：

$$\beta =\sum _x\min \left(q(x),p(x)\right)$$

3. ​​拒绝率α的定义​​

$$\alpha =1-\beta =1-\sum _x\min (p(x),q(x))x$$

• 假设每个token的接受事件独立且同分布（i.i.d.），草稿模型一次生成 K 个token：

• ​​首次拒绝发生在位置 r​​ 的概率为：

  $$P(r)=(1-\beta ){\beta }^{r-1}(1\le r\le K)$$

  所有token均被接受​​ 的概率为：${\beta }^K$

• 综上期望能生成的token数量为：

  $$\gamma =\underset{\text{拒绝前生成的token}}{\underbrace{\sum _{r=1}^{K}r\cdot P(r)}}+\underset{\text{全接受时生成K个token}}{\underbrace{K\cdot {\beta }^K}}$$

代入 $P(r)$ 后展开：

$$\gamma =\sum _{r=1}^{K}r\cdot (1-\beta ){\beta }^{r-1}+K{\beta }^K$$

4. 几何级数求和​

几何级数求和公式为：

对 ${\sum }_{r=1}^{K}r{\beta }^{r-1}$ 求和处理：

• ​令 $S={\sum }_{r=1}^K{\beta }^{r-1}$​：

$S=1+\beta +{\beta }^2+\cdots +{\beta }^{K-1}=\frac{1-{\beta }^K}{1-\beta }$

• ​​对 $S$ 求导​​：

${\sum }_{r=1}^{K}r{\beta }^{r-1}=\frac{d}{d\beta }\left({\sum }_{r=0}^K{\beta }^r\right)=\frac{d}{d\beta }\left(\frac{1-{\beta }^{K+1}}{1-\beta }\right)=\frac{1-(K+1){\beta }^K+K{\beta }^{K+1}}{(1-\beta {)}^2}$

• ​​代入γ表达式​​：

$\gamma =(1-\beta )\cdot \frac{1-(K+1){\beta }^K+K{\beta }^{K+1}}{(1-\beta {)}^2}+K{\beta }^K=\frac{1-(K+1){\beta }^K+K{\beta }^{K+1}}{1-\beta }+K{\beta }^K$

• 化简​​：

$\gamma =\frac{1-{\beta }^K}{1-\beta }$

​​物理意义​​：

• 当 $K\to \infty$时，$\gamma \to \frac{1}{1-\beta }=\frac{1}{\alpha }$（理想无限长草稿）。
• 例如 $\beta$ = 0.8` 时，${\gamma }_{\text{max}}=5$，即平均每次生成5个token。

得证

Walltime的时间优化

​​定理 3.8​​：算法 1 在总运行时间上的预期改进因子为
$‘\frac{1-{\alpha }^{\gamma +1}}{(1-\alpha )(\gamma c+1)}‘$

​​证明​​：
记运行目标模型 $M_p$ ​​单步​​的成本为 $T$。
算法 1 的​​单次运行成本​​为 $Tc\gamma +T$（其中 $c\gamma T$用于运行近似模型 $M_q$ $\gamma$ 次，$T$ 用于运行 $M_p$ 一次）。
根据单次算法期望能生成的token算法推导，单次运行​​平均生成 token 数量​​为 $\frac{1-{\alpha }^{\gamma +1}}{1-\alpha }$。
因此，使用算法 1 生成单个 token 的​​总体预期成本​​为：
$\frac{(c\gamma +1)(1-\alpha )}{1-{\alpha }^{\gamma +1}}T‘$
由于标准解码算法生成单个 token 的成本为 T，
比较可得上述改进因子。∎
（注：符号 “∎” 表示证明结束）

---

关键术语说明：

英文术语	中文翻译	符号	含义
walltime	总运行时间	-	算法从启动到结束的时钟时间
expected improvement factor	预期改进因子	-	优化后时间开销的缩减比例
cost per step	单步成本	$T$	目标模型 $M_p$ 推理一个 token 的时间
approximation model	近似模型	$M_q$	快速但低精度的草稿模型
tokens	标记（Token）	-	模型生成的基本文本单位
rejection rate	拒绝率	$\alpha$	草稿模型 $M_q$ 的 token 被目标模型 $M_p$ 拒绝的概率
$\gamma$	生成长度	$\gamma$	草稿模型单次运行的 token 生成数
cost ratio	成本比	$c$	$M_q$ 与 $M_p$ 的单步时间比值（$0<c<1$）

公式解析：

1. ​​改进因子​​
   $\frac{1-{\alpha }^{\gamma +1}}{(1-\alpha )(\gamma c+1)}$

• ​​分子​​ $1-{\alpha }^{\gamma +1}$：草稿模型连续生成 \gamma 个 token 均未被拒绝的概率补偿
• ​​分母​​ $(1-\alpha )$：单 token 接受率，$\gamma c+1$：草稿+验证的总时间成本

该值 ​​>1​​ 时表示加速，值越大加速效果越显著

2. ​​单 token 成本公式​​
   $\frac{(c\gamma +1)(1-\alpha )}{1-{\alpha }^{\gamma +1}}T$

• ​​分子​​ $(c\gamma +1)(1-\alpha )T$：草稿生成+验证的实际计算量
• ​​分母​​ $1-{\alpha }^{\gamma +1}$：有效 token 产出的概率加权

操作数计算

操作数的计算量也是类似的，直接贴结论了

$$\frac{(1-\alpha )(\gamma \hat{c}+\gamma +1)}{1-{\alpha }^{\gamma +1}}$$

采样和原分布的等价性证明

参考https://arxiv.org/pdf/2302.01318
其中需要一步代换证明下面两个公式等价：

原始公式

第一个公式：
$$=1-\sum _{x^{\prime }}\min \left(p\left(x^{\prime }\right),q\left(x^{\prime }\right)\right)$$

第二个公式：
$$=\sum _{x^{\prime }}\max \left(0,q\left(x^{\prime }\right)-p\left(x^{\prime }\right)\right)$$

推导步骤

步骤 1: 应用 min 函数的恒等式

对于任何两个实数 $a$ 和 $b$，都存在以下恒等关系：
$$\min (a,b)=a-\max (0,a-b)$$

令 $b=p(x^{\prime })$，$a=q(x^{\prime })$，得到：
$$\min (p(x^{\prime }),q(x^{\prime }))=q(x^{\prime })-\max (0,q(x^{\prime })-p(x^{\prime }))$$

步骤 2: 代入第一个公式

将恒等式代入原始公式：
$$\begin{array}{rl}& 1-\sum _{x^{\prime }}\min (p(x^{\prime }),q(x^{\prime }))\\ & =1-\sum _{x^{\prime }}\left[q(x^{\prime })-\max (0,q(x^{\prime })-p(x^{\prime }))\right]\end{array}$$

步骤 3: 拆分求和运算

将求和符号分配到表达式内部：
$$=1-\left[\sum _{x^{\prime }}p(x^{\prime })-\sum _{x^{\prime }}\max (0,p(x^{\prime })-q(x^{\prime }))\right]$$
$$=1-\sum _{x^{\prime }}q(x^{\prime })+\sum _{x^{\prime }}\max (0,q(x^{\prime })-p(x^{\prime }))$$

步骤 4: 应用概率分布性质

因为 $p$ 和 $q$ 都是概率分布函数，满足：
$$\sum _{x^{\prime }}p(x^{\prime })=1\text{和}\sum _{x^{\prime }}q(x^{\prime })=1$$

代入表达式：
$$=1-1+\sum _{x^{\prime }}\max (0,q(x^{\prime })-p(x^{\prime }))$$
$$=\sum _{x^{\prime }}\max (0,q(x^{\prime })-p(x^{\prime }))$$

得证

Reference

https://arxiv.org/pdf/2211.17192