Training Compute-Optimal Large Language Models

Jordan Hoffmann; Sebastian Borgeaud; Arthur Mensch; Elena Buchatskaya; Trevor Cai; Eliza Rutherford; Diego de Las Casas; Lisa Anne Hendricks; Johannes Welbl; Aidan Clark; Tom Hennigan; Eric Noland; Katie Millican; George van den Driessche; Bogdan Damoc; Aurelia Guy; Simon Osindero; Karen Simonyan; Erich Elsen; Jack W. Rae; Oriol Vinyals; Laurent Sifre

論文 · arXiv 2203.15556

Training Compute-Optimal Large Language Models

Jordan Hoffmann, Sebastian Borgeaud, Arthur Mensch, Elena Buchatskaya, Trevor Cai, Eliza Rutherford, Diego de Las Casas, Lisa Anne Hendricks, Johannes Welbl, Aidan Clark, Tom Hennigan, Eric Noland et al.

發表於 2022-03-29cs.CLcs.LG arXiv PDF

用 XYZNOT 閱讀（全文翻譯與 AI 講解）

三句話總結

🎯 在固定的運算預算下，找出模型參數量與訓練 token 數量之間的最佳權衡。
🔧 透過最小損失包絡線、IsoFLOP 切片與參數化損失配適，在超過 400 次 Transformer 訓練運行中實證配適擴展法則（scaling laws）。
📈 參數量與 token 數量應等比例擴展（ $a \approx 0.5$ ， $b \approx 0.5$ ）。以 1.4T token 訓練的 70B Chinchilla 模型符合 Gopher 的運算預算，並在 MMLU 上達到 67.6%。

摘要

我們探討了在給定運算預算下，訓練 Transformer 語言模型的最佳模型大小與 token 數量。我們發現當前的大型語言模型普遍存在明顯訓練不足的情況，這是由於近期研究過度聚焦於擴展模型大小，卻維持訓練資料量不變所導致。透過訓練超過 400 個參數量介於 7,000 萬至 160 億以上、訓練 token 介於 50 億至 5,000 億的語言模型，我們發現若要達到運算最佳化訓練，模型大小與訓練 token 數量的擴展比例應當均等：模型大小每增加一倍，訓練 token 數量也應增加一倍。我們藉由訓練一個預測達到運算最佳化的模型 Chinchilla 來驗證此假說。Chinchilla 使用與 Gopher 相同的運算預算，但僅具備 70B 參數並採用了 4 倍以上的資料量。在廣泛的下游評估任務中，Chinchilla 全面且顯著地超越了 Gopher (280B)、GPT-3 (175B)、Jurassic-1 (178B) 以及 Megatron-Turing NLG (530B)。這也意味著 Chinchilla 在微調和推論時所需的運算資源大幅降低，極大地便利了下游應用。值得一提的是，Chinchilla 在 MMLU 基準測試中達到 67.5% 的尖端平均準確率，相較於 Gopher 提升了超過 7%。

英文原文摘要

We investigate the optimal model size and number of tokens for training a transformer language model under a given compute budget. We find that current large language models are significantly undertrained, a consequence of the recent focus on scaling language models whilst keeping the amount of training data constant. By training over 400 language models ranging from 70 million to over 16 billion parameters on 5 to 500 billion tokens, we find that for compute-optimal training, the model size and the number of training tokens should be scaled equally: for every doubling of model size the number of training tokens should also be doubled. We test this hypothesis by training a predicted compute-optimal model, Chinchilla, that uses the same compute budget as Gopher but with 70B parameters and 4 $\times$ more more data. Chinchilla uniformly and significantly outperforms Gopher (280B), GPT-3 (175B), Jurassic-1 (178B), and Megatron-Turing NLG (530B) on a large range of downstream evaluation tasks. This also means that Chinchilla uses substantially less compute for fine-tuning and inference, greatly facilitating downstream usage. As a highlight, Chinchilla reaches a state-of-the-art average accuracy of 67.5% on the MMLU benchmark, greater than a 7% improvement over Gopher.

論文解讀

問題與研究動機

過去的研究（特別是 Kaplan et al., 2020）指出，當運算資源增加時，模型參數量應大幅擴展（ $N \propto C^{0.73}$ ），而 token 數量的增長速度則慢得多（ $D \propto C^{0.27}$ ）。因此，許多大型模型如 GPT-3 (175B)、Jurassic-1 (178B)、Gopher (280B) 以及 MT-NLG (530B) 皆普遍在約 3,000 億 token 上進行訓練。然而，這使得大型模型相對於其投入的運算資源呈現顯著的訓練不足。本研究重新審視了這項權衡：在固定的運算預算 $C$ 下，應如何配置參數量 $N$ 與 token 數量 $D$ ，以使測試損失最小化？

核心概念

等比例擴展（Equal scaling）：模型參數量每增加一倍，訓練 token 數量也應增加一倍（ $N_{o pt} \propto C^{0.5}$ ， $D_{o pt} \propto C^{0.5}$ ）。
餘弦排程對齊（Cosine schedule alignment）：學習率排程必須與實際訓練期程（horizon）相符。若使用在固定長週期（例如 1,300 億 token）衰減過程中擷取的中間點評估，會高估較小 token 數量下的損失，進而扭曲實證擴展估計。
下游任務效率：在預訓練運算量相當的情況下，以大幅增加的 token 訓練較小的模型，能取得相同甚至更佳的準確率，同時大幅降低記憶體佔用、微調成本與推論延遲。

研究方法（含核心公式）

作者使用三種互補的策略評估了超過 400 個模型（參數量介於 70M 至 16B，訓練 token 介於 5B 至 500B）：

訓練曲線最小值法（Minimum over training curves）：針對固定模型架構變化 4 種不同訓練時長。擷取每個 FLOP 的最小損失包絡線並配適冪律： $N_{o pt} \propto C^{a}, D_{o pt} \propto C^{b}$ 得出 $a = 0.50$ ， $b = 0.50$ 。
IsoFLOP 剖面分析（IsoFLOP profiles）：在 9 個離散預算（ $6 \times 1 0^{18}$ 至 $3 \times 1 0^{21}$ ）下固定總 FLOP 數，並將餘弦衰減與期程對齊。在參數量上配適拋物線以找出損失極小值，得出 $a = 0.49$ ， $b = 0.51$ 。
參數化損失配適（Parametric loss fitting）：使用風險分解模型對平滑損失進行建模： $\hat{L} (N, D) ≜ E + \frac{A}{N ^{α}} + \frac{B}{D ^{β}}$ 其中 $E$ 代表自然文本熵， $A / N^{α}$ 代表表徵次最佳性（容量不足），而 $B / D^{β}$ 則捕捉提早停止／有限資料的次最佳性。透過最小化對數殘差的 Huber loss，得出 $E = 1.69$ 、 $A = 406.4$ 、 $B = 410.7$ 、 $α = 0.34$ 、 $β = 0.28$ 。在限制條件 $FLOPs (N, D) \approx 6 N D$ 下： $a = \frac{β}{α + β} = 0.46, b = \frac{α}{α + β} = 0.54$

成果（含重要數據）

Chinchilla 70B：在與 Gopher 相同的 FLOP 預算（ $5.76 \times 1 0^{23}$ FLOPs，若採精確計算則為 $6.3 \times 1 0^{23}$ ）下訓練，包含 70B 參數與 1.4 兆 token。
MMLU 基準測試：Chinchilla 達到 67.6% 的平均 5-shot 準確率，超越 Gopher (60.0%)、GPT-3 (43.9%)，並突破了 2023 年 6 月的人類專家預測值 (63.4%)。在 4 個獨立的 MMLU 任務中達到 $> 90%$ 的準確率。
下游基準測試：在 BIG-bench (65.1% vs 54.4%)、LAMBADA (77.4% vs 74.5%) 以及 closed-book Natural Questions (5-shot: 31.5% vs 24.5%) 上均大幅超越 Gopher。
偏見與毒性（Bias and Toxicity）：未經提示引導的毒性程度在 Gopher 與 Chinchilla 間表現相當（中位數 0.064 vs 0.066），而整體 Winogender 準確率則從 71.4% 提升至 78.3%。

限制與開放研究問題

受限於訓練成本，本研究僅在邊界規模上直接驗證了兩個同運算量對等點（Gopher 280B 對比 Chinchilla 70B）。
分析假設模型在網路資料集上僅進行單一輪次（single-epoch）訓練；在多個 epoch 間重複資料的表現仍有待探索。
實證數據顯示在高運算預算下， $lo g (N_{o pt})$ 呈現輕微的負曲率（凹性），這表示冪律模型可能仍然些微高估了最佳模型大小。

關鍵術語

Chinchilla: 一個具備 700 億參數的自迴歸語言模型，使用與 Gopher 相同的運算預算，並於 1.4 兆 token 上訓練而成。
Gopher: 一個具備 2,800 億參數的 Transformer 語言模型，於 3,000 億 token 上訓練，作為本研究中運算量對等的基準模型。
IsoFLOP profiles: 一種在維持總訓練 FLOPs 不變的情況下，改變模型參數量以找出損失極小化參數規模的方法論。
Compute-optimal frontier: 連結可用訓練 FLOPs 與參數量及訓練 token 數量之最佳配置的數學曲線。
Parametric loss function: 將語言模型損失在統計上分解為不可約貝氏熵、參數容量缺失與 token 樣本最佳化誤差的函數。
Huber loss: 在方法 3 中使用的穩健損失函數，用於配適實證 LM 訓練成果，同時減弱小運算規模下殘差離群值的影響。
MassiveText: 包含網頁、書籍、新聞、程式碼與維基百科的多領域預訓練語料庫，用於訓練 Gopher 與 Chinchilla。
FLOPs calculation: 乘積累加運算預算，通常透過前向與反向傳播近似計算為 $6 N D$ ，或透過逐層計算精確求得。

常見問題

為什麼 Kaplan 等人會得出參數量擴展速度應快於 token 的結論？

Kaplan 等人採用了固定的餘弦學習率排程（例如衰減至 1,300 億 token），並在較小的 token 數據點上直接擷取訓練過程中尚未收斂的中間評估值。由於當循環長度大幅超過實際訓練步數時，餘弦學習率排程的效果會劣化，導致這些中間點的損失被非自然地抬高，進而低估了使用較少 token 訓練較小模型的效能。

根據 Chinchilla 的研究，運算最佳化擴展的經驗法則是甚麼？

模型參數量與訓練 token 數量應以大致相等的比例擴展：運算量每增加 $2 \times$ ，模型大小應增加 $\approx 2^{0.5} \approx 1.41 \times$ ，且訓練 token 數量也應增加 $\approx 2^{0.5} \approx 1.41 \times$ 。

要訓練一個達到運算最佳化的 175B 模型，建議需要多少 token？

根據論文的預測，一個 1,750 億參數的模型需要大約 3.7 兆至 4.3 兆個 token，以及 $3.85 \times 1 0^{24}$ 至 $4.54 \times 1 0^{24}$ FLOPs 才能達到運算最佳化，這遠遠超出了 GPT-3 通常採用的 3,000 億 token。

Gopher 與 Chinchilla 之間進行了哪些架構調整？

除了將層數與隱藏維度縮減（從 280B 降至 70B）之外，Chinchilla 還將 Adam 最佳化器替換為 AdamW、在分散式最佳化器狀態中保留 float32 主權重（master weights），並從 SentencePiece 標記器中移除了 NFKC 正規化。

本頁解讀由 AI 根據論文生成，可能有錯誤；細節請以原文為準。

用 XYZNOT 閱讀（全文翻譯與 AI 講解）