大模型推理效率無損提升3倍，滑鐵盧大學、北京大學等機構發布EAGLE
大語言模型（LLM）被越來越多應用于各種領域。然而，它們的文本生成過程既昂貴又緩慢。這種低效率歸因于自回歸解碼的運算規則：每個詞（token）的生成都需要進行一次前向傳播，需要訪問數十億至數千億參數的 LLM。這導致傳統自回歸解碼的速度較慢。
近日，滑鐵盧大學、加拿大向量研究院、北京大學等機構聯合發布 EAGLE，旨在提升大語言模型的推理速度，同時保證模型輸出文本的分布一致。這種方法外推 LLM 的第二頂層特征向量，能夠顯著提升生成效率。
- 技術報告：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 天內）和測試；
- 可以與 vLLM、DeepSpeed、Mamba、FlashAttention、量化和硬件優化等其他平行技術結合使用。
加速自回歸解碼的一種方法是投機采樣（speculative sampling）。這種技術使用一個更小的草稿模型，通過標準自回歸生成來猜測接下來的多個詞。隨后，原始 LLM 并行驗證這些猜測的詞（只需要進行一次前向傳播進行驗證）。如果草稿模型準確預測了 α 詞，原始 LLM 的一次前向傳播就可以生成 α+1 個詞。
在投機采樣中，草稿模型的任務是基于當前詞序列預測下一個詞。使用一個參數數量顯著更少的模型完成這個任務極具挑戰性，通常會產生次優結果。此外，標準投機采樣方法中的草稿模型獨立預測下一個詞而不利用原始 LLM 提取的豐富語義信息，導致潛在的效率低下。
這個局限啟發了 EAGLE 的開發。EAGLE 利用原始 LLM 提取的上下文特征（即模型第二頂層輸出的特征向量）。EAGLE 建立在以下第一性原理之上：
特征向量序列是可壓縮的，所以根據前面的特征向量預測后續特征向量比較容易。
EAGLE 訓練了一個輕量級插件，稱為自回歸頭（Auto-regression Head），與詞嵌入層一起，基于當前特征序列從原始模型的第二頂層預測下一個特征。然后使用原始 LLM 的凍結分類頭來預測下一個詞。特征比詞序列包含更多信息，使得回歸特征的任務比預測詞的任務簡單得多?？傊?，EAGLE 在特征層面上進行外推，使用一個小型自回歸頭，然后利用凍結的分類頭生成預測的詞序列。與投機采樣、Medusa 和 Lookahead 等類似的工作一致，EAGLE 關注的是每次提示推理的延遲，而不是整體系統吞吐量。
EAGLE——一種增強大語言模型生成效率的方法
上圖顯示了 EAGLE 與標準投機采樣、Medusa 以及 Lookahead 關于輸入輸出的區別。下圖展示了 EAGLE 的工作流程。在原始 LLM 的前向過程中，EAGLE 從第二頂層收集特征。自回歸頭以這些特征以及此前生成的詞的詞嵌入作為輸入，開始猜下一個詞。隨后，使用凍結的分類頭（LM Head）確定下一個詞的分布，使 EAGLE 能夠從這個分布中進行采樣。通過多次重復采樣，EAGLE 進行了類似樹狀的生成過程，如下圖右側所示。在這個例子中，EAGLE 的三次前向傳播 “猜” 出了 10 個詞組成的樹。
EAGLE 使用輕量級的自回歸頭來預測原始 LLM 的特征。為了確保生成文本分布的一致性，EAGLE 隨后驗證預測的樹狀結構。這個驗證過程可以使用一次前向傳播完成。通過這個預測和驗證的循環，EAGLE 能夠快速生成文本詞。
訓練自回歸頭代價很小。EAGLE 使用 ShareGPT 數據集進行訓練，該數據集包含不到 70,000 輪對話。自回歸頭的可訓練參數數量也很少。如上圖中的藍色部分所示，大多數組件都是凍結的。唯一要額外訓練的是自回歸頭，這是一個單層 Transformer 結構，具有 0.24B-0.99B 參數。即使是 GPU 資源不足的情況下，也可以訓練自回歸頭。例如，Vicuna 33B 的自回歸頭可以在 8 卡 RTX 3090 服務器上在 24 小時內完成訓練。
為什么使用詞嵌入來預測特征？
Medusa 僅使用第二頂層的特征來預測下一個詞，下下個詞......與 Medusa 不同，EAGLE 還動態地將當前采樣得到的詞嵌入作為自回歸頭輸入的一部分來進行預測。這額外的信息幫助 EAGLE 處理抽樣過程中不可避免的隨機性?？紤]下圖中的例子，假設提示詞是 “I”。LLM 給出了 “I” 后面跟著 “am” 或 “always” 的概率。Medusa 不考慮是抽樣了 “am” 還是 “always”，直接預測 “I” 下下個詞的概率。因此，Medusa 的目標是，在只給定 “I” 的基礎上，預測 “I am” 或 “I always” 的下一個詞。由于抽樣過程的隨機性，Medusa 的相同輸入 “I” 可能有不同的下下個詞輸出 “ready” 或 “begin”，導致輸入和輸出之間缺乏一致的映射。相比之下，EAGLE 的輸入包括了抽樣結果的詞嵌入，確保了輸入和輸出之間的一致映射。這種區別使 EAGLE 能夠考慮抽樣過程建立的上下文，進而更準確地預測后續詞。
樹狀生成結構
與投機采樣、Lookahead 和 Medusa 等其他猜測 - 驗證框架不同，EAGLE 在 “猜詞” 階段采用類似樹狀的生成結構，進而實現了更高的解碼效率。如圖所示，標準投機采樣和 Lookahead 的生成過程是線性或鏈式的。Medusa 的方法由于在猜測階段無法構建上下文，故通過笛卡爾積生成樹，導致相鄰層之間形成全連接圖。這種方法經常導致無意義的組合，例如 “I am begin”。對比之下，EAGLE 創建了一個更稀疏的樹結構。這種稀疏的樹結構防止形成無意義的序列，將計算資源集中在更合理的詞組合上。
多輪投機采樣
標準投機采樣方法在進行 “猜詞” 的過程中保持了分布的一致性。為了適應樹狀猜詞場景，EAGLE 將這種方法擴展成了多輪遞歸形式。下面呈現了多輪投機采樣的偽代碼。在樹狀生成過程中，EAGLE 記錄了每個抽樣詞對應的概率。通過多輪投機采樣，EAGLE 確保最終生成的每個詞的分布與原始 LLM 的分布保持一致。
更多實驗結果
下圖展示了 EAGLE 在 Vicuna 33B 上關于不同任務中的加速效果。涉及大量固定模板的 “編程”（coding）任務顯示出最佳的加速性能。
歡迎大家體驗 EAGLE，并通過 GitHub issue 反饋建議：https://github.com/SafeAILab/EAGLE/issues
