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

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

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

EAGLE 具有以下特點(diǎn)：

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

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

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

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

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

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

EAGLE——一種增強(qiáng)大語言模型生成效率的方法

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

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

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

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

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

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

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

多輪投機(jī)采樣

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

更多實(shí)驗(yàn)結(jié)果

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

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