
自回歸解碼已然成為大語(yǔ)言模型的推理標(biāo)準(zhǔn)。大語(yǔ)言模型每次前向計(jì)算需要訪問(wèn)它全部的參數(shù),但只能得到一個(gè) token,導(dǎo)致其生成昂貴且緩慢。
近日,EAGLE 團(tuán)隊(duì)的新作《EAGLE-3: Scaling up Inference Acceleration of Large Language Models via Training-Time Test》通過(guò)一系列優(yōu)化解鎖了投機(jī)采樣的 Scaling Law 能力,可以將大語(yǔ)言模型的推理速度提高 6.5 倍,同時(shí)不改變大語(yǔ)言模型的輸出分布,確保無(wú)損。同時(shí),隨著訓(xùn)練數(shù)據(jù)的增加,加速比還能進(jìn)一步提高。

- 論文標(biāo)題:EAGLE-3: Scaling up Inference Acceleration of Large Language Models via Training-Time Test
- arXiv 地址:https://arxiv.org/abs/2503.01840
- 項(xiàng)目地址:https://github.com/SafeAILab/EAGLE
- SGLang 版本:https://github.com/sgl-project/sglang/pull/4247
EAGLE-3 的加速效果(DeepSeek-R1-Distill-LLaMA 8B 在數(shù)學(xué)推理數(shù)據(jù)集 GSM8K 上測(cè)試,其他模型在多輪對(duì)話數(shù)據(jù)集 MT-bench 上測(cè)試):

不同方法的生成速度對(duì)比:
背景
投機(jī)采樣使用一個(gè)小的模型快速生成草稿,一次生成多個(gè) token。目標(biāo)大語(yǔ)言模型通過(guò)一次前向計(jì)算并行驗(yàn)證草稿的正確性,輸出正確的草稿,并確保無(wú)損。EAGLE 系列是投機(jī)采樣的最快實(shí)現(xiàn)。
EAGLE-1 在更有規(guī)律的特征層面而不是 token 層面進(jìn)行自回歸,同時(shí)輸入采樣結(jié)果(超前一個(gè)時(shí)間步的 token)消除了不確定性,明顯提升了草稿模型的準(zhǔn)確率。EAGLE-2 利用草稿模型的置信度近似接受率,據(jù)此動(dòng)態(tài)地調(diào)整草稿樹的結(jié)構(gòu),進(jìn)一步提升了投機(jī)采樣的效率。
機(jī)器之心之前已經(jīng)報(bào)道了 EAGLE-1 和 EAGLE-2 的工作:
- EAGLE-1:大模型推理效率無(wú)損提升3倍,滑鐵盧大學(xué)、北京大學(xué)等機(jī)構(gòu)發(fā)布EAGLE
- EAGLE-2:無(wú)損加速最高5x,EAGLE-2讓RTX 3060的生成速度超過(guò)A100
最新的大模型通過(guò)使用越來(lái)越多的訓(xùn)練數(shù)據(jù)以取得更好的性能。比如說(shuō),對(duì)于 LLaMA 系列 7B(8B)大小的模型,LLaMA 1、LLaMA 2 和 LLaMA 3 分別使用了 1T、2T、15T token 訓(xùn)練數(shù)據(jù),模型結(jié)構(gòu)和推理成本基本不變的前提下各項(xiàng)指標(biāo)取得了明顯提升。
EAGLE-3 的作者們希望通過(guò)增加草稿模型的訓(xùn)練數(shù)據(jù)量以提高接受率和加速比(藍(lán)色曲線)。遺憾的是, EAGLE-1 從訓(xùn)練數(shù)據(jù)增加中得到的提升非常有限(紅色曲線)。
EAGLE-1 和 EAGLE-3 的加速比隨著訓(xùn)練數(shù)據(jù)的增加而變化的趨勢(shì):

動(dòng)機(jī)
EAGLE 在特征層進(jìn)行自回歸,再使用目標(biāo)模型的分類頭得到草稿 token。這種方式利用了目標(biāo)模型的中間結(jié)果和信息,相比 token 層自回歸有更好的性能。因?yàn)椴莞迥P偷淖罱K任務(wù)是預(yù)測(cè)草稿 token,EAGLE 的損失函數(shù)包括兩部分,一部分是特征預(yù)測(cè)損失,另一部分是 token 預(yù)測(cè)損失。特征預(yù)測(cè)損失也可以被認(rèn)為是一種約束,限制了模型的表達(dá)能力。
EAGLE-3 的作者們進(jìn)行了實(shí)驗(yàn)驗(yàn)證,觀察第一個(gè)草稿 token 的接受率 0-α,在數(shù)據(jù)集較小時(shí),特征預(yù)測(cè)損失在訓(xùn)練數(shù)據(jù)較少時(shí)可以提高性能,但嚴(yán)重?fù)p害了草稿模型的 scaling up 能力。

不過(guò),不使用特征預(yù)測(cè)損失會(huì)導(dǎo)致生成后續(xù)草稿 token 時(shí)的輸入偏離訓(xùn)練分布,導(dǎo)致第二個(gè)草稿 token 的接受率 1-α 嚴(yán)重下降。為了解決這一問(wèn)題,EAGLE-3 使用 “訓(xùn)練時(shí)測(cè)試” 模擬多步生成,從而兼顧了模型的 scaling up 能力和生成多個(gè)草稿 token 的能力

EAGLE、Medusa 等投機(jī)采樣方法都重用目標(biāo)模型的最后一層特征作為草稿模型的提示,但 EAGLE-3 的作者們發(fā)現(xiàn)這存在缺陷。大語(yǔ)言模型的最后一層特征經(jīng)過(guò)線性變換就能得到下一個(gè) token 的分布。對(duì)于滿秩的分類頭,最后一層特征和下一個(gè) token 的分布一一對(duì)應(yīng)。最后一層特征只有下一個(gè) token 的信息,失去了目標(biāo)模型的全局性質(zhì)。
因此,EAGLE-3 不再使用目標(biāo)模型的最后一層特征作為輔助信息,而是混合目標(biāo)模型的低層、中層、高層信息來(lái)作為草稿模型的輸入。
方法
與其他投機(jī)采樣方法一致,EAGLE-3 交替執(zhí)行草稿階段和驗(yàn)證階段。
草稿階段以下圖為例。在預(yù)填充或者上一個(gè)驗(yàn)證階段,EAGLE-3 記錄模型的低、中、高層特征序列,分別記為 l、m 和 h,拼接 k 維向量 l、m 和 h 得到 3k 維向量,然后通過(guò)一個(gè)全連接層將其降維到 k 維,就得到融合了不同層級(jí)信息的特征 g,其中 k 為目標(biāo)模型的隱藏層維度。目標(biāo)是生成以 “How can I” 為前綴的草稿 token 序列,只輸入 g_how 和 g_can,草稿模型無(wú)法感知隨機(jī)的采樣過(guò)程。
因此 EAGLE-3 引入采樣結(jié)果 I 的詞向量嵌入 e_I。將 g 與對(duì)應(yīng)的超前一個(gè)時(shí)間步的嵌入 e 拼接,即 g_how 與 e_can 拼接,g_can 和 e_I 拼接。將拼接后的向量通過(guò)全連接層降到 k 維,并將其輸入進(jìn)一個(gè)單層的 transformer 得到 a。然后,將 a 輸入分類頭,采樣得到第一個(gè)草稿 token “do”。
在 step 1,以 “How can” 為前綴時(shí),EAGLE-3 重用了來(lái)自目標(biāo)模型的 g_how 和 g_can。
在 step 2,前綴是 “How can I”。理想的方式是重用來(lái)自目標(biāo)模型的 g_how、g_can 和 g_I。但這是不可能的,因?yàn)?token “I” 還沒(méi)有被目標(biāo)模型檢查,無(wú)法獲取 g_I。EAGLE-3 用上一個(gè) step 草稿模型的輸出 a_I 替代 g_I,拼接 a_I 與采樣結(jié)果 “do” 的詞向量嵌入作為草稿模型 step2 的輸入。
在 step 3,同樣無(wú)法獲得 g_do,所以使用 a_do 代替,拼接 a_do 與 e_it 作為草稿模型的輸入。之后的 step 同理。

實(shí)驗(yàn)
EAGLE-3 在多輪對(duì)話、代碼、數(shù)學(xué)推理、指令遵循、總結(jié)五項(xiàng)任務(wù)上分別使用 MT-bench、Humaneval、GSM8K、Alpaca、CNN/DM 數(shù)據(jù)集進(jìn)行了實(shí)驗(yàn),并與 7 種先進(jìn)的投機(jī)采樣方法(SpS、PLD、Medusa、Lookahead、Hydra、EAGLE、EAGLE-2)進(jìn)行了比較。
該實(shí)驗(yàn)分別在 Vicuna (V), LLaMA-Instruct 3.1 (L31), LLaMA-Instruct 3.3 (L33), DeepSeek-R1-Distill-LLaMA (DSL) 上進(jìn)行。

表格中的 Speedup 為加速比,τ 為平均接受長(zhǎng)度,也就是目標(biāo)模型每次前向計(jì)算能生成的 token 數(shù)。EAGLE-3 每次前向計(jì)算能生成大約 4-7 個(gè) token,而自回歸解碼每次生成 1 個(gè) token,因此 EAGLE-3 明顯加速了大語(yǔ)言模型的生成,加速比為3.1x-6.5x
在所有任務(wù)和模型上,EAGLE-3 的加速比和平均接受長(zhǎng)度都是最高的,明顯優(yōu)于其他方法。
應(yīng)用
EAGLE-3 發(fā)布第一天就被集成到 SGLang 中。在生產(chǎn)級(jí)框架中,EAGLE-3 也有數(shù)倍加速效果。以下實(shí)驗(yàn)由 SGLang 團(tuán)隊(duì)提供并以 LLaMA 3.1 8B(batch size=1, 1x H100)為例。

投機(jī)采樣往往被認(rèn)為在大 batch size 下會(huì)降低吞吐量。但是在 SGLang 這一生產(chǎn)級(jí)框架下,EAGLE-3 在 batch size 為 64 時(shí)仍可以提高 38% 的吞吐量,而 EAGLE 在 batch size 為 24 時(shí)就導(dǎo)致吞吐量下降。這里 1.00x 以 SGLang (w/o speculative) 的吞吐量作為基準(zhǔn)。以下實(shí)驗(yàn)由 SGLang 團(tuán)隊(duì)提供并以 LLaMA 3.1 8B(1x H100)為例。

作者介紹
李堉暉:北京大學(xué)智能學(xué)院碩士,滑鐵盧大學(xué)訪問(wèn)學(xué)者,受張弘揚(yáng)老師和張超老師指導(dǎo),研究方向?yàn)榇竽P图铀俸蛯?duì)齊。
魏芳蕓:微軟亞研院研究員,研究方向?yàn)榫呱碇悄?、圖像生成和 AI agents。
張超:北京大學(xué)智能學(xué)院研究員,研究方向?yàn)橛?jì)算機(jī)視覺(jué)和大模型加速。
張弘揚(yáng):滑鐵盧大學(xué)計(jì)算機(jī)學(xué)院、向量研究院助理教授,研究方向?yàn)榇竽P屯评砑铀?、AI 安全和世界模型。
熱門跟貼