Scaling Law with Learning Rate Annealing
Scaling Laws;Full Loss Curve Prediction;Learning Rate Schedule;LLM Pretraining
Yes
No ethics review needed.
It is an interesting paper overall, but I would like to see the authors' response to the two weaknesses I listed above.
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.
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.
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.
Yes
No ethics review needed.
See weakness
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.
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.
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.
Yes
No ethics review needed.
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?
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.
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.
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.
Yes
No ethics review needed.
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.
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).
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.
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.
A new scaling law formula with learning rate annealing that can fit and predict full loss curves.
@inproceedings{
anonymous2024scaling,
title={Scaling Law with Learning Rate Annealing},
author={Anonymous},
booktitle={Submitted to The Thirteenth International Conference on Learning Representations},
year={2024},
url={https://openreview.net/forum?id=o9YC0B6P2m},
note={under review}
}
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\cdot S_1^{-\alpha} - C\cdot S_2,$$
where $L(s)$ is the validation loss at step $s$, $S_1$ is the area under the LR curve, $S_2$ is the LR annealing area, and $L_0$, $A$, $C$, $\alpha$ are constant parameters.
This formulation accounts for two main effects: (1) power-law scaling over data size, and (2) the additional loss reduction during LR annealing. 
Unlike previous studies that only fit losses at final steps, our formulation captures the entire training curve, allowing for parameter fitting using losses from any training step.
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 under any learning rate scheduler (LRS).
This approach significantly reduces computational cost in formulating scaling laws while providing more accuracy and expressiveness.
Extensive experiments demonstrate that our findings hold across a range of hyper-parameters and model architectures and can extend to scaling effect of model sizes.
Moreover, our formulation provides accurate theoretical insights into 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 democratizing scaling laws, and it is helpful to guide both research and industrial participants in refining training strategies for further LLMs.
I certify that there is no URL (e.g., github page) that could be used to find authors’ identity.
I acknowledge that I and all co-authors of this work have read and commit to adhering to the ICLR Code of Ethics.
Scaling Laws Full Loss Curve Prediction Learning Rate Schedule LLM Pretraining
I certify that there is no acknowledgement section in this submission for double blind review.
/pdf/906d4803a09e95aa16140b6aea4998ea3708faa5.pdf
foundation or frontier models, including LLMs
I understand the reciprocal reviewing requirement as described on https://iclr.cc/Conferences/2025/CallForPapers. If none of the authors are registered as a reviewer, it may result in a desk rejection at the discretion of the program chairs. To request an exception, please complete this form at https://forms.gle/Huojr6VjkFxiQsUp6.
I certify that this submission complies with the submission instructions as described on https://iclr.cc/Conferences/2025/AuthorGuide.
Scaling Law with Learning Rate Annealing
ICLR 2025 Conference Submission
ICLR.cc/2025/Conference/Submission