Better & Faster Large Language Models via Multi-token Prediction

透過多 token 預測打造更好、更快的大型語言模型

原始論文：Better & Faster Large Language Models via Multi-token Prediction 作者：Fabian Gloeckle¹,²、Badr Youbi Idrissi¹,³、Baptiste Rozière¹、David Lopez-Paz⁺¹、Gabriel Synnaeve⁺¹ 機構：¹FAIR at Meta；²CERMICS, École des Ponts ParisTech；³LISN, Université Paris-Saclay arXiv ID：2404.19737v1 日期：2024-04-30 標籤：LLM Pretraining Multi-token Prediction Speculative Decoding Code Generation Meta Llama

摘要

像 GPT 和 Llama 這樣的大型語言模型 (Large Language Models, LLMs) 都是用 next-token prediction (下一個 token 預測) 損失來訓練的。本文提出，讓語言模型一次預測多個未來 token 可以帶來更高的樣本效率 (sample efficiency)。具體而言，在訓練語料的每個位置，我們要求模型用 $n$ 個獨立的輸出頭 (output head) 同時預測接下來的 $n$ 個 token，這些輸出頭共用同一個 model trunk (模型主幹)。

把 multi-token prediction (多 token 預測) 視為一種輔助訓練任務 (auxiliary training task)，我們在程式碼與自然語言模型上都觀察到下游能力的提升，且不增加任何訓練時間開銷。這個方法在較大模型尺寸下越發有用，並在多 epoch 訓練時仍保有優勢。增益在程式碼這類生成式 (generative) 評測上尤其顯著：我們的 13B 參數模型在 HumanEval 上多解了 12% 的題目，在 MBPP 上多解了 17%。在小型演算法任務上的實驗顯示，多 token 預測有利於 induction heads 與演算法推理 (algorithmic reasoning) 能力的形成。額外好處是：用 4-token 預測訓練的模型在推論時最多可達 3 倍加速，即使在大批次下也成立。

[Figure 1: Multi-token prediction 概觀。(上) 訓練時，模型透過共用 trunk 與 4 個專屬 output head 一次預測 4 個未來 token；推論時只用 next-token 頭，其他 3 個頭可選擇用來加速推論。(下) Multi-token prediction 在 MBPP code 任務上的 pass@1，隨模型規模增加而顯著提升。誤差棒為以 bootstrapping 計算的 90% 信賴區間。]

一、引言

人類將最具智慧的成果、令人驚奇的發現與美麗的創作都凝鍊成了文字。在這些語料上訓練的 LLM，能藉由實作一個簡單卻強大的非監督學習任務——next-token prediction——來抽取大量世界知識與基本推理能力。儘管近期成就令人驚豔 (OpenAI, 2023)，next-token prediction 仍是一種低效的語言、世界知識與推理能力學習方式。更具體地說，使用 next-token prediction 的 teacher forcing 容易抓住局部模式 (local pattern) 而忽略「困難的選擇 (hard decision)」。因此，SOTA 的 next-token 預測器仍需要比人類兒童多出好幾個數量級的資料才能達到相同的流暢度 (Frank, 2023)。

本文主張：訓練 LLM 一次預測多個 token 將推動模型走向更高的樣本效率。如圖 1 所示，multi-token prediction 指示 LLM 在訓練語料的每個位置「同時並行」地預測 $n$ 個未來 token (Qi et al., 2020)。

貢獻

1. 我們提出一個簡單的 multi-token prediction 架構，沒有訓練時間或記憶體開銷（第 2 節）。
2. 我們提供實驗證據，證明此訓練範式在規模上有益：13B 參數模型在程式碼問題上平均多解 ~15%（第 3 節）。
3. Multi-token prediction 啟用了 self-speculative decoding (自我推測解碼)，使模型在多種批次大小下推論最高可達 3 倍加速（第 3.2 節）。

雖然簡單且零成本，multi-token prediction 卻是一個有效的修改，可訓練出更強、更快的 transformer 模型。我們希望這個工作能激發人們對「next-token prediction 之外」之 LLM 輔助損失的興趣，以提升這些模型的表現、連貫性與推理能力。

二、方法

標準語言建模對大規模文本語料 $x_1, \ldots, x_T$ 學習一個 next-token prediction 任務，最小化交叉熵損失：

$$L_1 = -\sum_t \log P_\theta(x_{t+1} \mid x_{t:1}) \tag{1}$$

其中 $P_\theta$ 是訓練中的 LLM，目標是在過去 token 歷史 $x_{t:1} = x_t, \ldots, x_1$ 給定下，最大化下一個未來 token $x_{t+1}$ 的機率。

本文將其推廣為 multi-token prediction 任務：在訓練語料的每個位置，模型被指示一次預測 $n$ 個未來 token。對應的交叉熵損失為：

$$L_n = -\sum_t \log P_\theta(x_{t+n:t+1} \mid x_{t:1}) \tag{2}$$

為了讓問題可解，我們假設 LLM $P_\theta$ 用一個共用 trunk 從觀察到的上下文 $x_{t:1}$ 產生隱藏表徵 $z_{t:1}$，再餵給 $n$ 個獨立的頭並行預測 $n$ 個未來 token（圖 1）。這帶出 multi-token prediction 交叉熵的下列分解：

$$L_n = -\sum_t \log P_\theta(x_{t+n:t+1} \mid z_{t:1}) \cdot P_\theta(z_{t:1} \mid x_{t:1}) = -\sum_t \sum_{i=1}^{n} \log P_\theta(x_{t+i} \mid z_{t:1}) \cdot P_\theta(z_{t:1} \mid x_{t:1})$$

實作上，我們的架構由一個共用 transformer trunk $f_s$（從上下文 $x_{t:1}$ 產生隱藏表徵 $z_{t:1}$）、$n$ 個獨立的輸出頭 $f_{h_i}$（以 transformer layer 實作）、以及一個共用的 unembedding 矩陣 $f_u$ 組成。要預測 $n$ 個未來 token，我們計算：

$$P_\theta(x_{t+i} \mid x_{t:1}) = \mathrm{softmax}(f_u(f_{h_i}(f_s(x_{t:1}))))$$

對於 $i = 1, \ldots, n$，其中 $P_\theta(x_{t+1} \mid x_{t:1})$ 即是我們的 next-token 預測頭。

記憶體高效的實作　訓練 multi-token 預測器的一個大挑戰是降低 GPU 記憶體用量。原因在於：當代 LLM 的詞彙表大小 $V$ 遠大於潛在表徵維度 $d$，因此 logit 向量是 GPU 記憶體的瓶頸。若直接同時實體化所有 logit 與其梯度（形狀 $(n, V)$），會嚴重限制可用批次大小。為此，我們仔細地調整 forward/backward 的順序（圖 2）：在共用 trunk $f_s$ 的 forward pass 之後，循序對每個獨立輸出頭 $f_i$ 做 forward 與 backward，並在 trunk 上累積梯度。輸出頭 $f_i$ 的 logit 與梯度在進到下一個頭 $f_{i+1}$ 前就被釋放，只需要長期儲存 $d$ 維的 trunk 梯度 $\partial L_n / \partial f_s$。如此一來，峰值 GPU 記憶體用量從 $O(nV + d)$ 降為 $O(V + d)$，且沒有額外執行時間成本。

推論　推論時最基本的用法是 vanilla next-token autoregressive prediction，僅使用 next-token 頭 $P_\theta(x_{t+1} \mid x_{t:1})$，丟棄其他頭。然而，額外的輸出頭可用來加速 next-token 頭的解碼，方法包括 blockwise parallel decoding (Stern et al., 2018)、self-speculative decoding，以及與 Medusa-like tree attention (Cai et al., 2024) 並用的 speculative decoding。

[Figure 2: 在 $n=2$ 頭的 n-token 預測模型中 forward / backward 的順序。透過循序在頭上做 forward/backward，避免同時實體化所有 unembedding layer 的梯度，降低峰值 GPU 記憶體用量。]

三、真實資料的實驗

我們以七個大規模實驗展示 multi-token prediction 損失的有效性。為了讓 next-token 預測器與 n-token 預測器之間的比較公平，後續實驗皆採用參數量相等的模型——當我們在未來預測頭上多加 $n-1$ 層時，就從共用 trunk 上拿掉 $n-1$ 層。

3.1 收益隨模型規模擴大

我們從頭訓練 6 種尺寸（300M 到 13B 參數）的模型，至少 91B token 的程式碼資料。圖 3 顯示在 MBPP (Austin et al., 2021) 與 HumanEval (Chen et al., 2021) 上的評估結果：在完全相同的計算預算下，可以從 LLM 中擠出明顯更多效能。

我們相信「只在大規模才有用」是 multi-token prediction 過去長期被忽略的可能原因。

[Figure 3: n-token 預測模型在 MBPP 上的結果，按模型尺寸分組。模型規模從 300M 到 13B，以 pass@1, 10, 100 評估。Multi-token 預測模型在小尺寸時略差於 baseline，但在大規模上反超。]

3.2 加速推論

我們在 xFormers 上實作貪婪 self-speculative decoding，使用異質批次大小。對 7B 參數的最佳 4-token 預測模型測量解碼速度，在訓練時未見過的程式碼與自然語言提示上：

• 程式碼：3.0× 加速，平均 3 個建議中接受 2.5 個 token。
• 文本：2.7× 加速。
• 8-byte 預測模型：6.4× 加速。

用 multi-token prediction 預訓練讓額外的頭比「只是微調 next-token 模型」要準確得多，因此能完整釋放 self-speculative decoding 的潛力。

3.3 用 multi-byte 預測學習全域模式

為了證明 next-token 預測會抓住局部模式，我們走到極端：在 314B byte（約等於 116B token）上訓練一個 7B 參數的 byte-level transformer。8-byte 預測模型相對於 next-byte 預測獲得驚人提升——MBPP pass@1 多解 67%，HumanEval pass@1 多解 20%。

Multi-byte prediction 因此成為一個非常有前景的方向：解鎖 byte-level 模型的高效訓練。Self-speculative decoding 對 8-byte 預測模型可帶來 6 倍加速，足以完全抵消 byte-level 序列較長的推論成本，甚至比 next-token 預測模型快近兩倍。8-byte 預測模型是個強大的 byte-based 模型，儘管訓練資料少了 1.7×，仍逼近 token-based 模型的效能。

3.4 尋找最佳 n

為了理解預測 token 數的影響，我們在 7B 規模、200B token 程式碼上做完整消融，嘗試 $n = 1, 2, 4, 6, 8$。

Table 1：Multi-token 預測在 byte-level 訓練上的提升（7B 參數模型，從零訓練）

4 個未來 token 在 HumanEval 與 MBPP 的 pass@1, 10, 100 上一致超過所有其他模型。對 APPS/Intro 而言 $n=6$ 略佳。byte-level 模型的最佳窗口較一致為 8 byte。

3.5 多 epoch 訓練

即使在多 epoch 訓練同一資料的情況下，multi-token 訓練仍對 next-token prediction 保有優勢。增益縮小但仍有：MBPP pass@1 增加 +2.4%，HumanEval pass@100 增加 +3.2%。

3.6 微調 multi-token 預測模型

在 CodeContests (Li et al., 2022) 上微調 7B 模型，比較 4-token 預測模型與 next-token 預測 baseline。圖 4 顯示，兩種微調 4-token 預測模型的方式（保留 4-token 損失或剝離後僅以 next-token 微調）皆優於 baseline。在 4-token 預訓練之上做 next-token 微調是整體最佳方法，符合「以輔助任務預訓練 + 任務專屬微調」的經典範式。

3.7 自然語言上的 multi-token 預測

我們在 200B token 自然語言上訓練 7B 模型，比較 4-token、2-token、next-token 損失。

• 6 個標準 NLP 選擇題基準：2-token 模型表現與 baseline 相當，4-token 模型略差。我們認為多選題與基於 likelihood 的基準並不適合鑑別語言模型的生成能力。
• 摘要任務（8 個基準的 ROUGE-L F1 平均）：$n=2$ 與 $n=4$ multi-token 預測模型都優於 next-token baseline，差距隨資料量增加而縮小。
• GSM8K 自然語言數學：在 200B token 訓練下，$n=2$ 模型明顯超越 baseline；500B token 後形勢反轉，$n=4$ 一致較差。

[Figure 5/6: NLP 選擇題與摘要評測的訓練曲線。Multi-token 模型在選擇題上沒有改善甚至退步，但在摘要任務上一致勝過 next-token baseline。]

四、合成資料的消融

是什麼驅動 multi-token 預測模型的下游表現提升？我們透過小型受控實驗證明：multi-token 預測在能力與泛化行為上帶來質的改變。

4.1 Induction 能力

Induction 是一種簡單的推理模式，它根據最近的延續來補完部分模式 (Olsson et al., 2022)：若一段話含有「AB」，之後又出現「A」，induction 會預測下一個是「B」。我們設計受控設定來測量這個能力：在 100 個兒童故事的測試集中，把人物名換成隨機產生、用 tokenizer 切成兩個 token 的名字。預測這兩個 token 中的第二個（在名字至少出現過一次之後），可視為純粹的 induction 任務。

[Figure 7: n-token 預測模型在 induction 上的精度，依模型尺寸 1M~1B 變化。對小模型（30M nonembedding 參數以下），2-token 預測損失大幅改善 induction 能力的形成；100M 以上時優勢消失。]

我們的詮釋：multi-token 預測損失幫助模型學會跨序列位置的資訊轉移，這有助於 induction heads 與其他 in-context learning 機制的形成。一旦 induction 能力建立，學到的特徵就把 induction 任務轉化成「在當前 token 局部就能解決的任務」，next-token 預測也能勝任。

4.2 演算法推理

我們在環 $\mathbb{F}_7[X]/(X^5)$ 上的多項式算術任務（含一元負號、加、乘與多項式合成）訓練並評估模型，難度由表達式中操作數 $m$ 控制：訓練時 $m \in [1, 5]$，評估涵蓋 in-domain ($m \le 5$) 與 out-of-domain ($m > 5$)。

[Figure 8: 在多項式算術任務上，accuracy 隨表達式操作數變化。Multi-token 預測在所有難度上都提升精度，特別是 out-of-domain 泛化。]

Multi-token 預測在所有難度上都提升精度，特別是在 out-of-distribution 泛化上效果驚人。把模型從 30M 加大到 100M 的提升，還不如把 next-token 預測換成 multi-token 預測來得大。

五、為什麼有效？一些推測

為什麼 multi-token prediction 在程式碼評測與小型演算法推理上有更佳表現？我們的直覺是：multi-token prediction 緩解了訓練時 teacher forcing 與推論時 autoregressive 生成之間的分布差異。

5.1 Lookahead 強化選擇點 (choice points)

並非所有 token 決策對生成有用文本都同等重要 (Bachmann & Nagarajan, 2024; Lin et al., 2024)。有些 token 容許風格上的變化而不限制後續，有些則是與較高層語意性質連結的「選擇點 (choice points)」。

Multi-token prediction 隱含地依據每個訓練 token 與其後續者的相關度賦予權重。透過標記與計數損失項，我們發現 n-token prediction 對選擇點及其相關項賦予 $\frac{n(n+1)}{2}$ 的權重，對「無關緊要的點」則賦予較小的 $n$。我們相信文本生成的品質取決於在選擇點上做對決策，而 n-token 預測損失正是促進這種能力。

[Figure 9: Multi-token prediction loss 對「結果型」token 隱含賦予較高權重。在示例序列中除了「5 → A」之外的所有轉換都易預測；由於「5 → A」之後的結果同樣難預測，這個轉換透過其相關項「3 → A」、…、「5 → C」在整體損失中得到較高的隱含權重。]

5.2 資訊論論證

語言模型通常以 teacher-forcing 訓練，模型在訓練時收到每個未來 token 的真實值；測試時則是無引導 (unguided) 的 autoregressive 生成，誤差會累積。Teacher-forcing 鼓勵模型專注於極短期預測，可能犧牲對生成序列整體結構的長期依賴關係。

設 $X$ 為下一個未來 token、$Y$ 為再下一個。Vanilla next-token prediction 在預測 $X$ 時關注 $H(X)$，而 $n=2$ 的 multi-token prediction 關注 $H(X) + H(Y)$。我們將兩者分解為：

$$H(X) = H(X \mid Y) + I(X; Y)$$

$$H(X) + H(Y) = H(X \mid Y) + 2 I(X; Y) + H(Y \mid X)$$

捨棄 $H(Y \mid X)$（它會在下個位置的預測中再次出現），可以看到 2-token 預測將互資訊 $I(X; Y)$ 的重要性提高了 2 倍。也就是說，multi-token 預測器更善於預測那些對「後續文本」具相關性的 token $X$。

六、相關工作

語言建模損失　Dong et al. (2019) 與 Tay et al. (2022) 用混合多種注意力遮罩 (full / causal / prefix) 的 denoising 任務來縮小生成式任務上的差距；前者使用 BERT 風格 mask，後者使用 span corruption。Yang et al. (2019) 在保留原本位置嵌入的前提下訓練於 permuted 序列。然而這些任務都只在輸入文本的小比例 token 上訓練（平均約 15%）。我們的方法則在全部 token 上反向傳播。

語言建模中的 multi-token prediction　Qi et al. (2020) 主張 multi-token 預測鼓勵規劃、改善表徵、防止對局部模式的過擬合；但其技術上將 residual stream 複製 $n$ 倍，而我們允許 compute-matched 的比較且讓 residual 表徵更直接地參與輔助損失項。Stern et al. (2018) 與 Cai et al. (2024) 提出以 multi-token prediction 微調以加速推論，但未研究預訓練階段的效果。Pal et al. (2023) 用 probing 顯示 next-token 模型也能某種程度地預測後續 token，但不如我們專門訓練的模型。

Self-speculative decoding　Stern et al. (2018) 是最早提出 speculative decoding 的工作之一。我們的架構把他們的線性預測頭換成 transformer 層。Cai et al. (2024) 提出更精緻的 self-speculative decoding 方案，利用每個頭的 top-k 預測而非只用最佳的，可與我們訓練的 multi-token 預測模型結合。

Multi-target 預測　Multi-task learning (Caruana, 1997) 是讓神經網路在多個任務上聯合訓練以改善目標任務表現的範式。現代深度學習常以共用大型 trunk + 各任務獨立預測頭的形式實作。Multi-target 預測在許多領域（時序預測 (Vapnik & Vashist, 2009)、影片預測 (Mathieu et al., 2016)、未來幀表徵 (Vondrick et al., 2016)）都展現了成功。

七、結論

我們提出 multi-token prediction 作為訓練語言模型的改進，特別針對生成式或推理任務。實驗（最多 7B 參數、1T token）顯示這個方法在更大模型上越發有用，並對程式碼任務有明顯改善。我們認為這個方法降低了 teacher-forced 訓練與 autoregressive 生成之間的分布錯配。配合 speculative decoding，可使精確推論加速 3 倍。

未來工作包括：自動選擇 multi-token prediction 損失中的 $n$，例如使用損失尺度與損失平衡 (Défossez et al., 2022)；為 multi-token prediction 調整最佳詞彙量；發展在 embedding 空間中操作的改良輔助預測損失 (LeCun, 2022)。

影響聲明　本文的目標是讓語言模型在計算與資料上更高效。原則上這可降低訓練 LLM 的生態影響，但需注意反彈效應 (rebound effect)。

環境影響　所有模型訓練合計約需 500K GPU 小時（A100-80GB 與 H100），估計總排放約 50 噸 CO2eq，100% 由 Meta 永續計畫抵銷。

參考文獻

完整參考文獻請見原始論文。下列為主要引用：

• Qi, W., Yan, Y., Gong, Y., Liu, D., Duan, N., Chen, J., Zhang, R., & Zhou, M. (2020). ProphetNet: Predicting future n-gram for sequence-to-sequence pre-training.
• Stern, M., Shazeer, N., & Uszkoreit, J. (2018). Blockwise parallel decoding for deep autoregressive models.
• Leviathan, Y., Kalman, M., & Matias, Y. (2023). Fast inference from transformers via speculative decoding.
• Cai, T., Li, Y., Geng, Z., Peng, H., Lee, J. D., Chen, D., & Dao, T. (2024). Medusa: Simple LLM inference acceleration framework with multiple decoding heads.
• Olsson, C., Elhage, N., Nanda, N., et al. (2022). In-context learning and induction heads. Transformer Circuits Thread.
• Bachmann, G., & Nagarajan, V. (2024). The pitfalls of next-token prediction.
• Dong, L., Yang, N., Wang, W., et al. (2019). Unified language model pre-training for natural language understanding and generation.
• Tay, Y., et al. (2022). UL2: Unifying language learning paradigms.
• Yang, Z., et al. (2019). XLNet: Generalized autoregressive pretraining for language understanding.
• Frank, M. C. (2023). Bridging the data gap between children and large language models. Trends in Cognitive Sciences.
• Austin, J., et al. (2021). Program synthesis with large language models. (MBPP)
• Chen, M., et al. (2021). Evaluating large language models trained on code. (HumanEval)
• Hendrycks, D., et al. (2021). Measuring coding challenge competence with APPS.
• Li, Y., et al. (2022). Competition-level code generation with AlphaCode. Science.
• Cobbe, K., et al. (2021). Training verifiers to solve math word problems. (GSM8K)
• Goyal, S., et al. (2023). Think before you speak: Training language models with pause tokens.
• Caruana, R. (1997). Multitask learning. Machine Learning.
• LeCun, Y. (2022). A path towards autonomous machine intelligence.
• Lefaudeux, B., et al. (2022). xFormers: A modular and hackable transformer modelling library.

術語對照表