前言

如果給定了計算預算$C$，如何分配LLM的模型尺寸$N$和訓練的數據量$D$，才能使得模型的效果$L$最好呢？筆者在此介紹一篇經典的文章討論這個問題。如有謬誤請見諒并聯系指出，本文遵守CC 4.0 BY-SA版權協議，轉載請聯系作者并注明出處，謝謝。

• 關鍵字：最佳計算預算分配
• 發表信息：NIPS 2022

$?$ 聯系方式：

e-mail: FesianXu@gmail.com

github: https://github.com/FesianXu

知乎專欄: 計算機視覺/計算機圖形理論與應用

微信公眾號：機器學習雜貨鋪3號店

---

我們知道在大語言模型（Large Language Model, LLM）中，存在所謂的尺度擴展規律（Scaling Laws） [2]，如Fig 1所示，即是：

LLM的性能會隨著模型的參數量、模型的訓練量、模型的訓練數據量的增加而增加

我們也知道模型的參數量、模型的訓練量和模型的訓練數據量都會影響到最終的計算預算（可以用FLOPs計算），因此LLM的性能可以說和計算預算直接掛鉤，這也是Fig 1 左圖所表示的。我們不禁會有個疑問，給定了模型的計算預算$C$，我們應該怎么均衡模型參數量$N$和預訓練的Token數量$D$，才能使得模型的預訓練損失$L$最小化呢？我們期待得到最優的模型參數$N_{opt}$和最優的預訓練Token數量$D_{opt}$，可以使得預訓練損失最小，正如公式(1)所示。

$$N_{opt}(C),D_{opt}(C)=\underset{N,Ds.t.\mathrm{FLOPs}(N,D)=C}{\mathrm{argmin}}L(N,D)\text{\hspace{1em}(1)}$$

作者探索這個規律的方法論也很直接，作者步進遍歷了一遍不同的模型尺寸（從70M到16B參數量），也步進遍歷了一遍預訓練數據Token數量（從5B到400B），最終跑了超過400個組合的數據點，不得不說有算力真的可以為所欲為。從直觀上看越大尺寸的模型需要越多訓練的Token，當然我們需要研究具體的比例，作者采用了三種不同的方法去找這個比例關系。

固定模型尺寸下的性能分析

這種方法是分別固定住模型尺寸（從70M到10B多個模型尺寸都需要實驗），然后觀察訓練了不同數量的Tokens數量后，在每一個節點時哪一個模型尺寸能夠達到最小的訓練損失。如Fig 2 左圖 所示， 這里有些地方需要解釋。首先這里的橫坐標是浮點計算量FLOPs，在不同模型尺寸下，相同的FLOPs能訓練的Token數量是不同的，因此才會出現Fig 2左圖中同一個FLOPs中，大尺寸模型損失比小尺寸模型還大的情況。從Fig 2 左圖中，我們能發現在不同的FLOPs下，到達最小損失的模型尺寸是不一樣的（不太容易看出來，在左圖中是灰色點，它們形成了一個包絡線），不同的FLOPs在對應尺寸模型下能夠折算成訓練過的Token數量，因此能夠畫出Fig 2 中圖和右圖，橫坐標是FLOPs，縱坐標是達到最小損失（也就是左圖的灰色點）時的模型尺寸和過了的Tokens數。換句話說，Fig 2中圖和右圖就是給定計算預算$C$下的最佳模型尺寸$N_{opt}$和訓練數據量$D_{opt}$，我們發現有$N_{opt}\propto C^a,D_{opt}\propto C^b$，通過實驗可以算出$a=0.50,b=0.50$。

固定計算預算下的性能分析

第一種方法的計算量FLOPs沒有固定，在此方法中我們固定計算量$C$（也就是所謂的IsoFLOP），分析等量計算下的最佳模型參數量$N_{opt}$。同時，在知道了每個實驗固定的計算量，和在此之下的最佳模型參數量后，也就可以反推訓練Token數量。實驗如Fig 3 左圖所示，可以發現在不同的固定計算量下（從$6\times 10^{18}$到$3\times 10^{21}$ FLOPs），遍歷不同尺寸的模型能夠發現在某些尺寸處會存在明顯的低谷，這個低谷就是在固定計算預算情況下的最佳模型參數量，由此也能繪制出Fig 3 中圖和右圖，繪制邏輯如第一種方法所述。不難發現同樣有$N_{opt}\propto C^a,D_{opt}\propto C^b$這個規律，算出$a=0.49,b=0.51$。

對參數化損失函數進行擬合

在第1和2中方法中已經積累了很多最小損失$L$下的$FLOPs(N_{opt},D_{opt})=C$的數據點了，我們不妨把損失拆解為三大部分如公式(2)所示，其中第一項$E$為不可約損失，也就是自然文本的熵，是不可繼續減少的最基礎的損失項。第二項為（不完美的）參數量為$N$的Transformer模型訓練過程中產生的損失（因為參數量$N$總是有限，也就是不完美的，因此總是在理想損失$E$的基礎上有超額損失），第三項則是（不完美的）訓練數據量$D$下（因為訓練數據量$D$不可能是無限的）的產生的超額損失。
$$\hat{L}(N,D)?E+\frac{A}{N^{\alpha }}+\frac{B}{D^{\beta }}\text{\hspace{1em}(2)}$$

作者采用L-BFGS算法去最小化所謂的Huber loss（因為數據點只有400多個，這個loss作者說對離群點比較穩健）去進行估計$(A,B,E,\alpha ,\beta )$，筆者也沒細究，讀者有興趣的可以翻閱 [3] 和 [4]。最終估計出來的參數為：
$$E=1.69,A=406.4,B=410.7,\alpha =0.34,\beta =0.28\text{\hspace{1em}(3)}$$
在LLM Scaling Law的論文 [2] 中提出了一個估算: $FLOPs(N,D)\approx 6ND$，借此可以將公式(2)進行變形，得到公式(4)

$$\begin{array}{rl}{N}_{\mathrm{opt}}(C)& =G{\left(\frac{C}{6}\right)}^a,\\ D_{\mathrm{opt}}(C)& =G^{?1}{\left(\frac{C}{6}\right)}^b,\\ \text{其中}G& ={\left(\frac{\alpha A}{\beta B}\right)}^{\frac{1}{\alpha +\beta }},\\ a& =\frac{\beta }{\alpha +\beta },\\ b& =\frac{\alpha }{\alpha +\beta }\end{array}\text{\hspace{1em}(4)}$$

作者算得$a=0.46,b=0.54$，具體過程請自行參考原文。

給定計算量下的最優設計

Fig 4是將以上三種預測方法繪制成計算量——最佳模型尺寸估計曲線圖，其中那貼上了一些之前工作的估計 [2] 和一些模型的對比，如Gopher（280B參數量）、GPT-3（175B參數量）和Megatron-NLG （530B）參數量。從圖中能發現：

1. 方法1和方法2估計出來的曲線基本上貼合，方法3估計出的模型尺寸在計算預算小的時候和前兩者基本貼合，但在大計算預算下會偏小些，不過也不會差距特別大。
2. 主流的大模型，如Gopher、GPT3等在對應的計算預算下，模型尺寸明顯偏大，基本上是貼著 [2] 的曲線走的。

為了證明本文提出的估計方法更佳準確，作者在方法1和2中對齊Gopher的計算預算（大概是$5.76\times 10^{23}$ FLOPs），找到了最佳模型尺寸，約是70B，作者將這個訓練出來的模型稱之為Chinchilla，需要將這個模型的性能和Gopher進行公平對比。注意到在方法1和2中，從Fig 2和Fig 3的右圖中可以找出給定預算下的最佳訓練Token數量，對于Chinchilla來說是1.4-1.5T左右，因此$D_{opt}/N_{opt}\approx 20$。

作者在相當多語言下游任務的基準上進行了測試，都發現Chinchilla對比Gopher存在普遍優勢，在一些任務中甚至超過了Megatron-NLG 530B模型。這些實驗過于冗長，筆者就不展示細節了。

筆者讀后感

這篇論文的意義在于告訴我們，在給定了計算預算下，是存在一個最優的模型尺寸和訓練數據量的，他們存在一個比例（$D_{opt}\approx 20N_{opt}$），越大的模型就需要越多數據進行訓練，才能發揮出模型最優的性能。這篇論文的發表時間比較早，是2022年，現在已經有很多工作證實了在推理中進行復雜策略可以有效提高模型性能 [5,6]，并且這些推理策略同樣也存在Scaling Law。這意味著計算預算不僅可以花在預訓練上，而且可以花在推理時的Scaling，這也是這篇文章沒有考慮到的點。當然，在 [6] 中作者也承認，推理時的Scaling并非是萬能的，而是：

推理時計算與預訓練計算并非一對一“可互換”。對于模型能力范圍內的簡單和中等難度問題，或者在推理（實時性）要求較低的情況下，測試時計算可以輕松彌補額外的預訓練。然而，對于超出基礎模型能力范圍的具有挑戰性的問題，或者在推理（實時性）要求較高的情況下，預訓練可能更有效于提升性能。

也就是說預訓練的地位并不是通過推理時的Scaling就可以替代的，預訓練中分配一定量的計算預算對于全方面提高LLM的性能是必須的。結合了模型訓練、模型推理的更為綜合的最優配比，應該是值得去研究的。

Reference

[1]. Hoffmann, Jordan, Sebastian Borgeaud, Arthur Mensch, Elena Buchatskaya, Trevor Cai, Eliza Rutherford, Diego de Las Casas et al. “Training compute-optimal large language models.” arXiv preprint arXiv:2203.15556 (2022).

[2]. Kaplan, Jared, Sam McCandlish, Tom Henighan, Tom B. Brown, Benjamin Chess, Rewon Child, Scott Gray, Alec Radford, Jeffrey Wu, and Dario Amodei. “Scaling laws for neural language models.” arXiv preprint arXiv:2001.08361 (2020).

[3]. J. Nocedal. Updating Quasi-Newton Matrices with Limited Storage. Mathematics of Computation, 35(151):773–782, 1980. ISSN 0025-5718. doi: 10.2307/2006193. URL https://www.jstor.org/stable/2006193 aka L-BFGS

[4]. P. J. Huber. Robust Estimation of a Location Parameter. The Annals of Mathematical Statistics, 35 (1):73–101, Mar. 1964. ISSN 0003-4851, 2168-8990. doi: 10.1214/aoms/1177703732. URL
https://projecteuclid.org/journals/annals-of-mathematical-statistics/volume-35/issue-1/Robust-Estimation-of-a-Location-Parameter/10.1214/aoms/1177703732.full. aka Huber loss

[5]. https://fesianxu.github.io/2025/03/02/test-time-scaling-laws-20250302/， 《大模型推理時的尺度擴展定律》

[6]. Snell, Charlie, Jaehoon Lee, Kelvin Xu, and Aviral Kumar. “Scaling llm test-time compute optimally can be more effective than scaling model parameters.” arXiv preprint arXiv:2408.03314 (2024).