Scaling Law with Learning Rate Annealing

Q1: Some points in section 4.4 and 6 seem irrelevant.

In section 4.4, the first point, about loss decreasing suddenly at lr annealing

Thanks for your valuable feedback! $S_2$ is not the only factor that influence losses in LR annealing. Actually, when LR anneals, $S_1$ increases less and makes loss decrease less suddenly, while $S_2$ makes loss decrease more suddenly (see Fig. 22 for the specific trends). There is a trade-off between $S_1$ and $S_2$ and the more essential reason is that $S_2$ plays a more dominant role than $S_1$ in LR annealing. We agree that it is more empirical rather than real physics behind the phenomenon. We have revised our statements in our paper.

In section 6, the exploration of the question "why is there a momentum in loss while lr annealing" seems insufficient.

We are thankful for your thorough review and valuable comment on our work. We derive the momentum form based on the delay phenomenon in section 3.2. In section Appendix H.2, we analyzed three possible root reasons behind the delay phenomenon, including Adam optimizer, $S_1$ in our law, and a proposed simple intuition. We found that while these factors are related, they are not the root reasons. We believe this is a really important but unresolved issue that requires further investigation in the future (as we explicitly state this in the paragraph "Other Possible Reasons?" in Appendix H.2.).

Q2: Is it possible for given values of lambda and other parameters, to predict the optimal learning rate schedule for a given training duration?

Yes! We answer this question from two aspects: engineering and mathematics.

1. From a practical engineering aspect, we have indeed presented how to get the (nearly) optimal learning rate schedule in this work. Specifically, WSD LRS should be used to confirm the larger values for both $S_1$ and $S_2$, as stated in section 4.2 and Appendix G.2; WSD LRS annealing ratio can be determined by the method stated in section 4.3; Annealing function can be determined by the method stated in Appendix G.3. In practical engineering, it is quite easy to get the (nearly) optimal learning rate schedule based on the composition of the methods above.
2. From a mathematical aspect, it becomes an optimization problem, that is, searching a learning rate sequence $(\eta_1, \eta_2, \cdots, \eta_s)$, to optimize $\text{max }L(s)= L_0 + A\cdot S_1^{-\alpha} - C\cdot S_2$ . We believe that the solution can be built up by some non-linear optimization algorithms (e.g. neural networks + gradient decent). This is a quite interesting thing on which the following work can have deeper explorations. At this stage, we believe that the approach based on the practical engineering mentioned above is sufficient to obtain a (nearly) optimal learning rate schedule.

Q3: 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?

We are glad that you are interested in the MoE part in our work. Our scaling law formulation could not directly reflect the relationship between the number of experts and training durations. However, we believe that when the number of experts changes (and controlling other experimental settings), the training loss curves will change, too, and $L_0, A, C, \alpha$ in our scaling law will change accordingly. That is to say, it is quite possible to research this interesting question by researching how $L_0, A, C, \alpha$ changes over the number of experts. Also, it is feasible to substitute specific training durations into our scaling law to calculate the loss and compare which is better for different number of experts. We do not conduct the experiment in this work, but we believe that the our scaling law can be extended to many application scenarios (e.g. # experts v.s. training durations), and providing practical suggestions for researchers.

Finally, we thank again for your valuable feedback! If you have any other questions or concerns, we are happy to answer and continually improve our work based on your feedback.

Q1: Comparison to simple baselines. How good is the fit which averages the loss curves of all the schedules tried?

We deeply appreciate the time and effort you have put into reviewing our paper and providing valuable feedback. As you commented, we are the first the propose a scaling law to model the full loss curve. Therefore, there is no other baselines to compare with. Still, we would like to provide some simple baselines (ablation studies) to more comprehensively reflect the performance of our law. In each form, we re-fit based on 20K cosine + constant and predict the schedules in Fig. 3.

We have also included the results in our revised paper (Table 4).

Q2: Are the observations matched in practice in Section 4.3 ?

Thanks for the valuable feedback, we indeed validate our predictions empirically. Specifically, in Figure 24, we compare 10%, 20%, and 50% annealing ratios (in both ground-truth and prediction), The ground-truth and prediction coincide well, and we find that in our experimental setting, 50% annealing ratio is worse than 10-20% annealing ratio. Here we also present more explanations for the phenomenon you mentioned.

b. The optimal annealing fraction decreases with more training.

This is exactly the same conclusion made by Hägele' work [1]. Here we copy the statement from their paper (the caption of Fig. 6)

A long training run suggests that a small number of cooldown steps can match cosine for long training... we find that the required duration of cooldown to match the cosine loss decreases with longer training.

a. The difference between cosine and WSD (with optimal hyperparams is around .05)

We believe that the difference between cosine and WSD LRS is affected by the detailed experimental settings. For example, we present the difference between cosine and constant LRS in Fig. 23 in Appendix G.2. It is shown that if we set small total steps, the final loss of constant LRS could be even lower than cosine LRS, and vice versa.

Therefore, what really matters is that our scaling law is able to predict the final loss under all kinds of circumstances. We believe that this finding is Not against the evidence of Hägele' work [1] because we have totally different experimental settings (and thus we have different $L_0, A, C, \alpha$). Assuming we fit and get the parameters for their loss curves, we believe that the fitted equation can also lead to the same conclusion with theirs. In summary, our scaling law can uncover some general patterns (e.g., the optimal annealing fraction decreases with more training.) and provide specific predictive results under some particular experimental setups (e.g., the differences among various LRS at a specific step).

[1] Scaling Laws and Compute-Optimal Training Beyond Fixed Training Durations

Q3: Is this correlation for a single fit or across multiple model fits?

This is a single fit. We have labeled the $R^2$ value on all fitting plots, and each fitting plot represents a single fit, indicating that multiple loss curves are simultaneously fitted using exactly the same set of parameters.

Finally, we thank again for your valuable feedback! If you have any other questions or concerns, we are happy to answer and continually improve our work based on your feedback.

We are very grateful for the opportunity to improve our work through your valuable feedback. If you think our response has addressed your concerns and questions, we sincerely hope you might consider raising your score. Should you have any further inquiries or require additional clarification, we warmly welcome your questions and are eager to provide comprehensive responses.

Q4: Real applications of the scaling law.

We clarify the applications of the scaling law from two aspects:

1. Optimizing LR schedule: Many previous works have pointed out that LR schedule has a great impact on the final validation loss like [1], [2], [3]. We believe that Section 4, Application is well aligned with real and practical scenarios for tuning LR schedule. And reviewer rnjo, 5e85, up2N (you) all agree that our scaling law can be applied to study the optimal learning rate schedule. Specifically, Section 4.1 shows how to determine min lr and cycle length in cosine LRS; Section 4.2 shows how to use our scaling law to select LRS among cosine, and multi-step cosine WSD; Section 4.3 shows how to determine the annealing ratio in WSD LRS; Appendix G.3 shows how to determine the annealing function in WSD LRS ...

   Notably, Among all these applications, we judge which LRS is better or worse only through our scaling law prediction, which is cost-free and really efficient.

2. Doing exactly what Chinchilla scaling law does, but much more efficiently. Chinchilla scaling law, $L(D) = L_0 + A\cdot D^{-\alpha}$ ($D$ denotes dataset size) has been widely used to many scenarios like predicting the final loss when consuming much more tokens. In Section 5, we compare our scaling law with Chinchilla scaling law. We show that our scaling law is a generalized form of Chinchilla scaling law and can be reduced to Chinchilla scaling law. Simply speaking, data points that follow Chinchilla scaling law, also follow our scaling law. However, Chinchilla scaling law can only adopt the final loss as fitting points while our scaling law makes use of every points including intermediate losses. Therefore, as shown in Table 2, our scaling law can do what chinchilla do, but with only <1% cost for training to collect fitting points.

[1] MiniCPM: Unveiling the Potential of Small Language Models with Scalable Training Strategies

[2] Scaling Laws and Compute-Optimal Training Beyond Fixed Training Durations

[3] DeepSeek LLM: Scaling Open-Source Language Models with Longtermism

Q5: More theoretical justifications.

We are very glad that you are interested in the more theoretical justifications of our scaling law. We honestly admit that our scaling law is empirical, and we have claimed that our scaling law is empirical in the first line of the abstract. However, we would like to point out that even the Chinchilla and OpenAI scaling laws are, in fact, empirical, as explicitly stated in their respective papers. We fully understand your concern regarding the lack of theoretical justification, but we believe that in AI domain, theoretical proofs often follow empirical works. For example, [4] presents some reliable theoretical justifications for scaling law after three years since OpenAI scaling law paper is published. We are continuously working on the theoretical aspects and this is indeed a very important future work.

Actually, surprisingly, some related works have emerged very recently [5]. Their theoretical derivations are well matched with our scaling law. If you are interested in the theoretical justifications, we would greatly appreciate it if you could kindly take the time to read it.

Finally, we thank again for your valuable feedback! If you have any other questions or concerns, we are happy to answer and improve our work based on your feedback.

[4] The Quantization Model of Neural Scaling

[5] Understanding Warmup-Stable-Decay Learning Rates: A River Valley Loss Landscape Perspective

Q1: Adding more steps with zero learning rate always reduces the predicted final loss.

Thank your for your valuable comment!

We have indeed discussed this scenario in "Section 6: Discussion" (Line 527-528) and introduced a possible variant of our equation in appendix H.3. Specifically, we can add a LR-weighted coefficient to $S_2$, where $S_2=\sum \limits_{i=1}^{s}m_i\cdot \color{red} \eta_{i}^{\epsilon}$ and $\epsilon$ is a newly added parameter. This format can well address the issue that “adding more steps with zero learning rate always reduces the predicted final loss”, because it makes the new terms of $S_2$ become zero when learning rate is $0$. We encourage other researchers to try this format (see the footnote in Appendix H.3).

We choose $S_2=\sum \limits_{i=1}^{s}m_i$ as our main format for the following reasons (where we have briefly discussed in the Appendix H.3):

1. When adopting the variant format, $\epsilon$ is small after fitting in practice, which makes the whole result become very close to the original format's result (i.e. $\eta_i^{\epsilon}\approx 1$) when $\eta$ is not that extremely near zero, which shows limited use in practical effect as in most situations, the learning rate should not be (extremely near) zero.
2. Notably, for our main format, we have only one extra parameter, the coefficient of $S_2$ term, $C$, compared to Chinchilla scaling law. However, we get a really accurate prediction results (shown in Section 3) and show the main format is quite robust across very diverse experimental setups (e.g. training hyper-parameters, model architectures), acknowledged by reviewer rnjo, up2N and 5e85. In general, $S_2=\sum \limits_{i=1}^{s}m_i$ utilizes the fewest extra parameters but model the essential training dynamics in LR annealing process to the greatest extent.
3. Still, we do not deny the variant $S_2=\sum \limits_{i=1}^{s}m_i\cdot \eta_{i}^{\epsilon}$. When some extreme corner cases really happen (e.g., very small learning rates are really padded at last as assumed), we still recommend future researchers to try this variant format for fitting, and the results would be expected to get better.

Q2: Multiplying the entire schedule elementwise by a positive constant larger than 1 always decreases the predicted loss.

There might be a misunderstanding about our scaling law. We present the retrictions for our scaling in Line 279-282:

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, 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 S1 and S2 are defined in Eq. 1. $L_0$, $A$, $C$, $\alpha$ are positive constants.

Our scaling law cannot predict across different max learning rates using one single set of parameter. Loss curves with different max learning rates have totally different $L_0$, $A$, $C$, $\alpha$. When multiplying the entire schedule elementwise by a positive constant larger than 1, the max learning rate also becomes larger and the original parameters, $L_0$, $A$, $C$, $\alpha$ are not suitable anymore. Instead, we should substitute specific parameters into the different equations for comparison between different $\eta_{max}$.

For example, in our main experiments (Fig. 2 and Fig. 3), we use the equation $L(s) = 2.628 + 0.429 \cdot S_1^{−0.550} - 0.411 \cdot S_2$ for $\eta_{max}=2e-4$. We control other all conditions and only increase $\eta_{max}$ to $4e-4$ and re-train the models, then we get a different equation after fitting: $L(s) = 2.680 + 0.491\cdot S_1^{-0.639} - 0.253\cdot S_2$ ($R^2>0.999$).

We are very grateful for the opportunity to improve our work through your valuable feedback. If you think our response has addressed your concerns and questions, we sincerely hope you might consider raising your score. Should you have any further inquiries or require additional clarification, we warmly welcome your questions and are eager to provide comprehensive responses.

I thank the authors for their clarification on the fixed maximum learning rate. That resolves the concern raised in the original review on that point. However, I am still concerned about the undesired behavior of scaling law in case of appending zero or extremely small learning rates in the end. The discussion in the appendix doesn't really present a solution to this issue.

Regarding the application of the scaling law, I am referring to predicting phenomena that we do not know or proposing novel schedules that works well. For example, I think a basic sanity check the authors should include in the paper is to exactly solve the learning rate schedule that minimizes the predicted loss, and then test the performance of the optimal schedule predicted by your scaling law, e.g., in the context of Figure 2, $L(s) = 2.628 + 0.429\cdot S_1^{−0.550} − 0.411 · S_2$. Given that $S_2$ is a linear function of the learning rates per step and $S_1$ is a function of the sum of learning rates, solving the exact optimal schedule is very easy. I will increase my score to at least 6 if the performance of the optimal schedule predicted by the authors' scaling law yield comparable performance to WSD in context of Figure 2.

Regarding the corner case of appending zero learning rates, we suggest researchers to change the scaling law form to $S_2=\sum \limits_{i=1}^{s}m_i\cdot \eta_{i}^{\epsilon}$ (re-fitting and re-predicting) when appending zero learning rates, which can well avoid the undesired behavior you mentioned. May you specify the reason why you think the form $S_2=\sum \limits_{i=1}^{s}m_i\cdot \eta_{i}^{\epsilon}$ is not a solution to the issue?

Regarding the issue for the optimal learning rate schedule, it is not an easy thing to solve the exact optimal LR schedule, while we have presented a nearly optimal LR schedule (WSD + 10-20% annealing ratio + 1-sqrt annealing function) in this paper. We elaborate this issue from two aspects:

1. From a practical engineering aspect, we have indeed predicted the (nearly) optimal learning rate schedule in this work. Specifically, WSD LRS should be used to confirm the larger values for both $S_1$ and $S_2$, as stated in section 4.2 and Appendix G.2; WSD LRS annealing ratio can be determined by the method stated in section 4.3; Annealing function can be determined by the method stated in Appendix G.3. In practical engineering, it is quite easy to get the (nearly) optimal learning rate schedule based on the composition of the methods above. In Fig. 24, we compare the prediction and our empirical results under a better annealing function and annealing ratio, where the prediction and ground-truth curves coincide very well. Notably, we exactly use the equation as in Fig. 2, $L(s)=2.628+0.429⋅S_1^{−0.550}−0.411·S_2$, to predict and draw figures in all parts of section 4 and Appendix G.

2. From a mathematical aspect, it becomes an optimization problem, that is, searching a learning rate sequence $(\eta_1, \eta_2, \cdots, \eta_s)$, to optimize $\text{max }L(s)= L_0 + A\cdot S_1^{-\alpha} - C\cdot S_2$ . Given the large search space of for solving the high dimensional non-linear problem, this might be not that easy as you think. We can not find an exact optimal LR schedule solution based on mathematical deductions from our previous attempts. We would like to conduct some experiments if you can solve it and specify an optimal learning rate schedule. From another perspective, we believe that the solution can be built up by some non-linear optimization algorithms (e.g. neural networks + gradient decent) to get a locally optimal solution, but the minimization is local and the found LR schedule would be very like the schedule we mentioned above (i.e., WSD + 10-20% annealing ratio + 1-sqrt annealing function). Overall, this is a quite interesting thing on which the following work can have deeper explorations. At this stage, we believe that the approach based on the practical engineering mentioned above is sufficient to obtain a (nearly) optimal learning rate schedule, and the exprimental results coincide well with the predictions.

We look forward to further discussion if you have any more questions and concerns.

Thank you for the detailed analysis for our main format in eq.1.

First, we apologize for the possible confusion brought by the phase "corner case" in our response. We intend to underline the fact that few LLM participants will pad 0 LR steps at the end of their training. We start our study with the intention of covering the most common use cases. Our extensive experiments (acknowledged by Reviewer rnjo, up2N) have demonstrated the effectiveness of our equation in these most commonly used learning rate schedulers, such as WSD, cosine, constant, multi-step cosine, and cyclic scheduler (acknowledged by Reviewer rnjo, up2N, and 5e85).

We would like to point out that your math deduction has a set of typos. Specifically, it should be ${\partial{L}}/{\partial{\eta_t}}=-\alpha A S_1^{-\alpha-1} {\color{red}{-}} C \cdot (-\lambda^{T-t})$ for the partial derivative of $L$ (our simple main format). Besides this minor issue, your analysis (thanks a lot!) is correct: the solved optimal learning rate schedule has zero learning rates in the end.

As we have discussed in our paper (as well as in our responses to all reviewers), this problem is related to the issue ($\eta \approx 0$ case) that mentioned in Appendix H.3. Specifically, our main format does not well handle the case $\eta=0$ at last and the optimization leads to the unhandled case. Mathematical optimization strongly depends on an absolute accuracy, while our scaling law in such scenario does not achieve perfectly 100% accuracy (shown in Fig. 3). In comparison, our mentioned variant forms (e.g., Eq. 10, $S_2=\sum \limits_{i=1}^{s}m_i\cdot \eta_{i}^{\epsilon}$) with extra parameters should be more preferably used to solve the optimal LRS, because they cover more corner cases and have higher accuracy. It won't lead to the same optimal learning rate schedule as our main format if solving it using your mentioned method. We also believe that, in the future, there will be stronger and more parameterized specific forms, which are more suitable for directly mathematically solving the optimal LR schedule. At this stage, we believe that the approach based on the practical engineering elaborated in this paper is sufficient and more suitable to obtain a (nearly) optimal LR schedule.

We would like to emphasize our point of view about this issue (as we also stated in global response before):

• Studies of scaling laws are mostly empirically driven. Over-parameterizing the scaling law equation with more parameters essentially leads to more accurate fitting results. However, it will also complicate the final format and hinders us to focus on major factors for the training dynamics. For example, we compared our main format ( $L \propto S_2$) with $L \propto S_2^{\zeta}$ (adding a parameter $\zeta$ as power of $S_2$) in Fig. 29 of Appendix H.3, and the latter one is indeed better (marginally). We choose our main format not due to absolute prediction accuracy but to pursue the simplification (i.e., fewest extra parameters) to model the essential training dynamics of LLMs. Notably, in our main format in Eq.1, we only introduce one extra parameter compared to Chinchilla scaling law, i.e., the coefficient $C$ of the $S_2$ term. Moreover, it has been acknowledged by all the reviewers that our main scaling law format has a strong and robust capacity across many practical scenarios. We even validate our scaling law under random learning rate schedule in the response to another reviewer rnjo.

• We agree with you and another reviewer 5x7q that the current format might not be perfect. However, we want to highlight that our novel formulation (acknowledged by all reviewers) is the first to address the learning rate scheduling process in scaling law studies (acknowledged by reviewer 5x7q), and we believe and expect that there should be more powerful specific format (maybe with more parameters) after this work.

Dear Reviewer up2N,

Thank you once again for your thoughtful review and detailed analysis. Based on your valuable feedback, we have made several revisions to the paper.

Specifically, we have elaborated further on the reasons for selecting our main format in Appendix H.3, emphasizing its simplicity and effectiveness in most practical scenarios. Additionally, we have added a new subsection (Appendix H.4) to the Discussion section to discuss the issue regarding the optimal LR schedule that you raised. Here is the last paragraph in Appendix H.4:

From another perspective, one might adopt our scaling law to directly solve optimal LR schedule mathematically. It could be found that our Eq.1 leads to a collapsed LRS: some zero learning rates at last. This problem is related to the issue ($\eta \approx 0$ case) that we discussed above in Appendix H.3. Mathematical optimization strongly depends on an absolute accuracy, while our scaling law in such scenario does not achieve perfectly 100% accuracy (shown in Fig. 3). In comparison, our mentioned variant forms (e.g., Eq. 10) with extra parameters should be more preferably used to solve the optimal LRS, because they cover more corner cases and have higher accuracy. We believe that in the future, there will be stronger and more parameterized specific forms, which are more suitable for directly mathematically solving the optimal LR schedule. At this stage, we believe that the approach based on the practical engineering mentioned above is sufficient to obtain a (nearly) optimal LR schedule.

We hope these revisions adequately address your concerns, and we look forward to hearing your further thoughts on our updates. We are truly grateful for your insightful comments and remain committed to improving our work based on your suggestions.

Best regards,

Authors

1. The goal of this work is NOT to solve mathematically optimal LR schedule though we discuss much about this issue in rebuttal. Here is part of our abstract to present our motivations and contributions:

   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.

2. Our scaling law does have great generalization. In Fig. 3, these predicted full loss curves are all under unseen LR schedules including WSD, cyclic, etc. The training data (in Fig. 2) only includes constant and cosine LR schedules. The generalization ability of our scaling law is also acknowledged by reviewer 5e85, 5x7q.

   For those really uncommon LR schedules (e.g., padding zero LR at last), we recommend to use our variant form, which is explicitly mentioned in Line 286 and Line 527-528. The variant forms share the same core idea including $S_1$ and the momentum form in $S_2$, which we have demonstrated in Sec. 3.1 and Sec. 3.2. The variant forms are more accurate in uncommon LR schedules but introduce one extra parameter, for which we use the term "simplicity" to describe our main form.

Q1: An extremely large learning rate induces better result in the scaling law.

Thanks for your valuable feedbacks and there might be some misunderstandings about our scaling law.

1. We present the restrictions for our scaling law in Line 279-282:

   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, 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 S1 and S2 are defined in Eq. 1. $L_0$, $A$, $C$, $\alpha$ are positive constants.

   Our scaling law cannot predict across different max learning rates using one single set of parameter. Loss curves with different max learning rates have totally different $L_0$, $A$, $C$, $\alpha$. When setting a larger learning rate, the max learning rate also becomes larger and the orginal parameters, $L_0$, $A$, $C$, $\alpha$ are not suitable anymore. It is very likely that the newly fitted parameters under the condition of larger $\eta_{max}$, will lead to a worse prediction result compared with the condition of smaller $\eta_{max}$.

   • For example, in our main experiments (Fig. 2 and Fig. 3), we use the equation $L(s) = 2.628 + 0.429 \cdot S_1^{−0.550} - 0.411 \cdot S_2$ for $\eta_{max}=2e-4$. We control other all conditions and only increase $\eta_{max}$ to $4e-4$ and re-train the models, then we get a different equation after fitting: $L(s) = 2.680 + 0.491\cdot S_1^{-0.639} - 0.253\cdot S_2$ ($R^2>0.999$).

   Therefore, our scaling law never encourages to make learning rate as large as possible (e.g. 100). Conversely, When we set another max learning rate, the scaling law should be re-fitted and we get another set of parameters. Thus, we cannot directly compare the losses between larger and smaller learning rates only through the scaling law formulation. Instead, we should substitute specific parameters into the different equations for comparison between different $\eta_{max}$.

2. Our scaling law is used to fit and predict the loss curves when the curve is not divergent. It does not even follow a power law (Chinchilla and Openai scaling law) if the curve is divergent. Infinitely large LR would indeed "go up to the level of random guessing repeatedly", and inevitably lead to a divergent loss curve, which is out of scope of this work.

Thanks for your valuable feedbacks again! We will also emphasize this retrictions in our revised manuscript.

Q2: This law is invariant to padding zero learning rates at the end.

We are glad that you have noticed this issue and found that we have discussed this issue in Appendix H.3. Here we give some further clarifications for this issue.

1. We designed a variant of our scaling law by adding a LR-weighted coefficient to $S_2$, where $S_2=\sum \limits_{i=1}^{s}m_i\cdot \color{red} \eta_{i}^{\epsilon}$ and $\epsilon$ is a newly added parameter. This format can well address the issue that “adding more steps with zero learning rate always reduces the predicted final loss”, because it makes the new terms of $S_2$ become zero when learning rate is $0$. We also footnoted in the paper (in Appendix H.3) that we recommend future researchers to try this format for fitting if possible.

2. For clafication, we do not deny the variant $S_2=\sum \limits_{i=1}^{s}m_i\cdot \eta_{i}^{\epsilon}$. As you noticed, the fitting result $\epsilon \approx 0$ is based on our real data (i.e., cosine LRS + constant LRS as fitting data), which contains few learning rates that are extremely near $0$. In most practical situations, the learning rate will not be (extremely near) zero, making this variant unnecessary for practical use. Conversely, when some extreme corner cases really happen (e.g., very small learning rates are padded), we recommend future researchers to try this variant format, and the results would be expected to get different and much better because $\epsilon$ is no more near $0$.

3. Moreover, we would like to present some more explanations about why we choose our main format $S_2=\sum \limits_{i=1}^{s}m_i$. Extensive experiments have validated the effectiveness of our equation across many LLM pretraining settings (acknowledged by reviewer rnjo, up2N and 5e85). This demonstrates the robustness of our equation in practical applications. Moreover, we can also easily extend our formulation to model some extremely rare corner cases. We discuss and explain more about the form of $S_2$ in the following Q5, Q6, and Q7.

Q3: 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.

We state the issue about accuracy of our scaling law from the following three aspects:

1. We would like to point out that this prediction error is already very small. As shown in the middle sub-fig in Fig. 3(c) and 3(d), the predicted curve aligns very well with the ground truth. The seemingly large deviation in the zoomed-in view on the right is due to the vertical axis being scaled down to a very small range of 0.01. This indicates that the maximal absolute loss error in both cases is approximately only 0.01, which significantly outperforms many related scaling law papers. Here we list some famous and recognized scaling law works for comparison:

   • In OpenAI scaling law paper [1], in their Fig. 1 (medium), they use 0.3 as tick interval (30x of ours), and that is even a fitting result rather than prediction. Some points are obviously more deviated from the actual values than ours.

   • In OpenAI scaling law paper [1], in their Fig. 4 (left), they use 0.5 as tick interval (50x of ours), and shows significent deviated prediction results.

   • In GPT-4 paper [2], in their Fig. 1, they use 1.0 as tick interval (100x of ours) and get a less accurate prediction result than ours.

   • In Chinchilla scaling law paper [3], in their Fig. 3 (left), they use 0.2 as tick interval (20x of ours), and some points are obviously more deviated from the actual values than ours.

   • In Llama-3 paper [4], in their Fig. 2 (left), they use 0.05 as tick interval (5x of ours), and some data points are more deviated from the actual values than ours.

   • In Llama-3 paper [4], in their Fig. 4 (left), they use 0.025 as tick interval (2.5x of ours), and the prediction result is worse than ours.

   • In DeepSeek paper [5], in their Fig. 4(a), they use 0.2 as tick interval (20x of ours), and some points are obviously more deviated from the actual values than ours.

   • In Baichuan-2 paper [6], in their Fig. 4, they use 2.0 as tick interval (200x of ours), and get a much less accurate prediction result than ours.

​Overall, we would like to emphasize that our scaling law is quite accrate compared with other scaling law works. We set 0.01 as our tick interval to provide readers with details of the predicted curves. We believe that it is a more decent practice than hiding the error using some presentation tricks.

2. We did not include any WSD learning rate curves or learning rate curves with re-warmup in the fitting process. Instead, our current fitting data points are based solely on two different and separate schedulers (i.e. cosine and constant), and are limited to just 20K steps. Moreover, there could be more other factors influencing the validation loss. For example, some training samples near specific checkpoints are more similar to the validation set, and then it could make a temporarily lower validation loss. Overall, the task is quite difficult but our scaling law gives an accurate enough prediction.

3. Our scaling law equation can easily make more accurate predictions if more data points are used in fitting. For example, if we add loss curves obtained from WSD LRS (50K steps, 20% 1-sqrt annealing to 0), the equation becomes $L(s) = 2.611 + 0.449\cdot S_1^{-0.532} - 0.330 \cdot S_2$ ($R^2>0.999$). The new equation is more accurate in predicting Fig 3(c) and 3(d), and reducs the mean prediction error by almost 60%. Here we present a table below to show the prediction improvement by adding more fitting data.

   Fitting Points	Equation	Fig. 3(c) mean prediction error	Fig. 3(d) mean prediction error
   Cosine + Constant (in our paper)	$L(s) = 2.628 + 0.429 \cdot S_1^{−0.550} - 0.411 \cdot S_2$	0.235%	0.322%
   Cosine + Constant + WSD	$L(s) = 2.611 + 0.449\cdot S_1^{-0.532} - 0.330 \cdot S_2$	0.095%	0.273%

[1] Scaling Laws for Neural Language Models

[2] GPT-4 Technical Report

[3] Training Compute-Optimal Large Language Models

[4] The Llama 3 Herd of Models

[5] DeepSeek LLM: Scaling Open-Source Language Models with Longtermism

[6] Baichuan 2: Open Large-scale Language Models

Q4: 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?

This is a quite interesting experimental design! We conduct this experiment!

The conclusion is: our scaling law still works well, though the prediction error might be a bit larger.

For replication, we present the detailed experimental processes:

1. We first generate a random learning rate sequence using np.random.uniform(0, 2e-4, 30000) after we set the seed by np.random.seed(1442) (Numpy version: 1.24.4). The generated learning rate sequence is: [6.82985668e-05, 1.81342066e-04, 7.20995317e-05 ... ]. Please feel free to replicate and check the process above.

2. Then, we adopt setting B (shown in Table 3 and Appendix D.1, i.e. 293M model) for faster training, and the equation is ready-made: $L(s) = 2.761 + 0.517 \cdot S_1^{−0.491} - 0.458 \cdot S_2$. We launch a training for 30K steps according to the above random learning rate schedule after we warmup learning (100 steps) rate from 0 to 2e-4. Here we set U[0, 2e-4] because the max learnining rate of setting B is 2e-4.

3. Here is the prediction and ground-truth loss comparison table:

4. Finally, as the table above shows, the prediction results are really accurate, considering the task is quite challenging. We get a 0.605% mean prediction error, which is about 3x error compared to WSD LRS (Fig. 3c), and about 2x error compared to the Cyclic LRS (Fig. 3d). Notably, as we mentioned above, this mean prediction error can be also further reduced by adding more fitting data.

Interestingly, this random learning rate experiment can even serve as one another evidence for the power-law form of $S_1$, which is related to the following question you asked (Q5). We notice that in random learning rate schedule, the $S_2$ is quite stable over the training steps, while what really impacts the loss is $S_1$. The values of $S_1$ and $S_2$ are as follows (actually this can be easily replicated using several lines code based on the learning rate sequence generator above):

It is $S_1$ that is significantly different between random learning rate and constant learning rate ($S_1 = \eta_{max}\cdot s$ in constant LRS), and then finally makes a difference to the losses. We believe this supports our idea of using the powers of $S_1$ to predict the loss.

Thanks for designing such a convincing and interesting experiments, which helps to comprehensively show the performance of our scaling law.

Q5: Theoretical analysis of $S_1$ and $S_2$

We are very glad that you are interested in the more theoretical justifications of our scaling law. We honestly admit that the form of $S_1$ and $S_2$ are primarily based on our empirical observations (i.e. Section 3.1 and 3.2). And in Q4 above, we also use your designed experiments (random LRS) to validate the form of $S_1$. We have even claimed that our scaling law is empirical in the first line of the abstract.

However, we would like to point out that even the Chinchilla and OpenAI scaling laws are, in fact, empirical, as explicitly stated in their respective papers. As you said, proving any theorem about the loss curve in deep learning is hard, and we believe that in AI domain, theoretical proofs often follow empirical works. For example, [7] presents some reliable theoretical justifications for scaling law after three years since OpenAI scaling law paper is published. We are continuously working on the theoretical aspects and this is indeed a very important future work.

Actually, surprisingly, some related works have emerged very recently [8]. Their theoretical derivations are well matched with our scaling law. We would greatly appreciate it if you could kindly take the time to read it.

[7] The Quantization Model of Neural Scaling

[8] Understanding Warmup-Stable-Decay Learning Rates: A River Valley Loss Landscape Perspective

Q6: Some other possible $S_2$ form.

Thanks for your great and thoughtful specific form suggestions! We have the following clarifications:

1. For the sliding window form, it is actually a "hard" version of our momentum form. Sliding window form assigns the same weight to each one in the window and drops ones out of the window. Our momentum form assigns weight based on the historical distance. They are consistent in their underlying ideas.

2. Essentially, over-parameterizing the scaling law equation (like parameterizing $\lambda$ as a function of $\eta$ ) with more parameters indeed often leads to more accurate fitting results. However, it will also complicate the final format and hinders us to focuse on major factors for the training dynamics. For example, we compared our main format ($L \propto S_2$) with $L\propto S_2^{\zeta}$ (adding a parameter $\zeta$ as power of $S_2$) in Fig. 29 of Appendix H.3, and the latter one is indeed better (marginally). We choose our main format not due to absolute prediction accuracy but to pursue the simplification (i.e., fewest extra parameters) to model the essential training dynamics of LLMs. Notably, in our main format in Eq.1, we only introduce one extra paramter compared to Chinchilla scaling law, i.e., the coefficent $C$ of the $S_2$ term but extend to the full loss curve prediction. Moreover, it has been acknowledged by all the reviewers that our main scaling law format has a strong and robust capacity across many practical scenarios. We also validate our scaling law even under random learning rate schedule in the above response.

3. We have elaborated further on the reasons for selecting our main format in Appendix H.3, emphasizing its simplicity and effectiveness in most practical scenarios. The revisions are highlighed in red.

Q7: Ablation studies on $S_1$ and $S_2$

Thanks for your valuable feedback! In the paper, we show the importance of $S_1$ by considering training disount in annealing (Line 212), and $S_2$ is necessary because there must exist some term to make the final loss decrease when LR annealing.

We would like to also supplement some ablation studies on $S_1$ and $S_2$ to show that every component of the scaling law is important. Specifically, we compare some law forms without $S_1$ or $S_2$. We use the same setting as in Fig. 2 and Fig.3. For each format, we re-fit on 20K cosine + constant and re-predict on different LRS. The prediction results are shown as below:

The results suggest that the component of $S_1$ and $S_2$ is both important and Indispensable. We have also included the results in our revised paper (Table 4).

Finally, we thank again for your valuable feedback! If your have any other questions or concerns, we are happy to answer and continually improve our work based on your feedback.

We are very grateful for the opportunity to improve our work through your valuable feedback. If you think our response has addressed your concerns and questions, we sincerely hope you might consider raising your score. Should you have any further inquiries or require additional clarification, we warmly welcome your questions and are eager to provide comprehensive responses.

Dear Reviewer rnjo,

We are very grateful for the opportunity to improve our work through your valuable feedback. As the discussion period is near the end, this is a kind reminder for your further feedback.

If you think our response has addressed your concerns and questions, we sincerely hope you might consider raising your score. Should you have any further inquiries or require additional clarification, we warmly welcome your questions and are eager to provide comprehensive responses.

Best regards,

Authors

We are sincerely grateful for your recommendation to accept this work! In response to your concerns, we have the following answers.

Regarding your Q1: We have highlighted the following things updated in the latest version of our paper.

• We further emphasize the restrictions of our scaling law in Line 283-285. And we mention our proposed variants to counter the case of zeros learning rates at last in the main text including Line 286 and Line 527-528.
• We discuss why we choose our main format rather than our proposed two variant formats in Appendix H.3. Specifically, we prefer the simplification (i.e., the fewest parameters) to model the essential dynamics in LLM training.
• We add one subsection to explicitly discuss on solving optimal learning rate schedule in Appendix H.4.

Regarding your Q2: Thanks for your suggestions. We would like to clarify the point about common and corner LR schedules.

• Common LR schedules mean popular LR schedules that many previous researchers widely adopt, including constant, cosine, multi-step cosine, WSD, and cyclic LR schedule. These schedules are exactly the cases shown in Fig. 2 and Fig. 3.

• The "corner" LR schedules mean those LR schedules that few previous researchers adopt, including the case with many zero LR at last, and random learning rate schedule you mentioned (though we also validate our scaling law for random LR schedule in our previous discussion). We recommend our variant forms for these "corner" LR schedules if possible.

We would like to further emphasize this point in the next version of the paper, as the deadline for revisions during the discussion period has passed.

Thanks again for your recommendation for this paper. We are committed to continuously improving our work based on your valuable suggestions.

Q2: Adding more steps with zero learning rate always reduces the predicted final loss.

We have discussed this scenario in "Section 6: Discussion" (Line 527-528) and introduced a possible variant of our equation in appendix H.3. Specifically, we can add a LR-weight coefficient to $S_2$, i.e., $S_2=\sum \limits_{i=1}^{s}m_i\cdot \color{red} \eta_{i}^{\epsilon}$ where $\epsilon$ is a new parameter. This format can well address the issue that “adding more steps with zero learning rate always reduces the predicted final loss”, because it makes $S_2$ becomes zero when the learning rate $\eta_{i}$ is $0$. We encourage other researchers to try this format (see the footnote in Appendix H.3).

We choose to use $S_2=\sum \limits_{i=1}^{s}m_i$ as our main format for the following reasons (where we have briefly discussed in the Appendix H.3):

1. In our experiments, the absolute value of $\epsilon$ is extremely small after fitting. This makes the value of $\eta_{i}^{\epsilon}$ approaches to $1$ if the learning rate $\eta_{i}$ is within a normal scale, which is true in most practical training processes. The value of $\eta_{i}^{\epsilon}$ approaches to $0$ only when the learning rate $\eta_{i}$ is extremely small, which rarely happens in practical training.

2. Extensive experiments have validated the effectiveness of our equation across many LLM pretraining settings (acknowledged by reviewer rnjo, up2N and 5e85). This demonstrates the robustness of our equation in practical applications. Moreover, we can also easily extend our formulation to model some extremely rare corner cases.

3. Essentially, over-parameterizing the scaling law equation with more parameters leads to more accurate fitting results. However, it will also complicate the final format and hinders us to focus on major factors for the training dynamics. Notably, in our main format in Eq.1, we only introduce one extra parameter compared to previous works, i.e., the coefficient $C$ of the $S_2$ term. We intend to use a concise scaling law format to reveal more insights of the training dynamic to LLM practitioners, and shed light on (1) studying the optimal learning rate schedule in pretraining; (2) predicting losses like Chinchilla scaling law but with much less training cost to collect fitting data.

We agree with the reviewer 5x7q that the current format might not be perfect. However, we want to highlight that our novel formulation (acknowledged by all reviewers) is the first to address the learning rate scheduling process in scaling law studies (acknowledged by reviewer 5x7q). The more deep investigation of extreme corner cases is left in further studies.