[論文閱讀] Draft & Verify: Lossless Large Language Model Acceleration via Self-Speculative Decoding

Last Updated on 2024-11-14 by Clay

本篇論文重點

量化、剪枝、蒸餾同樣可以加速，但得面對與原始模型不同的輸出分佈、重新訓練的開銷等等問題
原先的 Speculative Decoding 面對的問題則為我們需要使用額外的記憶體空間去驅動 draft model（草稿模型），而 Self-Speculative Decoding 僅使用了自身部份神經網路作為 draft model
自適應草稿脫離機制（Adaptive Draft-Exiting Mechanism）可以基於自動調整信心分數閾值來自動調整草稿模型的推測 tokens 數量

Abstract – 摘要

研究者基於推測性解碼（Speculative Decoding），提出了另外一種名為自推測性解碼（Self-Speculative Decoding）的變體。

本來的 Speculative Decoding 可以分成兩個階段：草稿（drafting）和驗證（verification）。

在草稿階段，是由一個較小、但與我們真正想要加速的目標大模型擁有相同詞彙表的 draft model 進行生成。由於 draft model 生成的速度相較 target model 快上許多，所以可以快速解碼多個 tokens 並提供這一猜測的 tokens 在各自序列位置的機率分佈。

在驗證階段，則是由 target model 對於 draft model 猜測的序列，繼續預測下一個 token —— 在預測下一個 token 的同時，也就會把過往的 tokens 的 logits 一併得出，於是我們便可以拿這些 logits 經過採樣變換為機率分佈，然後與 draft model 預測的序列一一比較進行驗證。

驗證的詳細流程可以參考我之前寫的筆記：[論文閱讀] Fast Inference from Transformers via Speculative Decoding

而 Self-Speculative Decoding 就是不再使用額外的 draft model 了，而是使用自己的部份網路當作 draft model。

Self-Speculative Decoding 的實驗方法

我跳過了論文中提到的背景知識不紀錄，因為這其實與之前看的 Speculative Decoding 在很大程度上都是相同的，便不再贅述，僅專心紀錄 Self-Speculative Decoding 的亮點。

貝氏優化跳層策略

既然我們要讓模型跳過某些層不進行計算，並讓處於跳層模式的模型當作 draft model，那麼我們就很有必要『精挑細選』我們要跳過的層。

直覺上，很明顯如果我們一層也不跳，那麼 acceptance rate（AR）幾乎會是 100%；但反之若我們把模型拔到只剩一層，雖然草稿模式的模型會推理速度超級快，但是 acceptance rate 就會降低至趨近 0%。

順帶一提，所謂的跳層，基本上是會跳過 Attention Layer 以及 MLP Layer 兩種，雖然 Attention Layer 後面就接著 MLP，但兩者可以分開進行判斷，並不用一定需要綁定一起跳或不跳。

而論文中所採用的方法，則是貝氏優化（Bayesian Optimization），這是個非常恰當的選擇，因為貝氏優化就是可以用貝氏定理的假設來搜索要怎麼樣的參數組合作為輸入，才能最大化我們的『目標函數』。

在這裡的參數組合，指的當然是跳層的組合；而目標函數（Objective Function），則就是加速推理的平均 token 推理時間（越少越好）。

自適應草稿脫離機制（Adaptive Draft-Exiting Mechanism）

在 Speculative Decoding 的驗證機制中，如果 draft model 的推測序列，只要被拒絕一個 token，則從此 token 之後所推測的 token 也都會隨之廢棄，所以讓 draft model 推測太多步反而可能有損加速推理時間。

所以簡單的一個想法是，我們替 draft model 的解碼過程加上一個閾值 $\gamma$ ，只要 draft model 解碼時的信心分數低於 $\gamma$ ，我們就停止進行推測。

$P(x_{t+1} \mid x_{1}, \ldots, x_{t}) < \gamma$

然而，靜態的 $\gamma$ 值可能在實際應用中也不甚理想，因為 target model 終究也是基於自身產出的機率分佈去以一定機率拒絕 draft model 的解碼，所以實際上可能兩方都是處於沒有什麼把握的狀態，這時候我們可能會因為過於嚴苛而禁止 draft model 進行推測性解碼。

所以，自適應的草稿脫離機制就是把信心分數閾值 $\gamma$ 值設定為一個根據 acceptance rate（AR）調整的變數。

參數定義如下：

$\alpha$ : 是一個我們期望的接受率，若低於此期望接受率則會增加草稿模型的信心閾值，需要更高的信心分數才能進行推測
$\epsilon$ : $\gamma$ 更新的步長
$AR$ : 當前的接受率，不只是這一次，通常也會考慮到之前的接受率所定義出來的當前接受率
$\tilde{\gamma}$ : 計算出來的新閾值調整方向，之後還會透過其他參數平滑化來決定最終的閾值
$\beta_{1}$ : 這一次的接受率的權重，可以比較平滑地從前次的接受率過度到當前的接受率，不會讓接受率忽高忽低
$\beta_{2}$ : 這一次的信心程度閾值，同樣是比較平滑地從前次的閾值過度到當前的閾值，不會讓閾值忽高忽低

透過以上的方式，就可以實現自適應的草稿脫離機制。

實驗結果

主要加速結果

實驗結果可說是十分令人滿意，我想也跟大家的直覺相同：越是參數量大的模型，越是能接受跳層所帶來的影響，所以 70B 量級的模型加速的幅度是最大的，來到了 1.992x 的加速！

物極必反：跳層並不是跳越多越好

然後也有另外一張圖，實際地表明了並不是跳過越多層加速效果越好：潛台詞就是接受率實在太低了，到最後甚至還減速了。

自適應草稿脫離機制實驗結果

還有非常重要的是，他們證明自適應草稿脫離機制是確實有效的。無論是在固定產生 K 個推測 Token 還是固定信心程度閾值的對比上，自適應草稿脫離機制的效果也幾乎都是數一數二好。

結語

這篇論文對我個人來說，是一篇比較有著意義的論文。因為在我第一次看到 Speculative Decoding 時（2023年9月初），我腦海中就浮現了跟 Self-Speculative Decoding 類似的想法，不過當時腦海中想像的不是貝氏優化而是基因演算法去挑選跳層組合。

而後過了一段時間，我就看到了 Self-Speculative Decoding 的研究出來了。對我來說，意外地有一種自己的想法被承認的感覺，尤其是我深知自己還有如此多的不足之處，這種感受當真彌足珍貴。

有一種，論文看得多了，就算是不成熟的自己也終於在某種程度上開始與各位研究員使用相同的視角看待研究了。

以技術的角度來說，我也覺得這項研究的價值是非常重要的，因為維護一個 draft model 和 target model 終究還是一件準備麻煩的事情，若是將來的模型都可以無痛使用自己的部份神經網路做 draft model 進行加速推理，那豈不是非常便利？

或許將來的 AI 模型都會直接應用上這樣的技術說、反過來說，搞不好不使用才是比較奇怪的也說不定。

另外我也在想，以『層』（Layer）為單位的節省是否真的是最後的答案呢？還是我們其實可以應用各種剪枝、稀疏化、跳層、共享權重等等方式來從原始 target model 中建構模型？就像米開朗基羅所說的：「雕像就在石頭裡，我只是把不要的部份去掉」。

以上，僅僅只是想想，我還是會早日趕緊把 Self-Speculative Decoding 的實作解決掉。

如果感興趣，歡迎關注我的 GitHub 加速推理實作專案：https://github.com/ccs96307/fast-llm-inference

我會把看到的許多感興趣的加速推理實作都放在此專案上，會不斷更新的，歡迎隨時討論！

References

透過貝氏優化去搜索 LayerSkip 模型的最佳跳層策略

Self-Speculative Decoding 實現: 跳層 Transformer 模型實作筆記

[論文閱讀] Draft & Verify: Lossless Large Language Model Acceleration via Self-Speculative Decoding

本篇論文重點

Abstract – 摘要