李堉暉：北京大學(xué)智能學(xué)院碩士，受張弘揚(yáng)老師和張超老師指導(dǎo)，研究方向為大模型加速和對齊，正在尋找25屆工作機(jī)會

魏芳蕓：微軟亞研院研究員，研究方向為具身智能、圖像生成和AI agents

張超：北京大學(xué)智能學(xué)院研究員，研究方向為計算機(jī)視覺和機(jī)器學(xué)習(xí)

張弘揚(yáng)：滑鐵盧大學(xué)計算機(jī)學(xué)院、向量研究院助理教授，研究方向為LLM加速和AI安全

自回歸解碼已經(jīng)成為了大語言模型（LLMs）的事實標(biāo)準(zhǔn)，大語言模型每次前向計算需要訪問它全部的參數(shù)，但只能得到一個token，導(dǎo)致其生成昂貴且緩慢。

今日，一篇題為《EAGLE-2: Faster Inference of Language Models with Dynamic Draft Trees》的論文提出了動態(tài)草稿樹投機(jī)采樣，依據(jù)草稿模型的置信度動態(tài)調(diào)整草稿樹的結(jié)構(gòu)，最高可以將大語言模型的推理速度提高5倍，同時不改變大語言模型的輸出分布，確保無損。

• 論文鏈接：https://arxiv.org/pdf/2406.16858
• 項目鏈接：https://github.com/SafeAILab/EAGLE
• Demo鏈接：https://huggingface.co/spaces/yuhuili/EAGLE-2

EAGLE-2在多輪對話數(shù)據(jù)集MT-bench上的加速效果（上圖為貪婪生成，下圖為采樣生成）：

使用EAGLE-2,2張RTX 3060（$300）的推理速度可以超過A100（$10000）。

背景

投機(jī)采樣使用一個小的模型快速生成草稿，原始的大語言模型可以通過一次前向計算驗證草稿的正確性，將正確的草稿作為輸出，從而一次生成多個token，并確保無損。EAGLE是投機(jī)采樣的一種改進(jìn)。它在更有規(guī)律的特征層面而不是token層面進(jìn)行自回歸，同時輸入采樣結(jié)果（超前一個時間步的token）消除了不確定性，明顯提升了草稿模型的準(zhǔn)確率。

到目前為止，EAGLE在第三方測試Spec-Bench（https://github.com/hemingkx/Spec-Bench/blob/main/Leaderboard.md）中排名第一。

思路

EAGLE和Medusa等方法使用靜態(tài)的草稿樹，隱式地假設(shè)草稿token的接受率和上下文無關(guān)，下面是一個簡單的例子

上文是“10+2”時，下一個token難以預(yù)測，EAGLE在這個位置添加兩個候選token以增加草稿命中率，“10+2=”和“10+2+”有一個正確即可。當(dāng)上文是“10+2=”時，下一個token明顯是“1”，但是EAGLE使用靜態(tài)的草稿結(jié)構(gòu)，仍然添加兩個候選“1”和“3”，“10+2=3”不可能通過大語言模型的檢查，存在浪費。EAGLE-2旨在解決這一問題，如下圖所示，當(dāng)上文是“10+2=”時，EAGLE-2只增加一個候選token“1”，將節(jié)約出的token用于讓草稿樹更深，這樣“10+2=12”通過大語言模型的檢查，EAGLE-2可以一次生成更多的token。

EAGLE-2的作者們在Alpaca數(shù)據(jù)集上進(jìn)行了簡單的測試，下圖顯示了不同位置的草稿token的接受率，左圖中的P1-P6代表位置，與右圖的橫軸坐標(biāo)對應(yīng)。實驗結(jié)果顯示，在相同的位置上的草稿token的接受率也有較大的差異，這說明了使用動態(tài)草稿樹可能取得比靜態(tài)草稿樹更好的效果。

上述例子中，EAGLE-2根據(jù)預(yù)測草稿token的難易程度決定草稿樹的結(jié)構(gòu)，精確計算難易程度（接受率）需要原始大語言模型的計算結(jié)果，這違背了投機(jī)采樣減少對原始大語言模型訪問的初衷。幸運的是，EAGLE的草稿模型的置信度與接受率（難易程度）高度正相關(guān)。下圖顯示了草稿模型不同置信度區(qū)間的草稿token的平均接受率，紅色虛線連接（0,0）和（1,1）。由此可見，草稿模型的置信度可以作為接受率的有效近似。

方法

EAGLE-2包括兩個階段，擴(kuò)展和重排，擴(kuò)展階段加深加大草稿樹，重排階段修剪草稿樹，丟棄部分節(jié)點（token）。

為了保證無損，一個草稿token被接受的前提是它的祖先節(jié)點都被接受，所以EAGLE-2將一個節(jié)點的價值定義為它和它祖先的接受率的乘積，用置信度的乘積來近似。

在擴(kuò)展階段，EAGLE-2選擇草稿樹最后一層價值最高的m個節(jié)點（token）進(jìn)行擴(kuò)展。這些token被送入草稿模型，然后將草稿模型的輸出作為子節(jié)點連接到輸入節(jié)點，加深加大草稿樹。在重排階段，EAGLE-2按照價值對整棵草稿樹進(jìn)行重排序，保留前n個節(jié)點（token）。草稿token的置信度在0-1之間，兩個節(jié)點價值相同時優(yōu)先保留淺層節(jié)點，因此重排后保留的草稿樹一定是連通的，保證了語義上的連貫性。重排后草稿樹變小，降低了原始大語言模型驗證的計算量。為了保證計算結(jié)果的正確性，還需要調(diào)整attention mask，確保每一個token只能看到它的祖先節(jié)點，不受其他分支的影響。下面是一個簡單的例子。

擴(kuò)展（Expand）階段的黃色框表示被選中進(jìn)行擴(kuò)展的節(jié)點，綠色框為以這些節(jié)點為輸入時草稿模型的預(yù)測。重排（Rerank）階段的藍(lán)色框表示被保留的節(jié)點，之后它們被展平成一維作為原始大語言模型的輸入。EAGLE-2根據(jù)樹的結(jié)構(gòu)調(diào)整attention mask，比如，”a”只能看到它的祖先“It”和“is”，看不到另一個分支的“has”。EAGLE-2也同時調(diào)整位置編碼，確保和標(biāo)準(zhǔn)自回歸解碼的一致性。

實驗

EAGLE-2在多輪對話、代碼、數(shù)學(xué)推理、指令遵循、問答、總結(jié)六項任務(wù)上分別使用MT-bench、Humaneval、GSM8K、Alpaca、CNN/DM、Natural Questions數(shù)據(jù)集進(jìn)行了實驗，與6種先進(jìn)的投機(jī)采樣方法（SpS、PLD、Medusa、Lookahead、Hydra、EAGLE）進(jìn)行了比較。

表格中的Speedup為加速比，τ 為平均接受長度，也就是原始大語言模型每次前向計算能生成的token數(shù)。EAGLE-2每次前向計算能生成大約4-5個token，而自回歸解碼每次生成1個token，因此EAGLE-2明顯加速了大語言模型的生成，加速比為2.5x-5x。加速比和接受長度在代碼生成任務(wù)（Humaneval數(shù)據(jù)集）上最高，這是因為代碼中存在大量確定性的模板，草稿更容易命中。在所有任務(wù)和大語言模型上，EAGLE-2的加速比和平均接受長度都是最高的，明顯優(yōu)于其他方法。

應(yīng)用

EAGLE-2也在工業(yè)界得到應(yīng)用，集成至Intel/intel-extension-for-transformers等。