A Multi-Power Law for Loss Curve Prediction Across Learning Rate Schedules

Authors: Kairong Luo, Haodong Wen, Shengding Hu, Zhenbo Sun, Zhiyuan Liu, Maosong Sun, Kaifeng Lyu, Wenguang Chen
Primary area: foundation or frontier models, including LLMs
Keywords: Large language model, Learning rate scheduler, Scaling Law, Hyperparameter optimization

Abstract:
Training large models is both resource-intensive and time-consuming, making it crucial to understand the quantitative relationship between model performance and hyperparameters. In this paper, we derive an empirical law that predicts pretraining loss for large language models for every intermediate training step across various learning rate schedules, including constant, cosine, and step decay schedules. Our proposed law takes a multi-power form, combining a power law based on the sum of learning rates and additional power laws to account for a loss reduction effect as learning rate decays. We validate this law extensively on Llama-2 models of varying sizes and demonstrate that, after fitting on a few learning rate schedules, it accurately predicts the loss curves for unseen schedules of different shapes and horizons. Moreover, by minimizing the predicted final pretraining loss across learning rate schedules, we are able to find a schedule that outperforms the widely-used cosine learning rate schedule. Interestingly, this automatically discovered schedule bears some resemblance to the recently proposed Warmup-Stable-Decay (WSD) schedule (Hu et al, 2024) but achieves a slightly lower final loss. We believe these results could offer valuable insights for understanding the dynamics of pretraining and for designing learning rate schedules to improve efficiency.

Decision: accept (poster)
Reviewer scores: 6, 5, 5, 8, 6

Reviews and discussion:
<thread>
creation date:
1730600835000

thread_type:
official_review

<thread_object>
creator: reviewer_ic14

<content>
summary:
The paper presents an empirical scaling law for loss prediction while considering learning rate schedules. The formula is in a multi-power form, which includes a power law that is applied on the sum of learning rates and another power law related to learning rate decay. The rationale behind the approximation has been explained. Loss curves with both predicted and actual values on 25M, 100M, and 400M parameter models have been obtained from experiments and shown in the paper. An approach for finding an optimal learning rate scheduled based on the empirical law is also presented.

strengths:
- The inclusion of learning rate schedules in scaling law. 
- Experiments were conducted to study various aspects in this work.

weaknesses:
- The scaling law expression in (2) is quite complicated, making it difficult to understand the intuitive meaning behind the formula. 
- To apply the result, the initial learning rate still needs to be specified. 
- One way of using the result, as presented in the paper, is to find the optimal learning rate schedule. However, the usefulness of such optimization in practice is not very clear. From Figure 1, the loss difference between the different learning rate schedules seems to be quite small. Only the loss is shown throughout the paper, and no benchmark evaluation result is presented, so it is not clear whether such small changes in loss would actually cause a significant difference in the model accuracy or not. 
- The consideration of models with 25M, 100M, and 400M parameters, adhering to the Llama-2 architecture, may be due to limitations in the GPU resource, which is understandable. However, most practical LLMs have at least 1B parameters, so the current results are insufficient for showing that this scaling law actually works for models with practical sizes. 
- Some figures, such as Figure 2, are quite hard to read and understand, due to the many curves in the figure without a clear labeling of them.

questions:
- Is it possible to analyze the sensitivity of different terms in the expression (2), and try to simplify the formula while still giving a similar prediction?
- Could you add benchmark evaluation results, to see whether the loss improvement with optimized learning schedule indeed improves the accuracy?
- For figures with two y axes in the same plot, such as Figure 2, could you put each of them into two separate plots with one y axis each, to improve the readability?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730659906899

thread_type:
official_review

<thread_object>
creator: reviewer_ytrq

<content>
summary:
The authors formulate and fit an intricate multi-power law (MPL)
function to estimate a subset of learning rate schedules using only
very simple staged LR schedules, i.e., those with constant LR over
different sections of time, for its derivation. The main contribution
is the inclusion of learning rate decay in the formulation of the
scaling law using the empirically motivated MPL term. It is important
to note that the empirical law was fit with extremely few examples,
requiring only 9 training runs for strong fits on unseen data at test
time.

The method is used to propose an optimised WSD-like schedule with a
smoother transition between the constant ("stable") and decay phase.

strengths:
The paper's findings and the design of the empirical law are extremely
exiting, even with limitations such as only being able to work with a
fixed maximum learning rate. Similarly, being able to extrapolate to
longer horizons is both important in practice but a testament to the
design of the MPL.

weaknesses:
Despite highlighting the flexibitily of the method, it is only tested
on few unseen LR schedules.

It would be nice to quantify how well/badly extrapolation to other
model sizes works.

questions:
### Please address

Page 1, line 052f:  
I do not understand: the statement is made here (and in section 5.1)
that only monotonously decreasing schedules ("we restrict our
attention to the class of LR schedules that decay the LR over time")
are used. However, in sections 2 and 3 and appendix C.1, you state
that warmup was included in every experiment and taken into
consideration in practice.

Page 10, line 516f:  
You state that the final learning rate for the "tuned" WSD schedule
was fixed to 1/10 of its peak. I don't think this makes for a fair
comparison, since your proposed optimized version decays to a far
lower final learning rate. This is consistent with results on learning
rate schedules in the LLM setting, which found that a decay to very
low learning rates (more specifically, zero) is best.

### Minor comments

Page 1, line 053:  
Should this be "$\eta_{t} >= \eta_{t + 1}$" instead of "$\eta_{t} >=
\eta_{t - 1}$"?

Page 3, line 123:  
I believe this should this be "as illustrated in Figure 2(b)", not "as
illustrated in Figure 1(b)".

Page 10, line 512f:  
"So the optimization of the WSD schedule is still disturbing."  
Please rephrase this sentence, it is unclear to me what you mean here.
Do you want to imply that optimizing the WSD schedule is difficult?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730679198769

thread_type:
official_review

<thread_object>
creator: reviewer_h6ae

<content>
summary:
This paper introduces a multi-power law for loss curve prediction of large language models. By using power law to fit loss curve of different learning rate in each step, it can accurately predict the loss curves for unseen schedules of different shapes and horizons in different sizes of Llama-2 models. Moreover, by minimizing the predicted final loss across learning rate schedules, it is able to find a schedule that outperforms the widely-used cosine learning rate schedule.

strengths:
1.	It achieves good loss curve prediction performance for 25M, 100M and 400M Llama-2 models and unseen learning rate schedules. 
2.	It can also help to find a better learning rate schedule than cosine learning schedule.

weaknesses:
1.	It is empirical law which is not well-proofed as well as fully verified. It would be better to include more kinds of LLMs and larger model sizes since it can benefit more for large model training and new model structure exploration.
2.	It only considers the relationship between learning rate and the loss curve. However, many other hyper-parameters will also affect the loss curve, such as model size, batch size, initialization seeds, etc. It would be better to take a comprehensive view so that the prediction could be more practical.
3.	It is interesting that the multi-power law can also be used to tune optimal learning rate schedule for each case. It would be better to give more results on different kinds of models and larger model sizes.
4.	It seems the parameters A, B, C, \alpha, \beta, \gama,and L0 need to be tuned per model (or even per combination of other hyper-parameters). The tuning process needs to collect data from different learning rates. Each optimization takes over 5 × 10^4 steps. The tuning cost is a little high.

questions:
Please address the questions in weaknesses.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730702205105

thread_type:
official_review

<thread_object>
creator: reviewer_uu5m

<content>
summary:
This paper introduces a multi-power law for predicting loss curves in LLMs. The law is empirically derived, with all experiments conducted using Llama-2 models of various sizes (25M, 100M, 400M). The proposed law predicts training loss curves under decaying learning rates across different iterations. Building on Chinchilla's scaling law, it incorporates additional terms to account for the impact of learning rate decay, which Chinchilla’s law does not explicitly address. The authors fit the multi-power law to three loss curves under distinct learning rate schedules, demonstrating that it can predict loss curves for unseen schedules, validated using WSD. Additionally, they applied the law to optimize learning rate schedules, identifying one that resembles WSD but achieves faster convergence.

strengths:
1. The proposed multi-power law shows potential for predicting loss curves for unseen learning rate schedules, though it has not been sufficiently validated.

2. The proposed law could be leveraged to optimize learning rate schedules, potentially accelerating LLM training.

3. The paper provides valuable insights into pretraining dynamics, offering guidance for designing learning rate schedules that can simplify the complex task of learning rate tuning.

weaknesses:
1. The theoretical justification for the proposed multi-power law is limited. The formulation of the law in Equation 2 lacks rigorous derivation or justification, raising doubts about its generalizability to unseen learning schedules.

2. The paper does not clearly state the fundamental assumptions underlying the proposed law. It is unclear whether the additional terms in Chinchilla's law (the second and third terms in Equation 2) should be model-independent, batch-size-independent, or optimizer-independent. The paper appears to assume that the impact of learning rate schedules remains independent of these factors, a point with which the reviewer may disagree.

3. The experimental section does not explore other types of learning rate schedules, such as cyclic or increasing schedules, leaving it unconvincing that the proposed law applies universally to various schedulers.

4. It remains unclear if the learning rate scheduler derived from the proposed law consistently outperforms baseline schedulers. While the paper demonstrates reduced training loss at specific iterations with the proposed scheduler, it does not clarify if this advantage holds at convergence.

5. The scope of application for the proposed law appears restrictive. The parameters of the law need to be fitted under the same task, model, batch size, and optimizer before yielding an improved learning rate schedule. It seems that the model, batch size, and optimizer must remain unchanged to apply the law effectively.

6. There is a typo on line 491, and Figure 8 lacks clarity. The curves for the baseline methods are difficult to distinguish due to missing labels.

questions:
1. How does the proposed law generalize to other model types, such as high-parameter models like Llama-7B? Currently, it has only been tested on models with 25M, 100M, or 400M parameters.

2. The paper claims that the proposed approach is effective for long-horizon predictions, but the validation is limited to training curves with cosine and constant learning rate schedulers. Could you add plots supporting this claim for additional learning rate schedulers?
</content>
</thread_object>

</thread>
<thread>
creation date:
1731112377480

thread_type:
official_review

<thread_object>
creator: reviewer_g7pt

<content>
summary:
In this paper authors introduce an empirical power law that predicts the pretrianing loss for LLMs as a given training step. Method is tested across the set of learning rate schedules: constant, cosine and step decay. Using the proposed power law authors were able to find the schedule that outperforms widely used cosine learning rate schedule.

strengths:
Improving the predictability of LLMs at different scales has been a very important task for many years. This paper attempts to make the projections very granular, in particular, authors propose an empirically tuned power law that allows to predict the training loss for a given model at every individual step across the set of learning rates. Within the range of tested LR schedules, authors also demonstrated that their method improves training loss compared to the popular cosine learning rate schedule.

weaknesses:
I think this work has a lot of potential, but it is important to conduct additional testing and analysis of the proposed method:

1. Given that the training loss demonstrates an improvement for the provided method, it is important to understand how the gain in training loss quality effects validation loss. In case the proposed method has a lower generalization capacity, it might not be beneficial to train with it even provided the improvement in the training loss. 
2. Several prior works tried to control training dynamics and make LLMs more predictable with various levels of granularity. For example, MuP paper: https://arxiv.org/abs/2203.03466 that proposes a list of adjustments in the architecture to improve predictability of the LLMs at scale. It would be very interesting to understand the benefits of the proposed method in the scenario when combined with techniques like MuP. 
3. During optimization of LLMs it is usually very important to look at the batch size together with learning rate. Here is the paper that proposes how the critical batch size should be estimate assuming the perfectly tuned learning rate: https://arxiv.org/abs/1812.06162. I think it is important to understand the temperature of the training in the proposed method and how the benefits would be affected once the batch size is tuned. 
4. One other thing that I struggled with was understanding what is the actionable outcome from this work. If the authors propose a new learning rate schedule derived from the power law equations, it should be rigorously tested across a set of downstream tasks, or at least authors should demonstrate improvements in terms of the validation loss and not just the training loss.  
5. Finally, I would recommend authors testing their methodology against the existing scaling laws that allow us to understand benefits of this method when scaling the pretraining FLOP budget.

questions:
It would be good to see the quality change on the validation or test set in addition to the training set.
If you have quality on the downstream tasks, that would be very important to look at as well for the proposed method. 
From these, I would also want to see how marginal the improvements are by looking at the stddev values for the validation loss due to the randomness.
</content>
</thread_object>

</thread>
<thread>
creation date:
1732204736083

thread_type:
official_comment

<thread_object>
creator: area_chair_6m1b

<content>
title:
No author response yet

comment:
Dear Submission13754 Authors,

ICLR encourages authors and reviewers to engage in asynchronous discussion up to the 26th Nov deadline. It would be good if you can post your responses to the reviews soon.
</content>
</thread_object>

</thread>
<thread>
creation date:
1732380445032

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
General Response

comment:
We sincerely thank all the reviewers for taking the time to review our paper and for engaging with the content of our work. We appreciate your interest in the multi-power law (MPL), and your insights and feedback have been crucial in helping us improve the paper. We have slightly reorganized the paper and added additional experiments based on the reviewers' suggestions. 

We notice that reviewers `Uu5m`, `H6ae`, `ic14`, and `ytrq` have a strong concerning in verifying MPL on larger model sizes, while reviewers `g7Pt`, `ic14`, and `H6ae` share a common focus on evaluating the optimized LR schedule on downstream tasks. We conduct additional experiments from three perspectives—model size, model type, and downstream tasks—to provide a more comprehensive verification of MPL and address the reviewers' concerns.

**Q: Verification of Multi-Power Law in larger model size. (`Uu5m`; `H6ae`; `ic14`)**

**A: We have added experiments with 1B model size, and MPL still holds in the setting with 1B model size and induces a schedule outperforming the Cosine schedule.** In these experiments, we trained models using cosine and constant LR schedules of 24,000 steps and used the resulting loss curves to fit MPL. Then we showed the MPL could predict the longer-horizon (72,000 steps) loss curve with high accuracy. We also conduct experiments to test the final validation loss of the optimized schedule of long horizon (72,000 steps), consistently outperforming the Cosine schedules. The results demonstrate that MPL retains strong generalization capabilities for 1B models. (Figure 8 \& Section 5, third paragraph \&  Appendix I, second paragraph in revision)

**Q: Downstream task evaluation of 1B Model trained with optimized LR schedule. (`g7Pt`; `ic14`; `H6ae`)**

**A: The optimized LR schedule consistently outperforms the cosine LR schedule across multiple downstream tasks**. In the experiment of 1B model, we derived the optimized LR schedule based on its corresponding MPL. Following the evaluation criteria outlined in the Mamba paper[1], we selected the first four downstream tasks to compare the performance of the cosine and optimized LR schedules. The comparison results are shown in the table below, demonstrating that the optimized LR schedule performs better than the cosine schedule. (Table 2 \& Section 5, third paragraph \& Appendix I, second paragraph in revision)

| Downstream Dataset| LAMBADA     |HellaSwag  |PIQA       |ARC-E     |
| -----------       | ----------- |-----------|-----------|----------|
| cosine schedule            | 46.54       |37.12      |**65.13**  |43.56     |
| optimized schedule       | **48.71**($\uparrow$ 2.17%)   |**37.74**($\uparrow$ 0.62%) | 65.07($\downarrow$ 0.06%)     |**44.09**($\uparrow$ 0.53%) 

**Q: Verification of Multi-Power law in Different Model Types. (`Uu5m`; `H6ae`)**

**A: MPL still holds across different model types including GPT2[2] and OLMo[3] (model size 7B)**. In the previous experiments, we used Llama2 architecture. We additionally conduct two sets of experiments to validate that the multi-power law holds across the model architecture. (Section 5, second paragraph \& Appendix I, first paragraph in revision) 
1. We pretrain GPT2 models and fit the law over 24000-step Cosine and Constant schedules. Then we validate the accuracy of unseen schedules and longer horizons and the superiority of the optimized schedules. MPL outperforms other baselines, consistently match loss curves in longer horizon with unseen LR schedule as shown in Figure 21.
2. We fit the multi-power law over the open-source training curve of 7B OLMo which uses a linear decay schedule. As a comparison, we also show the fitting of vanilla Chinchilla Law, which shows significant deviations, while the MPL shows better alignment with the OLMo curve. See Figure 7, right side.

[1] [Mamba: Linear-Time Sequence Modeling with Selective State Spaces](https://arxiv.org/pdf/2312.00752)

[2] [Language Models are Few-Shot Learners](https://arxiv.org/abs/2005.14165)

[3] [OLMo: Accelerating the Science of Language Models](https://arxiv.org/abs/2402.00838)
</content>
</thread_object>

</thread>
<thread>
creation date:
1732540973439

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
Summary of Revisions

comment:
We sincerely thank all reviewers for their valuable feedback and constructive suggestions. Based on your comments, we have enhanced the paper and clarified our contributions as follows (section references correspond to the revised version):

- **Section 1**: We reorganized the introduction for clarity, explicitly defined our problem setting, and summarized our contributions in a clear, enumerated list.  
- **Section 3.1**: We extended the experiments to include **complex and non-monotonic schedules** of **long horizon** such as cyclic, random polyline schedules (Figures 6). (`Uu5m`, `ytrq`)  
- **Section 3.2**: We added a detailed **comparison with Chinchilla Scaling Laws**, demonstrating the superior sample efficiency and prediction accuracy of the multi-power law (MPL) compared to the Chinchilla Law (Figure 7). (`g7Pt`, `Uu5m`, `H6ae`)  
- **Section 4.2**: We provided extensive validation of the **superiority of the optimized schedules**, including comparisons with zero-ending-LR variants of baseline schedules (Figures 1, 8, 17, 18, 19). To facilitate practical usage, we also introduced and evaluated approximate optimized schedules (Figures 8, 10, 21). (`g7Pt`, `ytrq`)  
- **Section 5**: We validated the robustness of the MPL across different hyperparameter settings:  
    - **Model Types** (Figure 7, 21): Added experiments with GPT2 and open-source 7B OLMo models. (`Uu5m`, `H6ae`)  
    - **Model Size** (Figure 8): Validated the law and optimized schedule using a 1B model trained on 144B tokens. (`Uu5m`, `H6ae`, `ic14`, `ytrq`)  
    - **Downstream Task Evaluation** (Table 2): Showed that 1B models trained with optimized schedules achieve improved downstream task performance compared to baseline schedules. (`g7Pt`, `H6ae`, `ic14`)  
    - **Peak Learning Rate** (Figure 16): Validated the MPL's applicability with peak learning rates of $4 \times 10^{-4}$ and $6 \times 10^{-4}$.  
    - **Batch Size** (Figure 22): Demonstrates that the MPL consistently holds across varying batch sizes of $64$ and $256$. (`g7Pt`, `Uu5m`, `H6ae`)
    - **Random Seed** (Figure 20): Demonstrated that random seed variations have a negligible effect on results. (`g7Pt`, `H6ae`)  
    - **Theoretical Discussion**: We included additional discussion on the quadratic case to address reviewer comments. (`Uu5m`)  

We believe these revisions address your concerns and strengthen the paper. Thank you for your thoughtful reviews!
</content>
</thread_object>

</thread>
<thread>
creation date:
1733043367337

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
Request for Assistance in Facilitating Reviewer Feedback

comment:
**Dear Area Chair,**  

We hope this message finds you well. We sincerely seek your assistance regarding our submission. We believe we have thoroughly addressed all concerns raised by the reviewers and are eager to engage in further discussion to ensure any remaining issues are resolved.

However, we regret that we have not yet received follow-up feedback from the majority of reviewers. We would appreciate it if you could let us know how to handle this situation or assist with further feedback from reviewers.

Thank you for your time, effort, and consideration.  

Best regards,
  
The Authors
</content>
</thread_object>

</thread>


Metareview:
metareview:
The paper extends scaling laws to account for learning rates (LRs), with special attention paid to how LRs change over time according to a schedule. Using the new scaling laws, the paper proposes a new LR schedule that approaches recent (2024) state of the art results.

Reviewers agreed that accounting for LRs in scaling laws was a valuable problem to solve. While the initial version of the paper only presented results on 400M-parameter or smaller models (which was a major concern shared by all reviewers), the authors have since added results on 1-B and 7B-parameter models.

Because the proposed scaling laws are more complex than previous ones, multiple reviewers requested additional analyses to validate the reasoning behind these choices. Not all concerns were fully addressed; at the same time, some requests (e.g. theoretical analysis) are also atypical for scaling law papers.

Overall, the paper solves an interesting problem and presents results at adequate scale, even if the analysis is imperfect in places.