LLM推理加速新范式！推測(cè)解碼（Speculative Decoding）最新綜述

1. Background

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

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

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

圖1: 自回歸解碼（左），推測(cè)解碼（右）

2. Speculative Decoding

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

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

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

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

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

圖2: 推測(cè)解碼研究思路的演化

3. Key Facets of Speculative Decoding

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

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

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

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

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

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

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

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

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

4. “推測(cè)”的高效性和準(zhǔn)確性

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

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

4.1 Independent Drafting

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

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

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

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

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

4.2 Self-Drafting

然而，采用一個(gè)獨(dú)立的“推測(cè)”模型也有缺點(diǎn)：

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

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

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

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

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

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

5. 驗(yàn)證策略的選擇

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

讓我們重新回顧推測(cè)解碼的驗(yàn)證過程：

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

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

圖6：recap of Speculative Decoding

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

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

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

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

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

6. 總結(jié)

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

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