無損加速最高5x，EAGLE-2讓RTX 3060的生成速度超過A100
李堉暉：北京大學智能學院碩士，受張弘揚老師和張超老師指導，研究方向為大模型加速和對齊，正在尋找25屆工作機會
魏芳蕓：微軟亞研院研究員，研究方向為具身智能、圖像生成和AI agents
張超：北京大學智能學院研究員，研究方向為計算機視覺和機器學習
張弘揚：滑鐵盧大學計算機學院、向量研究院助理教授，研究方向為LLM加速和AI安全
自回歸解碼已經成為了大語言模型（LLMs）的事實標準，大語言模型每次前向計算需要訪問它全部的參數，但只能得到一個token，導致其生成昂貴且緩慢。
今日，一篇題為《EAGLE-2: Faster Inference of Language Models with Dynamic Draft Trees》的論文提出了動態草稿樹投機采樣，依據草稿模型的置信度動態調整草稿樹的結構，最高可以將大語言模型的推理速度提高5倍，同時不改變大語言模型的輸出分布，確保無損。
- 論文鏈接：https://arxiv.org/pdf/2406.16858
- 項目鏈接：https://github.com/SafeAILab/EAGLE
- Demo鏈接：https://huggingface.co/spaces/yuhuili/EAGLE-2
EAGLE-2在多輪對話數據集MT-bench上的加速效果（上圖為貪婪生成，下圖為采樣生成）：
使用EAGLE-2,2張RTX 3060（$300）的推理速度可以超過A100（$10000）。
背景
投機采樣使用一個小的模型快速生成草稿，原始的大語言模型可以通過一次前向計算驗證草稿的正確性，將正確的草稿作為輸出，從而一次生成多個token，并確保無損。EAGLE是投機采樣的一種改進。它在更有規律的特征層面而不是token層面進行自回歸，同時輸入采樣結果（超前一個時間步的token）消除了不確定性，明顯提升了草稿模型的準確率。
到目前為止，EAGLE在第三方測試Spec-Bench（https://github.com/hemingkx/Spec-Bench/blob/main/Leaderboard.md）中排名第一。
思路
EAGLE和Medusa等方法使用靜態的草稿樹，隱式地假設草稿token的接受率和上下文無關，下面是一個簡單的例子
上文是“10+2”時，下一個token難以預測，EAGLE在這個位置添加兩個候選token以增加草稿命中率，“10+2=”和“10+2+”有一個正確即可。當上文是“10+2=”時，下一個token明顯是“1”，但是EAGLE使用靜態的草稿結構，仍然添加兩個候選“1”和“3”，“10+2=3”不可能通過大語言模型的檢查，存在浪費。EAGLE-2旨在解決這一問題，如下圖所示，當上文是“10+2=”時，EAGLE-2只增加一個候選token“1”，將節約出的token用于讓草稿樹更深，這樣“10+2=12”通過大語言模型的檢查，EAGLE-2可以一次生成更多的token。
EAGLE-2的作者們在Alpaca數據集上進行了簡單的測試，下圖顯示了不同位置的草稿token的接受率，左圖中的P1-P6代表位置，與右圖的橫軸坐標對應。實驗結果顯示，在相同的位置上的草稿token的接受率也有較大的差異，這說明了使用動態草稿樹可能取得比靜態草稿樹更好的效果。
上述例子中，EAGLE-2根據預測草稿token的難易程度決定草稿樹的結構，精確計算難易程度（接受率）需要原始大語言模型的計算結果，這違背了投機采樣減少對原始大語言模型訪問的初衷。幸運的是，EAGLE的草稿模型的置信度與接受率（難易程度）高度正相關。下圖顯示了草稿模型不同置信度區間的草稿token的平均接受率，紅色虛線連接（0,0）和（1,1）。由此可見，草稿模型的置信度可以作為接受率的有效近似。
方法
EAGLE-2包括兩個階段，擴展和重排，擴展階段加深加大草稿樹，重排階段修剪草稿樹，丟棄部分節點（token）。
為了保證無損，一個草稿token被接受的前提是它的祖先節點都被接受，所以EAGLE-2將一個節點的價值定義為它和它祖先的接受率的乘積，用置信度的乘積來近似。
在擴展階段，EAGLE-2選擇草稿樹最后一層價值最高的m個節點（token）進行擴展。這些token被送入草稿模型，然后將草稿模型的輸出作為子節點連接到輸入節點，加深加大草稿樹。在重排階段，EAGLE-2按照價值對整棵草稿樹進行重排序，保留前n個節點（token）。草稿token的置信度在0-1之間，兩個節點價值相同時優先保留淺層節點，因此重排后保留的草稿樹一定是連通的，保證了語義上的連貫性。重排后草稿樹變小，降低了原始大語言模型驗證的計算量。為了保證計算結果的正確性，還需要調整attention mask，確保每一個token只能看到它的祖先節點，不受其他分支的影響。下面是一個簡單的例子。
擴展（Expand）階段的黃色框表示被選中進行擴展的節點，綠色框為以這些節點為輸入時草稿模型的預測。重排（Rerank）階段的藍色框表示被保留的節點，之后它們被展平成一維作為原始大語言模型的輸入。EAGLE-2根據樹的結構調整attention mask，比如，”a”只能看到它的祖先“It”和“is”，看不到另一個分支的“has”。EAGLE-2也同時調整位置編碼，確保和標準自回歸解碼的一致性。
實驗
EAGLE-2在多輪對話、代碼、數學推理、指令遵循、問答、總結六項任務上分別使用MT-bench、Humaneval、GSM8K、Alpaca、CNN/DM、Natural Questions數據集進行了實驗，與6種先進的投機采樣方法（SpS、PLD、Medusa、Lookahead、Hydra、EAGLE）進行了比較。
表格中的Speedup為加速比，τ 為平均接受長度，也就是原始大語言模型每次前向計算能生成的token數。EAGLE-2每次前向計算能生成大約4-5個token，而自回歸解碼每次生成1個token，因此EAGLE-2明顯加速了大語言模型的生成，加速比為2.5x-5x。加速比和接受長度在代碼生成任務（Humaneval數據集）上最高，這是因為代碼中存在大量確定性的模板，草稿更容易命中。在所有任務和大語言模型上，EAGLE-2的加速比和平均接受長度都是最高的，明顯優于其他方法。
應用
EAGLE-2也在工業界得到應用，集成至Intel/intel-extension-for-transformers等。
