Scaling Law with Learning Rate Annealing

## Abstract

Abstract We find that the cross-entropy loss curves of neural language models empirically adhere to a scaling law with learning rate (LR) annealing over training steps: L ⁢ ( s ) = L 0 + A ⋅ S 1 − α − C ⋅ S 2 , 𝐿 𝑠 subscript 𝐿 0 ⋅ 𝐴 superscript subscript 𝑆 1 𝛼 ⋅ 𝐶 subscript 𝑆 2 L(s)=L_{0}+A\cdot S_{1}^{-\alpha}-C\cdot S_{2}, italic_L ( italic_s ) = italic_L start_POSTSUBSCRIPT 0 end_POSTSUBSCRIPT + italic_A ⋅ italic_S start_POSTSUBSCRIPT 1 end_POSTSUBSCRIPT start_POSTSUPERSCRIPT - italic_α end_POSTSUPERSCRIPT - italic_C ⋅ italic_S start_POSTSUBSCRIPT 2 end_POSTSUBSCRIPT , where L ⁢ ( s ) 𝐿 𝑠 L(s) italic_L ( italic_s ) is the validation loss at step s 𝑠 s italic_s , S 1 subscript 𝑆 1 S_{1} italic_S start_POSTSUBSCRIPT 1 end_POSTSUBSCRIPT is the area under the LR curve, S 2 subscript 𝑆 2 S_{2} italic_S start_POSTSUBSCRIPT 2 end_POSTSUBSCRIPT is the LR annealing area, and L 0 subscript 𝐿 0 L_{0} italic_L start_POSTSUBSCRIPT 0 end_POSTSUBSCRIPT , A 𝐴 A italic_A , C 𝐶 C italic_C , α 𝛼 \alpha italic_α are constant parameters.
This formulation takes into account two factors: (1) power-law scaling over data size, and (2) the additional loss reduction during LR annealing.
Therefore, this formulation can describe the full loss curve at each step, rather than the single loss point at the end of training.
Applying the scaling law with LR annealing and fitting only one or two training curves, we can accurately predict the loss at any given step across any learning rate scheduler (LRS).
This approach significantly reduces computational cost in formulating scaling laws while providing more accuracy and expressiveness for training dynamics.
Extensive experiments demonstrate that our findings hold across a range of hyper-parameters and model architectures, and our equation can extend to scaling effect of model sizes.
Moreover, our formulation provides accurate theoretical verification and explanation for empirical results observed in numerous previous studies, particularly those focusing on LR schedule and annealing.
We believe that this work is promising to enhance the understanding of LLM training dynamics while greatly democratizing scaling laws, and it can guide researchers in refining training strategies (e.g. critical LRS) for further LLMs 1 1 1 We welcome any feedback, comment, and discussion at AlphaXiv or HuggingFace , where we also provide our implementation codes. .

## 1 Introduction

In recent years, large language models (LLMs) have garnered significant academic and industrial attention (Brown et al., ; Touvron et al., ). The scaling law suggests that the validation loss of language models follow a power-law pattern as model and data sizes increase (Hestness et al., ; Kaplan et al., ; Henighan et al., ). This law provides a powerful framework for forecasting LLM performances before large scale training by fitting losses at smaller scales (OpenAI, ; DeepSeek-AI, ; Dubey et al., ). Numerous studies have explored on the formulation to model the scaling effect of LLMs under various different settings (Bahri et al., ; Hernandez et al., ; Caballero et al., ; Michaud et al., ; Muennighoff et al., ).

However, typical scaling law formulations focus only on the final loss at the end of training (Hoffmann et al., ). Specifically, previous approaches generally rely on a set of training runs and fit the scaling law curve solely on the final loss from each run. Essentially, the middle points with different degrees of LR annealing fail to follow typical scaling laws, which do not consider local loss drop brought by LR annealing.
The previous approach under-utilizes the training compute and fails to capture the *training dynamics* within each run. Further, the application of scaling laws in LLM developments is limited since the loss curve through the whole training process is not modeled. An expressive formulation that models full loss curves enables prediction of future training dynamics and also offers insights on understanding the learning process of LLMs.

In this study, we propose a scaling law that models the full loss curve within a complete LLM training run.
Specifically, we dive deeper into the training dynamics during LR annealing, and incorporate a LR annealing factor into the traditional scaling law formula to formulate the process.
This design is motivated by the observed correlation between LRS and loss curves, where loss gradually decreases as we consume more training steps (^2^22In this paper, we use training steps to quantify the amount of consumed data, as they are typically proportional, with data amount calculated as training steps multiplied by batch size.) and then sharply declines when the LR undergoes significant annealing (Loshchilov & Hutter, ; Smith et al., ; Ibrahim et al., ; Hu et al., ).
Fig. 4 depicts how loss curves change over different learning rate schedules.
Overall, we discover that the validation loss of a language model at any step is determined by two factors: the forward area $S_{1}$ under the LR curve and the degree of LR annealing $S_{2}$ at that step. Formally, the expectation of loss $L$ at step $s$ of a language model follows::

$$ $\displaystyle L(s)$ $\displaystyle=L_{0}+A\cdot S_{1}^{-\alpha}-C\cdot S_{2},$ (1) $\displaystyle S_{1}$ $\displaystyle=\sum\limits_{i=1}^{s}\eta_{i},$ $\displaystyle S_{2}$ $\displaystyle=\sum\limits_{i=1}^{s}\sum\limits_{k=1}^{i}(\eta_{k-1}-\eta_{k}) \cdot\lambda^{i-k},$ $$

where $\eta_{i}$ is the learning rate at step $i$, and $\lambda$ is a hyper-parameter representing the decay factor for LR annealing momentum (see Sec. 3 in detail),which typically ranges from $0.99$ to $0.999$. $L_{0}$, $A$, $C$, $\alpha$ are undetermined positive constants. $S_{1}$ is also known as the summed learning rate (Kaplan et al., ), and $S_{2}$ represents the LR annealing area. A visualization of $S_{1}$ and $S_{2}$ is provided as Fig. 1.

Eq. 1 describes how loss changes for each step in a full loss curve during training. In Eq. 1, the term $L_{0}+A\cdot S_{1}^{-\alpha}$ represents a rectified scaling law that captures the expected loss decreases as a power-law function of the number of training steps. The new term $-C\cdot S_{2}$ accounts for the further loss drop due to learning rate annealing. Remarkably, this simple formulation accurately describes the validation loss at any training step across various LRS and even allows us to predict the loss curve for unseen LRS. For example, we can fit Eq. 1 to the full loss curve of constant and cosine LRS with 20K total steps (Fig. 2), and then predict the full loss curve for various unseen LRS with longer total steps (e.g. 60K) (Fig. 3).

Figure: Figure 1: Visualization of $S_{1}$ and $S_{2}$ at the 20-th step of a cosine LR scheduler. $S_{1}$ is the forward area, i.e., sum of red grid areas, which can be approximately regarded as the total amount of movement for neural network parameters; $S_{2}$ is the decayed annealing area, i.e., weighted sum of blue grid areas, where lighter shades indicate smaller weights. Both $S_{1}$ and $S_{2}$ contribute to loss reduction, and balancing their values is crucial for achieving the lowest possible final loss.
Refer to caption: x1.png

Figure: (a) LR curves.
Refer to caption: x2.png

Figure: (a) Full curve prediction of constant LRS.
Refer to caption: x5.png

We validate our proposed equation through extensive experiments and find that: (1) Our formulation performs consistently well across various hyper-parameters and model architectures; (2) Eq. 1 can be extended to incorporate other scaling factors, such as model sizes; (3) Our proposed equation accurately fits the loss curves of open-sourced models;
(4) Our formulation can be used to verify and explain numerous previous findings regarding LR annealing and scheduling.

In Sec. 3, we derive the scaling law formulation with LR annealing and elucidate the potential theory underpinning our formulation. Extensive experiments are conducted to validate the formulation.
In Sec. 4, we apply our formulation to verify and explain the empirical results from various previous studies.
Our approach offers theoretical insights into the crux of loss drop, LR schedule, and LR annealing,
enabling LLM participants to better understand training dynamics of LLM and select optimal training recipes in advance.
In Sec. 5, we compare our approach to typical scaling law formula, such as the Chinchilla scaling law (Hoffmann et al., ). We show that
our formulation is more general and requires significantly less compute (less than 1%) to fit,
which greatly democratizes the development of LLMs and scaling laws.

## 2 Preliminary

### 2.1 Scaling Laws

Cross-entropy loss of language models on the validation set is a reliable indicator of LLMs’ performance on downstream tasks (Caballero et al., ; Du et al., ). Kaplan et al. () empirically discovered a power-law relationship between validation loss $L$ and three factors: model size $N$, dataset size $D$, and training compute.
As an application of scaling law, Hoffmann et al. () developed Chinchilla, a compute-optimal LLM, by balancing model size and dataset size. They used a simplified and intuitive scaling law equation:

$$ $\displaystyle L(D,N)=L_{0}+A\cdot D^{-\alpha}+B\cdot N^{-\beta},$ (2) $$

where $L_{0}$, $A$, $B$, $\alpha$, $\beta$ are positive constants.
Traditional scaling law formulations fit only the loss at the final training step, while ignoring losses from other steps. Note that collecting a new loss value of data size requires launching a another training run with the same LRS, which is resource-intensive.
Previous works have conducted some preliminary studies on the impact of the learning rate on the scaling laws.
For example, OpenAI and chinchilla scaling laws both report that the choice of learning
rate schedule does not influence the power-law format (Kaplan et al., ; Hoffmann et al., ).
Also, OpenAI’s experiments suggest that the specific choice of learning rate schedule has minimal impact on the final validation loss, provided that the total summed learning rate is adequately large and the schedule incorporates both a warmup stage and a final annealing stage, reducing the learning rate to nearly zero at the end of training (Kaplan et al., ).

### 2.2 Learning Rate Annealing

Learning rate annealing is a widely-used technique in training neural networks, where the learning rate is progressively reduced from a maximum to a minimum value following a pre-defined LRS. Various LRS schemes have been proposed to improve the performance and stability of model training. For example, the popular cosine LRS (Loshchilov & Hutter, ) reduces the LR in a cosine-like pattern over full training steps. WSD LRS (Hu et al., ) keeps a constant LR for the majority of training, and applies annealing only in the final (e.g. $10\%\sim 20\%$) steps. During the training of LLMs, it has been widely observed
that a more pronounced decrease in the learning rate often results in a more precipitous drop in the
validation loss (Loshchilov & Hutter, ; Ibrahim et al., ; DeepSeek-AI, ; Hu et al., ). However, to the best of our knowledge, all previous studies end with providing a rough and
qualitative description of how loss changes during LR annealing, while our work provides an accurate equation to
quantitatively formulate the loss changes during LR annealing.

## 3 Theory

In this section, we elaborate the origin, the intuition, and the experimental basis behind Eq. 1. We then validate our formula through extensive experiments.

### 3.1 Similarity between Learning Rate, Gradient Norm, and Loss

Figure: Figure 4: The shapes of LR (top), gradient norm (medium), and validation loss (bottom) curves exhibit high similarity across various LRS (labeled as different colors).
Refer to caption: x20.png

The first key observation is that the shapes of LR curve, gradient norm curve, and validation loss curve are quite similar across various LRS when training LLMs (Fig. 4). This suggests an implicit connection between learning rate and loss, where gradient norm could be the bridge.

#### Scaling Laws for Constant LRS.

A constant LRS is a special LRS, where every training step can be viewed as an endpoint of the LRS. Notably, the Chinchilla scaling law (Hoffmann et al., ) exactly fits losses of last steps, i.e., LRS endpoints.
Therefore, the expectation of validation loss of all steps under in constant LRS adheres to a power-law over training step $s$.

#### Extra Loss Changes in LR Annealing.

Unlike a constant LRS, LR annealing (or re-warmup) brings significant local changes in the loss (see Fig. 4), causing the full loss curve to deviate from the traditional power-law formulation that consider only the training steps $s$. We hypothesis that such loss changes can be captured by an additional LR ($\eta$) related term, i.e.,

$$ $\displaystyle L(s)={\color[rgb]{0,1,1}\underbrace{L_{0}+A\cdot s^{-\alpha}}_{ \text{Traditional scaling law}}}\quad{\color[rgb]{1,0,0}\underbrace{-f(\eta)}_ {\text{LR annealing term}}},$ (3) $$

where the first two terms (blue part) follow traditional scaling laws, while the last term (red part) denotes the extra loss change brought by LR annealing. Recall the similarity between learning rate and loss curves, we can form a naive guess for $f(\eta)$ as $f(\eta)=C\cdot\eta_{s}$, where $C$ is a positive constant.

#### Training Discount in Annealing.

The form of Eq. 3 is still imperfect.
Note that the gradient norm $\|\mathbf{g}\|$ decreases along with LR during the annealing process (shown in Fig. 4).
Thus, the amount of parameter movement (approximately $\eta\cdot\|\mathbf{g}\|$ per step) in the LR annealing stage declines at an almost quadratic rate compared to stages before annealing.
As the parameter movement become smaller, the change in loss also slows down accordingly. Therefore, the loss drop brought by the power law term (i.e., the first two terms in Eq. 3) should also diminish during LR annealing.
This consideration leads to an improved equation:

$$ $\displaystyle L(s)$ $\displaystyle={\color[rgb]{0,1,1}L_{0}+A\cdot S_{1}^{-\alpha}}{\color[rgb]{ 1,0,0}-f(\eta)}$ (4) $\displaystyle S_{1}$ $\displaystyle=\sum\limits_{i=1}^{s}\eta_{i},$ $$

where $S_{1}$ is the forward area, i.e., the area under the LR curve (as visualized in Fig. 1), which could be approximately interpreted as the total amount of parameter updates.

### 3.2 LR Annealing Momentum

Figure: (a) Different delay steps in the annealing process associated with different annealing steps (0.1K, 0.5K, 1K and 2K).
Refer to caption: x21.png

Another key observation is that LR annealing has momentum.
To refine the formulation of $f(\eta)$,
we design a special LRS where the LR decreases linearly from $\eta_{max}$ to $\eta_{min}$ and then increases.
The increasing stage always has a fixed slope, reaching the maximum value in 5K steps, while the slope of the decreasing stage is varied, with durations of 0.1K, 0.5K, 1K, and 2K.
Symmetrically, we design another LRS where the LR increases linearly from $\eta_{min}$ to $\eta_{max}$ and then decreases. Fig. 5 shows the corresponding LR and loss curves.

We observe a delay phenomenon between the LR and the validation loss. Firstly, the turning point of the validation loss curve consistently lags behind the turning point of the LR curve, indicating that the validation loss continuous along its previous trajectory for some steps even after the LR changes direction. Secondly, the steeper the slope of the LR annealing (or re-warmup), the more pronounced the delay phenomenon becomes.
Thirdly, given the same LR slope, the left figure (where LR decreases then increases) consistently shows a longer delay compared to the right figure (where LR increases then decreases). We discuss some possible root reasons of delay phenomenon in Sec. 6.2.

Interestingly, this phenomenon closely resembles the physical experiment of a small ball rolling down a slope.
The steeper the slope, the faster the ball accelerates.
When the ball lands, the accumulated momentum causes the ball to slide further.
Inspired by this delay phenomenon, we hypothesize that $f(\eta)$, the loss reduction induced by LR annealing, has cumulative historical formation so that the past change of learning rate will affect the following
loss curve for a few steps.
In summary, *learning rate annealing exhibits momentum*.
To capture this, we define $f(\eta)=C\cdot S_{2}$, where $S_{2}$ is calculated as:

$$ $\displaystyle m_{i}$ $\displaystyle=\lambda\cdot m_{i-1}+(\eta_{i-1}-\eta_{i}),$ (5) $\displaystyle S_{2}$ $\displaystyle=\sum\limits_{i=1}^{s}m_{i}=\sum\limits_{i=1}^{s}\sum\limits_{k=1 }^{i}(\eta_{k-1}-\eta_{k})\cdot\lambda^{i-k},$ $$

where $m_{i}$ is the LR annealing momentum at step $i$ ($m_{1}=0$), and $\Delta\eta=\eta_{i-1}-\eta_{i}$ denotes the LR annealing amount at step $i$. $\lambda$ is the decay factor that signifies how much historical information is retained.
We find that $\lambda$ values between $0.99$ and $0.999$ generally works well.
In contrast, $\lambda=0$ implies no momentum effect, reducing $f(\eta)$ to $C\cdot\eta_{s}$, which degenerate to the initial form mentioned above.
Note that $S_{2}$ applies not only to LR annealing ($S_{2}>0$), but also to LR re-warmup ($S_{2}<0$). This means that our equation is applicable to scenarios like continual pre-training, where LR re-warmup serves
as an important factor for better outcomes (see Sec. 4).
More intuitively, the definition of $S_{2}$ can be visualized
in Fig. 1, as the weighted sum of blue grid areas.

### 3.3 Final Formulation

#### Scaling Law with LR Annealing.

*Given the same training and validation dataset, the same model size, the same training hyper-parameters such as warmup steps, max learning rate $\eta_{max}$ and batch size, the language modeling loss at training step $s$ empirically follows the equation $L(s)=L_{0}+A\cdot S_{1}^{-\alpha}-C\cdot S_{2}$, where $S_{1}$ and $S_{2}$ are defined in Eq. 1. $L_{0}$, $A$, $C$, $\alpha$ are positive constants.*

Our formulation describes the loss of each training step across different LRS. It allows fitting based on a simpler LRS with shorter training steps and enables the prediction of validation losses for more complex LRS with longer training steps.

Figure: Figure 6: Loss surface of language models as a *slide* after simplification. Optimization direction could be decomposed into two directions: power-law scaling direction ($S_{1}$, sliding down) and annealing direction ($S_{2}$, inner height of the slide).
Refer to caption: x23.png

#### Only One Extra Parameter.

Notably, we add only one parameter, the coefficient of $S_{2}$ term $C$, compared to Chinchilla scaling law (Hoffmann et al., ).
That is, we utilize the fewest extra parameters but model the essential training dynamics in LR annealing to the greatest extent.

#### Loss Surface as a Slide.

To better understand our formulation, we view the loss surface of language models as a slide in Fig. 6. The optimization process can be seen as sliding down the slide according to the power-law scaling ( orange line), while oscillating on the inner wall ( blue dashed line). When the learning rate anneals ( red line), the amplitude of the oscillation decreases, resulting in a reduction in loss.

#### Balance between S 1 subscript 𝑆 1 S_{1} italic_S start_POSTSUBSCRIPT 1 end_POSTSUBSCRIPT and S 2 subscript 𝑆 2 S_{2} italic_S start_POSTSUBSCRIPT 2 end_POSTSUBSCRIPT .

Note that in Eq. 1, $\frac{\partial L}{\partial S_{1}}<0$ and $\frac{\partial L}{\partial S_{2}}<0$ always hold, indicating that increases in both $S_{1}$ and $S_{2}$ help to reduce the loss.
However, as shown intuitively in Fig. 1, there exists delicate balance between $S_{1}$ and $S_{2}$. When LR begins to anneal and $S_{2}$ starts to increase, the forward area $S_{1}$ of subsequent steps starts to diminish instead.
Our equation aptly describes this delicate balance.
In Sec. 4, we elaborate this topic in detail.

### 3.4 Experiments

#### LR Warmup.

Different warmup steps can result in different loss curves in training from scratch.
During the warmup stage, neural networks are prone to random optimization, resulting in *unpredictable* outcomes (Hestness et al., ).
Various studies, along with our own pilot experiments (Appendix A), show that LR warmup significantly accelerates model convergence.
High gradient norms are usually observed during the LR warmup stage, especially in the initial steps of training (see Fig. 4). This indicates that model parameters undergo substantial updates during this stage.
Therefore, in all our experiments, we linearly warmup LR to reach $\eta_{max}$ but compute $S_{1}$ and $S_{2}$ assuming a constant LR value $\eta_{max}$ in the warmup stage.

#### Experimental Setups.

We use standard experimental setups for LLM pre-training. In our main experiments,
the training dataset is Fineweb (Penedo et al., ) and the validation
dataset is RedPajama-CC (Computer, ). We train a $594$M non-embedding parameters LLAMA
architecture-like model (Touvron et al., ) from scratch. We use AdamW optimizer (Loshchilov & Hutter, ) with $\beta_{1}=0.9$ and $\beta_{2}=0.95$. The weight decay is set as $0.1$ and gradient clip is set as $1.0$. We set maximal learning rate as $2\times 10^{-4}$ and batch size as $4$M tokens.
To verify the robustness of our formulation across different experimental settings, we have four other distinct experimental setups (see Appendix B).
We adopt the decay factor or learning rate annealing $\lambda=0.999$ in our all experiments and discuss the impact of $\lambda$ in Sec. 6.1.

#### Fitting Details.

Given a LRS of total step $s$, i.e., a learning rate value sequence $\{\eta_{1},\eta_{2}\,\cdots,\eta_{s}\}$, the value of $S_{1}$ and $S_{2}$ for each step can be computed using Eq. 1 in advance.
To estimate $(L_{0},A,C,\alpha)$, we adopt a similar fitting method as used by Hoffmann et al. (). Specifically, we minimize the Huber loss (Huber, ) between the predicted and the observed log loss values using the L-BFGS algorithm (Nocedal, ):

$$ $\displaystyle\min_{L_{0},A,C,\alpha}\quad\sum_{\mathrm{Step~{}}i}\mathrm{Huber }_{\delta}\left(\log\hat{L}(i)-\log L(i)\right).$ (6) $$

We use the implementation of the minimize method provided by the scipy library. Huber loss is to enhance the robustness of the fitting process and we set the $\delta$ value of Huber loss to $1.0\times 10^{-3}$. To address the potential issue of local minima during fitting, we select the optimal fit by testing across a range of initial conditions. Note that in practice, we can also fit the loss curves generated by multiple LRS using the same set of parameters $(L_{0},A,C,\alpha)$. In this situation, we sum the Huber losses from Eq. 6 of all fitted LRS.

#### Fitting and Prediction Results.

We fit Eq. 1 on the loss curves under constant and cosine LRS with 20K total steps (see Fig. 2), and then predict the full loss curves under several unseen LRS with 60K total steps (see Fig. 3). The results show an almost perfect fit, achieving a coefficient of determination ($R^{2}$) greater than $0.999$. This underscores the robust capability of our equation to accurately fit loss curves across diverse LRS using a single parameter tuple.

The prediction results in Fig. 3 indicate that our formulation is broadly applicable and generalizes robustly across four unseen LRS, with a mean prediction error as low as 0.2%. Moreover, our equation can accurately predict losses even for complex LRS that include multiple LR re-warmup stages (Fig. 3(e)), despite that the loss curves used for fitting do not contain any LR re-warmup stages.

#### Extensive Experiments on Different Setups.

To demonstrate the broad applicability of our proposed equation, we conduct additional fitting and prediction experiments using various setups. (1) We use an alternative set of training hyper-parameters (Appendix C.1); (2) We test our equation on the Mixture of Experts (MoE) architecture (Appendix C.2); (3) We apply our equation to predict loss curves for a much longer ($10\times$) training run involving a 1.7B parameter model trained on 1.4T tokens (Appendix C.3).
(4) We fit the loss curves of open-sourced models, including BLOOM-176B trained on 300B tokens (BigScience, ) and OLMo-1B trained on 2T tokens  (Groeneveld et al., ) (Appendix C.4).
All experiments produce excellent results, indicating that our equation is effective across diverse experimental setups, including different training hyper-parameters, architecture, model sizes, and dataset scales.

Figure: (a) The loss drop brought by LR annealing for different model size $N$. Dashed lines represent trends over steps. The loss drop brought by LR annealing scales with data size and model size.
Refer to caption: x24.png

### 3.5 Extension to Model Size Scaling

#### Loss Drop During Annealing Scales with Model Size N 𝑁 N italic_N .

We explore the effect of model size $N$ on the loss drop during the annealing stage. Specifically, we compare the final losses obtained with a constant LRS and a WSD LRS (10% cosine annealing to $\eta_{min}=0$) to estimate the loss drop due to LR annealing. We conduct this experiment on different total steps and different model sizes. The experimental results are shown in Fig. 7(a). It suggests that the loss drop from LR annealing scales with both annealing steps and model sizes. This implies that the annealing area $S_{2}$ in our equation should also increase as the model size $N$ increases. We suppose there is a simple relationship of $S_{2}\propto N^{\gamma}$ where $\gamma$ is a positive constant.

#### Model Size Scaling.

Building on the experiments and analysis above, we extend our proposed Eq. 1 to incorporate model size scaling, based on traditional scaling laws (Eq. 2):

$$ $\displaystyle L(s,N)=L_{0}+A\cdot S_{1}^{-\alpha}+B\cdot N^{-\beta}-C\cdot S_{ 2}\cdot N^{\gamma},$ (7) $$

where $N$ is the number of non-embedding model parameters, and $B$, $\beta$, $\gamma$ are positive constants related to $N$.
We parameterize $S_{2}\propto N^{\gamma}$ via a multiplier $N^{\gamma}$ to the original annealing term $-C\cdot S_{2}$.

#### Fitting and Prediction with Model Size.

We validate Eq. 7 by fitting the full loss curves of models with varying sizes.
We then apply the obtained equation to predict full loss curve on the unseen largest model size.
Results in Fig. 7(b) show an almost perfect fit ($R^{2}>0.998$) and prediction for entire training dynamics of larger-scale models. This indicates the effectiveness and robustness of our proposed $N$-extended equation.
Additional $N$-extended experiments with other setups further confirm the robustness of our formulation (see detail in Appendix C.5).

## 4 Takeaways: Experimental Findings Verification and Explanation

We apply our proposed formulation to validate and provide a theoretical explanation for numerous existing experimental findings regarding the training dynamics of language models. These key insights also guide researchers in selecting critical LRS before initiating model training. An interesting summary is that

> The art of learning rate schedule lies in the delicate balancing act between forward area and annealing area.

### 4.1 It verifies and explains why loss drops more sharply when LR anneals.

Our equation helps researchers understand why loss values drop more sharply when LR anneals. This phenomenon has been widely observed in many previous studies.
Without loss of generality, consider the fitted equation from Fig. 2.
We analyze how the $S_{1}$-item ($A\cdot S_{1}^{-\alpha}$) and the negative $S_{2}$-item ($-C\cdot S_{2}$) impact the loss values (notated as “$\Delta$Loss”) obtained across a WSD scheduler (see Fig. 8). It can be seen that in the annealing stage, the negative $S_{2}$-item has a much more significant impact on the overall loss than the $S_{1}$-item. Therefore, the loss drops more sharply compared to the stage with constant LR.
In conclusion, LR annealing increases the annealing area, which further leads to a drastic decrease for the validation loss.

Figure: Figure 8: How S1-item and negative $S_{2}$-item changes in a WSD LRS. The upper and lower shaded area indicate the loss drop brought by $S_{1}$ and $S_{2}$, respectively, in the annealing stage.
Refer to caption: x26.png

### 4.2 Determining cosine cycle length and minimum LR in cosine LRS.

Many papers have found that in LLM pre-training using cosine LRS, setting the cosine cycle length $T$ as the total steps $S$, and setting min LR as nearly 0 (rather than 10% max LR) can lead to the optimal loss (Hoffmann et al., ; Hu et al., ; Hägele et al., ; Parmar et al., ).
We theoretically validate this observation using our equation in Fig. 9. The predicted loss curve with $T=S$ and a minimum LR of 0 indeed achieves the optimal loss in the final step.
Moreover, our equation gives a quite intuitive explanation: setting $T>S$ leads to incomplete annealing, while $T<S$ leads to a small forward area $S_{1}$ due to early annealing. Thus, the optimal configuration is to set $T$ equal to $S$. Also, setting the minimum LR to 0 0 maximizes the annealing amount, thereby increasing the annealing area $S_{2}$, which facilitates lower final loss.

Figure: Figure 9: Predicted loss curves of different cycle length $T$ and min LR in cosine LRS. The results well align with previous studies.
Refer to caption: x27.png

### 4.3 It verifies and explains the phenomenon, where constant LRS gets a lower loss than cosine LRS if setting small total steps, and vice versa.

Figure: (a) The predicted loss curve of 10K total steps.
Refer to caption: x28.png

In our experiments shown in Fig. 2, we observe that the constant LRS can sometimes yield lower final loss values than the cosine LRS.
To validate this phenomenon, we use our equation to predict the loss curve of a constant and cosine LRS for 10K and 100K total steps in Fig. 10. It can be seen that with more training steps, the cosine LRS outperform the constant LRS, while with less training steps, the constant LRS is better.
Moreover, Fig. 10(c) shows the predicted final loss of different total steps using constant and cosine LRS. It further convincingly suggests that constant LRS indeed gets a lower loss if setting small total steps, but the scaling slope is smaller than cosine LRS’s, resulting in higher loss in more steps.

We can further analyze this phenomenon by taking the derivative of $S_{1}$ and $S_{2}$ of Eq. 1. Specifically, $|\frac{\partial L}{\partial S_{1}}|=\alpha A\cdot S_{1}^{-\alpha-1}$ is a power-law decreasing function while $|\frac{\partial L}{\partial S_{2}}|=C$ is a constant.
When the training step $s$ is small, $|\frac{\partial L}{\partial S_{1}}|$ is larger than $|\frac{\partial L}{\partial S_{2}}|$ since we have a small $S_{1}$. Therefore $S_{1}$ plays a dominant role over $S_{2}$. In this case, increasing $S_{1}$ by maintaining a large LR, i.e. using a constant LRS rather than a cosine LRS, leads to more sharply decrease of the loss value.

When the training step $s$ is large,
$|\frac{\partial L}{\partial S_{1}}|$ becomes smaller than $|\frac{\partial L}{\partial S_{2}}|$. Therefore, $S_{2}$ plays a dominant role over $S_{1}$. In this case, increasing $S_{1}$ does not contribute much for decreasing the loss. It is time to start LR annealing to increase $S_{2}$.
Interestingly, this formulation aligns with the idea of WSD LRS (Hu et al., ):
In the early stages, the neural network is exploring globally and it is a suitable time to use a larger LR; In the later stages, the neural network is exploring locally and it is a suitable time to use a smaller LR.
We will delve further into WSD LRS in the following subsections.

### 4.4 It verifies and explains WSD and multi-step cosine LRS have matched or even lower loss than cosine LRS.

Figure: (a) Learning rate curves of three types of LRS.
Refer to caption: x31.png

Recent studies have shown that WSD LRS (Hu et al., ) and multi-step cosine LRS (DeepSeek-AI, ) result in lower loss compared to the traditional cosine LRS.
Our experiments also support this finding (refer to the ground-truth loss in Fig. 3(b), 3(c), 3(d)).
We validate and elucidate this finding using our proposed equation. Fig. 11 shows the learning rate curve (left) and the predicted loss drop (right) for different LRS.
The results suggest that for WSD and multi-step cosine LRS, the negative $S_{2}$-item ($-C\cdot S_{2}$) is slightly larger than that of the cosine LRS, whereas the $S_{1}$-item ($A\cdot S_{1}^{-\alpha}$) is significantly lower.
Specifically, both the WSD LRS and multi-step cosine LRS unintentionally employ strategies that marginally reduces $S_{2}$ but substantially increases $S_{1}$, leading to an overall decrease in validation loss.

### 4.5 Determining Optimal Annealing Ratio of WSD LRS.

In the case of WSD LRS, it is crucial to ascertain the optimal annealing ratio for training steps.
Hägele et al. () found that there is an optimal annealing ratio for WSD LRS, i.e. both excessively high or low annealing ratios lead to sub-optimal model performance.
This phenomenon can be further elucidated through our proposed equation.
Specifically, a high annealing ratio results in a significant reduction of the forward area $S_{1}$ while a low annealing ratio leads a diminished annealing area $S_{2}$.
Our scaling law equation describes the trade-off between the forward area $S_{1}$ and the annealing area $S_{2}$ about the annealing ratio.

Fig. 12 depicts the final loss predicted by our equation for various annealing ratios and total training steps.
The predictions form parabola-like curves, and align well with the actual experimental results reported in previous studies.
This suggests that a moderate annealing ratio, typically around 10% to 20%, is optimal, as it balances $S_{1}$ and $S_{2}$ to maximize their combined effect, thereby minimizing the overall validation loss.
Moreover, our equation can directly predict the optimal annealing ratio for different total steps without large-scale experiments, which saves a lot of resources.

Figure: (a) The relationship between the predicted final loss and the ratio of annealing steps under the condition of different total steps.
Refer to caption: x33.png

### 4.6 Determining the optimal annealing function in WSD LRS.

Figure: (a) The predicted loss curve of cosine and 1-sqrt annealing method of different annealing ratio.
Refer to caption: x35.png

For the WSD LRS, the selection of the annealing function is another important factor for obtaining further lower loss values. Hägele et al. () conclude that the 1-sqrt annealing (see Appendix 10 for the detailed formula) yields lower final losses compared to the other annealing methods (e.g. cosine).
They claim that this conclusion is true across different annealing ratios.

However, our equation suggests different results (see Fig. 13(a)). The 1-sqrt annealing approach does get a lower loss than the cosine annealing approach when using small annealing ratios (e.g. 10%), but it performs much worse than the cosine annealing approach when using 50% annealing ratio.

To verify our prediction, we conduct experiments by training models using different annealing methods and ratios within 50K total steps. As illustrated in Fig. 13(b), at a 10% annealing ratio, the 1-sqrt annealing method obtains lower final losses compared to the cosine annealing method, whereas at a 50% annealing ratio, the cosine annealing method obtains a lower final loss compared to the 1-sqrt annealing method. The experimental results align quite well with our prediction, which also overturns some of the claims made by previous works. We conclude that the optimal annealing function in WSD LRS depends on the annealing ratio.

Our scaling law equation provides an explanatory framework for these empirical observations. We draw the LR curves associated with the 1-sqrt and cosine annealing approaches in Appendix 10.
When using an small annealing ratio, the forward area $S_{1}$ of the cosine annealing method is slightly larger than that of the 1-sqrt annealing method, while the annealing area $S_{2}$ of the 1-sqrt method plays a more dominate role in decreasing the final loss. As the annealing ratio increases, the difference of $S_{1}$ between two annealing methods becomes larger and larger. Therefore the forward area $S_{1}$ gradually takes the dominate role. The delicate balance between $S_{1}$ and $S_{2}$ breaks at 50% annealing ratio, resulting in a lower final loss for the cosine annealing function.

The above analysis underscores the importance of carefully selecting the annealing strategy for the WSD LRS to optimize model training outcomes. Our equation can help predict a better annealing function without experiments, which saves a lot of resources.

### 4.7 It verifies and explains that in continual pre-training, a higher max learning rate during re-warmup leads to a higher initial peak in the loss, followed by a more rapid decrease.

In continual pre-training (CPT) settings, the LRS typically involves re-warmup to a new max LR at the start of training.
Through numerous experiments, Gupta et al. () concludes that a higher max LR during re-warmup results in a higher initial peak loss, followed by a more rapid decrease.

According to our scaling law formulation (^3^33Strictly speaking, continual pre-training process often include LR re-warmup as well as data distribution shift. Here we primarily research on the setting that there is no distribution shift between two training stages. The conclusions transfer across most cases because the loss change brought by LR re-warmup is significantly larger than the loss change brought by data distribution shift (Gupta et al., ; Ibrahim et al., ).), in the LR re-warmup process, the annealing area $S_{2}$ will reduce to a negative value ($S_{2}<0$) and thus the validation loss increases. A higher max LR in re-warmup lead to a lower annealing area $S_{2}$, and thus there would be a higher peak loss.
Moreover, a higher max LR also leads to a faster growth of the forward area $S_{1}$ after re-warmup. We use the fitted equation to predict the continual pre-training process with different max LR as shown in Fig. 14(a). The predicted loss curves reproduce a quite similar phenomenon with previous works (Gupta et al., ).

There is a more profound strategy using our equation in CPT. As shown in Fig. 14(a), after the total steps of CPT is determined, we can apply our equation to predict a better max LR and scheduler to get the lowest final loss without experiments, which saves a lot of resources.

### 4.8 It verifies and explains that in continual pre-training, the steps of re-warmup have little impact on the final loss.

Meanwhile, how many steps to re-warmup is another important issue in the continual pre-training.
Gupta et al. () find that longer re-warmup steps smooth the transition of loss curve but the number of re-warmup steps does not significantly influence the final validation loss.
We use the fitted equation to predict the continual pre-training dynamics with different re-warmup steps. The results, shown in Fig. 14(b), present a good alignment with previous observations (Gupta et al., ).

Based on our theory, given a fixed max LR, longer re-warmup steps lead to a slower decreases of the annealing area, resulting in a smoother raise in the loss curve. However, both the final values of the forward area ($S_{1}$) and the annealing area ($S_{2}$) remain relatively stable across different re-warmup steps. The annealing area ($S_{2}$) corresponding to different re-warmup steps are very close since the max and min LR remain the same.
Besides, although different re-warmup steps lead to temporary distinct loss fluctuations, re-warmup only accounts for a small portion of the whole training process. Thus, the forward area $S_{1}$ is also close across different re-warmup steps, resulting in the close overall loss across different steps of re-warmup.

Figure: (a) The predicted validation loss with different re-warmup max LR in the continual pre-training process. All the re-warmup steps are 500 steps.
Refer to caption: x37.png

## 5 Comparison with Chinchilla Scaling Law

### 5.1 Reduction to Chinchilla Scaling Law

Our scaling law equation can predict the full loss curve across any given LRS. In this section, we show
that our equation has no contradiction with traditional scaling laws, and it is a generalized form of the Chinchilla scaling law (Hoffmann et al., ). Specifically, all the final loss values for different total training steps following our equation should also follow a power-law relationship. We prove this by dividing two conditions: (1) constant LRS, and (2) other LRS.

Constant LRS. In the case of a constant LRS, the annealing area $S_{2}$ is always zero and the forward area $S_{1}=\eta_{max}\cdot s$, where $s$ is the step, and $\eta_{max}$ is the constant maximal LR. Thus, the whole loss curve formula becomes:

$$ $L(s)=L_{0}+(A\cdot\eta_{max}^{-\alpha})\cdot s^{-\alpha}=L_{0}+A^{\prime}\cdot s ^{-\alpha}$ (8) $$

which aligns with the Chinchilla scaling law equation.

Other LRS. For non-constant LRS, we use a statistical approach to show that our equation can be reduced to the Chinchilla scaling law.
Specifically, we verify whether the Chinchilla scaling law adequately fits the endpoints of loss curves predicted by our equation.
The parameter tuple of our equation is $(L_{0},A,C,\alpha)$. We then randomly sample 1000 sets of parameter tuples from some uniform distributions:
$L_{0}\sim U(1,3)$, $A\sim U(0.3,0.5)$, $C\sim U(0.2,0.6)$, $\alpha\sim U(-0.6,-0.4)$.
Each parameter tuple could be seen as the fitting result of a distinct set of experimental setups (^4^44It’s worth noting that some of these sampled parameter tuples might not be reasonable or likely to be observed in real experiments, but we choose to keep them nonetheless.). (e.g. dataset, batch size, model size, etc.).
For each generated parameter tuple, we apply our equation to predict the final loss of different total training steps on two LRS including cosine and WSD (10% annealing ratio),
ranging from 5K to 60K steps.
The predicted loss values on final steps are used to fit the chinchilla equation.
The fitting results are shown in Fig. 15.

**Table 1: Mean and standard deviation of $R^{2}$, as well as the mean Huber loss of 1,000 randomly generated parameter tuples.**
| LR Scheduler | mean($R^{2}$) $\uparrow$ | std($R^{2}$) $\downarrow$ | Huber Loss $\downarrow$ |
| --- | --- | --- | --- |
| Cosine | 0.972 | 0.056 | 0.00017 |
| WSD | 0.979 | 0.053 | 0.00013 |

Figure: (a) The predicted loss of different total steps with cosine LRS and the fitted chinchilla curve.
Refer to caption: x39.png

Each parameter tuple represents a synthetic fitting result corresponding to a distinct set of experimental setups (e.g. dataset, model size, etc.).
For each sampled parameter tuple, we apply our equation to predict the final loss for different total training steps with both cosine and WSD LRS, and then employ the predicted losses to fit the Chinchilla scaling law.
We calculate the mean and standard deviation of $R^{2}$ values for each fit.
The results in Table 1 demonstrate that Chinchilla scaling law fits well on the data predicted by our scaling law equation. Thus, our equation can be considered a generalization that can be reduced to the Chinchilla scaling law.

### 5.2 Scaling Law Fitting and Prediction Democratization

Our scaling law equation allows us to utilize all loss values from a full loss curve during training, while traditional scaling laws can only collect a single data point from the full loss curve. This feature allows us to fit scaling laws with much less cost.
For a direct comparison, we compare the computational efficiency of our approach and the Chinchilla scaling law (Hoffmann et al., ).
Specifically, we assume and evaluate the computational cost of obtaining 100 fitting points for each scaling law equation with an step interval of $K$:

- •
Adopting Chinchilla scaling law, typical cosine LRS requires total steps of at least $1K+2K+3K+\cdots+100K=5050K$;
- •
Adopting Chinchilla scaling law, WSD LRS (notating annealing ratio as $r$) requires total steps of at least $(1K+2K+3K+\cdots+100K)r+100K(1-r)=(100+4950r)K$.
- •
Adopting our scaling law, all we need is only one or two training curves with moderate total steps, such as one curve with $50K$ steps under cosine LRS. Further, the number of loss values we can get for fitting our equation is far more than $100$. Note that more fitting points generally achieve more accurate fitting results. Our equation can also utilize loss curves from various different LRS (e.g. $20K$ constant + $30K$ cosine)

**Table 2: The comparison of computational cost for fitting different scaling law equations.**
| Equation | LRS | Computational cost | Applicable to other LRS? |
| --- | --- | --- | --- |
| Chinchilla | Cosine | 100% | No |
| Chinchilla | WSD (20% annealing) | 21.6% | No |
| Chinchilla | WSD (10% annealing) | 11.8% | No |
| Ours | Any (except constant) | $<$1.0% | Yes |

We present a comparison of the computational costs associated with different laws and LRS in Table 2. The results indicate that our proposed equation uses less than 1% of the computational cost required by the Chinchilla scaling law. Further, our scaling law with LR annealing, can be universally applied to predict full loss curves for unseen LRS, thus conserving even more computational resources. This approach significantly democratizes the study of scaling laws in LLM pre-training, paving the way for a more environmentally friendly and carbon-efficient methodology.

## 6 Discussion

### 6.1 The impact of Decay Factor λ 𝜆 \lambda italic_λ

Figure: (a) Learning rate curves.
Refer to caption: x41.png

Figure: Figure 17: The predicted loss in different annealing ratios of WSD LRS for $\lambda=0.99$ and $\lambda=0.999$.
Refer to caption: x44.png

The decay factor $\lambda$ in our equation indicates the information retaining degree in LR annealing. We set $\lambda$ as $0.999$ in our all experiments. We explore the difference from another decay factor $\lambda=0.99$. We fit and get different equations for different $\lambda$. We compare them (1) on the predicted loss curves for 1-square and 1-sqrt annealing methods, and (2) on the predicted loss curves in different annealing ratios of WSD LRS (cosine annealing).

The results, illustrated in Fig. 16 and 17, reveal several key insights into the impact of decay factor:

#### Delay Steps.

A larger decay factor results in longer delay steps. Comparing
Fig. 16(b) and Fig. 16(c), $\lambda=0.999$ introduces a more obvious delay phenomenon, which is consistent across both the 1-square and 1-sqrt annealing methods.
The root reason is simple: larger $\lambda$ can retain more LR historical momentum, causing longer delay steps after LR finish annealing.

#### Optimal Annealing Ratio.

a larger decay factor tends to favor a higher annealing ratio. As shown in Fig. 17,
The optimal annealing ratio of $\lambda=0.999$
is larger than that of $\lambda=0.99$. Meanwhile, due to the similar reason, $\lambda=0.999$ favors 1-sqrt annealing method while $\lambda=0.99$ favors 1-square annealing method, as shown in Fig. 16.

#### Balance Point between S 1 subscript 𝑆 1 S_{1} italic_S start_POSTSUBSCRIPT 1 end_POSTSUBSCRIPT and S 2 subscript 𝑆 2 S_{2} italic_S start_POSTSUBSCRIPT 2 end_POSTSUBSCRIPT .

More essentially, the selection of $\lambda$ decides the balance point of $S_{1}$ and $S_{2}$.
For example,
$\lambda=0.999$ means that LR annealing only retain the information of previous approximately $\frac{1}{1-\lambda}=1000$ steps, which can be seen as the window size of LR annealing momentum. The window size could be very close to the optimal annealing steps if LR changes smoothly over steps. After reaching window size, $S_{2}$ increases very slowly, with the cost of large decrease of $S_{1}$.

The analyses above highlights the importance of selecting a decay factor that aligns closely with empirical data to ensure the accuracy of predictions.
We recommend that the future developers try different $\lambda$ for their own setups (^5^55Actually, $\lambda$ can be fitted as a parameter, instead of a hyper-parameter requiring manual tuning. We regard $\lambda$ as a hyper-parameter because $\lambda=0.999$ performs well in our all experiments. Besides, fitting with $\lambda$ could bring in additional time complexity due to the recomputation of $S_{2}$.).

### 6.2 Possible Root Reasons of Delay Phenomenon in Learning Rate Annealing

In Sec. 3, we discover the delay phenomenon, which proves that LR annealing has momentum.
We discuss possible root reasons of the phenomenon in this section.

#### Adam Optimizer? No.

We notice that Adam optimizer (Kingma & Ba, ) also has the first-order momentum decay factor $\beta_{1}$ and the second-order momentum decay factor $\beta_{2}$, which presents the possible connection to the the delay phenomenon.

We keep $\beta_{1}=0.9$, and conduct delay experiments
on different $\beta_{2}\in\{0.95,0.99,0.999\}$ (default: $0.95$) of AdamW optimizer (Loshchilov & Hutter, ) to observe whether larger $\beta_{2}$ causes a more longer delay steps.
The learning rate and ground-true loss curve are shown in Fig. 18(a). It suggests that the ground-truth loss curves of different $\beta_{2}$ almost coincide with each other, and their delay steps are also the same. Therefore, we believe that Adam optimizer has little to do with the delay phenomenon, despite its momentum form seeming very related to our experiments. Speaking of which, we even once tried to mimic the form of Adam Optimizer to describe LR annealing momentum, attempting to discover a connection between them, but the fitting results were a mess.

Figure: (a) The comparison of true loss curve with setting different $\beta_{2}$ of Adam optimizer.
Refer to caption: x45.png

#### Forward Area S 1 subscript 𝑆 1 S_{1} italic_S start_POSTSUBSCRIPT 1 end_POSTSUBSCRIPT ? Not Really.

No matter how LR changes, $S_{1}$ is always increasing over steps, resulting in consistently reducing the validation loss brought from $S_{1}$. Therefore, the forward area, $S_{1}$ would lengthen delay steps in LR annealing then re-warmup, but would shorten delay steps in LR re-warmup then annealing. The delay phenomenon is indeed related to $S_{1}$.

But still, $S_{1}$ is not all the reasons of delay phenomenon. We prove this by Fig. 5(b), which suggests that even though in LR re-warmup then annealing, the delay phenomenon, while not that pronounced, still exists. Moreover, we conduct delay experiments by adjusting the slope of LR after tuning point of LR. As shown in Fig. 18(b), We find that more smooth slope of LR re-warmup, with smaller $S_{1}$, but still causes longer delay steps. Therefore, we conclude that $S_{1}$ indeed influences the specific delay length, but is not the root reason.

#### Other Possible Reasons?

The delay phenomenon could be intuitive in some cases. For example, suppose that learning rate decreases directly from 2e-4 to 2e-5 in one step, and then maintains 2e-5. In this case, although the loss would decrease to a lower value but the parameter changes in one step is too small in neural networks. Given a sudden low LR, neural networks still require some steps to gradually optimize to a local minimum, incurring delay phenomenon.
But still, analysis above still ends with a rough description, and we have not figured out the root reasons and look forward to completing this part in future work.

### 6.3 Other Possible Scaling Law Formats with LR annealing

#### Adding a LR-weighted Coefficient to S 2 subscript 𝑆 2 S_{2} italic_S start_POSTSUBSCRIPT 2 end_POSTSUBSCRIPT ?

Imagine that when LR anneals to nearly 0 0, the neural network’s parameters almost do not change and the validation loss should not change, either.
However, as defined in our equation, Eq. 1, $S_{2}$ still has historical momentum even if LR is nearly 0 0, making the loss continue to decrease and misalign with observed training dynamics.

To cover this corner case, we try a revision to our equation and add a LR-weighted coefficient to $S_{2}$. Specifically, we adjust $S_{2}$ to more approach 0 0 when $\eta$ is close to 0 0, counteracting the original formulation’s tendency to overestimate loss reduction when $\eta\approx 0$.

The revised equation for the annealing area $S_{2}$ in our scaling law function is as follows:

$$ $\displaystyle m_{i}$ $\displaystyle=\lambda\cdot m_{i-1}+(\eta_{i-1}-\eta_{i})$ (9) $\displaystyle=\sum\limits_{k=1}^{i}(\eta_{k-1}-\eta_{k})\cdot\lambda^{i-k},$ $\displaystyle S_{2}$ $\displaystyle=\sum\limits_{i=1}^{s}m_{i}\color[rgb]{1,0,0}\cdot\eta_{i}^{ \epsilon},$ $$

where the red part is the added LR-weighted coefficient and $\epsilon$ is a undetermined positive constant. $\epsilon$ could be very small in practice.

We have tried the revised function to fit data. We find that the fitting results are quite similar and $\epsilon$ is very close to 0, showing little use in practical effect. Hence, we adopt the original format in our experiments. However, we still recommend future developers to try this format if possible.

#### L ∝ S 2 ζ proportional-to 𝐿 superscript subscript 𝑆 2 𝜁 L\propto S_{2}^{\zeta} italic_L ∝ italic_S start_POSTSUBSCRIPT 2 end_POSTSUBSCRIPT start_POSTSUPERSCRIPT italic_ζ end_POSTSUPERSCRIPT rather than L ∝ S 2 proportional-to 𝐿 subscript 𝑆 2 L\propto S_{2} italic_L ∝ italic_S start_POSTSUBSCRIPT 2 end_POSTSUBSCRIPT ?

Figure: (a) LR of cosine and constant.
Refer to caption: x47.png

Actually, all we know is that $L$ and $S_{2}$ have a positive correlation. Thus $L\propto S_{2}^{\zeta}$ rather than $L\propto S_{2}$ might be a more reasonable assumption.
That is, our equation would be changed to
$L(s)=L_{0}+A\cdot S_{1}^{-\alpha}-C\cdot S_{2}^{{\color[rgb]{1,0,0}\zeta}}$.
Theoretically, the introduction of $\zeta$ as an additional fitting parameter is expected to provide a more nuanced control over how changes in the learning rate annealing affect validation loss, potentially leading to improve the accuracy of our equation.

However, the empirical results, as depicted in Fig. 19, demonstrate that the fitting improvement with the inclusion of $\zeta$ is quite marginal when compared to the version without this parameter. This slight enhancement does not justify the additional complexity introduced by managing negative values of $S_{2}$. Furthermore, the empirical observation that $\zeta$ tends to converge close to $1$ (e.g. $1.125$ in Fig. 19(c)) reinforces the notion that the original formulation of the function, without the power term $\zeta$, is adequately robust. This finding suggests that the direct influence of the learning rate annealing area, as initially modeled, sufficiently captures the essential dynamics without the need for this additional complexity. Another additional complexity lies in that $S_{2}^{\zeta}$ becomes incalculable when $S_{2}<0$ in LR re-warmup.

## 7 Future Works

### 7.1 Tuning Scaling Law Format

We have tested a variety of equation formats to enhance the accuracy of the entire training process. As a result, the final equation format, as presented in Eq. 1, proves to be optimal so far
across a range of scenarios.
We add only one extra parameter but obtain a very good fitting and prediction result.
The formulation has achieved a level of practicality that enables the prediction of future loss when scaling training steps and model sizes. We expect more following researches to explore the format of the scaling law with learning rate annealing.

### 7.2 More Applications via Our Scaling Law

In Sec. 4, we present many instances to apply our scaling law with LR annealing, to predict future training dynamics with a cost-free manner.
We believe that our equation can help analyze and select more training recipes in specific scenarios.

### 7.3 Extension to Post-training

In this work, we research primarily on the scope of pretraining of LLM. We also show how to apply our equation to guide the LR re-warmup strategy in continual pre-training. We will continue researching on how to extend our equation to post-training, which might include data distribution shift, data mixture, model alignment, and specific downstream evaluations.

## 8 Conclusion

In conclusion, we propose that the loss curves of neural language models empirically adhere to a scaling law with learning rate annealing over training steps $s$: $L(s)=L_{0}+A\cdot S_{1}^{-\alpha}-C\cdot S_{2}$.
This equation can accurately predict full loss curves across unseen learning rate schedulers.
We present the underlying intuition and theory for deriving our equation and demonstrate that our approach can be extended to capture the scaling effect of model sizes.
Extensive experiments demonstrate that our proposed scaling law has good accuracy, scalability, and holds under various experimental setups.
It also offers accurate theoretical insights to the training dynamics of LLMs,
and explains numerous phenomena observed in previous studies. We believe that the scaling law with LR annealing is promising to reshape the understanding of researchers for LLM training and scaling laws.

## Appendix A Impact of Warmup Steps

We investigate the impact of LR warmup steps on loss curves. Specifically, different values of warmup steps are experimented. As shown in Fig. 20, we find that a warmup step value of 500 accelerates convergence and achieves the lowest validation loss compared to 100 or no warmup.
This finding is aligned with previous works Liu et al. (); Kosson et al. ().

Based on these findings,
we use 500 warmup steps in our main experiments when training models from scratch.
Note that LR warmup when training a randomly initialized model is different from LR re-warmup in continual training. In fact, the re-warmup process can be seen as a special types of LRS where $S_{2}<0$ in our equation.

Figure: Figure 20: Validation loss curves corresponding to different warmup steps. We use cosine LRS with 20K total steps.
Refer to caption: x50.png

## Appendix B Experimental Setups

In this work, we use multiple experimental setups to validate the effectiveness of our equation across a variety of conditions. For clarity, we summarize all experimental setups in Table 3. The majority of our experiments use Setting A. Additionally we also successfully fit our equation to the loss curves of BLOOM (BigScience, ) and OLMo (Groeneveld et al., ). The experimental settings used to produce these loss curves are significantly different from ours. Please refer to their papers for details on their specific experimental setups.

**Table 3: Experimental settings adopted in this work. Model size denotes the number of non-embedding parameters. Our datasets include Fineweb (Penedo et al., ) and RedPajama-CC (Computer, ). * denotes pre-training multilingual datasets including mixture of sources such as common crawls, books, arxiv, code, etc. We use AdamW Optimizer (Kingma & Ba, ; Loshchilov & Hutter, ), denoted as AO. Most experiments adopt Llama-3’s tokenizer (Dubey et al., ). Ext Llama-2’s is extended from Llama-2’s tokenizer (Touvron et al., ) by adding vocabulary.**
| Setups | Setting A (main) | Setting B | Setting C |
| --- | --- | --- | --- |
| Model Size | $594$M | $293$M | multiple |
| Train Dataset | Fineweb | Finweb | Mixture-train* |
| Val Dataset | RedPajama-CC | RedPajama-CC | Mixture-valid* |
| Total Steps | $60$K | $120$K | $143$K |
| Maximal LR | $2\times 10^{-4}$ | $2\times 10^{-4}$ | $1.381\times 10^{-3}$ |
| Warmup Steps | $500$ | $100$ | $500$ |
| Batch Size (tokens) | $4$M | $2$M | $4$M |
| Sequence Length | $4096$ | $4096$ | $4096$ |
| Tokenizer | Llama-3’s | Llama-3’s | Ext Llama-2’s |
| $\beta_{1}$,$\beta_{2}$ in AO | $0.9$, $0.9$5 | $0.9$, $0.95$ | $0.9$, $0.95$ |
| Weight Decay | $0.1$ | $0.1$ | $0.1$ |
| Gradient Clip | $1.0$ | $1.0$ | $1.0$ |
| Setups | Setting D (MoE) | Setting E (1.4T tokens) |  |
| Model Size | $8\times 106$M | $1704$M |  |
| Train Dataset | Fineweb | Mixture-train* |  |
| Val Dataset | RedPajama-CC | Mixture-valid* |  |
| Total Steps | $60$K | $350$K |  |
| Maximal LR | $2\times 10^{-4}$ | $6\times 10^{-4}$ |  |
| Warmup Steps | $500$ | $1000$ |  |
| Batch Size (tokens) | $4$M | $4$M |  |
| Sequence Length | $4096$ | $8192$ |  |
| Tokenizer | Llama-3’s | Llama-3’s |  |
| $\beta_{1}$,$\beta_{2}$ in AO | $0.9$, $0.9$5 | $0.9$, $0.95$ |  |
| Top-$k$ Experts | $2$ | - |  |
| Auxiliary Loss | $0.01$ | - |  |
| Weight Decay | $0.1$ | $0.1$ |  |
| Gradient Clip | $1.0$ | $1.0$ |  |

## Appendix C Results on Extensive Experiments Setups

### C.1 Another Set of Training Hyper-parameters

Fig. 2 and Fig. 3 show that our equation can work very well on our main experimental setup detailed in Table 3.
To validate the effectiveness of our scaling law formulation on different experimental settings, we report the results on Setting $B$ (see Table 3). The fitting results are shown in Fig. 21, and the prediction results are shown in Fig. 22. The results suggest that our scaling law with LR annealing works well across different experimental setups.

Figure: (a) Learning rate.
Refer to caption: x51.png

Figure: (a) Full curve prediction of constant LRS.
Refer to caption: x54.png

### C.2 Experiments on the Mixture of Experts Architecture

We validate our equation on the Mixture of Experts (MoE) architecture.
The experimental setting is shown as setting $D$ in Table 3.
We add the widely-used auxiliary loss for load balancing among experts (Fedus et al., ). Moreover, we change the LRS and total steps to 60K WSD with 10K annealing steps in fitting, testing whether our scaling law is effective
under various circumstances.
The fitting results shown in Fig. 23 and the prediction results shown in Fig. 24 indicates that our equation still works well on the MoE architecture.

Figure: (a) Learning rate.
Refer to caption: x63.png

Figure: (a) Full curve prediction of multi-step cosine LRS (80% + 10% + 10%) (DeepSeek-AI, )
Refer to caption: x66.png

### C.3 Scaling Up for Much Longer Steps

In this section, we validate the effectiveness of our equation in predicting loss curves over significantly long total steps. This scalability is particularly useful in large-scale training.
Specifically, we fit our equation on loss curves yield by the constant and cosine LRS with 30K steps and predict the loss curve for the WSD LRS with 350K steps. The model has 1.7 billion parameters, and the training uses 1,400 billion tokens. More detailed experimental setup is shown as the setting $E$ in Table 3.

The fitting and prediction results are shown in Fig. 25 and Fig. 26 respectively. It shows that we accurately predict the loss curve for the annealing stage after 10x longer steps in advance.

Figure: (a) Learning rate.
Refer to caption: x75.png

Figure: (a) Learning rate.
Refer to caption: x78.png

### C.4 Open-sourced Models

To further validate our proposed scaling law, we apply our equation on modeling the loss curves of open-sourced language models, including BLOOM-176B (BigScience, ) and OLMo-1B (Groeneveld et al., ). As shown in Fig. 27, our equation fits very well on loss curves of open-sourced models, even when the model size scales up to 176B (e.g. BLOOM) and token number scales up to 2000B over 740K steps (e.g. OLMo).

Figure: (a) Full loss curve fitting on BLOOM-176B.
Refer to caption: x81.png

### C.5 Our N 𝑁 N italic_N -extended Scaling Law on Another Experiments Setups

Fig. 7(b) shows that our $N$-extended equation (Eq. 7) works well on our main experimental setup.
Similarly, to validate the effectiveness of our $N$-extended scaling equation on various different experimental settings, we change our setup from setting $A$ to setting $C$ (refer to Table 3). The fitting results are shown in Fig. 28. The results suggest that our $N$-extended scaling law with LR annealing works well across different experimental setups.

Figure: Figure 28: Curve fitting on cosine LRS (143K steps to $\eta_{min}=0$) for various different model sizes using our scaling law extended to model size $N$. Refer to setting C in Table 3 for detailed experimental setups.
Refer to caption: x83.png

## Appendix D WSD Scheduler and Annealing Functions

Figure: (a) LR curve of WSD (20% 1-sqrt/cosine annealing).
Refer to caption: x84.png

Hu et al. () proposes a warmup-stable-decay (WSD) LRS including three learning rate stages, which could help get a lower validation loss compared to the typical cosine LRS. The format is like

$$ $\displaystyle WSD(s)=\begin{cases}&\frac{s}{T_{\text{warmup}}}\eta_{max},\quad s \leq T_{\text{warmup}}\\ &\eta_{max},\quad T_{\text{warmup}}<s\leq T_{\text{stable}}\\ &\eta_{min}+f(s)\cdot(\eta_{max}-\eta_{min}),\quad T_{\text{stable}}<s\leq T_{ \text{total}}\end{cases}$ (10) $$

Where $0\leq f(s)\leq 1$ is typically a decreasing function about step $s$, and $\eta_{max}$ is the maximal learning rate. Hägele et al. () consolidate the effectiveness of the WSD scheduler with numerious empirical experiments. Moreover, Hägele et al. () also find that using 1-sqrt annealing and a moderate
annealing ratio (e.g. 20%) can further decrease the final loss.
The 1-sqrt annealing is defined as:

$$ $\displaystyle f(s)=1-\sqrt{\frac{s-T_{\text{stable }}}{T_{\text{total}}-T_{ \text{stable }}}}$ (11) $$

Also,  Hägele et al. () mention the 1-square annealing method as a baseline, which is defined as:

$$ $\displaystyle f(s)=1-\left({\frac{s-T_{\text{stable }}}{T_{\text{total}}-T_{ \text{stable }}}}\right)^{2}$ (12) $$

We draw the LR curve of the WSD LRS proposed by Hu et al. () (20% and 50% 1-sqrt annealing) along with the WSD LRS with a cosine annealing in Fig. 29.

## Peer Reviews

confidence: 4
contribution: 3
initial_rating: 6
presentation: 2
questions: It is an interesting paper overall, but I would like to see the authors' response to the two weaknesses I listed above.

=============

Post-rebuttal update:

I'm satisfied with the authors' response and would like to thank them for the additional experiments and clarification.

Assuming that the authors will further clarify to what schedules their scaling laws are applicable, I'm willing to raise my score to 7.
* Finding a good scaling law to capture LR schedules is challenging, and I don't think we should reject a paper because it has some failure cases. But the paper has to make enough efforts to search for failure cases and report them clearly for future exploration.
* I would encourage the authors to think more deeply about the precise definitions of "common" and "corner" cases. What properties are shared by widely adopted schedules that make the proposed scaling law work in practice?

As only 6 and 8 options are available, I have to keep my score 6 in the system.
soundness: 3
strengths: 1. Deriving a more accurate scaling law for LLM pretraining is an important topic. This paper presents a novel scaling law for loss curve prediction.
2. The proposed scaling law is validated in many LLM pretraining settings.
3. The proposed scaling law can be applied to guide practitioners in choosing hyperparameters when setting the LR schedules.
summary: This paper proposes a scaling law that predicts the full loss curve of LLM pretraining: $L(s) = L_0 + A \cdot S_1^{-\alpha} - C \cdot S_2$. Extensive experiments demonstrate the law's effectiveness across many LLM pretraining settings. It can also be used to guide practitioners in choosing hyperparameters when setting the LR schedules.
weaknesses: 1. It is not very clear what kind of LR schedules this scaling law is applicable to and how large errors we will get for different schedules. I understand the authors' point that the law can be applied to various practical learning rate schedules. Still, it is easy to find or imagine many failure cases.
    * If we use an extremely large learning rate (say, 100) once in a while, the pretraining loss will go up to the level of random guessing repeatedly. Near the end of training, we turn to use very small learning rates. The proposed scaling law in the paper would suggest that doing this will lead to a much smaller loss than training without the extremely large learning rates, but it is hard to believe it to be true in practice. 
    * This law is invariant to padding zero learning rates at the end. For a training run with some learning rate schedule, adding more steps with zero learning rate should not change the final training loss. However, if we pad zeros at the end of the learning rate, the scaling law will predict a much smaller loss as $S_2$ will gradually increase to its maximum possible value $\eta_0$. There is a related discussion on this in Appendix H.3, and the authors proposed to add a multiplicative factor at some places of the formula. But this does not really lead to a different loss formula because the factor $\approx 1$ after fitting real data (i.e., $\epsilon \approx 0, \eta_i^\epsilon \approx 1$).
    * For the learning rate schedules in Figures 3(c) and 3(d), the scaling law's predictions deviate from the actual values by quite a lot.
    * If at each step, I randomly choose the current learning rate from a uniform distribution (e.g., U[0, 1e-4]), will the scaling law still work?
2. While the title of Section 3 is called “Theory,” there is no real theoretical analysis in the paper to justify the proposed scaling law. The power law form for $S_1$ is obtained through speculation. The momentum form of $S_2$ is obtained from the observation of delay phenomenon, but it doesn't exclude the possibility that other forms could be better, e.g., sliding window average, or the same momentum form but the momentum coefficient $\lambda$ depends on learning rates. I acknowledge that proving any theorem about the loss curve in deep learning is hard, but I would like to see more evidence that every component of the scaling law is important (and perhaps irreplaceable). Doing some ablation studies could also be informative.

---

confidence: 3
contribution: 2
initial_rating: 5
presentation: 3
questions: See weakness
soundness: 2
strengths: This paper is well-written and easy to understand. The topics studied in this paper,  scaling law for training loss when annealing learning rates, is very related to the topic of ICLR conference and is of great significance, both theoreitcally and empirically.  The specific scaling law (Equation (1)) seems novel to me. The authors also provide extensive experimental justification for the accuracy of the scaling law, e.g. for cosine or WSD schedules. Overall I feel like this scaling law is quite non-trivial and will shed some light on studying the optimal learning rate schedule in pretraining, at least under certain conditions.
summary: This paper proposes a novel scaling law for training loss of language models when annealing learning rates. The new formulation
accounts for two main effects: (1) power-law scaling over data size, and (2) the additional loss reduction during LR annealing.  Applying the scaling law with LR annealing and fitting only one or two training curves, the authors show that their scaling law can accurately
predict the loss at any given step under any learning rate scheduler, under certain restrictions. The authors provide extenisive experimental verification for their scaling law for a range of hyper-parameters and model architectures, which can even be extended to scaling effect of model sizes.
weaknesses: It is not completely clear under what assumptions the scaling law holds. Without necessary restrictions or premises, the scaling law seems problematic for the following obvious reasons.

1. Adding more steps with zero learning rate always reduces the predicted final loss. Any reasonable scaling law should be invariant to such operations like adding steps with zero learning rate which does not change the parameters at all.
2. Multiplying the entire schedule elementwise by a positive constant larger than 1 always decreases the predicted loss. This suggests that the scaling law must break when we use a sufficiently large peak learning rate. This is a very important limitation and I could not find related discussion on this.
3. Combing 1 and 2, if I solve for optimal LR schedule, what I get is at the first step I use an infinitely large LR, then I use 0 LR for many steps. This will give me a negative loss.


Besides the above main concern, I would like to see real applications of the scaling law, e.g., how we can make training or LR tuning more efficient by using this scaling law. I also would like to see more theoretical justifications for the scaling law.

---

confidence: 4
contribution: 3
initial_rating: 8
presentation: 3
questions: 1. Is it possible for given values of lambda and other parameters, to predict the optimal learning rate schedule for a given training duration?

2. It would be really interesting to study the tradeoff in MOEs for the number of experts with different training durations - should the number of experts be increased for longer training durations?
soundness: 4
strengths: The paper proposes a new scaling law for modeling the complete training loss curve for any schedule. It shows impressive results on predicting the loss curves of various lr schedules such as WSD, re-warmup, by training only on a single constant and cosine schedule trajectory. It shows results across various architectures such as MOEs as well as various model scales.
summary: The paper proposes a new scaling law which captures the complete loss curve of a training run, instead of only fitting the final loss as done by traditional scaling laws. It shows that just using a constant and cosine schedule curve, it can predict the loss of any other lr schedule, including re-warmup periods.
weaknesses: Some points in section 4.4 and 6 seem irrelevant. In section 4.4, the first point, about loss decreasing suddenly at lr annealing, the given equation has an inherent term S2 modeling the same. It feels a bit redundant to state this point, and this definitely cannot be stated as 'the reason behind why loss drops suddenly'. Similarly, in section 6, the exploration of the question 'why is there a momentum in loss while lr annealing' seems insufficient.

---

confidence: 4
contribution: 3
initial_rating: 6
presentation: 3
questions: Some questions are present in the previous section of the review.

“The results show an almost perfect fit, achieving a coefficient of determination (R2) greater than 0.999.” Is this correlation for a single fit (i.e. “Given the same training and validation dataset, the same model size, the same training hyper-parameters such as warmup steps, max learning rate ηmax and batch size) or across multiple model fits? Because fitting across multiple models would not be statistically sound, if that is the case the authors should instead report the average of R^2 across different models.
soundness: 2
strengths: Learning rate scheduling is an important problem in the training of LLMs, which had not yet been addressed in general (i.e. given a schedule predict the final loss), this paper addresses that problem. While I do not believe the results in the paper are perfect, they are an important first step in addressing this problem. As an example of their predictive power they are able to approximately predict the behavior of loss even when it is non-montonic (Fig 3d).
summary: This paper gives an analytical expression for the loss a a function of the learning rate schedule (after fixing “training and validation dataset, the same model size, the same training hyper-parameters such as warmup steps, max learning rate ηmax and batch size”). Their common expression accounts for both the benefit in loss due to more training duration as well the benefit due to learning rate annealing.
weaknesses: Weakness:
1. The results in the paper would have been much easier to interpret had there been a comparison to simple baselines. For example, how good is the fit which averages the loss curves of all the schedules tried?  (for a fixed eta_max, batch size etc.)
2. One of the main application the results is comparing different schedule. In Section 4.3 “Determining optimal annealing ratio of wsd scheduler” the authors claim to explain the commonly observed durations of annealing in WSD (10-20%). It would have been more convincing if the authors had validated their curves empirically. For example in the predicted curves 
   
   a. The difference between cosine and WSD (with optimal hyperparams is around .05)
   
   b. The optimal anneleaning fraction decreases with more training. 
   
    Are these two observations matched in practice? As an evidence against a., in Hägele 2024 they observe that after decaying cosine to 0 they can match the performance of WSD.

Meta-review: This paper proposes a scaling law for predicting the full loss curve of language model training, accounting for both data size scaling and learning rate annealing effects. The authors claim their formulation can predict loss at any training step under various learning rate schedules using only one or two training curves for fitting, demonstrating results across different model architectures and hyperparameters.

The paper addresses an important aspect of LLM training by incorporating learning rate scheduling effects into scaling laws and shows potential for optimizing learning rate schedules. However, it suffers from significant weaknesses that ultimately lead to its rejection. The lack of robust theoretical justification for the proposed scaling law form is a major concern. Additionally, the paper fails to adequately address issues with corner cases, such as padding with zero learning rates, where the law breaks down. The inability to derive an optimal learning rate schedule from the proposed law further limits its practical utility.

The primary reason for rejection is the unresolved concerns raised by Reviewer up2N, particularly regarding the handling of corner cases and the lack of a clear optimal learning rate schedule derivation. The paper does not sufficiently explore the limitations of the proposed scaling law or provide a solid theoretical foundation. These issues indicate that the work requires substantial further development before it can be considered for publication.

Decision: Reject