The Surprising Agreement Between Convex Optimization Theory and Learning-Rate Scheduling for Large Model Training

We thank the reviewer for the feedback and suggestions, and address each question below in detail. We hope that after clarifying all concerns we can improve the reviewer's rating of our submission.

• First, training loss is not an important criterion for LLMs: For all LLM experiments in our paper, we already plot the validation set loss. We also comment on the relationship between train and test/validation loss in the last paragraph of the Limitation section.
• the gap between strict assumptions and real-world neural network training is not well-addressed: The limitations of these assumptions are explicitly addressed in Section 6. As the title suggests, one of the insights of this paper is that theory and practice match very well, even though the assumptions from theory are not satisfied.
• the loss function is extremely smooth: Can you explain what is meant with the term 'extremely smooth'? In fact, our theoretical results hold for non-smooth functions with bounded gradients (e.g. the absolute value function would satisfy this, but it’s not smooth). Thus we are confused by this comment. We also remark that we set $D=1$ and $G=1$ in the simulation purely for simplicity (see details further below). The bound holds for any other values. (In fact, for the LLM training setup, the gradient norms are below one except for the first ~hundred steps; we will provide a plot of the empirical gradient norms in the final version.)
• is achieved by approximating a linear function of the hyper-parameters in the theorem: We did not not understand this comment. Could you explain which linearization is meant here?
• Fig 4, where a larger decay rate leads to under-training (large loss): we are confused about this comment, as Fig 4 shows that larger decays lead to lower losses if the learning rate is tuned. Could the reviewer expand on this, so that we can hopefully answer the question?
• The alignment is problematic: We kindly disagree: our theoretical findings are verified on multiple schedules+model sizes+time horizons, and we show that multiple aspects of the theory match the experiments (e.g. schedule adaptation for longer runs, effect of cooldown length, LR transfer across cooldown lengths etc.).
• LLMs are less than 1B, which is not sufficient: Our experiments require a high number of individual runs, for example to sweep the learning rate (e.g. only Figure 12b requires around 20 runs). With our computational budget it would not be feasible to execute this amount of runs on a >10x larger scale. Besides that, scaling law research has repeatedly shown that loss improvements transfer to larger scales highly predictably in this regime (see https://arxiv.org/pdf/2203.15556). We also point out that for the central message of our paper (“convex bounds match empirical behaviour in deep learning”) any scale would be sufficient to prove the point.

On Questions:

1. In Equation (9), $\gamma$ is arbitrary, and not necessarily equal to $\gamma^\star$.
2. The simulations simply evaluate the bound $\Omega_t$ for different schedules and learning rates. For this it is not necessary to specify $f$, but only constants $D$ and $G$ which we set to 1 in the simulations (purely for simplicity). The choice of $D$ and $G$ only affects the scale of the bound and learning rates, but not its shape or scaling in $T$.
3. As the caption of Fig 8 explains, T_1=400 and we test several values of T_2. Indeed the goal is to reuse the training up to iteration T_1, as it is explained in the beginning of Section 5.1
4. The right plot is only the iterate path, and one can not exactly read off the timescale of it. However, for (i) cosine at the step where its path makes a sharp left-wards turn and (ii) for WSD before cooldown are roughly similar and this matches the contours in the right plot.

We thank the reviewer for acknowledging the theoretical and empirical contributions of our paper. We are delighted that the reviewer considers the paper to be interesting to the optimization and ML community. Regarding the questions:

• Convergence rate of cosine: unfortunately, due to the form of the cosine schedule it is quite difficult to explicitly derive the bound. We tried several simplifications and have not reached a nice result yet. However, from the simulations we conduct, the bound of cosine is almost the same as the one for linear-decay, for which we can explicitly derive the bound (see lines 1240 ff.).

We thank the reviewer for their positive feedback, especially for the assessment that our experimental results are strongly supporting the claims/contributions.

• We ran additional experiments with SGD on Imagenet as suggested by the reviewer (see details here: https://anonymous.4open.science/r/icml25-additional-experiments-lr-schedules/). The results confirm the close match between the bound and actual loss curves.
• Regarding additional schedules, we would like to point to Fig 18 and 23 where we compare many different schedules beyond cosine and WSD, including constant and 1/sqrt schedule as well as the 1-sqrt schedule proposed in Hägele et al 2024. We also test a piecewise constant schedule with cooldown for the continual learning experiment.
• Regarding other optimizers, we have not done a systematic comparison. However, from results reported in the literature, it is known that the sudden drop of the loss also appears for other optimizers, for example Muon, Shampoo and SOAP (see https://github.com/KellerJordan/modded-nanogpt/tree/master/records/102924_Optimizers) or Ademamix (see Fig 3b in https://arxiv.org/pdf/2409.03137). This makes us confident that the characteristic schedule behaviours that we describe also generalize to other optimizers; investigating this in detail would be an interesting direction for followup work.

We thank the reviewer for their thoughtful review and detailed feedback. We address all questions in detail below. We hope that this will clarify all concerns, and allows for a higher scoring of our submission.

• limited empirical evaluation, one real dataset: We agree that it is beneficial to validate our findings on additional datasets, and for the rebuttal we performed additional experiments on Imagenet and OpenWebText; the results confirm our previous findings (see detail here https://anonymous.4open.science/r/icml25-additional-experiments-lr-schedules/). However, even though we had only considered one dataset, our empirical evaluation is not limited: we verify our findings on multiple schedules+model sizes+time horizons, and find that multiple aspects of the theory match the experiments (e.g. schedule adaptation for longer runs, effect of cooldown length, LR transfer across cooldown lengths etc.). Besides the dataset choice (which we address in the rebuttal), are there any other reasons why the reviewer considers our empirical evaluation to be limited?

• Related to the above, the characteristic behaviour of WSD during cooldown has been reported in multiple papers (which all train in slightly different settings and on different datasets). For example (see also references in our paper): https://arxiv.org/pdf/2410.05192, https://arxiv.org/abs/2502.15938. These independent experiments further justify that the bound indeed reflects real-world behaviour, and our paper is the first one to highlight and exploit this connection between theory and practice.

• lack of significant theoretical novelty: One interesting theoretical insight is that the cooldown of WSD achieves to remove a log-factor in the convergence bound (compared to the constant schedule). Moreover, this improvement is obtained exactly by the drop during cooldown, as Figure 20 (left) shows. More broadly, the main goal of our paper is not to develop new proof techniques, but to show how existing theory can be used in order to (i) design/tune the learning rate (schedule) and (ii) explain the benefit of cooldown observed in practice (and the behaviour of LR-schedules in general).

• Correct, Assumption A3 is implied when assuming Lipschitzness of the objective. However, as the bound only depends on the expected gradient norms on the SGD trajectory (and the proof does not explicitly use the Lipschitz condition), we wanted to formulate the assumption as tight as possible - bounded gradient norms over a finite number of steps could be satisfied with much weaker assumptions than Lipschitz continuity.

• We thank the reviewer for pointing out the work by Liu and Zhou; we were not aware of this paper and will discuss it in the related work section. From our understanding the main differences are:

  • the results by Liu and Zhou do not hold for arbitrary schedules
  • In the Lipschitz-smooth case (M=0 in Liu and Zhou), instead of the Lipschitz constant, we have the noise constant in the bound, and additionally the step size needs to be smaller than $\mathcal{O}(1/2L)$. However, as we can decompose the bound $G_t$ in our notation in a noise part plus the Lipschitz constant of the objective $f$, this appears to result in a very similar structure. Extending our analysis to the smooth case is definitely interesting and we are actually working actively in this direction.