Scaling Law with Learning Rate Annealing

W1: Theoretical justification and related works

Yes, our scaling law is empirical, as noted in our limitations section. We appreciate your suggestion and will revise certain expressions in the paper accordingly. Additionally, we will incorporate more related work to strengthen the paper. We will make this revision in the next version of the paper. Additionally, we appreciate your suggestion to include a more detailed discussion of related work to better contextualize our findings. We will expand the related works to provide a clearer comparison with prior studies and to highlight how our work fits into the broader research landscape. This will help strengthen the paper’s contribution and provide a more comprehensive perspective.

Q1: How feasible would it be to consider generalization bounds under this new scaling law?

Thank you for your insightful question regarding the feasibility of deriving generalization bounds under the new scaling law. We believe that exploring generalization bounds is feasible but would require additional theoretical work. We would like to emphasize some strong generalization effects demonstrated in our experiments:

1. In our experiments, the learning rate schedule (LRS) used for prediction was unseen during fitting, and the error was controlled within a very small range (<0.35%).
2. In our experiments, we performed data fitting using only a short number of training steps (e.g., 20K) and extrapolated generalization to longer training steps (e.g., 60K).
3. As you also noted, our work is a generalized form of the Chinchilla scaling law. We can achieve significantly more fitting points under the same computational resources, leading to better prediction and generalization performance. Moreover, compared to the Chinchilla scaling law, our model introduces only one additional parameter, and this simplicity contributes to greater robustness.

We will incorporate the paper you mentioned in the next version of our manuscript. We are committed to addressing these points in the revised manuscript and hope our responses clarify our approach and intentions. Please let us know if you have any further recommendations or require additional clarification. We greatly appreciate your time and expertise in reviewing our work.

I thank the authors for clarifying my question. This paper is outside of my area of expertise, but having considered the other reviews, I'm happy to maintain my original score and look forward to the upcoming discussion with the other reviewers.

We sincerely appreciate your thorough engagement and the numerous insightful questions you’ve raised, which have greatly enriched our work.

W1: The $L_0+A\cdot S_1^{-\alpha}−C\cdot S_2$ is not very grounded.

We present in detail how we derive $S_2$ in Section 3.2. Basically, we observe the delay phenomenon in LR annealing and we find that there exists momentum in LR annealing. Therefore, the final $S_2$ is momentum-like (or EMA-like) form. As noted by reviewer EhjS comments, the formulation "is interpretable, and illustrates the trade‐off between rapid progress at high learning rates and the benefits of annealing". To further substantiate our findings, we conduct extensive experiments in Section 3 to validate the proposed scaling law.

W2: The hyper-param $\lambda$ in $S_2$ is not interpretable.

1. $\lambda$ is the momentum coefficient, which is quite common in momentum-like form which means the . We adopt $\lambda = 0.999$ in our all experiments and it shows that the choice of $\lambda$ is not that sensitive and $\lambda=0.999$ is quite robust across many experimental setups.
2. In Appendix I.1, we have fully discussed and investigated the essence and the impact of decay factor $\lambda$ in detail. We show that $\lambda$ influences three aspects: (1) delay steps; (2) optimal annealing ratio; and (3) balance point between $S_1$ and $S_2$.

W3: Concerns about Figure 2b

This is an actually great question, and we fully discuss this in Appendix H.2. Simply speaking, the loss curves of constant/cosine LR, and how different they are, depend on some important hyper-parameters such as training steps. And our curve is actually not contradictory to the paper you mentioned.

In Figure 24, we show that the curves of cosine and constant could be quite similar if we train 10K steps (shown in Figure 24a), but become quite different if training for 100K steps (Figure 24b). Moreover, we show that in small training steps, constant LR could lead to better performance than cosine LR. When training for more steps, the final performance of cosine LRS surpasses that of constant LRS, and this gap gradually widens (Figure 24c).

Moreover, the max learning rate can also influence the curves of constant and cosine LR schedules. And larger max learning rate often leads to larger gap between them. We set max learning rate as 2e-4 in Figure 2, and it is not large. Our results are all based on rigorous and standardized experiments.

Q1: Figure 3

No, we did extrapolation prediction rather than interpolation. We fit only 20K training steps (Figure 2) and predict the loss for 60K training steps under many unseen different LRS (Figure 3).

In our additional experiments in our Appendix E, we also did extrapolation prediction. For example, in Appendix E.1, We adopt another experimental setup and fit only 30K training steps (Figure 10) and predict the loss for over 100K training steps under unseen LRS.

Q2: Figure 7b

Our proposed scaling law calculates the expected loss for each training step. In empirical experiments, 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. Additionally, larger batch sizes typically result in lower variance and makes the loss curve smoother. Furthermore, prior scaling laws, such as OpenAI and Chinchilla scaling laws, also compute expected losses with varying degrees of variance. Therefore, we believe this does not suggest that our scaling law lacks robustness.

Q3: Theoretical grounding of the annealing term

We fully discussed all your mentioned concerns about $\lambda$ in Appendix I.1 and I.2 of our paper.

Selection of $\lambda$: We adopt $\lambda=0.999$ in all of our experiments (with many different setups, see Table 3), and the results are all good, which means that $\lambda=0.999$ is quite robust across different setups. Moreover, we even mention in the footnote of page 27 that one can fit $\lambda$ as a parameter like $L_0,A,C,\alpha$, which is also reasonable and totally feasible.

How $\beta_1$ and $\beta_2$ in Adam Optimizer impacts: We have discussed this issue in our paper. We conduct an experiment in Figure 29a of Appendix I.2. In conclusion, we find that changing $\beta_2$ does not nearly influence the delay steps as well the momentum term (including $\lambda$).

How $\lambda$​ impacts: In Appendix I.1, we discussed and investigated the essence and the impact of decay factor $\lambda$ in detail. We show that $\lambda$ influences three aspects: (1) delay steps; (2) optimal annealing ratio; and (3) balance point between $S_1$ and $S_2$.

Overall, we believe that we have sufficiently discussed $\lambda$ in our paper.

Q4 & Q5 & Q6 & Q7

In this section, we aim to demonstrate that the experimental conclusions from previous works can be also validated and explained using our predicted scaling law.

Claim on LRS Superiority: We cited the conclusion drawn from [1] and [2], and used our scaling law to verify and explain their conclusions. We provide a perspective about the balance between $S_1$ and $S_2$. We believe there are still cases where WSD underperforms compared to cosine (as I mentioned above regarding the relationship between cosine and constant LRS). In summary, we believe what really matters is that our scaling law can predict the final loss under various circumstances, allowing for case-specific evaluations.

Lack of Empirical Validation: Similarly, we cited the conclusion drawn from [3], and used our scaling law to verify and explain their conclusions. Their paper has already shown that 10%~20% annealing ratio in WSD LRS is quite good in most situations, which is consistent with the results predicted by our scaling law.

Speculative Applications: Section 4 demonstrates how we apply our scaling law to validate and explain insightful findings, and these results align with the conclusions of prior work. In fact, we have already conducted extensive experiments in Section 3 to verify our scaling law, including the main experiments in Figure 2 and Figure 3, as well as four additional validation experiments documented in Lines 209–219 and Appendix E.1–E.4.

About Section 5: Section 5 presents a comparison between our scaling law and the Chinchilla scaling law. First of all, in Table 4, we have validated the performance comparison (in real experiments) between our and Chinchilla’s scaling laws. In Section 5, we focus more on formal reduction and analysis of resource consumption.

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

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

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

Thank you again for your valuable feedback and the many thought-provoking questions that have helped refine our work. We look forward to any further insights you may have!

We sincerely appreciate your valuable feedback and the opportunity to address your comments.

Q1: Empirical observations rather than rigorous theory

We believe that theoretical analysis can strengthen our claims. However, it’s important to recognize that in real-world LLM training, theoretical analysis is extremely difficult. Almost all scaling law papers are proposed empirically. The paper you mentioned [1] also addresses this issue and simplifies the problem to a highly toy scenario to derive its results: "we consider a simple setting where SGD optimizes a quadratic loss function with noisy gradients."

We believe that advancing theoretical analysis is an important research direction, and we also believe that our conclusions are sufficiently supported by the extensive experiments in Section 3 we have conducted.

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

Q2: Optimizing learning rate schedule

Optimizing learning rate schedule: In Appendix I.4 of the paper, we discuss the optimization of the learning rate schedule in detail. Specifically, when optimizing Eq. 10 (as opposed to Eq. 1), the problem becomes quite difficult to solve, as it is no longer convex (the Hessian matrix of this problem is no longer positive definite). From another practical engineering aspect, it is feasible and efficient to select better LRS from many candidates based on the prediction of the scaling law. For example, we show how to choose between WSD and cosine LRS and how to choose the optimal annealing ratio in Section 4.

The scope of our scaling law: Besides, based on your concern "The law can compare configurations only within certain families of schedulers". We have the response below:

When the LR schedule becomes even stranger and irregular, the loss becomes more difficult to predict, which is quite intuitive and inevitable. In this case, our scaling law will still make accurate enough predictions and try to keep the error within a usable range.

As the Table 4 shows, when the predicted LR schedule becomes more complicated, the error of our scaling law becomes indeed larger from 0.159% to 0.322%. The errors are all quite small so that we can adopt our scaling law on the schedules.

More extremely, we also perform an extra experiment on irregular and random LR, where we set LR of each step uniformly sampled from 0~2e-4, and S2 term indeed approaches zero during the whole training run. The mean error in such extreme situation is 0.605%, indeed larger but within a moderate range. Your mentioned "constant‐then‐zero" is a quite extreme situation. Overall, our scaling law applies within a specific scope (e.g., consistent maximum learning rate and standard learning rate schedules). However, for commonly used learning rate schedules, it offers high accuracy, simplicity, and strong interpretability.

Q3: About $N$-extended scaling law

Yes and the experiment in Figure 7b shows that we can derive a model size extended scaling law from smaller model sizes. The fitted equation is $1.504 + 0.440\cdot S_1^{-0.546} + 3.067\cdot N^{-0.158} - 0.273\cdot S_2\cdot N^{0.052}$. We didn't use this equation to select the better learning rate schedule but we validated this fitted equation by predicting the full loss curve for a larger model. Actually, if we replace $N$ with specific value (e.g. 594.6M), and the equation becomes $2.622 + 0.440\cdot S_1^{-0.546} - 0.380\cdot S_2$, which is quite similar with our fitted equation in Figure 2: $L(s) = 2.628 + 0.429 · S_1^{−0.550} − 0.411 · S_2$. This means that our equation can also incorporate certain values of N to make specific predictions for a given model size.

We hope this clarification addresses your concerns and thank you again for your thoughtful review.

Thank you for your thoughtful recognition and for sharing such detailed insights!

W1: Deriving an optimal rate schedule

Yes and we discussed this issue in detail in Appendix I.4. Simply speaking, when optimizing LR schedule mathematically, one should use Eq. 4 rather than Eq. 1, which is more precise but quite difficult to solve in math. From another practical engineering aspect, it is feasible and efficient to select better LRS from many candidates based on the prediction of the scaling law. For example, we show how to choose between WSD and cosine LRS and how to choose the optimal annealing ratio in Section 4.

W2: How our parameters change

We answer this question one by one.

(a) Indeed, the parameters (L0, A, C, $\alpha$) would change under those conditions. And we emphasized this in Section 3.3.

(b) Great question! It is quite possible to research this interesting question by researching how $L_0,A,C,\alpha$ changes over the pre-set max learning rate. For instance, as the learning rate increases, these parameters may follow a consistent pattern, potentially leading to a parabolic-like trend in final performance, which could help predict the optimal learning rate. While we haven’t conducted this experiment, we believe it offers a promising approach to determining the optimal learning rate.

Thank you again for your valuable input and for sparking such an interesting discussion! Please let us know if there’s anything further we can assist with.