機器之心報道
機器之心編輯部
近年來,語言模型的顯著進展主要得益于大規模文本數據的可獲得性以及自回歸訓練方法的有效性。
在這種自回歸的訓練框架中,每一個 token 都作為其前文上下文的預測目標。
這一方法無需顯式標注,使得自回歸模型在訓練階段具有明顯優勢,因而成為主流范式。
然而,在推理階段,自回歸生成本質上是順序執行的,計算開銷大,因為每一步解碼都需要完整地運行一次模型。相比之下,人類在表達時常常會先在句子層面組織思路,再逐詞輸出。
那么在大模型推理階段,能否讓大模型擺脫單個 token 順序預測的瓶頸,也和人類的邏輯一樣,從更大層面考慮輸出內容呢?
假如能夠打破自回歸模型的順序執行本質,就產生了如網友所說的「具有時間跳躍能力的 LLM」
受到啟發,蘋果的研究人員為此展開思考,借助類似策略,開發了一個框架,使預訓練的自回歸大型語言模型能夠執行多 token 預測,在保持生成質量的同時,為代碼和數學任務提供高達 5.35 倍的推理加速,以及為一般任務提供約 2.5 倍的推理加速
AI 工程師 Jackson Atkins 認為這是一項具有開創性的工作,「最令人興奮的是這一切都可以通過對現有模型進行 LoRA 微調實現」。
想象一下,將 AI 運行成本削減數倍,或是在輕量設備上流暢運行強大的實時助手,這對推理優化而言是一次顛覆性的變革。
如果你手中有一個運行速度提升 5 倍的大語言模型,第一件事會做什么?
- 論文標題:Your LLM Knows the Future: Uncovering Its Multi-Token Prediction Potential
- 論文鏈接:https://www.alphaxiv.org/abs/2507.11851
首先探討一個基本問題:語言模型能否在單個推理步驟中生成多個 token?
令人鼓舞的是,答案是肯定的。
現有的推測解碼研究已經探索了這個方向以加快生成速度。推測解碼方法利用一個草稿模型生成多個 token,然后通過一個驗證器檢查它們與標準自回歸輸出的一致性。雖然這種方法提供了加速,但它仍然根本上依賴于自回歸生成。
在這項工作中,研究者們提出了一個更深層次的問題:我們能否訓練真正非自回歸的語言模型?
研究者們通過設計完全非自回歸的訓練算法來探討這一問題,例如基于擴散機制的語言模型。然而,此類方法通常需要構建全新的建模與訓練流程。
于是研究者們進一步提出一個問題:是否可以在盡可能少改動的前提下,適配現有的自回歸訓練與推理框架?目標是實現在保留自回歸模型核心優勢的同時,實現高效的多 token 生成。
為進一步論證這一點,研究者們首先觀察到,自回歸模型雖然并未明確針對未來 token 進行訓練,但實際上已經在一定程度上編碼了關于未來 token 的信息。
例如,給定提示詞 “what is two plus two?”,一個預訓練模型在標準自回歸解碼過程中通常會生成 “two plus two equals four”。為了檢驗模型是否具備對未來 token 的感知能力,研究者們在提示后添加占位 token(圖中以 ? 表示),并分析輸出的 logits,如圖 1 (左)所示。
令人驚訝的是,未來 token 的正確序列出現在前 200 個 logits 中,說明模型在某種程度上已隱含地掌握了即將生成的 token 信息
圖 1:自回歸模型能夠隱式預判未來的 token。
基于上述觀察,研究者們進一步探究能否引導模型發揮其潛在能力,從而更好地組織對未來 token 的預測。
為此,研究者在提示詞末尾引入若干mask token,并對模型進行訓練,使其能夠直接預測這些 token。如圖 1 (中)所示,經過微調后的模型能夠將正確的 token 提升至前 10 個 logits 中。
最后,為生成連貫的多 token 輸出,研究者們引入了一個輕量級的采樣模塊:一個兩層感知機(two-layer perceptron),該模塊在預測每個 token 時,結合先前已采樣出的 token 作為條件,具體結構如圖 1 (右)所示。
與之前的一些方法不同,本文訓練模型填充 mask token,以預測未來 token。該模型在推理這些 token 時,充分利用了自身的全部深度與表示能力,并結合整個序列的上下文信息,從而在效果上顯著優于現有的多 token 預測方法。
此外,得益于一種簡單而有效的技術 ——門控 LoRA 適配(gated LoRA adaptation),本方法在生成質量上無任何下降
在較為輕量的監督微調下,研究者們相較于傳統回溯式生成模型實現了加速效果。如圖 2 所示,在 Tulu3-8B 模型基礎上微調以預測 8 個額外 token,即可獲得顯著性能提升。
圖 2:在通過監督微調訓練門控 LoRA 和采樣頭后獲得的加速效果。
方法
為在盡量減少重新訓練的前提下實現多 token 生成,引入了一類特殊的 token,稱為 mask。設原始序列為 X = [x1, …, xn],基本思路是在該序列末尾附加 k 個唯一的 mask token,從而構造出擴展序列 Xm = [x1, …, xn, m1, …, mk]。其中,mask token [m1, …, mk] 的表示向量被初始化為隨機向量,并添加至模型的嵌入表中。
在本文中,將模型對標準下一個 token 的預測稱為 NTP(Next Token Prediction),而對 mask token 的預測稱為 MTP(Mask Token Prediction)。
本文所提出的 MTP 模型的整體架構如圖 3 所示,其中展示了在微調時使用 k=2 個 mask 的模型在推理階段的工作流程。在圖中 box-1(左上角),擴展后的序列 Xm 被輸入至解碼器,模型生成的潛在表示中,[z1, …, zn] 對應于 NTP token 的表示,而 [zn+1, …, zn+k] 則對應于 MTP token 的表示。
圖 3:MTP 模型的組成部分。框 1(左上)顯示了帶有門控 LoRA 參數的自回歸模型。框 2(左下)說明了采樣頭。框 3(右)展示了門控 LoRA 模塊的框圖。
圖中 框 2(左下)展示了采樣器頭(sampler head)的結構。第一個(NTP)token 通過標準的反嵌入層(unembedding layer)以自回歸方式生成,即由 z_n 預測出 y_n+1。而剩余的(MTP)token 則由采樣器模塊依次生成。在每一步中,采樣器根據 z_n+k+1 與前一個已生成的 token y_n+k 來預測 y_n+1+k,確保每個生成的 token 同時融合模型的潛在表示與先前已采樣的 token 信息。
為了在微調時保留預訓練模型的行為,研究者們在解碼器層中引入了門控 LoRA 模塊(gated LoRA)。微調過程中,僅更新 LoRA 參數與采樣器頭參數,原始解碼器權重保持凍結狀態。門控 LoRA 模塊通過對 NTP 與 MTP token 分別采取不同的計算路徑,確保微調過程不會影響 NTP token 的生成行為,如圖中 框 3(右側)所示。該差異通過引入二進制 mask 實現,二進制 mask 被作為額外輸入傳入解碼器層。
詳細的模型訓練過程請參閱原論文。
實驗
研究者們在 Tulu3-8B SFT 模型上開展了實驗。該模型屬于 LLaMA-3 系列,并基于 Tulu3 數據集進行了有監督微調。
生成質量評估
研究者們首先在微調過程中跟蹤了模型的準確率變化。由于所用模型為 SFT 模型,其對額外訓練非常敏感。圖 6 (a) 展示了模型在 ARC-Challenge 基準上的 zero-shot 準確率,該評估通過 Harness 庫完成。
如圖 6 (a) 中的虛線所示,gated LoRA 能夠保持準確率穩定,原因在于其保證了 NTP token 的輸出不會受到微調影響。
如圖 6 (b) 中所示,在標準 LoRA 與 gated LoRA 兩種訓練方式都實現了有效的收斂。
接下來,研究者們在圖 6 (c) 中分析了 NTP token 上的交叉熵損失,該指標與生成質量直接相關。如圖所示,使用標準 LoRA 訓練的模型,其 NTP 損失呈上升趨勢,表明生成質量逐漸下降。相比之下,采用 gated LoRA 的模型在整個訓練過程中保持了幾乎恒定的 NTP 損失,這是因為在 gated LoRA 機制下,梯度不會在 NTP token 上傳播,從而有效避免了對原始生成能力的干擾。
圖 6:使用標準 LoRA 和門控 LoRA 適配器訓練的模型收斂性分析。
加速能力分析
為了評估所提方法在生成速度上的提升,研究者們采用了自推測解碼算法(self-speculative decoding)。
實驗最多運行 100 個生成步驟,若模型在第 100 步之前生成了句子結束符(end-of-sentence token),則提前終止生成過程。設總生成步驟數為 T,總生成 token 數為 G,則計算接受率(acceptance rate)為 G/T。該指標反映了模型平均每一步生成多少個被接受的 token,進而衡量所提方法帶來的加速效果。
接受率的理論最小值為 1,因為每一步至少會生成一個 token(即標準的下一個 token 預測)。最大理論值為 k+1=9,其中 k=8 是訓練時使用的 mask token 數量。
表 1 報告了模型在五個不同任務領域下的接受率,包括:知識問答、數學、編程、對話和安全。
如表所示,所提出的多 token 生成算法在不同任務上實現了約 1.5 到 5.2 倍不等的加速,具體取決于任務類型及插入的 mask 數量。值得注意的是,在編程和數學任務中,加速效果更為顯著,可能是因為這些領域中未來 token 的可預測性較高。
表 1:Tulu-3 在不同領域中通過多標記預測獲得的加速。
消融實驗
表 1 中的結果對應于本方法的最優配置,包含以下三個關鍵組件:
(1)使用采樣器 MLP 頭(sampler MLP head);
(2)訓練過程中引入 LCM 損失(LCM loss);
(3)生成階段采用二次解碼算法(quadratic decoding algorithm)。
研究者們對各個組件的貢獻進行了消融實驗。
圖 7: MTP 模型的簡單配置和最先進配置所達到的平均加速效果。基礎版本使用線性解碼,不包含 LCM 損失或采樣頭,而高級版本則包含 LCM 損失、采樣頭和二次解碼。
研究者們訓練了多個不同 LoRA 秩的模型,以進一步探索其影響。圖 8 總結了相關結果:左側與中間子圖分別展示了在未使用與使用采樣器頭的條件下的接受率;右側子圖則展示了采樣器頭與 LoRA 參數所帶來的顯存開銷。
圖 8:LoRA 秩的影響:無采樣頭加速(左),帶采樣頭加速(中),以及 LoRA 和采樣頭的內存開銷(右)。
總結
研究者們評估了自回歸模型在語言模型有監督微調階段對多 token 預測任務的適應能力。未來值得探索的一個方向,是在預訓練階段或下游任務自適應階段引入該方法,以進一步檢驗其適用性與效果。另一個具有前景的研究方向是將基于擴散的生成方法應用于多 token 預測任務。研究者們認為,多 token 預測位于完全自回歸生成與完全擴散生成之間,能夠在兩者之間取得優勢的平衡,兼具效率與質量的潛力。
更多信息,請參閱原論文。
特別聲明:以上內容(如有圖片或視頻亦包括在內)為自媒體平臺“網易號”用戶上傳并發布,本平臺僅提供信息存儲服務。
Notice: The content above (including the pictures and videos if any) is uploaded and posted by a user of NetEase Hao, which is a social media platform and only provides information storage services.