1. Background

近年來，隨著LLM (Large Language Model) 規(guī)模的逐漸增大（200M->7B->175B），LLM的推理加速技術(shù)正逐步引起NLP學界的廣泛關(guān)注。尤其是像ChatGPT[1]，Bard[2]這種線上實時交互的應(yīng)用，LLM的inference latency（推理耗時）極大程度地影響了用戶的使用體驗。那么，LLM的Latency主要來自哪里呢？

相關(guān)研究表明，LLM推理主要是受內(nèi)存帶寬限制的（memory-bandwidth bound）[3][4]-- LLM每個解碼步所用的推理時間大部分并不是用于模型的前向計算，而是消耗在了將LLM巨量的參數(shù)從GPU顯存（High-Bandwidth Memory，HBM）遷移到高速緩存（cache）上（以進行運算操作）。也就是說，LLM推理下的GPU并不是一個合格的打工人：他把每天大多數(shù)的時間都耗費在了早晚高峰堵車上，在公司沒干啥實事兒（可不就是我摸魚仙人:P）。

這個問題隨著LLM規(guī)模的增大愈發(fā)嚴重。并且，如下左圖所示，目前LLM常用的自回歸解碼（autoregressive decoding）在每個解碼步只能生成一個token。這導(dǎo)致GPU計算資源利用率低下（->每個token的生成都需要重復(fù)讀寫LLM的巨量參數(shù)），并且序列的生成時間隨著序列長度的增加而線性增加。

圖1: 自回歸解碼（左），推測解碼（右）

2. Speculative Decoding

那么，如何更好地利用GPU資源，讓它成為一個合格的打工人呢？相信大家心里已經(jīng)有答案了：把公司當作家，減少通勤次數(shù)，就可以少摸魚多打工了（淚目）。

推測解碼（Speculative Decoding），作為2023年新興的一項LLM推理加速技術(shù)，正是提出了一種類似的解決方案：通過增加每個解碼步LLM計算的并行性，減少總的解碼步數(shù)（即減少了LLM參數(shù)的反復(fù)讀寫），從而實現(xiàn)推理加速。

如上右圖所示，在每個解碼步，推測解碼首先高效地“推測”target LLM（待加速的LLM）未來多個解碼步可能生成的token，然后再用target LLM同時驗證這些token。通過驗證的token作為當前解碼步的解碼結(jié)果。如果“推測”足夠準確，推測解碼就可以在單個解碼步并行生成多個token，從而實現(xiàn)LLM推理加速。并且，使用target LLM的驗證過程可以在理論上保證解碼結(jié)果和target LLM自回歸解碼結(jié)果的完全一致[5][6]。

也就是說，推測解碼在實現(xiàn)對target LLM推理加速的同時，不損失LLM的解碼質(zhì)量。這種優(yōu)異的性質(zhì)導(dǎo)致推測解碼受到了學界和工業(yè)界的廣泛關(guān)注，從2023年初至今涌現(xiàn)了許多優(yōu)秀的研究工作和工程項目（如Assisted Generation[7]，Medusa[8]，Lookahead Decoding[9]等等）。

考慮到推測解碼領(lǐng)域2023年以來飛速的研究進展，我們撰寫了一篇系統(tǒng)性的survey，給出推測解碼的統(tǒng)一定義和通用算法，詳細介紹了推測解碼研究思路的演化，并對目前已有的研究工作進行了分類梳理。在下文中，我們將文章內(nèi)容凝練為太長不看版——分享一些關(guān)于推測解碼關(guān)鍵要素的看法，以及目前常用的研究思路，歡迎感興趣的小伙伴一起討論～

圖2: 推測解碼研究思路的演化

3. Key Facets of Speculative Decoding

首先，我們總結(jié)推測解碼的定義：

推測解碼是一種“先推測后驗證” (Draft-then-Verify) 的解碼算法：在每個解碼步，該算法首先高效地“推測”target LLM未來多個解碼步的結(jié)果，然后用target LLM同時進行驗證，以加速推理。

也就是說，所有符合在每個解碼步“高效推測->并行驗證“模式的推理算法，都可以稱為是推測解碼（或其變體）。推測解碼實現(xiàn)加速的關(guān)鍵要素，主要在于如下三點：

相比于生成單一token，LLM并行計算額外引入的latency很小，甚至可以忽略；

“推測”的高效性&準確性：如何又快又準地“推測”LLM未來多個解碼步的生成結(jié)果；

“驗證“策略的選擇：如何在確保質(zhì)量的同時，讓盡可能多的“推測”token通過驗證，提高解碼并行性。

如上文所述，LLM推理的主要latency瓶頸在于推理過程中參數(shù)的反復(fù)讀寫。在只考慮一個解碼步的情況下，decoder-only LLM的forward latency主要和decoder層數(shù)有關(guān)——層數(shù)越深，推理時間越長。相比于這兩者，LLM運算并行性帶來的額外latency很小，這一點在非自回歸解碼的多個相關(guān)工作中有所討論[10][11]。

因此，推測解碼的算法設(shè)計主要考慮如下兩點：“推測”（Drafting）的高效性和準確性，以及“驗證“策略（Verification）的選擇：

圖3: 推測解碼相關(guān)研究的歸納分類

4. “推測”的高效性和準確性

“推測“階段（Drafting）的目的是精準地“預(yù)測”LLM未來多個解碼步的生成結(jié)果，且不引入過多的latency。

因此，“推測”階段的設(shè)計聚焦在“推測精度（accuracy）”和“推測耗時（latency）“的權(quán)衡上。一般來說，用以推測的模型越大，推測精度越高（即通過驗證的token越多），但是推測階段的耗時越大。如何在這兩者之間達到權(quán)衡，使得推測解碼總的加速比較高，是推測階段主要關(guān)注的問題。

4.1 Independent Drafting

最簡單的Drafting思路是，拿一個跟target LLM同系列的smaller LM進行“推測”[12][13]。比如OPT-70B的加速可以用OPT-125M進行推測，T5-XXL可以用T5-small。這樣的好處是可以直接利用現(xiàn)有的模型資源，無需進行額外的訓練。而且，由于同系列的模型使用相近的模型結(jié)構(gòu)、分詞方法、訓練語料和訓練流程，小模型本身就存在一定的和target LLM之間的“行為相似性“（behavior alignment），適合用來作為高效的“推測“模型。

圖4: https://huggingface.co/blog/assisted-generation

這一思路由Google和Deepmind同時提出[12][13]。作為Speculative Decoding的早期探索，這種“推測”思路易于實踐和部署。并且，這兩篇工作同時在理論上證明了推測解碼不僅支持greedy decoding，還支持nucleus sampling的無損加速（我們下文會講到）。這兩種解碼策略涵蓋了LLM應(yīng)用的大多數(shù)場景。因此，這兩篇工作極大地促進推測解碼在LLM推理加速中的應(yīng)用，吸引了工業(yè)界和學術(shù)界的大量關(guān)注。

然而，同系列小模型的“推測”精度還有提升空間嗎？

顯然是有的。最直接的思路，就是去增強小模型和大模型之間的“行為相似性”（behavior alignment），讓小模型模仿得“更像”一些。目前在這方面的研究進展集中在知識蒸餾（knowledge distillation）上：將target LLM作為教師模型，小模型作為學生模型，通過知識蒸餾讓小模型更加趨向于target LLM的預(yù)測行為[14][15]。并且，知識蒸餾還可以有效地增強小模型的生成質(zhì)量，通過減少低級的預(yù)測錯誤，增加通過驗證的token數(shù)量。

4.2 Self-Drafting

然而，采用一個獨立的“推測”模型也有缺點：

首先，并不是所有的LLM都能找到現(xiàn)成的小模型，比如LLaMA-7B。重新訓練一個小模型需要較多的額外投入。

另外，引入一個額外的小模型增加了推理過程的計算復(fù)雜度，尤其不利于分布式部署場景。

因此，相關(guān)研究工作提出利用target LLM自己進行“高效推測”。比如Blockwise Decoding[5]和Medusa[8]在target LLM最后一層decoder layer之上引入了多個額外的FFN Heads（如下所示），使得模型可以在每個解碼步并行生成多個token，作為“推測”結(jié)果。

圖5: https://sites.google.com/view/medusa-llm

然而，這些FFN Heads依然需要進行額外的訓練。除了這兩個工作，還有一些研究提出利用Early-Existing或者Layer-Skipping來進行“高效推測“[16][17]，甚至僅僅是在模型輸入的最后插入多個[PAD] token，從而實現(xiàn)并行的“推測”[18]。然而，“部署的便捷性”和“推測精度”之間依然存在一定的權(quán)衡關(guān)系。如何選擇合適的“推測”策略，達到令人滿意的加速效果，就見仁見智了。

感興趣的友友可以移步具體論文查看細節(jié)，我們后續(xù)也準備提供一個公平的加速評測，給大家提供一個參考～

5. 驗證策略的選擇

“驗證“階段（Verification）的首要目的是保證解碼結(jié)果的質(zhì)量。

讓我們重新回顧推測解碼的驗證過程：

如下圖所示，在給定“草稿”（即推測結(jié)果）時，LLM的并行驗證其實和訓練階段teacher-forcing的形式是一致的——在生成每個token時，都假設(shè)LLM的前綴輸入是正確的。比如，在驗證第三個“推測”token時，LLM以綠色前綴和兩個黃色的"推測“token作為前綴輸入。以貪婪解碼（greedy decoding）為例，以該前綴作為輸入時，LLM會自己生成一個概率最大的token。如果這個token（綠色）和第三個“推測”token相同，就說明第三個“推測”token通過了“驗證”——這個token本來就是LLM自己會生成的結(jié)果。

因此，第一個沒有通過驗證的“推測”token （圖中的紅色token）后續(xù)的“推測”token都將被丟棄。因為這個紅色token不是LLM自己會生成的結(jié)果，那么前綴正確性假設(shè)就被打破，這些后續(xù)token的驗證都無法保證前綴輸入是“正確”的了。

圖6：recap of Speculative Decoding

由此可見，推測解碼是可以保證最終解碼結(jié)果和target LLM原先的貪婪解碼結(jié)果完全一致的。因此，貪婪解碼經(jīng)常被用于推測解碼的demo展示[8]，用以清晰直觀地表示推測解碼在保持和target LLM解碼結(jié)果等同的前提下，實現(xiàn)了數(shù)倍的推理加速。

然而，嚴格要求和target LLM解碼結(jié)果完全匹配（exact-match）是最好的策略嗎？

顯然，并不是所有概率最大的token都是最合適的解碼結(jié)果（比如beam search）。當推測模型的性能較好時，嚴格要求和target LLM結(jié)果匹配會導(dǎo)致大量高質(zhì)量的“推測”token被丟棄，僅僅是因為它們和target LLM top-1解碼結(jié)果不一致。這導(dǎo)致通過驗證的“推測”token數(shù)量較小，從而影響推測解碼的加速比。

因此，有一些工作提出可以適當?shù)胤潘伞膀炞C”要求，使得更多高質(zhì)量的“推測”token被接受，增大每個解碼步通過驗證的“推測”token數(shù)量，進一步提升加速比[12][14][15]。

除了支持貪婪解碼，推測解碼還可以在理論上保障和target LLM nucleus sampling的分布相同[12][13]，具體證明感興趣的朋友可以查看相關(guān)paper～。另外，相比于只驗證單一的“推測”序列，相關(guān)研究還提出可以讓target LLM并行驗證多條“推測”序列，從而進一步增大通過驗證的“推測”token數(shù)量[19]。

6. 總結(jié)

表1: 推測解碼算法總結(jié)

在上表中，我們給出目前常用的推測解碼算法的總結(jié)～。作為一種新興的推理加速算法，推測解碼在實現(xiàn)對target LLM推理加速的同時保障了解碼結(jié)果的質(zhì)量，具有廣闊的應(yīng)用前景和極大的科研潛力，個人比較看好～。然而，推測解碼研究本身也存在許多尚未解答的問題，比如如何更好地實現(xiàn)target LLM和“推測”模型之間的行為對齊、如何結(jié)合具體任務(wù)的特點設(shè)計相應(yīng)的推測解碼策略（比如多模態(tài)模型加速），都是值得思考的問題。