[Paper review] Scaling Laws for Neural Language Models

Summary

• The test loss of a language model follows a power law with respect to each of model size $N$, dataset size $D$, and training compute $C$, when the other two are not the bottleneck. In contrast, architecture shape such as depth, width, and the number of heads has almost no effect on loss as long as the non-embedding parameter count $N$ is the same.
• When $N$ and $D$ are scaled together, the degree of overfitting is determined by a single ratio, $N^{0.74}/D$. In other words, if the model is made 8 times larger, the data only needs to be increased by about 5 times.
• Under a fixed compute budget, the optimal strategy is to make the model much larger ($N \propto C^{0.73}$) and stop training well before convergence.

1. Introduction

The performance of language modeling depends on the model architecture, the number of parameters, the amount of compute, the amount of available data, and so on. This paper targets Transformer language modeling and experimentally measures which of these factors are the main ones that actually determine performance. According to the results of several experiments conducted under varying conditions, performance scales as a power law with respect to training time, context length, dataset size, model size, and compute budget.

Since the paper compares results across many conditions, it defines and uses several notations.

$L$: Cross-entropy loss
$N$: Number of model parameters
$D$: Dataset size
$C$: Total non-embedding training compute
$C_{min}$: Estimate of minimum amount of non-embedding compute to reach a given loss

2. Background & Methods

2.1. Parameter and Compute Scaling of Transformers

The hyperparameters of the Transformer architecture are set as follows.

• $n_{layer}$: number of layers
• $d_{model}$: dimension of the residual stream
• $d_{ff}$: intermediate dimension of the feed-forward layer
• $d_{attn}$: dimension of the attention output
• $n_{head}$: number of attention heads per layer

① Model size $N$

A single layer consists of the following parameters.

• Q, K, V projections and output projection of attention: $4 \times d_{model} \times d_{attn}$
• Two linear layers of the FFN: $2 \times d_{model} \times d_{ff}$

Therefore, the total number of parameters excluding embeddings and biases is

② Compute

The amount of computation (FLOPs) for the forward pass of a single token is

Here, the second term is a context-length-dependent term arising from the attention score computation, and when $d_{model} \gg n_{ctx}/12$ it becomes much smaller than $N$, so it can be ignored. In addition, since the compute of the backward pass is about twice that of the forward pass, the total training compute per token can be approximated as follows.

3. Empirical Results and Basic Power Laws

3.1. Transformer Shape and Hyperparameter Independence

To check the effect of the Transformer's shape on performance, the loss was measured while fixing the non-embedding parameter count $N$ and changing one of $n_{layer}$, $n_{head}$, and $d_{ff}$ at a time. The resulting loss is as follows.

• Even when the feed-forward ratio, aspect ratio ($d_{model}/n_{layer}$), and head dimension are varied over a wide range, the change in loss stays within a few percent, and even architectures whose aspect ratios differ by a factor of 40 achieve similar performance.
• In other words, once $N$ is fixed, the performance of the Transformer does not depend much on its shape.

3.2. Performance with Non-Embedding Parameter Count $N$

The value used to represent the size of a model is the number of parameters in the model. Here, parameters can be viewed from two main perspectives, depending on whether only the parameters trained for the actual objective function are counted, or whether embedding parameters are also included. Therefore, to express the relationship between parameter count and model performance more accurately, it must be decided which of the two values to use as the parameter count, and the following experiment was conducted as the verification process.

• Left: When embedding parameters are included in the count, the loss appears to depend not only on the number of parameters but also on $n_{layer}$.
• Right: When embedding parameters are excluded from the count, models with different depths converge onto a single line (power law). Only models with just one layer or with extreme depth/width ratios were exceptions.

The reason it looks like the left plot is that the smaller the model, the larger the share of embedding in the total parameters. The parameter count looks large, but the part actually used for computation is small, so the loss comes out poor. For this reason, the paper uses the non-embedding parameter count as $N$ in all subsequent analyses.

With this definition, the loss can be written as a function of $N$.

3.3. Performance with Dataset Size and Compute

Dataset size $D$
To make the amount of data the bottleneck of the model, a large model ($n_{layer}=36$, $d_{model}=1280$) is trained on subsets of WebText2, and the lowest loss reachable with a given amount of data is measured. To compare the lowest losses, training is stopped once the test loss no longer decreases. The resulting relationship between loss and dataset size is as follows.

Compute $C$

The training compute is $C = 6NBS$. With $C$ fixed, $N$ is varied to find the model that achieves the lowest loss within that compute. Connecting these optimal points gives the following relationship.

The values written as $X_c$ in each equation are obtained through fitting to express the relationship on a log-log graph. The relationships of loss with respect to parameters, dataset, and compute are summarized as follows.

All three graphs are linear on log-log axes. The light blue curves in the left graph are the learning curves of models of different sizes, and their lower envelope (black line) forms the power law with respect to compute, $L = (C_{\min}/2.3 \cdot 10^8)^{-0.050}$. However, this power law holds only when the other two factors are not the bottleneck.

4. Charting the Infinite Data Limit and Overfitting

Up to this point, only one variable was made the bottleneck at a time in order to check the independent effect of each variable. This section deals with how the loss behaves when the two values $N$ and $D$ are changed simultaneously, and in particular with the overfitting that occurs when data is insufficient.

4.1. Proposed $L(N, D)$ Equation

This loss as a function of $N$ and $D$ was chosen to satisfy the following principles.

1. Rescaling: If the vocabulary size or tokenization changes, the overall loss is expected to change by a constant factor, so the equation should naturally allow for such rescaling.

2. A limit imposed by the other variable even when one variable grows infinitely: If $D$ is fixed and $N \to \infty$, the loss should converge to $L(D)$, and if $N$ is fixed and $D \to \infty$, it should converge to $L(N)$.

3. $L(N, D)$ is analytic at $D = \infty$: It should be expandable as a series in integer powers of $1/D$.

4.2. Results

Training was performed while varying $N$ and $D$, stopped once the test loss no longer decreased, and then the parameters of the above equation were fitted.

Overfitting
How much is lost compared to the loss with infinite data, $L(N, \infty)$, is defined as $\delta L$.

• The larger $\delta L$ is, the more severe the overfitting.
• From the equation, $\delta L$ depends only on $N^{\alpha_N/\alpha_D}/D \approx N^{0.74}/D$.
• Since the variation in loss due to the random seed is about 0.02, there is considered to be no overfitting when $\delta L < 0.02$. The condition that satisfies this is

In other words, when increasing model size, overfitting can be avoided by increasing data only sub-linearly. In this way, not only can the loss equation be fitted with data $D$ and parameters $N$, but the relationship between $N$ and $D$ needed to prevent overfitting can also be identified.

5. Scaling Laws with Model Size and Training Time

Now, the number of training steps $S$ is also taken into account, and the loss as a function of $N$ and training time is combined into a single equation. Before that, the critical batch size is introduced to correct for the step count and compute that vary with batch size.

5.1 Critical Batch Size $B_{crit}(L)$

Training time (number of steps) and compute vary with batch size $B$.

• $B \ll B_{crit}$: Increasing the batch reduces the number of steps almost proportionally. Compute efficiency is good.
• $B \gg B_{crit}$: Increasing the batch further barely reduces the number of steps. The number of steps is close to the minimum, but compute is wasted.

The number of steps $S$ needed to reach a target loss and the amount of data processed $E = BS$ have the following relationship. Here, $E$ is the number of tokens actually processed, which differs from the dataset size $D$.

• $S_{\min}$: minimum number of steps needed to reach the target loss (when $B \to \infty$)
• $E_{\min}$: minimum amount of data needed to reach the target loss (when $B \to 0$)

The critical batch size is defined as the ratio of the two.

• Since $E_{\min}$ and $S_{\min}$ are determined by the target loss, $B_{crit}$ is also a function of the loss.
• Training with $B = B_{crit}$ gives $S = 2S_{\min}$ and $E = 2E_{\min}$, which is the balance point between time and compute.

Empirically, $B_{crit}$ depends only on the loss, independent of model size, and follows the power law below.

The critical batch size measured as a function of loss is shown below.

Every time the loss decreases by 13%, $B_{crit}$ roughly doubles. Using this, the results ($S$, $C$) of training with an arbitrary batch size $B$ can be converted into the minimum number of steps when $B \gg B_{crit}$ and the minimum compute when $B \ll B_{crit}$.

5.2. Performance with Model size and Compute

Using $S_{\min}$, the loss can be written as a function of $N$ and training steps.

The first term is the limit set by model size when trained infinitely, and the second term is the loss caused by insufficient training.

• Left: When compute is fixed, there exists an optimal model size that minimizes the loss. A model that is too small lacks capacity, and a model that is too large cannot be trained sufficiently with the same compute.
• Right: When the number of steps is fixed, the larger the model, the lower the loss.
• The dotted lines are the $L(N, S)$ equation, which fits the experimental values well except early in training (small $S$).

5.3. Lower bound on Early Stopping Step

The learning curves for finite and infinite data are nearly identical up to around $S_{\min} \approx S_{stop}$. Therefore, the lower bound on the number of steps at which training must be stopped due to overfitting can be estimated as follows.

The smaller the loss gap caused by insufficient data, $L(N,D) - L(N,\infty)$, the longer training can continue before overfitting begins.

6. Optimal Allocation of the Compute Budget

6.1. Optimal Performance and Allocations

Using $C_{\min}$ corrected with $B_{crit}$, the loss can be rewritten as a function of $C_{\min}$, $L(C_{\min})$. Then, finding the optimal model size $N(C_{\min})$ that minimizes the loss for a given compute gives

Since $C_{\min} = 6NB_{crit}S$ and $B_{crit} \propto C_{\min}^{0.24}$, the remaining number of steps is

In other words, the conclusion is as follows.

When scaling up a language model with an optimal allocation of compute, one should mainly increase the model size $N$, increase the batch size along with $B_{crit}$, and barely increase the number of steps.

6.2 Predictions from $L(N, S_{\min})$

By combining the $L(N, S_{\min})$ equation from Section 5.2 with the $B_{crit}(L)$ equation, the exponent of $L(C_{\min})$ can be predicted theoretically.

This agrees well with the directly measured value of 0.050. This is evidence that the three independent power laws are consistent with one another.

6.3 Contradictions and a Conjecture

Continuing to apply these scaling laws leads to a contradiction.

• To avoid overfitting, data must grow as $D \propto N^{0.74} \propto C_{\min}^{0.54}$.
• However, the data actually seen when training compute-efficiently is

which grows much more slowly than the required rate. Therefore, even without reusing data, training will eventually enter the overfitting regime.

• $L(C_{\min})$ (orange dotted line) decreases quickly as compute increases. To the left of the intersection, compute is the bottleneck, so increasing compute reduces the loss.
• $L(D(C))$ (red line) is the lower bound on the loss achievable with the amount of data seen at that compute. This line decreases more slowly.
• To the right of the intersection, $L(C_{\min})$ falls below $L(D(C))$, which would mean the loss drops below the limit possible with that data even though data is insufficient, which is impossible. Therefore, the scaling laws must break down before the intersection.

The intersection is at roughly $C^* \sim 10^4$ PF-days, $N^* \sim 10^{12}$ parameters, $D^* \sim 10^{12}$ tokens, and $L^* \sim 1.7$ nats/token. The paper conjectures that the loss at this point, $L^*$, may be an estimate of the per-token entropy of natural language. However, since the location of the intersection is highly sensitive to the power-law exponents, it cannot be taken as an exact value.

7. Discussion

• The loss of a language model follows consistent power laws in $N$ (model size), $D$ (data), and $C_{\min}$ (compute), and depends very little on architecture or hyperparameters.
• Large models are far more sample-efficient than previously thought. They need fewer steps and less data to reach the same loss. Therefore, given a fixed compute budget, it is more efficient to train a large model and stop before convergence than to train a small model until convergence.

8. My Take

1. What the loss at the intersection means
The loss $L^*$ at the intersection where the two power laws meet in the compute-data tradeoff corresponds, if the paper's conjecture is right, to the per-token entropy of natural language. If so, no matter how much model size or compute is increased, the loss cannot be pushed below this value. However, as the paper itself notes, this is a conjecture, and the location of the intersection is highly sensitive to the exponent values.

2. What if $D$ were defined as information content rather than token count?
Following this interpretation, I think what ultimately matters is the amount of information each token carries when the data is encoded. This paper simply defines dataset size $D$ as the number of tokens and expresses $L$ as a function of $D$. If information content were defined as something like the total dimension of the encoded data and used in place of $D$, I wonder whether a more accurate relationship could be obtained than when using the token count.

3. Which to scale: the model or the data
The LLaMA paper (based on the results of Chinchilla (Hoffmann et al., 2022)) holds that, for the same cost, training on more new data yields better performance than only increasing model size. In contrast, comparing the exponents in this paper, as compute increases, the optimal model size grows as $N \propto C^{0.73}$ while data grows only as $D \propto C^{0.26}$, so the model grows much more than the data. It is interesting that the two papers reach nearly opposite conclusions on the same question.

References

• Kaplan, J., McCandlish, S., Henighan, T., Brown, T. B., Chess, B., Child, R., Gray, S., Radford, A., Wu, J., & Amodei, D. (2020). Scaling Laws for Neural Language Models. arXiv:2001.08361
• McCandlish, S., Kaplan, J., Amodei, D., & OpenAI Dota Team (2018). An Empirical Model of Large-Batch Training. arXiv:1812.06162
• Hoffmann, J. et al. (2022). Training Compute-Optimal Large Language Models. arXiv:2203.15556
• Touvron, H. et al. (2023). LLaMA: Open and Efficient Foundation Language Models. arXiv:2302.13971

---

This post was translated from the original Korean version with the help of AI, so some errors may remain.