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～p(x)$。

給定通過在條件前綴上運行 $M_q$ 獲得的分布 $q(x)$，我們可以采樣一個標記 $x_1～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?P(r)}}+\underset{\text{全接受時生成K個token}}{\underbrace{K?{\beta }^K}}$$

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

$$\gamma =\sum _{r=1}^{K}r?(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+?+{\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 )?\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