大語言模型（LLM）被越來越多應(yīng)用于各種領(lǐng)域。然而，它們的文本生成過程既昂貴又緩慢。這種低效率歸因于自回歸解碼的運算規(guī)則：每個詞（token）的生成都需要進行一次前向傳播，需要訪問數(shù)十億至數(shù)千億參數(shù)的 LLM。這導(dǎo)致傳統(tǒng)自回歸解碼的速度較慢。

近日，滑鐵盧大學、加拿大向量研究院、北京大學等機構(gòu)聯(lián)合發(fā)布 EAGLE，旨在提升大語言模型的推理速度，同時保證模型輸出文本的分布一致。這種方法外推 LLM 的第二頂層特征向量，能夠顯著提升生成效率。

• 技術(shù)報告：https://sites.google.com/view/eagle-llm
• 代碼（支持商用 Apache 2.0）：https://github.com/SafeAILab/EAGLE

EAGLE 具有以下特點：

• 比普通自回歸解碼（13B）快 3 倍；
• 比 Lookahead 解碼（13B）快 2 倍；
• 比 Medusa 解碼（13B）快 1.6 倍；
• 可以證明在生成文本的分布上與普通解碼保持一致；
• 可以在 RTX 3090 上進行訓練（1-2 天內(nèi)）和測試；
• 可以與 vLLM、DeepSpeed、Mamba、FlashAttention、量化和硬件優(yōu)化等其他平行技術(shù)結(jié)合使用。

加速自回歸解碼的一種方法是投機采樣（speculative sampling）。這種技術(shù)使用一個更小的草稿模型，通過標準自回歸生成來猜測接下來的多個詞。隨后，原始 LLM 并行驗證這些猜測的詞（只需要進行一次前向傳播進行驗證）。如果草稿模型準確預(yù)測了 α 詞，原始 LLM 的一次前向傳播就可以生成 α+1 個詞。

在投機采樣中，草稿模型的任務(wù)是基于當前詞序列預(yù)測下一個詞。使用一個參數(shù)數(shù)量顯著更少的模型完成這個任務(wù)極具挑戰(zhàn)性，通常會產(chǎn)生次優(yōu)結(jié)果。此外，標準投機采樣方法中的草稿模型獨立預(yù)測下一個詞而不利用原始 LLM 提取的豐富語義信息，導(dǎo)致潛在的效率低下。

這個局限啟發(fā)了 EAGLE 的開發(fā)。EAGLE 利用原始 LLM 提取的上下文特征（即模型第二頂層輸出的特征向量）。EAGLE 建立在以下第一性原理之上：

特征向量序列是可壓縮的，所以根據(jù)前面的特征向量預(yù)測后續(xù)特征向量比較容易。

EAGLE 訓練了一個輕量級插件，稱為自回歸頭（Auto-regression Head），與詞嵌入層一起，基于當前特征序列從原始模型的第二頂層預(yù)測下一個特征。然后使用原始 LLM 的凍結(jié)分類頭來預(yù)測下一個詞。特征比詞序列包含更多信息，使得回歸特征的任務(wù)比預(yù)測詞的任務(wù)簡單得多。總之，EAGLE 在特征層面上進行外推，使用一個小型自回歸頭，然后利用凍結(jié)的分類頭生成預(yù)測的詞序列。與投機采樣、Medusa 和 Lookahead 等類似的工作一致，EAGLE 關(guān)注的是每次提示推理的延遲，而不是整體系統(tǒng)吞吐量。

EAGLE——一種增強大語言模型生成效率的方法

上圖顯示了 EAGLE 與標準投機采樣、Medusa 以及 Lookahead 關(guān)于輸入輸出的區(qū)別。下圖展示了 EAGLE 的工作流程。在原始 LLM 的前向過程中，EAGLE 從第二頂層收集特征。自回歸頭以這些特征以及此前生成的詞的詞嵌入作為輸入，開始猜下一個詞。隨后，使用凍結(jié)的分類頭（LM Head）確定下一個詞的分布，使 EAGLE 能夠從這個分布中進行采樣。通過多次重復(fù)采樣，EAGLE 進行了類似樹狀的生成過程，如下圖右側(cè)所示。在這個例子中，EAGLE 的三次前向傳播 “猜” 出了 10 個詞組成的樹。

EAGLE 使用輕量級的自回歸頭來預(yù)測原始 LLM 的特征。為了確保生成文本分布的一致性，EAGLE 隨后驗證預(yù)測的樹狀結(jié)構(gòu)。這個驗證過程可以使用一次前向傳播完成。通過這個預(yù)測和驗證的循環(huán)，EAGLE 能夠快速生成文本詞。

訓練自回歸頭代價很小。EAGLE 使用 ShareGPT 數(shù)據(jù)集進行訓練，該數(shù)據(jù)集包含不到 70,000 輪對話。自回歸頭的可訓練參數(shù)數(shù)量也很少。如上圖中的藍色部分所示，大多數(shù)組件都是凍結(jié)的。唯一要額外訓練的是自回歸頭，這是一個單層 Transformer 結(jié)構(gòu)，具有 0.24B-0.99B 參數(shù)。即使是 GPU 資源不足的情況下，也可以訓練自回歸頭。例如，Vicuna 33B 的自回歸頭可以在 8 卡 RTX 3090 服務(wù)器上在 24 小時內(nèi)完成訓練。

為什么使用詞嵌入來預(yù)測特征？

Medusa 僅使用第二頂層的特征來預(yù)測下一個詞，下下個詞......與 Medusa 不同，EAGLE 還動態(tài)地將當前采樣得到的詞嵌入作為自回歸頭輸入的一部分來進行預(yù)測。這額外的信息幫助 EAGLE 處理抽樣過程中不可避免的隨機性。考慮下圖中的例子，假設(shè)提示詞是 “I”。LLM 給出了 “I” 后面跟著 “am” 或 “always” 的概率。Medusa 不考慮是抽樣了 “am” 還是 “always”，直接預(yù)測 “I” 下下個詞的概率。因此，Medusa 的目標是，在只給定 “I” 的基礎(chǔ)上，預(yù)測 “I am” 或 “I always” 的下一個詞。由于抽樣過程的隨機性，Medusa 的相同輸入 “I” 可能有不同的下下個詞輸出 “ready” 或 “begin”，導(dǎo)致輸入和輸出之間缺乏一致的映射。相比之下，EAGLE 的輸入包括了抽樣結(jié)果的詞嵌入，確保了輸入和輸出之間的一致映射。這種區(qū)別使 EAGLE 能夠考慮抽樣過程建立的上下文，進而更準確地預(yù)測后續(xù)詞。

樹狀生成結(jié)構(gòu)

與投機采樣、Lookahead 和 Medusa 等其他猜測 - 驗證框架不同，EAGLE 在 “猜詞” 階段采用類似樹狀的生成結(jié)構(gòu)，進而實現(xiàn)了更高的解碼效率。如圖所示，標準投機采樣和 Lookahead 的生成過程是線性或鏈式的。Medusa 的方法由于在猜測階段無法構(gòu)建上下文，故通過笛卡爾積生成樹，導(dǎo)致相鄰層之間形成全連接圖。這種方法經(jīng)常導(dǎo)致無意義的組合，例如 “I am begin”。對比之下，EAGLE 創(chuàng)建了一個更稀疏的樹結(jié)構(gòu)。這種稀疏的樹結(jié)構(gòu)防止形成無意義的序列，將計算資源集中在更合理的詞組合上。

多輪投機采樣

標準投機采樣方法在進行 “猜詞” 的過程中保持了分布的一致性。為了適應(yīng)樹狀猜詞場景，EAGLE 將這種方法擴展成了多輪遞歸形式。下面呈現(xiàn)了多輪投機采樣的偽代碼。在樹狀生成過程中，EAGLE 記錄了每個抽樣詞對應(yīng)的概率。通過多輪投機采樣，EAGLE 確保最終生成的每個詞的分布與原始 LLM 的分布保持一致。

更多實驗結(jié)果

下圖展示了 EAGLE 在 Vicuna 33B 上關(guān)于不同任務(wù)中的加速效果。涉及大量固定模板的 “編程”（coding）任務(wù)顯示出最佳的加速性能。

歡迎大家體驗 EAGLE，并通過 GitHub issue 反饋建議：https://github.com/SafeAILab/EAGLE/issues