Scaling Laws for Neural Language Models — part 3/5
Cleaned-text code points [31205, 46825). Bibliography omitted here; source.md includes it.
A simple empirical theory for the batch size dependence of training was developed in [ MKAT18 ] (see also [ SLA+18 , ZLN+19 ] ). It was argued that there is a critical batch size B crit B_{\rm crit} for training; for B B up to B crit B_{\rm crit} the batch size can be increased with very minimal degradation in compute-efficiency, whereas for B > B crit B>B_{\rm crit} increases in B B result in diminishing returns. It was also argued that the gradient noise scale provides a simple prediction for B crit B_{\rm crit} , and that neither depends directly on model size except through the value of the loss that has been attained. These results can be used to predict how training time and compute will vary with the batch size. To utilize both training time and compute as effectively as possible, it is best to train with a batch size B ≈ B crit B\approx B_{\rm crit} . Training at B ≫ B crit B\gg B_{\rm crit} minimizes the number of training steps, while B ≪ B crit B\ll B_{\rm crit} minimizes the use of compute.
More specifically, it was demonstrated that for a wide variety of neural network tasks, the number of training steps S S and the number of data examples processed E = B S E=BS satisfy the simple relation ( S S min − 1 ) ( E E min − 1 ) = 1 \left(\frac{S}{S_{\rm min}}-1\right)\left(\frac{E}{E_{\rm min}}-1\right)=1 (5.1) when training to any fixed value of the loss L L . Here S min S_{\rm min} is the minimum number of steps necessary to reach L L , while E min E_{\rm min} is the minimum number of data examples that must be processed.
We demonstrate the relation ( 5.1 ) for Transformers in Figure 18 in the appendix. This relation defines the critical batch size B crit ( L ) ≡ E min S min B_{\rm crit}(L)\equiv\frac{E_{\rm min}}{S_{\rm min}} (5.2) which is a function of the target value of the loss. Training at the critical batch size makes a roughly optimal time/compute tradeoff, requiring 2 S min 2S_{\rm min} training steps and processing E = 2 E min E=2E_{\rm min} data examples.
In Figure 10 we have plotted the critical batch size and gradient noise scale 4 4 4 Although the critical batch size roughly matches the gradient noise scale, we are using a direct measurements of B crit B_{\rm crit} from Figures 18 and 10 for all our later analyses. as a function of training loss for two different models. We see that B crit ( L ) B_{\rm crit}(L) is independent of model size, and only depends on the loss L L . So the predictions of [ MKAT18 ] continue to hold for Transformer language models. The critical batch size can be fit with a power-law in the loss B crit ( L ) ≈ B ∗ L 1 / α B B_{\rm crit}(L)\approx\frac{B_{*}}{L^{1/\alpha_{B}}} (5.3) where B ∗ ≈ 2 × 10 8 B_{*}\approx 2\times 10^{8} and α B ≈ 0.21 \alpha_{B}\approx 0.21 .
We have chosen this parameterization for B crit ( L ) B_{\rm crit}(L) because as the loss approaches its minimum value L min L_{\rm min} , the gradient noise scale is expected to diverge, and we expect B crit B_{\rm crit} to track this noise scale. We do not know L min L_{\rm min} , as we see no sign that our models are approaching it, but L min > 0 L_{\rm min}>0 since the entropy of natural language is non-zero. Since apparently L min L_{\rm min} is much smaller than the values of L L we have achieved, we used a parameterization where B crit B_{\rm crit} diverges as L → 0 L\to 0 .
We will use B crit ( L ) B_{\rm crit}(L) to estimate the relation between the number of training steps S S while training at batch size B = 2 19 B=2^{19} tokens and the number of training steps while training at B ≫ B crit B\gg B_{\rm crit} . This is simply S min ( S ) ≡ S 1 + B crit ( L ) / B ( minimum steps, at B ≫ B crit ) S_{\rm min}(S)\equiv\frac{S}{1+B_{\rm crit}(L)/B}\qquad(\text{minimum steps, at }B\gg B_{\rm crit}) (5.4) for any given target value L L for the loss. This also defines a critical value of the compute needed to train to L L with a model of size N N if we were to train at B ≪ B crit ( L ) B\ll B_{\rm crit}(L) . This is C min ( C ) ≡ C 1 + B / B crit ( L ) ( minimum compute, at B ≪ B crit ) C_{\rm min}(C)\equiv\frac{C}{1+B/B_{\rm crit}(L)}\qquad(\text{minimum compute, at }B\ll B_{\rm crit}) (5.5) where C = 6 N B S C=6NBS estimates the (non-embedding) compute used at batch size B B .
### 5.2 Results for L ( N , S min ) L(N,S_{\rm min}) and Performance with Model Size and Compute
Now we will use S min S_{\rm min} defined in Equation ( 5.4 ) to obtain a simple and universal fit for the dependence of the loss on model size and training time in the infinite data limit. We will fit the stable, Adam-optimized training runs using Equation ( 1.6 ), repeated here for convenience: L ( N , S min ) = ( N c N ) α N + ( S c S min ) α S L(N,S_{\rm min})=\left(\frac{N_{c}}{N}\right)^{\alpha_{N}}+\left(\frac{S_{c}}{S_{\rm min}}\right)^{\alpha_{S}} (5.6) for the loss. We include all training steps after the warmup period of the learning rate schedule, and find a fit to the data with the parameters:
With these parameters, we obtain the learning curve fits in Figure 4 . Though the fits are imperfect, we believe they are quite compelling given the simplicity of Equation ( 5.6 ).
The data and fits can be visualized in a different and more interesting way, as shown in Figure 11 . There we study the test loss as a function of model size while fixing either the total non-embedding compute C C used in training, or the number of steps S S . For the fits we use Equation ( 5.5 ) and ( 5.4 ) along with the parameters above and Equation ( 5.6 ).
The power-law dependence of the loss on S min S_{\rm min} reflects the interplay of optimizer dynamics and the loss landscape. Since the fits are best late in training, when the loss may be approximately quadratic, the power-law should provide information about the spectrum of the Hessian of the loss. Its universality suggests that the Hessian eigenvalue density is roughly independent of model size.
### 5.3 Lower Bound on Early Stopping Step
The results for L ( N , S min ) L(N,S_{\rm min}) can be used to derive a lower-bound (and rough estimate) of the step at which early stopping should occur when training is data limited. It is motivated by the idea that finite and infinite D D learning curves for a given model will be very similar until we reach S min ≈ S stop S_{\rm min}\approx S_{\rm stop} . Thus overfitting should be proportional to the correction from simply ending training at S stop S_{\rm stop} . This will underestimate S stop S_{\rm stop} , because in reality the test loss will decrease more slowly when we have a finite D D , and therefore we will require more training steps to reach the optimal test loss at finite D D . This line of reasoning leads to the inequality S stop ( N , D ) ≳ S c [ L ( N , D ) − L ( N , ∞ ) ] 1 / α S S_{\rm stop}(N,D)\gtrsim\frac{S_{c}}{\left[L(N,D)-L(N,\infty)\right]^{1/\alpha_{S}}} (5.7) where L ( N , ∞ ) L(N,\infty) is the converged loss, evaluated with infinite available data. This inequality and its comparison to the empirical data is displayed in Figure 16 in the appendix. In that figure, the values of S stop S_{\rm stop} and L ( N , D ) L(N,D) are empirical (though S stop S_{\rm stop} is adjusted to mimic training at B ≫ B crit B\gg B_{\rm crit} ), while L ( N , ∞ ) L(N,\infty) is computed from the fit to L ( N , D ) L(N,D) evaluated at D = ∞ D=\infty .
## 6 Optimal Allocation of the Compute Budget
We displayed the empirical trend of performance as a function of the computation used during training in the top-right of Figure 1 . However, this result involved training at a fixed batch size B B , whereas we know that in fact we could train more efficiently 5 5 5 One might ask why we did not simply train at B crit B_{\rm crit} in the first place. The reason is that it depends not only on the model but also on the target value of the loss we wish to achieve, and so is a moving target. by training at the batch size B crit B_{\rm crit} discussed in Section 5.1 . Large and small values of the loss could have been achieved with fewer samples or fewer steps, respectively, and correcting for this inefficiency by standardizing to the critical batch size results in cleaner and more predictable trends.
In this section we will adjust for this oversight. More importantly, we will use the results of Section 5 to determine the optimal allocation of compute between model size N N and the quantity of data processed during training, namely 2 B crit S min 2B_{\rm crit}S_{\rm min} . We will determine this allocation both empirically and theoretically, by using the equation for L ( N , S min ) L(N,S_{\rm min}) , and we will demonstrate that these methods agree.
### 6.1 Optimal Performance and Allocations
Let us first study the loss as a function of the optimally allocated compute from Equation ( 5.5 ). The result is plotted in Figure 13 , along with a power-law fit. We see that as compared to the compute plot of Figure 1 , the new fit with C min C_{\rm min} is somewhat improved.
Given L ( C min ) L(C_{\rm min}) , it is natural to ask for the optimal model size N ( C min ) N(C_{\rm min}) that provides the minimal loss with a given quantity of training compute. The optimal model size is shown in Figure 14 . We observe that N ( C min ) N(C_{\rm min}) can be fit very well with a power-law N ( C min ) ∝ ( C min ) 0.73 . N(C_{\rm min})\propto(C_{\rm min})^{0.73}. (6.1) In Figure 12 , we show the effect of training models of sub-optimal sizes (see Appendix B.4 ).
By definition C min ≡ 6 N B crit S C_{\rm min}\equiv 6NB_{\rm crit}S , and so we can use N ( C min ) N(C_{\rm min}) to extract further results. In particular, since prior fits show B ∝ L − 4.8 B\propto L^{-4.8} and L ∝ C min − 0.05 L\propto C_{\rm min}^{-0.05} , we can conclude that B crit ∝ C min 0.24 B_{\rm crit}\propto C_{\rm min}^{0.24} . This leads us to conclude that the optimal number of steps will only grow very slowly with compute, as S min ∝ ( C min ) 0.03 , S_{\rm min}\propto(C_{\rm min})^{0.03}, (6.2) matching the empirical results in Figure 14 . In fact the measured exponent is sufficiently small that our results may even be consistent with an exponent of zero.
Thus we conclude that as we scale up language modeling with an optimal allocation of computation, we should predominantly increase the model size N N , while simultaneously scaling up the batch size via B ∝ B crit B\propto B_{\rm crit} with negligible increase in the number of serial steps. Since compute-efficient training uses relatively few optimization steps, additional work on speeding up early training dynamics may be warranted.
### 6.2 Predictions from L ( N , S min ) L(N,S_{\rm min})
The results for L ( C min ) L(C_{\rm min}) and the allocations can be predicted from the L ( N , S min ) L(N,S_{\rm min}) equation obtained in Section 5 . Given our equation for L ( N , S min ) L(N,S_{\rm min}) , we can substitute S min = C min 6 N B S_{\rm min}=\frac{C_{\rm min}}{6NB} and then find the minimum of the loss as a function of N N , while fixing the training compute. We carry out this procedure in detail in Appendix B , where we also provide some additional predictions.
For the loss as a function of training compute, we predict that L ( C min ) = ( C c min C min ) α C min L(C_{\rm min})=\left(\frac{C_{c}^{\rm min}}{C_{\rm min}}\right)^{\alpha_{C}^{\rm min}} (6.3) where α C min ≡ 1 1 / α S + 1 / α B + 1 / α N ≈ 0.054 \alpha_{C}^{\rm min}\equiv\frac{1}{1/\alpha_{S}+1/\alpha_{B}+1/\alpha_{N}}\approx 0.054 (6.4) in excellent agreement with the exponent of Figure 13 . We also predict that N ( C min ) ∝ ( C min ) α C min / α N ≈ ( C min ) 0.71 N(C_{\rm min})\propto(C_{\rm min})^{\alpha_{C}^{\rm min}/\alpha_{N}}\approx(C_{\rm min})^{0.71} (6.5) which also matches the scaling of Figure 14 to within a few percent. Our scaling laws provide a predictive framework for the performance of language modeling.
### 6.3 Contradictions and a Conjecture
We observe no signs of deviation from straight power-law trends at large values of compute, data, or model size. Our trends must eventually level off, though, since natural language has non-zero entropy.
Indeed, the trends for compute-efficient training described in this section already contain an apparent contradiction. At scales several orders of magnitude above those documented here, the performance predicted by the L ( C min ) L(C_{\rm min}) scaling law decreases below what should be possible given the slow growth in training data with compute. This implies that our scaling laws must break down before this point, but we conjecture that the intersection point has a deeper meaning: it provides an estimate of the point at which Transformer language models reach maximal performance.
Since the amount of data used by compute-efficient training grows slowly with the compute budget, the performance predicted by L ( C min ) L(C_{\rm min}) eventually hits a lower bound set by the L ( D ) L(D) power law (see Figure 15 ). Let us work this out in more detail.
To keep overfitting under control, the results of Section 4 imply that we should scale the dataset size as D ∝ N 0.74 ∝ C min 0.54 D\propto N^{0.74}\propto C_{\rm min}^{0.54} (6.6) where we have used the compute-efficient N ( C min ) N(C_{\rm min}) from Figure 14 .
Let us compare this to the data requirements of compute-efficient training. If we train at the critical batch size (i.e. C = 2 C min C=2C_{\rm min} ) and never re-use data during training, we find that data usage grows with compute as D ( C min ) = 2 C min 6 N ( C min ) ≈ ( 4 × 10 10 tokens ) ( C min / PF - Day ) 0.26 D(C_{\rm min})=\frac{2C_{\rm min}}{6N(C_{\rm min})}\approx\left(4\times 10^{10}\ {\rm tokens}\right)(C_{\rm min}/\mathrm{PF}{\text{-}}\mathrm{Day})^{0.26} (6.7) This is the maximum rate at which the dataset size can productively grow with compute, since it means that we are only training for a single epoch. But it grows the dataset much more slowly than in Equation ( 6.6 ). It appears to imply that compute-efficient training will eventually run into a problem with overfitting, even if the training process never re-uses any data!
According to Figure 1 , we expect that when we are bottlenecked by the dataset size (ie by overfitting), the loss should scale as L ( D ) ∝ D − 0.095 L(D)\propto D^{-0.095} . This implies that the loss would scale with compute as L ( D ( C min ) ) ∝ C min − 0.03 L(D(C_{\rm min}))\propto C_{\rm min}^{-0.03} once we are data-limited. Once again, we have a contradiction, as this will eventually intersect with our prediction for L ( C min ) L(C_{\rm min}) from Figure 13 , where we found a scaling L ( C min ) ∝ C min − 0.050 L(C_{\rm min})\propto C_{\rm min}^{-0.050} .
The intersection point of L ( D ( C min ) ) L(D(C_{\rm min})) and L ( C min ) L(C_{\rm min}) occurs at C ∗ ∼ 10 4 PF - Days N ∗ ∼ 10 12 parameters , D ∗ ∼ 10 12 tokens , L ∗ ∼ 1.7 nats/token C^{*}\sim 10^{4}\penalty\ \mathrm{PF}{\text{-}}\mathrm{Days}\quad N^{*}\sim 10^{12}\penalty\ \text{parameters},\quad D^{*}\sim 10^{12}\penalty\ \text{tokens},\quad L^{*}\sim 1.7\penalty\ \text{nats/token} (6.8) though the numerical values are highly uncertain, varying by an order or magnitude in either direction depending on the precise values of the exponents from the power-law fits. The most obvious interpretation is that our scaling laws break down at or before we reach this point, which is still many orders of magnitude away in both compute and model size.