AIxiv專欄是機器之心釋出學術、技術內容的欄目。過去數年，機器之心AIxiv專欄接收報導了2000多篇內容，覆蓋全球各大高校與企業的頂級實驗室，有效促進了學術交流與傳播。如果您有優秀的工作想要分享，歡迎投稿或者聯絡報導。投稿郵箱：liyazhou@jiqizhixin.com；zhaoyunfeng@jiqizhixin.com
李堉暉：北京大學智慧學院碩士，受張弘揚老師和張超老師指導，研究方向為大模型加速和對齊，正在尋找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等。