Temporal Scaling Law for Large Language Models

Pretraining Language Models. We mainly discuss the generation process of decoder-based generative language models. A generative language model is trained via next-token prediction. During pretraining, the text corpus is commonly segmented into token sequences. The token sequences are then fed to the model for calculating the prediction loss on each token position $i$. Formally, given a training text sequence $T_{n}$ consisting of $n$ tokens, i.e., $T_{n}=[t_{1},\cdots,t_{i-1},t_{i},\cdots,t_{n}]$, when predicting the token $t_{i}$, the language model takes the previous tokens $T_{i-1}=[t_{1},\cdots,t_{i-1}]$ as input, and generates the probability distribution $\mathbf{p}_{i}$ for the token $t_{i}$. The loss $\mathcal{L}_{i}$ w.r.t. the token $t_{i}$ is commonly calculated via the cross-entropy loss function. In decoder-based transformers Vaswani et al. (2017), a look-ahead mask is applied in the multi-head attention (MHA) module to ensure that each token can only attend to previous tokens, preserving the causal order of the generated sequence. Then, all tokens’ $\mathbf{p}_{i}$ can be calculated in a single forward pass in parallel during pretraining. The loss $\mathcal{L}_{i}$ on each token’s $\mathbf{p}_{i}$ is then averaged as the loss for sequence $T_{n}$:

We use experiment results to derive the temporal scaling law and conduct predictions after deriving parameters of the temporal scaling law from an early training period. Please refer to appendix B for more experiment details.

Train Dataset. We train LLMs on the Pile Gao et al. (2021), consisting of 22 English text domains. For all experiments, we tokenize it using the LLaMA tokenizer Touvron et al. (2023a) with a 32k vocabulary. We apply the domain weights in Su et al. (2023) for LLM pretraining.

Validation Dataset. We apply two validation datasets for test loss calculation. For the in-distribution (ID) dataset, we randomly sample 800 sequences (1024 consecutive tokens for each) for each text domain from the validation set in the Pile. For the out-of-distribution (OOD) dataset, we simply adopt the validation split of the RealNews-Like domain in the C4 dataset Raffel et al. (2020), as it is claimed to be “distinct from the Pile” Gao et al. (2021). We refer to both validation datasets as ID-Val and OOD-Val, respectively. When calculating test loss, we forward all sequences in the validation dataset and average the results. It is worth noting that calculating test loss is equivalent to calculating the perplexity (PPL) for the validation dataset since PPL can be directly acquired via test loss:

Models. We train two LLaMA-structure Touvron et al. (2023a) generative language models with 9.8M and 58M parameters to illustrate our temporal scaling law. The architectures of both are identical to the 14M and 70M parameter models in Biderman et al. (2023). The differences between model parameters can be ascribed to the different vocabulary sizes and different activation functions (i.e., SwiGLU v.s. GeLU). When utilizing the temporal scaling law for predictions and further applications, we adopt the architectures of the 410M and 1.0B parameter models in Biderman et al. (2023), which leads to 468M and 1.2B models, respectively. To further scale up and generalize the validation of temporal scaling law, we also pre-train: (1) a 6.7B model with the same architecture of LLaMA-7B, (2) a 468M GPT-NeoX model, and (3) an 2B MoE model. See appendix F for results on (2) and (3).

Training. Following LLaMA Touvron et al. (2023a), we apply the AdamW optimizer Loshchilov and Hutter (2019) with a learning rate of 3$e$-4. Following most open-source LLMs Touvron et al. (2023a); Biderman et al. (2023), we use the cosine learning rate decay schedule Loshchilov and Hutter (2017) for all experiments to guarantee the broad applicability of our work. All models are pretrained with 400B tokens and $1k$ total warmup steps.

Metric for Fitting Model Evaluation. We propose the temporal scaling law to fit the test loss during pretraining. Following the common practical choice Cheng et al. (2014), we choose the coefficient of determination ($R^{2}$) to evaluate the quality of the fit. $R^{2}$ demonstrates the proportion of the variability in the ground truth that the proposed fit could explain. Formally, we have:

where $y_{i}$ is the ground-truth test loss w.r.t. the $i$-th step, and $\overline{y}$ is the average of $y_{i}$. $f_{i}$ is the fit test loss w.r.t. the $i$-th step. $R^{2}\in(-\infty,1]$ and a perfect fit has $R^{2}=1$. A larger $R^{2}$ indicates a better fit.

The original scaling law Kaplan et al. (2020) has validated that the final test loss for training with different model sizes and data sizes follows a power-law. To validate the effectiveness of the power-law for the temporal evolution of test loss, we pretrain 9.8M and 70M models following Sec. 2.2 and fit the power-law function ($\mathcal{L}=(p_{1}x)^{p_{2}}+p_{3}$, $\{p_{i}\}$ are fitting parameters) to the test loss evolution curve (i.e., treating test loss as a whole to fit in this case) using non-linear least squares. As shown in Tab. 2, though somehow effective, applying the power-law is not precise enough to portray the temporal evolution for the test loss (follow-up experiments in Fig. 3 also demonstrate that directly using the power-law in predicting the future test loss results in larger errors). Such results motivate us to find a better method to understand the test loss evolution during LLM pretraining.

The Dynamic Hyperbolic-Law of Loss on Different Token Positions. To fit and predict the test loss more precisely, we propose to break it down and investigate the temporal scaling law from a finer granularity and look into the inherent patterns of the loss on every token position. Language models are statistical models trained to model the probabilistic process of next-token prediction. Formally, for a partial sequence of $n$ tokens $T_{n}=(t_{1},t_{2},\cdots,t_{n})$, the probability of the following token being $t_{n+1}$ given the preceding tokens is denoted as $P(t_{n+1}|T_{n})$. Intuitively and statistically, in a consecutive sequence:

since $t_{n+1}$ has a longer context than $t_{n}$. To illustrate this pattern, we plot the results of test loss for tokens in different positions, as validated on our ID-Val, using both our 9.8M and 58M models after training for 100B, 200B, 300B, and 400B tokens. We show the results for 400B tokens and leave the rest in the Appendix. As shown in  Fig. 1, both models, although varying in scales, follow the same pattern that test loss for tokens with shorter context is commonly higher than that for tokens with longer context. For both models, the loss values for tokens near the end of the sequence seem to converge at a certain value. We find that this trend can be well fit with a hyperbolic relation:

where $\mathcal{L}_{i}$ is the loss on token position $i(1\leq i\leq n)$, and $a_{0}$, $a_{1}$, and $a_{2}$ are fitting parameters solved by non-linear least squares. The fit for data from our 9.8M and 58M model are presented in Fig. 1. By applying this fit to test loss results of all checkpoints across the training phase, we find that over 99% of them can be fit with such hyperbolic relation with $R^{2}>0.95$, well demonstrating its generality. We term this phenomenon as the dynamic hyperbolic-law. See appendix C for more insights.

We also tested Kaplan’s power-law Kaplan et al. (2020) and other functions that conform to Eq. 4 to fit the curve, but the results are suboptimal. For example, using the power-law yields $R^{2}<0.75$ for over 95% checkpoints. The conclusion still holds on the OOD-Val and on larger-scale models. See detailed results in appendix E.

From the Dynamic Hyperbolic-law to the Overall Temporal Pattern. In the dynamic hyperbolic-law, $a_{2}$ is the converging factor denoting the converged value of loss on each tail token as context length increases. $a_{0}$ denotes the loss gap between the first token and the last token, and $a_{1}$ is the scaling factor along sequence length. To analyze how the derived values of $a_{0}$, $a_{1}$, and $a_{2}$ vary as training steps increase, we collect the derived value tuples of $(a_{0},a_{1},a_{2})$ for all evaluated small model checkpoints during pretraining: $[(a_{0}^{N_{1}},a_{1}^{N_{1}},a_{2}^{N_{1}}),\cdots,(a_{0}^{N_{tot}},a_{1}^{N_{tot}},a_{2}^{N_{tot}})]$, where $N_{i}$ is the number of trained tokens at the $i$-th checkpoint, and $N_{tot}$ is the total number of tokens used for pretraining. Then we plot the collected $(a_{0},a_{1},a_{2})$ values on ID-Val w.r.t. the number of tokens trained (denoted as $N$) for the 58M model in Fig. 2. It can be seen that the evolution of $a_{0}$, $a_{1}$, and $a_{2}$ presents an evident pattern that can be fit by common functions. Therefore, it is possible for us to firstly fit the temporal evolution of $a_{0}$, $a_{1}$, and $a_{2}$, and then fill them into Eq. 5 to depict the temporal evolution of the overall test loss.

First, for $a_{0}$ and $a_{1}$, we use a series of common functions to fit their temporal evolution, and eventually choose the following functions that yield the lowest fitting errors:

where the $\{\alpha_{i}\}$ and $\{\beta_{i}\}$ fitting parameters are solved by non-linear least squares. Observed from the original data in Fig. 2, we find that the value of $a_{0}$ and $a_{1}$ generally converges after a period of training, but suffer from fluctuations after converging due to the uncertainty of the first prediction position in a sequence, which has no previous context to make prediction. To mitigate the fluctuations, we define a separation point ${N}_{sep}$:

And we manually stabilize $a_{0}$ and $a_{1}$ after ${N}_{sep}$:

Here we empirically set $\epsilon=10^{-4}/\text{1B tokens}$.

For parameter $a_{2}$, we observe that ${N}_{sep}$ is also a separation point as it marks the shift of its decreasing pattern. Before ${N}_{sep}$, we apply the same function form of the loss gap factor $a_{0}$ as in Eq. 6. After ${N}_{sep}$, we find that its temporal pattern holds strong correlations with the cosine learning rate scheduler, and in turn apply a cosine relation. Similar to Eq. 6, we search the common functions and choose the best fitting function among them for $a_{2}$:

where $\{\gamma_{i}\}$ are fitting parameters solved by non-linear least squares. Interestingly, from the fit for $a_{2}$, we find that $\gamma_{5}\approx\frac{\pi}{{N}_{tot}}$, indicating that the training schedule resembles half of a cosine period, consistent with the cosine scheduler. It further validates our speculation that $a_{2}$ has a strong correlation with the cosine learning rate decay.

We plot the fit of $a_{0}$, $a_{1}$, and $a_{2}$ in Fig. 2. Our fit captures the primary patterns of parameter evolution and ignores the insignificant fluctuations. The conclusion holds on the OOD-Val, and for larger models and models with other structures. See detailed results in appendices E and F.

After acquiring $a_{0}^{N}$, $a_{1}^{N}$, and $a_{2}^{N}$, we could in turn measure the pattern of the total test loss $\mathcal{L}_{N}$ by averaging the loss on each token:

As listed in Tab. 2, the fit for all ground-truth test losses during the whole training process achieves $R^{2}>0.99$ across different settings, demonstrating the validity of our temporal scaling law. Following Gao et al. (2021), We also validate the fitting results on the LAMBADA Paperno et al. (2016) and the WikiText Merity et al. (2017) datasets, both achieving $R^{2}>0.98$. See appendix E for detailed results.

Summary. Overall, we summarize our temporal scaling law as follows:

• •

  For LLM pretraining, in each training step, the loss for tokens in different positions follows a dynamic hyperbolic-law, in which the values of parameters $(a_{0},a_{1},a_{2})$ will change as training steps increase.

• •

  The temporal patterns (i.e., values at different training steps) of dynamic hyperbolic-law parameters (i.e., $a_{0}$, $a_{1}$, and $a_{2}$) can be captured individually with their corresponding coefficients (i.e., $\{\alpha_{i}\}$, $\{\beta_{i}\}$, and $\{\gamma_{i}\}$) in the fitting functions of fixed mathematical forms.

• •

  The temporal pattern of the overall test loss is derived by averaging losses at all token positions acquired by the dynamic hyperbolic-law, i.e., filling Eq. (6-9) into Eq. (10).

The primary value of scaling laws lies in their capability of predicting training outcomes Kaplan et al. (2020). From the temporal perspective, accurate predictions of the upcoming training periods can assist in a direct selection of hyperparameters based on the target LLM’s performance, early stopping, etc. In this section, we derive practical methods from our temporal scaling law to accurately predict the overall loss in the upcoming training periods, using data from an early training period to derive parameters of the temporal scaling law, during which only $N_{train}$ tokens are trained.

To predict the overall loss, we first predict the temporal trajectory of $a_{0}^{N}$, $a_{1}^{N}$, and $a_{2}^{N}$. We could fit for $a_{0}^{N}$ and $a_{1}^{N}$ with the functions in Eq. 6:

where $\{\tilde{\alpha}_{i}\}$ and $\{\tilde{\beta}_{i}\}$ are fit using all collected $\{\mathcal{L}_{i}^{N}\}(N<N_{train})$ with non-linear least squares. Following Eq. 7, we could estimate the separation point as $\tilde{N}_{sep}$. Based on the relationship of $N_{train}$ and $\tilde{N}_{sep}$, the test loss is predicted as:

Situation #1 (In most occasions): $N_{train}\leq\tilde{N}_{sep}$, i.e., in an early training stage where the separate point is not reached yet. For $\tilde{a}_{0}^{N}$ and $\tilde{a}_{1}^{N}$, we apply the fit Eq. 11 for $N\leq\tilde{N}_{sep}$. For $\tilde{a}_{2}^{N}$, we first fit the first segment in Eq. 9:

where $\{\gamma_{i}\}$ are fit using non-linear least squares. Due to the absence of test loss data for $N_{train}>\tilde{N}_{sep}$, we could not acquire the value for $\tilde{a}_{0}^{N}$, $\tilde{a}_{1}^{N}$, and $\tilde{a}_{2}^{N}$ at that period directly. To estimate the patterns after the separation point, we propose general boundary conditions for estimation:

According to the definition in Eq. 7, we simply stabilize the values for $\tilde{a}_{0}^{N}$ and $\tilde{a}_{1}^{N}$ as $\tilde{a}_{0}^{\tilde{N}_{sep}}$ and $\tilde{a}_{1}^{\tilde{N}_{sep}}$ correspondingly for $N>\tilde{N}_{sep}$, fulfilling the derivative condition in Eq. 13 approximately. For $\tilde{a}_{2}^{N}$ at $N\geq\tilde{N}_{sep}$, to eliminate the number of unknown variables, we directly adopt the discovery from fitting Eq. 9 that $a_{2}$ resembles half of a cosine period, and directly set $\tilde{\gamma}_{5}=\frac{\pi}{N_{tot}}$. We set $\tilde{\gamma}_{6}=-\frac{\pi}{N_{tot}}N_{w}$, where $N_{w}$ is the number of tokens used for warmup training, to ensure that $N=N_{w}$ marks the start point of a cosine period. We in turn solve for the rest parameters $\tilde{\gamma}_{4}$ and $\tilde{\gamma}_{7}$ in the following function based on Eq. 13:

Situation #2 (Otherwise): $N_{train}>\tilde{N}_{sep}$, i.e., in a middle stage where the separate point is already reached: Apart from steps in Situation #1, we additionally calibrate $\tilde{\gamma}_{4}$ and $\tilde{\gamma}_{7}$ with the $N_{train}>\tilde{N}_{sep}$ data by estimating $\epsilon_{4}$ and $\epsilon_{7}$:

After acquiring the complete estimations of $\tilde{a}_{0}^{N}$, $\tilde{a}_{1}^{N}$, and $\tilde{a}_{2}^{N}$, we could calculate the test loss prediction $\tilde{\mathcal{L}}^{N}$ for each training step with Eq. 10.

Experiments. Predicting test losses in upcoming training periods is a time series modeling problem. Following standard practice Wu et al. (2021); Liu et al. (2022), we choose MSE for evaluation:

Apart from the power-law, we also choose the reciprocal ($\mathcal{L}=\frac{a_{0}}{1+a_{1}N}+a_{2}$) and the logarithmic ($\mathcal{L}=\log(a_{1}+a_{2}N)+a_{3}$) relations as baselines that model the test loss as a whole. We predict the subsequent test loss after deriving parameters of the temporal scaling law from, respectively, 10%, 20%, 30%, and 40% of the training process. Note that here we apply the larger models with 468M, 1.2B, and 6.7B parameters for validation. As shown in LABEL:{fig:fit_exp_result}, on different model scales and proportions of training, while the baselines hardly generate reliable results, our method always yields accurate results with $\text{MSE}<10^{-2}$. Our method yields far better $R^{2}$ as well. For example, when using 10% data for 1.2B model on ID-Val, our method yields $R^{2}=0.87$, while the power-law, reciprocal, and logarithmic baselines yield $R^{2}=-1.26,-3.47,-1.02$, correspondingly. Those prediction results provide strong evidence for the reliability of our temporal scaling law. The conclusion holds on the OOD-Val. See appendix G for detailed results.