The Surprising Agreement Between Convex Optimization Theory and Learning-Rate Scheduling for Large Model Training
Authors: Fabian Schaipp
Alexander Hägele
Adrien Taylor
Umut Simsekli
Francis Bach
Venue: ICML 2025 poster (2025)
Keywords: Learning rate schedules
convex optimization theory
large model training
continual training
Primary area: optimization
TL;DR: Learning rate schedules in deep learning behave strikingly similar to what convex optimization theory predicts

Abstract: We show that learning-rate schedules for large model training behave surprisingly similar to a performance bound from non-smooth convex optimization theory. We provide a bound for the constant schedule with linear cooldown; in particular, the practical benefit of cooldown is reflected in the bound due to the absence of logarithmic terms.
Further, we show that this surprisingly close match between optimization theory and practice can be exploited for learning-rate tuning: we achieve noticeable improvements for training 124M and 210M Llama-type models by (i) extending the schedule for continued training with optimal learning-rate, and (ii) transferring the optimal learning-rate across schedules.

The Surprising Agreement Between Convex Optimization Theory and Learning-Rate Scheduling for Large Model Training
Fabian Schaipp Alexander Hägele? Adrien Taylor'
Umut Şimşekli
Francis Bach
Abstract
0.00
CUCOA PM
We show that learning-rate schedules for large model training behave surprisingly similar to a performance bound from non-smooth convex op- timization theory. We provide bound for the constant schedule with linear cooldown; in par- ticular, the practical benefit of cooldown is re- flected in the bound due to the absence of loga- rithmic terms. Further, we show that this surpris- ingly close match between optimization theory and practice can be exploited for learning-rate tuning: we achieve noticeable improvements for training 124M and 210M LI ama-type models by (i) extending the schedule for continued training with optimal learning-rate, and (ii) transferring the optimal learning-rate across schedules.
Volebetee Lowe
(aal 233. 441
Sehepitrealt Aeceed
3.00
0.0
0:00
10331
20000
10000
30000
-3801
30000
+3033
Brnisoo
(a) Real Loss Curves
(b) Theoretical Bound
Figure f. Strikingly similar: Validation loss for a 210M Llama model trained with AdamW (left) and the theoretical suboptimality bound (6) from convex optimization (right). Both plots show wad and cosine schedule with different training lengths T, and with base learning-rate cosine heing twice as large as for wad.
by an arguably much simpler schedule, that combines a con- stant part with a cooldown period in the end (Hu et al., 2024; Hägele et al., 2024). This alternative schedule established under the name wsd (warmup-stan sapp-sttbb-edcya)) One distin- guishing feature of wsd is the drastic decrease of the loss shortly after initiating the cooldown.
1. Introduction
Large-scale machine learning requires a fine-tuned training recipe. In particular, the choice of appropriate learning-rate schedules is a crucial step for classical optimization meth- ods. This usually decomposes into the choice of a schedule, determining the shape of learning rates over time, and the tuning of a multiplicative base learning-rate, determining the magnitude of the step sizes.
However, the recent advancements in learning- rrating-rateee schedul- ing have emerged almost exclusively from empirical rather than from theoretical considerations (Loshchilov & Hutter, 2017; Goyal et al., 2017: Hoffmann et al., 2022: Hägele et al., 2024). We do not yet have a fundamental understand- ing that could explain the features of the above-mentioned schedules and why they perform better worse at given task, restraining the tuning procedure to a trial-and-error approach.
Over the years, the cosine schedule (Loshchilov & Hut- ter, 2017) has emerged among the most commonly used schedules in large (language) model training (Brown et al., 2020; Touvron al., 2023). The standard practice is set the frequency of the cosine half of the total number of training steps (Hoffmann et al., 2022); as consequence, the entire schedule depends the length of training, which makes it unsuitable for continued training. Recently, it has been shown that the performance of cosine can be matched
Summary and contributions. In this paper, we show that several empirical findings on scheduling can be reproduced with a suboptimality bound for SGD on convex problems that was introduced by Defazio al. (2023). Among others, this theory allows to reproduce
(i) the matching performance of
'Inria, Departement d' d'Informatique de l'Ecole Normale Superieure, PSL Research University, Paris, France "EPFL, Lausanne, Switzerland. Correspondence Fabian Schaipp Sabian.schaipp@innia.t
cosine and wsd; (ii) the necessity and optimal length of the cooldown period in the end of training (see Section 4).
In a second step, we take the reverse direction and show how the theoretical bound can be exploited in practice (see Section 5): for continued training of a 124M and 210M Llama-style model, using the theoretically optimal schedule
Proceedings of the 42nd International Conference Machine Learning. Vancouver, Canada. PMLR 267, 2025. Copyright 2025 by the author(s).

---

Learning Rate Scheduling for Large Model Training
notably improves performance compared continuing with the same learning rate; it also allows to transfer the optimal learning-rate across schedules (Figs. 10 and 12). This leads us to the perhaps surprising conclusion that the empirical behavior of learning-rate schedules in (non-convex) deep learning can be described precisely with a theoretical bound from non-smooth stochastic convex optimization.
warmup, the wsd schedule is constant up to iteration To S T, then decays linearly zero. Formally, we have
< To: To 515TT
(3)
1
T-TO T+1-16
The cosine schedule is given by T/2 = (1 + cos( 1-1 for S S T + Note that for both schedules we have MT+1 = 0 (we choose T + 1 in order to ensure that 1/+ >> 0 for S 7). It is also common to decay the cosine to a factor of 0.1 of the peak learning-rate instead of O.
T T))
particular focus of this paper is put on the wsd schedule: we derive a convergence result for this schedule (without warmup) in the non-smooth stochastic convex setting, see Section 3.2. Most importantly, the cooldown period of wsd leads to vanishing log-terms in the bound, which provides an explanation of the benefit of cooldown observed in practice. Second, we show that the sudden drop during cooldown can be observed in upper and lower bounds of the suboptimality, as well for non-smooth convex toy problem. Code for all experiments is available at https: / //ggtthubbb.. com/ fabian- sp/ sese-rssheeellnnn
Notation and naming convention. We will use wsd in the paper as it is the most established abbreviation in the literature; however, similar to Zhai et al. (2022); Hägele et al. (2024), we will refer to the phase where the schedule decays to zero as cooldown instead of decay, in order to avoid confusion with other terminology (e.g., weight decay). Unless explicitly stated otherwise, || || and (.,-) denote the standard Euclidean norm and its scalar product.
Setup. We consider the training problem
min f(x), [() =- E. [f s)].
(I)
=E3d
2. Related Work
In the above, I E Idd are the learnable parameters of a machine learning model, and f is a loss function. The ex- pectation is taken over the distribution of a random variable that maps to the space or set S (typically the training set). We assume that f(,,s) has a suitable subdifferential for ev- ery= $ E S (for example, see Rockafellar (1970); Clarke (1983); Bolte& & Pauwels (2021)). We denote elements of the subdifferential as g E (z s).
Learning-rate schedules. The cosine schedule (Loshchilov & Hutter, 2017) can be considered the de-facto default in large-scale deep learning. Convergence results for SGD with cosine schedule have been shown by Li et al. (2021). Recently, the wed schedule (short for warmup-stable--ecaay also called trapezoidal schedule) has been proposed as an alternative (Zhai et al., 2022; Hu et al., 2024; Hägele et al., 2024). Hägele et al. (2024) show that wsd matches the performance of cosine on LLM pretraining, while largely reducing the compute needed for scaling-law experiments, as the constant part of the schedule can be reused.
We study the iterates of SGD with a learning-rate schedule, given by
Te+1 = It - M,... 9. E (f( EE
(2)
Here, > is called the base learning-raa Prateg---eeaaeeeee and > is called the schedule. While it might seem redundant to sepa- rate y from (ne). this reflects the standard practice in deep learning libraries such as Pytorch. Most importantly, for different schedules (constant, cosine, wsd ), the optimal value of is in general different.
Last-iterate convergence. We will see that it is crucial to use a bound for the last-iterate in order to closely match empirical loss curves. This is in contrast to many stan- dard convergence results that prove an upper bound on the quantity min,=1....T ([f())) - fx=))] Due to convexity and Jensen's 's inequality, the same bound usually holds for E(f( - f(x.)) )[, where ëy is some (weighted) average over {=1.... FT Last-iterate results, that is, bounds on Ei (xr) - f(x.)), ). are less standard: convergence of SCD has been proven for constant step sizes (Zhang, 2004), and for decreasing step sizes in bounded domains (Shamir & Zhang, 2013). Other results are restricted to a specific choice of schedule (Jain et al., 2021: Zamani & Glineur, 2023). The backbone of this article will be a result from Defazio et al. (2023), which proves a last-iterate bound for general sched- ules; compared to previous work (Orabona, 2020) it has the advantage that the bound remains meaningful if the last step size or is very small.
We remark that the most commonly used optimizer for train- ing in practice is Adam (W) (Kingma& & Ba, 2015; Loshchilov & Hutter, 2019), and all empirical results we present or refer to in this paper are obtained with Adam (W). However, the theoretical results apply to SCD; we address this limitation in detail in Section 6.
Cosine and wsd schedules. We now formally introduce the two running examples cosine and wsd. Without
In case the reader is uncomfortable with the notion of subdif- ferentials, the entire article can be read with ge being the gradient V f (Te, 8t) instead.
2

---

Learning Rate Scheduling for Large Model Training
Understanding cooldown. For the wsd schedule, one can consistently observe sudden drop in train/validation loss after the start of the cooldown phase (Hägele et al., 2024). Hu et al. (2024) find that the curvature of the loss increases during cooldown; Hägele et al. (2024) expand this and con- clude that "the cooldown phase is a smooth transition to a basin in the loss landscape". More recently, Wen et al. (2025) hypothesize that the sudden drop caused by a river- valley loss landscape, that arises from "heterogeneity in the stochasticity of different tokens". In this work, we will offer an additional (and potentially much simpler) model: the drop of the loss can be observed in upper and lower bounds of the suboptimality, based on first-order convex optimiza- tion theory. In particular, this phenomenon happens for a toy instance of Lo -nom regression.
Remark 3.2. In general, the choice of y will affect the iterates (are) and therefore the gradient norm bounds (G+). Thus, the following Corollary can be applied only if we apply the same bound G, independent of y. This is the case for the standard assumption of f(;,)) being Lipschitz with constant G > O; in that case, choose G. = G for all t E N.
Let T/ = (m. .r)), and G == (G1, ,,r, For convenience, we define the quantities
1 D².
Ti Ti,.T, = 2 1/1
T
Ta(mi,, GITTT) =
2 2 T7t
1
(
m
G
(7)
r=1
T + 2 k=
T
mk
1
Σ
n
3. Convergence Results
2t=k41 Tht
Ei= Tt
t=k
Let us assume convexity of the objective and recall the definition of the iterates.
Corollary 3.3. In the setting of Theorem 3.1, under (A3), for any T E N it holds B[f(xT) - f(x))] S Str with
ΩT = T1 + TT(nnr GIT,T). T)- (8)
(Al) For each E S the function s(),s) : 12d + R is convex, that is, for all x. y Rd and g E (fxx)) 8)
7
For given (G)), D and T, minimizing the right-hand side of (8) with respect to 7 > 0 gives the solution
f(y 8) fxx,s) 2 (g. x)-
(4)
>* = V Taimrr. ..T)
Talmar G... Piugging " back into (8), we have E(f(xr) fxx)) S 2 Ti D, TYTa (mhr, G T).
(A2) Let > 0 anddny ne > 0. For & N. consider the iterates
T++1 = Zt - Mi9., ge E Of(zz, st).
(5)
Next, we plug in the cosine and wsd schedule into The- orem 3.1. Applying" Corollary 3.3 with T + t, we get
Let x, E Rd be an arbitrary point of interest, for example the (local) minimum of f that is closest to 21. We do not make any other assumption on x. for now.
B(f(x)) - f(x.)) S 24 for t E (7) with
Ω = (Tma, D,1) + TT(Ma, GLa,t).
Theorem 3.1 (cf. Thm. 10 from Defazio et al. (2023)). Let (xz) be given by (A2), with nr > for = ,....,T T and
(9)
> O. Let F. E Rd and define D = || - x. and
3.1. Comparison of cosine and wsd
Tr= Et Mt- Under (AI), for any T E N it holds
For a training horizon E N, we define both schedules (n) (mp)geeT++2 +1 such that they reach 7r+1 = 0. For a formal definition of wsd and cosine see (3) and thereafter. For each different training horizon T, and for both schedulers. we pick the optimal base learning-rate 7 given by Corol- lary 3.3 and plot the bound 82 in Fig. 2 (with G, = D = for all i E N). We plot sweep of y in Fig. 3a.
T
E[f(xT)- - f(xx))
D2
+
x
Σ
27TT
T-1
T
E
99.
12
km Ei=++1
Etzk 17t
The above result is essentially the same as (Defazio et al., 2023, Thm. 10); the only difference is that we explicitly separate y and (nx) which will be convenient subsequently. We refer to Appendix E for a proof.
Perhaps surprisingly, the shape of the theoretical bound 2 (for the convex case) matches closely the empirical loss curves of (the non-convex problem of) language model pre- training in Hägele al. (2024); see Fig. 1 for bide-by-side comparison. This is especially visible in the sudden drop of the loss for the wsd schedule during cooldown. However, us- ing the last-iterate result is crucial for this: we demonstrate
Our next goal is to compute the base learning-rate -rae 7, given schedule Th, that minimizes the bound in (6). To do so, we assume bound on the expected gradient norms:
FFrr Corollary 3.3 we require T) > 0 for f = 1,..., T, which is why we construct the schedule such that u instead of rr = 0.
(A3) Assume that there exists (Gr)esT > such that Ellg. 2| z G for all & T.
3

---

Learning Rate Scheduling for Large Model Training
o T + Hr-1l Therefore, Corollary 3.3 yields
Trant korting Nortingttte
DG - (((=T)-((aa)))) V VT V1+ HT-1-
0/2
0.00
2.1
The wsd schedule. We will now compute a suboptimality bound for the wsd schedule (without warmup). We will show that if higher-order terms are ignored, the improve- ment of wsd over a constant schedule is essentially due to the absence of the logarithmic term Hr-1- In Theorem 3.4, we surpress some terms due to space constraints. The full version is given in Theorem G.I.
231
J00 0:3
200
400
500
Iteraton
Figure 2. Schedule (left) and theoretical bound (right) for cosine and wad, and various T, with base lcarning-rate y
105
VoC DMA LDA DILEXO
0.00
moke of wit ginae
ved. of
Theorem 3,4. Assume that (A3) holds with G. G for some G > 0 for all E N. Let 7 = - from Corollary 3.3. Then, for the wsd schedule (3) with S To < T we get
Finil bouze Dp
0.00
0.10-
0.04
0.00
Elf(rr) f(xx)) S DG 7+% A1 + A2 As [1]]
01000
10
Box: Anzaiza-nte
where A := T+2To A2 == Hr+1--2 -1++
As 3(T (T-To) 3(T+Ta)'
(a) Learning-rate sweep
(b) Optimal base learning-ratt
(T and o(1) summarizes terms that
Figure 3. Optimal base learning-rate decays with inverse square- root of training horizon T (right). For cosine, it is roughly twice large for wed (as 0.92/0.47 2). The dashed curve the right-hand side plot is obtained with a least-squares fit.
go to zero as T +00.
Assume that the cooldown length is proportional to T. that is, To = 3T for B E (0, 1). For large T, we have Ag R 3(1-8)(1+]) we can estimate H(++8)T-2 S 1 In((1 + 3)T) and
(1-9)8T2 = 3(1-8)- Using Lemma D.3,
this with an ablation study that uses a standard bound on the minimum suboptimality instead; there, the theoretical bound does not resemble empirical loss curves (cf. Appendix A).
HQL-9r+11 2 In((1-3)T). This yields A2 S 1+In + Altogether,
DG
The last-iterate bound in Corollary 3.3 matches the shape of the loss curves in Hägele et al. (2024). In particular, the sudden drop for wad during cooldown can be observed.
Takeaway:
B[f(xT) f(x.){
VT
I+8
A4 + In (1))]
31+33 In total, the term in the square-root does not contain logarithmic terms in T. This is the main difference to the constant schedule (where in the square-root have 1 HT-1 for visualization. We defer additional remarks and the proof of Theorem 3.4 to Appendix G.
where A4 =
1+28
3(1+3)
The optimal base learning-rate from Corol- lary 3.3 scales 1/VT with the training horizon, and roughly twice as large for cosine as for wsd (Fig. 3b). This matches empirical observations (Fig. 4 in Shen et al. (2024) and Fig. 3 in Hägele et al. (2024)).
Takeaway:
+ In())). See Fig. 20
The wsd schedule improves over the con- stant schedule by a logarithmic term. This improve- ment in the bound happens during the cooldown pe- riod (cf. Figs. 2 and 20).
Takeaway:
3.2. Bound for wsd Schedule
We now derive the bound in Corollary 3.3 for (n.) being the wad schedule. To the best of our knowledge, this schedule has not been analyzed theoretically before. For this section, assume that G; = G >> 0 for all 1. & N. A useful notation will be the harmonic number H,, defined as = for & E N. and Ho := O. We recall that H, behaves like lt(t) in the limit. As baseline, we first compute the bound for the constant schedule.
4. Theoretical Simulations
In the following, we simulate the bound from Theorem 3.1 in order to analyze its dependence on the cooldown length for wsd, and on the gradient norm magnitude. Additional ex- periments (e.g.. on the cosine cycle length, and a compari- son of classical schedules) and supplementary information are deferred to Appendix B. Unless explicitly mentioned, we set G, = for all t E N and D 1 for the entire simulation. We do not use warmup for neither schedule.
Constant schedule. If nt for all t N. is easy to compute To (m D, T) = 2T: Ta(mer: GL., T) =
D²

---

Learning Rate Scheduling for Large Model Training
1.0
10-
10
Oodioo tranties
C se
(7)
0.6
0.20
0:15
10'
10-1
10-1
300
200
500
-10
100
200
2000
430
Iteratton 8
Irrastion
Figure (Left) Optimal base learning-rate increases with cooldown fraction. (Right) For fixed y, the optimal cooldown fraction can smaller than 1. The analogous curves for real ex- periments with similar parabola shapes are in Fig. 21.
Figure 6. Assumed gradient shape (leff) and theoretical convereg- nce (right). Only with a = (constant Go), the sudden drop for wad is clearly visible.
4.2. Gradient Norm
0.35
00.4
0.30
We now analyze how the bound of the expected gradient norms G1:T influences the shape of 92g. In this section only, assume that G, 10, ar {0, -0.5, -1). We sweep the base learning-rate 7 by computing the minimal 2r from (9) for the above choice of GI.T (Gi,... ....... GT). We set
0.00
D2 Gar
0..35
0.00
0.35
0.10
190
200
322
212
T= 400, and the cooldown fraction to 0.2 for wsd. Fig, 6
The 2
shows that the sudden drop in loss for wsd is only visible G, does not go to zero as xC.
Figure 5. Schedule (left) and theoretical convergence (right) for varying cooldown fraction. With optimal base learning-rate y", starting the cooldown at To = 1 is optimal. Fig. 21 shows the analogous plot for real experiments with the same behavior.
Takeaway:
The sudden drop during cooldown is most pronounced if the gradient norms do not go to zero.
Interestingly, if the gradient norms go to zero, the wed sched- ule also obtains a significantly better bound than cosine. So far we have observed that non-vanishing gradient norms lead to the characteristic drop in the upper bound Ω4 Next, we show that the same phenomenon can be observed for (i) suboptimality lower bound and (ii) for the loss of the iterates of SCD on a simple non-smooth convex problem.
4.1. Cooldown Length
Previously, we have set To 0.8 ' for wad. In Figs, 4 and 5, we vary the cooldown fraction, defined as 2-2h Specifically, we vary from To = T to To = 1 (constant schedule to lincar-decay schedule similar to Defazio et al. (2023, Corollary 2)).
T
4.3. Lower Bounds and Convexity
The simulation suggests that if the base learning-rate is fully tuned, then the optimal cooldown fraction is 1 (linear decay). For fixed 7, the optimal cooldown fraction can smaller than one.
Takeaway:
In all previous sections we analyzed an upper bound Ωg of [x[((;)) - f=x+)) How tight is this upper bound? To answer this, we compute lower bounds of (z,) - f(x))] using the PEP framework: for given function class and algorithm, worst-case example can be constructed by solv- ing semidefinite program (Drori & Teboulle, 2014; Taylor et al., 7a;b; Goujaud et al., 2024). Additional details are provided in Appendix B.4.
The first observation is in line with empirical observations from Defazio al. (2023) that compares many different schedules across several machine learning tasks, and find that the lincar-decay schedule performs best on average. Further, it is known that the linear-decay schedule matches the exact lower-bound convergence bound for the (stochas- tic) convex, Lipschitz case (Defazio al., 2023: Zamani & Glineur, 2023); see Appendix for detailed comments. The second observation matches the finding of Hägele al. (2024): in Appendix B.6, Fig. 21 we show the analogous figure on real training data (also see Fig. 5 in Hägele et al. (2024)). For small base learning-rate % we obtain the same parabola shape; however, for large enough 7, the parabola turns into a monotonically decreasing curve.
Takeaway: The sudden drop during cooldown appears well for the PEP lower bound (Fig. 7). The worst- case suboptimality value at T = 60 is very similar for cosine and wsd.
In Section 4.2, we have shown that non-vanishing gradient norms are characteristic for the sudden drop (of the upper bound Ω22) during cooldown of wsd. We observe the same behavior for the actual loss when running gradient descent with the wed schedule for the 2-dimensional, convex, non-
5

---

Learning Rate Scheduling for Large Model Training
smooth problem mingeR2 ||
|rr bllso. The experimental details and plots are deferred to Appendix B.1.
hus becosdeer
0.003
adayed tollioss
The sudden drop of the loss for wad is not specific to model architecture or even to non-convexity. It can be observed when minimizing simple non- smooth, convex objective function (Fig. 17).
Takeaway:
f 0.000
27
00000
breatiie 10000
10000
5. Applications
Figure 8. (Left) Transferring the wsd schedule from horizon Ti 1000 to T2 E [1.57], ]. (Right) Not adapting the cooldown length leads to significant suboptimality. Dashed horizontal lines mark bound for the linear-decay schedule with tuned y
So far, we have shown that the bound from Theorem 3.1 matches very closely empirical loss curves. However, the bound in Theorem 3.1 contains quantities that are unknown in practice, such as the gradient norm bounds G2 and D. Thus, the question arises how to convert the theoretical result into practical upplications. The following two scenar- ios demonstrate that using the optimal schedule and base learning-rate predicted from theory improves pretraining of a 124/210M Ll styyll transformer (Vaswani al., 2017: Touvron et al., 2023).
(B2) Alternatively, we can keep the same cooldown fraction, but decrease the learning-rate in the steps from Ti to T2: assume piecewise constant schedule with 7t- = 1 for up to the start of cooldown of the short run, and p for l up to the start of cooldown of the long run. How do we need to set p. such that y remains the optimal base-learning rate for this schedule?
5.1. Schedule Construction for Continued Training
The first application is to construct learning-rate schedules for longer horizons: for example, assume we have trained model for Ti steps, but later want to continue training up to T2 > Ti steps. The main benefit of the wed schedule is that the training steps up to the cooldown phase can be reused, thus reducing the amount of additional compute required for continual training (Hägele et al., 2024). This is not true neither for the linear-decay nor for the cosine schedule, as the value of the schedule in each step depends on the total number of steps.
Simulation. We simulate both options in Figs. 8 and 9. Here, we set Ti = 4000 and T22 ranging from 1.57; to 47,. We construct the extended schedule by sweeping e₁ for (BI) and p for (B2), and picking the value where the optimal base learning-rate according to Corollary 3.3 is closest to y" The cooldown phase of the short run is set to 20%. Specifically, our analysis suggests to decrease the schedule by p = 0.525 for T2 = 2T, and by 0.375 for T. = 4T, (see Fig. 22, left). We verified that changing the values of G. D. or Ti do not affect the result (plots not shown): the values might be different for other cooldown fractions than 20%.
Assume we have tuned the base learning-rate >" of wsd for the short run Tj. We have seen in Fig. 3b that 7 decreases with T; thus, continuing training at y until T2 would use a suboptimal learning rate. We present two solutions:
For (B2) (Fig. 9), conclude that by decreasing the sched- ule by the correct factor we can reuse the entire constant part of the short run, while obtaining a bound f2 close to the bound for a tuned lincar-decay schedule. Importantly, keep- ing the same base learning-rate for the entire long run would result in a significantly worse bound l4. For (B1) (Fig, 8), the required increase in cooldown fraction is large, and hence for long extensions, only small parts of training can be reused. When doubling the training length (T2 = 2T the adapted cooldown fraction is roughly c1 = 0.6. As an alternative, one could use the 1/sqrt schedule, defined by 1/t : 1/V1, combined with cooldown (Zhai et al., 2022). Fig. 23 shows that for 1/sqrt the cooldown fraction can roughly stay the same, which however comes at the cost of larger gap to linear-decay. From a theoretical and prac- tical perspective, we conclude that the approach (B2) is preferable, as it allows to reuse the entire short run with no drawbacks terms of the bound, and - in a similar fashion as before - allows iteratively continuing from the newly obtained checkpoints of a constant learning-rate phase.
(B1) We have seen that 7* increases with the cooldown frac- tion e (Section 4. We can increase the cooldown fraction e₁ for the long training run to T2. to compen- sate for the decrease in y due to Ti- + T2.
Jeee teex!
tuppor
10P
10F
10
231
201
31
33
10
301
1
(a) wsd with 20% cooldown
(b) cosine
Figure 7. PEP lower bound matches the upper bound Zr in shape.
6

---

Learning Rate Scheduling for Large Model Training
Experiments. Based on the above, we extend the train- ing of 124M and 210M Llama-style transformer (Tou- vron et al., 2023) on the SlimPajama dataset (Soboleva et al., 2023). For details on model, dataset and training procedure see Appendix B.5. We set 71 = 50k and T2 E {100K, 200k]; a sweep over 50k steps gives Yook R 0.001. As a baseline, we use wed schedule that continues with the same Y over the extended training length. For the adapted schedule from (B2), we decrease the learning-rate after 40k steps linearly over 1000 steps (e.g., from 10 to 5.25-10-1) as a precautionary measure; however, we did not observe that a decrease in-one-go results in significantly different performance. We use a cooldown of 20% for all runs.
Twor done
hus broodderr
0.003
alayed ateccll
f 0.000
27
0.000
10000
10000
Figure 9. Transfering the learning-rate schedule from horizon Ti 1000 to T2 E [1.57], [1T]] (see also Fig. 22, left). De- creasing the learning rate (green) after the short run (at iteration 3200) leads to significant better bound Pe as keeping it constant (grey). Dashed horizontal lines (blue) mark bounds for lincar- decay schedule with tuned y'
Considering the results in Fig. 10, we conclude that the schedule adaptation suggested by theory leads a slight but noticeable improvement in validation loss for both extended horizons. The improvement is more pronounced for the larger 210M model. Moreover, we observe a sudden drop in loss after decreasing the schedule at 40k steps, analogous to what the theoretical bound predicts, albeit the loss decrease thereafter is slower than expected (ef. Fig. 9). We also test adapting the cooldown length as described in (BI): for a total length of 100k steps, if cooldown is initialized immediately after 40k steps (cooldown 60%), we observe even larger improvements as previously (see Fig. 24).
:12
2012
3.1
Valihion bee
(ae, 2011
2.0
2.7
20
100
150
500
380 150
200
Thrration geea Tol)
Dretatio a grvit 30")
(a) 124M model
(b) 210M model
Figure I0. Transferring the learning- raae schedule from horizon Ti 50 000 to T2 4T]. Decreasing the base learning- rate (green) after 40k steps lcads to small improvements in vali- dation loss compared to keeping it the same (grey). We discuss the significance of the difference in loss values of (around 0.01) in Section 5.1 and Appendix B.5. See Fig. 22 for schedules.
From Figs. 10 and 24, we see that the improvement in loss of the adapted wad schedule over a naive continuation is in the range of 0.01. This raises the natural question of the rele- vance of such an improvement. To answer this, we estimate the slope of our loss curves?: we find that for T2- 100k, decrease of 0.01 takes roughly 6k steps in the constant learning-rat rate phase; for T-- = 200k, it takes roughly 14.5k steps. This translates to 0.6B and 1.5B tokens, respectively. Notably, to match the adapted wsd schedule, this would re- quire a substantial amount of 6% and 7.25% longer training. Another way to reason about the significance of the loss improvement is through the use of scaling laws, which leads to very similar estimates (see Appendix B.5).
better model, we might want to retrain with the linear-decay schedule. However, we do not yet know (1)). Can we compute > (1) from 7 (c) based on the theoretical bound?
Simulation. In Fig. 11 we show the quantity In In(?"(!) for e E (0. We simulate both the linear cooldown (3), and the 1-sqrt cooldown which has the form 7)- =
(c()
Tn (Hägele al., 2024). Across several orders of T, V T+1-1 the results are consistent; for example, knowing >" for 20 % of linear cooldown, we can compute 7 (1) R e0 (0.2). For Fig. 11, we set G = D 1; the resulting curve looks the same if we vary D or G (plots not shown).
5.2. Learning-Rate Transfer Across Schedules
One insight from Corollary 3.3 is that if G. = G, then the dependence of on G' and D is multiplicative. This implies that if we know - for given practical problem, any multiplicative transfer can be realized. For example, assume we know the optimal base learning-rate for the wad schedule with cooldown fraction e E [0, 1], and let us denote the tuned value as >*(c). As we have seen, the linear-decay schedule (c = 1) attains the optimal bound; thus, to obtain
Experiments.
We now analyze the quantity In( (1))
x- (c)
with real data (training a 124M Llama-style model for 50k steps), with linear cooldown. We estimate (c) from grid of base learning-rates y and cooldown fractions e (see Fig. 12b and Appendix B.5 for details on this step). We plot
*We do linear regression on the loss values the baseline run between [64000, 84 000] for the 100k run, and between [144 000, 164 000) for the 200k run. This accounts proportion- ally for the cooldown.
'For example, assume that (c) has been made public on Github or obtained it from sweep that used the wad schedule due to practical constraints.
7

---

Learning Rate Scheduling for Large Model Training
lincor
I-spr
is based on theoretical result for SGD, while most of the empirical evidence use is obtained with Adam (W). More generally, the result in Theorem 3.1 can not explain any performance differences that stem from the optimization algorithm. However, we believe that this gap can be closed in future work for several reasons: (i) by showing similar theoretical results for the methods used in practice; as a first step, we provide a proof for mirror descent (an entire family of methods) in Appendix F. It has been shown that for diagonal networks, the iterates of SGD are equivalent to mirror descent on a convex problem formulation (Even et al., 2023). (ii) Several recent variants of SGD close the gap to Adam on transformer problems (Kunstner et al., 2023; Xu et al., 2024). (iii) It has been shown that most of the parameters of language models can be equally well trained with SGD (Zhao et al., 2025).
pan Sta 1900 30
LG
LA toa cao 2000 HUMOX)
1.0
3-y011
0.5
0.5
11.0
Figure/ 11. Transferring the optimal base learning-rate rate from cooldown fraction to linear-decay (c = 1): for linear cooldown (left) and 1-sqrt cooldown (right). Dashed lines are fitted poly- nomial of degree 6.
100
10t -
1.4
30hool tobs)
$,00
Blual vel
2:00
The second obvious limitation of Theorem 3.1 is the con- vexity assumption, while modern deep learning problems are non-convex. At this point we have no explanation for why the convex theory is still closely matching (some) real- world observations. However, it has been shown that the landscape of neural network optimization problems might be reasonably close to being convex (Hardt al., 2018; Liu et al., 2023; Islamov et al., 2024). The sudden performance increase during cooldown is not restricted to language mod- eling and has also been reported for image problems, e.g., training ResNets with SGD (Sandler et al., 2023) or ViTs (Zhai et al., 2022). We verify this through additional experi- ments for SGD on Imagenet in Appendix C, which further contains experiments OpenWebText2.
1.01
3.8
002
203
Coollonn Exction.
Thee leorskg
(a) Learning rate transfer
(b) Learning rate sweep
Figure 12. (Left) Re-analysis of learning- rate transfer (Fig. 11) for 124M model. Yo denotes the best performing hase learning-rate for cooldown fraction e, estimated from a sweep (right). We observe that the learning-rate transfer (black line) almost perfectly matches the predictions by theory (e.g.. (0.9944 z ed.7 (0.2)). Note that the maximal cooldown fraction is 0.994 due to warmup and corresponds to a full linear schedule.
In(2 (n(21()) in Fig. 12a; it matches almost perfectly with the (-te)
Finally, the empirical quantity we compare to the theoretical bound is the test loss. This is limited to situations where the generalization gap between training and test loss is neg- ligible; that being said, the current practice of single-past training for large models falls within this category (Aitchi- son, 2024: Xiao, 2024).
curve predicted from theory in Fig. 11. This implies that knowing the optimal base learning-rate for 20% cooldown, one can immediately transfer the learning-rate to linear- decay (100% cooldown), without any additional sweeps; for the setup we consider, the linear-decay run obtains a final validation loss of 2.9535 vs. the best run with 20% cooldown obtaining a final loss of 2.9660.
7. Conclusion
6. Limitations
In this paper, we show that learning-rate schedules in prac- tice behave surprisingly similar to what convex optimization theory predicts. This spans across the necessity and optimal length of the cooldown period at the end of training as well as the optimal learning-rate transfer. Notably, our experi- ments suggest that the theoretical bounds can be used as testbed for schedule design before training: we have shown that theoretically inspired schedules achieve notable im- provements in practical scenarios. More broadly, our results suggest that one key characteristic underlying the observed behavior is gradient norms that do not go to zero; in practice, this could be due to non-smoothness (of the objective) or due to the problem-inherent gradient noise. We leave it as future work to explain this phenomenon.
We have shown that the empirical performance of vari- ous learning-rate schedules for large model training re- flects closely the theoretical suboptimality for non-smooth stochastic convex optimization. We want to stress that we can not expect the bound from Theorem 3.1 to match train- ing curves perfectly: first, it is an upper bound of the loss for convex problems only, and in practice many other fac- tors (e.g., randomness, architecture choices, data mixture) and training techniques (e.g., loss function, weight decay) will impact convergence and stability of training (Wortsman et al., 2024).
The perhaps most glaring limitation of our work is that it
8

---

Learning Rate Scheduling for Large Model Training
Acknowledgments
Defazio, A., Cutkosky, A., Mehta, H., and Mishchenko, K. Optimal linear decay learning rate schedules and further refinements. arXiv:2310.07831, October 2023.
A. Taylor is supported by the European Union (ERC grant CASPER 101162889). U. Şimşekli is supported by the Eu- ropean Union (ERC grant DYNASTY 101039676). Views and opinions expressed are however those of the author(s) only and do not necessarily reflect those of the European Union or the European Research Council Executive Agency (ERCEA). Neither the European Union nor the granting authority can be held responsible for them. The French government partly funded this work under the manage- ment of Agence Nationale de la Recherche as part of the "France 2030" program, reference ANR-23-IACL-0008 (PRIAI]RIE-PSAD.
Drori, Y. and Teboulle, M. Performance of first-order meth- ods for smooth convex minimization: a novel approach. Mathematical Programming. 145(1-2):451-422 2014.
Even, M., Pesme, S., Gunasekar, S., and Flammarion, N. (S)GD over diagonal linear networks: Implicit bias, large stepsizes and edge of stability. In Advances in Neural Information Processing Systems, volume 36, pp. 29406- 29448, 2023.
Goujaud, B., Moucer, C., Glineur, F., Hendrickx, J. M., Tay- lor, A. B., and Dieuleveut, A. PEPit: computer-assisted worst-case analyses of first-order optimization methods in Python. Mathematical Programming Computation, 16 (3):3373367, 2024.
Impact Statement
This paper presents work whose goal is to advance the field of Machine Learning. There are many potential societal consequences of our work, none which we feel must be specifically highlighted here.
Goyal, P., Dollár, P., Girshick, R., Noordhuis, P., Wesolowski, L., Kyrola, A., Tulloch, A., Jia, Y., and He, K. Accurate, large minibatch SGD: Training ImageNet in hour. arXiv:1706.02677 June 2017.
References
Aitchison, L. Why you don't 't ovvrffit, and don't need Bayes if you only train for one epoch. arXiv:241 144478, Novem- ber 2024.
Hägele, A., Bakouch, E.,, Kosson, A., Ben allal, L., Von Werra, L., and Jaggi, M. Scaling laws and compute- optimal training beyond fixed training durations. In Ad- vances in Neural Information Processing Systems, vol- 2024.
Beck, A. and Teboulle, M. Mirror descent and nonlinear projected subgradient methods for convex optimization. Operations Research Letters, 3173)1667177, May 2003.
ume 37, pp.
Hardt, M., Ma. T., and Recht, B. Gradient descent learns linear dynamical systems. Journal of Machine Learning Research (JMLR), 19(1): January 2018.
Besiroglu, T., Erdil, E., Barnett, M., and You, Chinchilla scaling: A replication attempt. arXiv:2404100102, April 2024.
He, K., Zhang, X., Ren, S., and Sun, J. Deep residual learning for image recognition. In 2016 IEEE Conference on Computer Vision and Pattern Recognition (CVPR), pp. 770-778, 2016.
Bolte, J. and Pauwels, E. Conservative set valued fields, automatic differentiation, stochastic gradient methods and deep learning. Mathematical Programming, 188(1, Ser. A):19-51, 2021.
Hoffmann, J., Borgeaud, S., Mensch, A., Buchatskaya, E., Cai, T., Rutherford, E., de Las Casas, D.,, Hendricks, L. A., Welbl, J., Clark, A., Hennigan, T., Noland, E., Millican, K., van den Driessche. G., Damoc, B., Guy, A., Osindero, S., Simonyan, K., Elsen, E.,, Vinyals, O., Rae, J., and Sifre, L. An empirical analysis of compute- optimal large language model training. In Advances in Neural Information Processing Systems, volume 35, pp. 50016-30034 2022.
Brown, T., Mann, B., Ryder, N., Subbiah, M., Kaplan, J. D.,
Dhariwal, P., Neelakantan, A., Shyam, P., Sastry, G.,
Askell, A., Agarwal, S., Herbert-Voss. A., Krueger, G., Henighan, T., Child, R., Ramesh, A., Ziegler, D., Wu, J., Winter, C., Hesse, C., Chen, M., Sigler, E., Litwin, M., Gray, S., Chess, B., Clark, J., Berner, C., McCandlish, S., Radford, A., Sutskever, I., and Amodei, D. Language models are few-shot learners. In Advances in Neural Information Processing Systems, volume 33, pp. 1877 1901, 2020.
Hu, S., Tu, Y., Han, X., Cui, G., He, C., Zhao, W., Long, X., Zheng, Z., Fang, Y., Huang, Y., Zhang, X., Thai, Z. I.., Wang, C., Yao, Y., Zhao, C., Zhou, J., Cai, J., Zhai, Z., Ding, N., Jia, C., Zeng, G., dahai li, Liu, Z., and Sun, M. MiniCPM: Unveiling the potential of small language models with scalable training strategies. In First Conference on Language Modeling, 2024.
Clarke, F. H. Optimization and nonsmooth analysis. Cana- dian Mathematical Society Series of Monographs and Advanced Texts. John Wiley & Sons, Inc., New York, 1983.

---

Learning Rate Scheduling for Large Model Training
Islamov, R., Ajroldi, N., Orvieto, A., and Lucchi, A. Loss landscape characterization of neural networks without over-parametrization. arXiv:2410. 12455, October 2024.
scheduler: A batch size and token number agnostic leamn ing rate scheduler. MiXiv2240..13355 August 2024.
Soboleva, D., Al-Khateeb, F., Myers, R., Steeves, J. R.,
Jain, P., Nagaraj, D. M., and Netrapalli, Making the last it- erate of sgd information theoretically optimal. SIAM Jour- nal on Optimization. 3 31(2):1 (100::100--11111 -130, January 2021.
Hestness, J., and Dey, N. SlimPajama: A 627B token cleaned and deduplicated version of RedPajama, 2023.
Taylor, A. B., Hendrickx, J. M., and Glineur, F. Smooth strongly convex interpolation and exact worst-case per- formance of first-order methods. Mathematical Program- ming, 161(1-2)3377345, 2017a.
Kingma, D. P. and Ba. J. Adam: A method for stochastic optimization. In International Conference on Learning Representations, 2015.
Kunstner, F., Chen, J., Lavington, J. W., and Schmidt, M. Noise is not the main factor behind the gap between sgd and adam on transformers, but sign descent might be. In International Conference on Learning Representations, 2023.
Taylor, A. B., Hendrickx, J. M., and Glineur, F. Exact worst- case performance of first-order methods for composite convex optimization. SIAM Journal on Optimization, 27 (3):1283-1313, 2017b.
Touvron, H., Lavril, T., Izacard, G., Martinet, X., Lachaux, M.-A., Lacroix, T., Rozière, B., Goyal, N., Hambro, E., Azhar, F., Rodriguez, A., Joulin, A., Grave, E., and Lam- plc, G. LLaMa: Open and efficient foundation language models. arXiv:2302. 13971, February 2023.
Li, X., Zhuang, Z., and Orabona, F. A second look at exponential and cosine step sizes: Simplicity, adaptivity, and performance. In Proceedings of the International Conference on Machine Learning, volume 139, pp. 6553- 6564, July 2021.
Vaswani, A., Shazeer, N., Parmar, N., Uszkoreit, J., Jones, L., Gomez, A. N., Kaiser, L.,, and Polosukhin, I. Atten- tion is all you need. In Advances in Neural Information Processing Systems, volume 30, 2017.
Liu, C., Drusvyatskiy, D., Belkin, M., Davis, D., and Ma, Y. Aiming towards the minimizers: fast convergence of sgd for overparametrized problems. In Advances in Neural Information Processing Systems, volume 36, pp. 60748-60766 2023.
Wen, K., Li, Z.,, Wang. J. S., Hall, D. L. W., Liang, P., and Ma, T. Understanding warmup-stable-decay learning rates: A river valley loss landscape view. In International Conference on Learning Representations, 2025.
Loshchilov, I. and Hutter, F. SGDR: stochastic gradient descent with warm restarts. In International Conference on Learning Representations, 2017.
Wightman, R. Pytorch image models. bttps://githwn
Loshehilov, and Hutter, F. Decoupled weight decay reg- ularization. In International Conference on Learning Representations, 2019.
con/rvightman/pyye wrigghttmmmpytrrchhtt http//yytttohchhee iea---mmaeemmm models, 2019.
Wortsman, M., Liu, J., Xiao, L., Everett, K. E., Alemi, A. A., Adlam, B., Co-Reyes, J. D., Gur, I., Kumar, A., Novak, R., Pennington, J., Sohl-Dickstein, J., Xu, K., Lee, J., Gilmer, J., and Kornblith, S. Small-scale prox- ies for large-scale transformer training instabilities. In International Conference on Learning Representations, 2024.
Orabona, Last iterate of SGD converges (even in unbounded domains), 2020. [Online; accessed 08- November-2024),
Rockafellar, R. T. Convex Analysis. Princeton Mathematical Series, No. 28. Princeton University Press, Princeton, N.J. 1970.
Xiao, L. Rethinking conventional wisdom in machine learn- ing: From generalization to scaling. arXiv:2409.55156, September 2024.
Sandler, M., Zhmoginov, A., Vladymyrov, M., and Miller, N. Training trajectories, mini-batch losses and the curious role of the learning rate. trXiiv2301.02311 January 2023.
Xu, M., Xiang, L.,, Cai, X., and Wen, H. No more Adam: Learning rate scaling at initialization is all you need, arXiv:2412.11766 December 2024.
Shamir, O. and Zhang, T. Stochastic gradient descent for non-smooth optimization: Convergence results and opti- mal averaging schemes. In Proceedings of the Interna- tional Conference Machine Learning, volume 28, pp. 71-79, June 2013.
Zamani, M. and Glineur, F. Exact convergence rate of the last iterate in subgradient methods. arXiv:2307.11133: 11134, July 2023.
Shen, Y. Stallone, M., Mishra, M., Zhang, G., Tan, S., Prasad, A., Soria, M., Cox, D. D., and Panda, R. Power
Zhai, X., Kolesnikov, A., Houlsby, N., and Beyer, L. Scaling vision transformers. In Proceedings of the IEEE/CVE
10

---

Learning Rate Scheduling for Large Model Training
Conference on Computer Vision and Pattern Recognition (CVPR), pp. 12104-12113, June 2022.
Zhang, T. Solving large scale linear prediction problems using stochastic gradient descent algorithms. In Interna- tional Conference on Machine Learning, pp. 116, 2004.
Zhao, R., Morwani, D., Brandfonbrener, D., Vyas, N., and Kakade, S. M. Deconstructing what makes a good op- timizer for autoregressive language models. Interna- tional Conference on Learning Representations, 2025.
Zinkevich, M. Online convex programming and generalized infinitesimal gradient ascent. In Proceedings of the Inter- national Conference on Machine Learning, pp. 928 935, 2003.

---

Learning Rate Scheduling for Large Model Training
Appendix
The supplementary material organized as follows:
Appendix A: ablation on bounds on the minimal suboptimality bounds.
Appendix B: supplementary information our experiments.
Appendix C: additional experiments for training language and vision models
Appendix D: technical lemmas
Appendices E and G: missing proofs of our theoretical results
Appendix F: additional analysis for mirror descent
A. Ablation: Min-Suboptimality Bounds
Standard convergence results for the SGD method (2) make statements on the suboptimality of average iterate, of the best objective value (in expectation) found up to iteration T. We state one of the standard results for the non-smooth (stochastic) convex setting below (Zinkevich, 2003):
Theorem A.1. Assume that each f(,, 8) is convex. Let (It) be the iterates given by (A2), with the > 0 and > 0. Let I, E Id and define D = [ x. Under (A3), we have
min B[f(x) - {(x))] t=......
[D2
+ Σ n
G.
(10)
27 Et= lt
The right-hand side of the above is minimized by VET Gful
D
Plugging in 7 yields
D DVEL Gini
min Blf(x)) 11fx))) f(x))
Et=i
We now repeat the theoretical simulations, but, instead of 24 from (9), using
Ω 27 E1 s==1 77
D
+
2
n G
(11)
A.1. Ablation of Section 3.1
0.12
0.4
0.10
(200,300, 400, 500, 600)
coeine
Thned learnnng rate
0.08
T' E (200,300, 300, 400, 500, 600)
Bound Qr (for 3*)
0.3
0.06
0.2
0.04
0.02
0.1
0.00
200
400
600
200
400
600
Iteration f
Iteration f
Figure 13. Same as Fig. 2, hut with Q2 from (11)
The bound on the best-so-fan far bound has a very different shape of the last-iterate bound. This shows that standard bounds such in Theorem not capture the real-world convergence observed Hägele al. (2024).
12

---

Learning Rate Scheduling for Large Model Training
10a
Wood
Optimal base learning- +t
cosine, 7" 1.61 7
(200, 100, 600]
wsd, >* 1.07 -700.000
cosine
Final bound Or
E (200, 300, 200, 500, 800)
0.10
10-1
0.06
0.04
10
10"
200
400
600
Base learning-rate
Training horizon T
(a) Leaming-rate sweep
(b) Optimal base learning-rate vs. T
Figure 14. Same Fig. 3. but with Q2 from (11)
A.2. Ablation of Section 4.1
1.0
0.07
0.04
0.12
0.05
0.02
Final bound Or
Cooldown fraction
Final bound for
0.10
0.5
0.08
10
0.06
0.0
10-2
10
0.2
0.4
0.6
Base learning-rate
Cooldown fraction
(a) Leamiin- rate sweep
(b) Final bound vs. cooldown fraction
Figure 15. Same as Fig. 4, but with Q4 from (11)
0.35
0.08
Toned leariin ? 7/
0.30
Bound f (for )*
0.06
0.25
0.20
0.04
0.15
0.02
0.00
0.05
200
300
100
200
400
Itcration t
Iteration
Figure 16. Same as Fig. 5. but with Qt from (11)
13

---

Learning Rate Scheduling for Large Model Training
B. Experiments: Supplementary Material
B.1. Non-smooth Convex Example
Here, we provide details for the non-smooth convex toy example of min || - blloo mentioned in Section 4.3. We set = 2 and pick A E 200dd uniformly random from 1-1.11 1]. We generate an oracle x, E R and set Az.. We then gradient descent (GD) for T 400 iterations with the wad schedule (cooldown fraction 0.2 and y 0.02). As baseline, plot the constant schedule with 0.02 and cosine schedule with y 0.04. We pick zero starting point, except for the constant schedule, where we pick (10-3, 10- to obtain visually distinguishable path.
The objective function and iterate paths are shown in Fig. 17.
0.10
4.90
wed
0.08
constant
cosine
0.06
4.85
f(m))
0.04
0.02
4.80
0.00
4.75
Iteration
1.22
1.23
1.24
1.25
Figure 17. (Left) Sudden drop of the loss for wsd schedule for a convex, non-smooth problem. (Right) Iterate path for the three schedules. For wad, the cooldown period indicated with the dashed line. Star marks solution.
B.2. Schedule Comparison
We compare the upper bound Q4 from (9) for various schedules:
wsd with cooldown fraction 0.2,
cosine,
constant schedule,
linear-decay schedule, that is, wsd with cooldown fraction of 1,
1/sqrt schedule, where T) 1/Vt,
1-sqrt schedule, where 17:
7
We assume D = G, = 1 and set T 400. For each schedule we sweep the base learning-rate y and plot the bound Ωr for > - obtained from Corollary 3.3.
B.3. Cosine Cycle Length
For the cosine schedule, an important hyperparameter is its cycle length, that is, the amount of training where the schedule first reaches zero. Originally, it was proposed in Loshchilov & Hutter (2017) to use multiple warm restarts (a cycle length less than one). Later, Hoffmann et al. (2022) show empirically that the best performance in language modeling tasks is obtained by setting the cycle length to one (the half-cosine matches exactly the training duration).
14

---

Learning Rate Scheduling for Large Model Training
0.30
wsd 0.2
wsd 0.2
Tineed leaening tate y11p
linser
0.25
l/ggrt cunstaut cosine
Find boand fbr
1/rrrt
(Ex
constant
costne
0.01
0:20
I-agre
I-gge
0.02
0.15
0.10
10-1
300
10"2
10
10"
Iteratoon
lirtttoon
Biee leatning-rate
(a) Convergence
(b) Learning-ran sweep
Figure 18. Comparison various learning-rate schedules, Convergence plotted with the optimal base learning-rate 7 (chosen individually for each schedule).
To the best of our knowledge, this recommendation is based mostly on empirical insights. Using the bound obtained in Theorem 3.1, our analysis shows that cycle length of obtains the lowest bound Ωr. Thus, the theoretical bound is in accordance to the empirical conclusion from Hoffmann et al. (2022).
Note that Hoffmann et al. (2022) choose the base learning-rate equally for all cycle lengths. To match the setting of their experiment, we pick y" for cycle length of one, and use this for all other cycle lengths as well. Picking > for each cycle length individually yields qualitatively the same result (the optimal cycle length being one), but with slightly less pronounced differences (plots not shown). contrast to previous simulations, the final value of the schedule is chosen as 0.1 of the peak learning-ral (instead of zero), again order match the setting of Hoffmann et al. (2022).
0.35
0.04
0.30
Learning rate M/
0.03
Bound Ω
0.25
0.20
0.02
0.15
0.01
0.10
100
200
300
400
100
200
300
400
Iteration
Iteration
Figure 19. Comparison of cycle lengths for the cosine schedule. Compare Figure Al in Hoffmann al. (2022).
B.4. Details Lower Bound Computation
We provide additional details for the simulation in Section 4.3. We compute the lower bounds with the PEPit package (Goujaud et al., 2024). For our purpose, use the class of convex G-Lipschitz functions and gradient descent (GD) with the step sizes
In PEPit, we use the MOSEK solver. As the size of the semidefinite program grows with I', we choose a rather small T = 60, and compute the lower bound of (x, - f(x,) for all + that are multiple of 5. Note that the specific worst-case function ff constructed by PEPit can be different for each t (as maximizes the suboptimality exactly at iteration t). We set 7 to 7 minimizing the upper bound 22 (ef. Corollary 3.3). We set G D
15

---

Learning Rate Scheduling for Large Model Training
B.5. Details Experiments in Fig. la and Section
Training details. The loss curves in Fig. la are an exact reproduction of the curves in (Hägele et al., 2024, Fig. 3); they obtained from training 210M Ll Llama-style transformer (Touvron et al., 2023). The base Cearning-rate of cosine is 0.001, and for wsd it is 0.0005.
All of the following applies to the training runs used in the experiments in Section 5: we use exactly the same model architecture as in Hägele et al. (2024), which is a L1 ammasstyle transformer with 12 (24) layers and 12 attention heads for the 124M (210M) model. The dataset used for training is SlimPaji j ama (Soboleva et al., 2023). Specifically, for runs with 50 000 steps (5B tokens), we use the SlimPaj ama-6B subset available on Hugging Face (link below). For the extension runs with 100 and 200 000 steps (approximately 10B and 20B tokens), randomly sample 550M documents (roughly 5% of full corpus) from the full SlimPaj jaaa-aaaaa--ol aaaajammm62777 627B arrive corpus of 30B tokens.
We train for 50000 steps, where the first 300 steps are reserved for linear warmup. We use Adamä (Loshchilov & Hutter, 2019) with weight decay of 0.1. For all further details we refer to Hägele et al. (2024, App. A. Note that all training curves show the validation loss computed over a subset of 32 batches, while the final validation loss is computed over approx. 6 000 batches; hence, the final value of the loss curve might not be identical to the final loss computed over the full validation set. One single rum over 50 000 steps takes roughly 2 hours on two Nvidia H100 GPUs.
The training runs be reproduced with the following repositories:
Training code from Hägele et al. (2024): Dataset: https::/ httg:///ggggggggeeeeeee co/ /ttaaaaatts////ooo///iii/////e//
Fitting procedure. We execute training runs on grid of base learning-rates E {0.0005, 0.001, 0.002, 0.003) and cooldown fractions с E {0.1, 0.2, 0.4, 0.6, 0.8, 0.9, 0.994J. Note that the largest cooldown fraction is slightly smaller than 1 as the remaining 0. 0.6% percent of steps constitute warmup. The final validation set loss (after 50k steps) for all runs is displayed in Fig. 12b (every dot marks one single training run).
We then fit, for each cooldown fraction separately, a function of the form he(?) A- + Bey + Co. where Ae, Bey Ce are fittable parameters. The resulting function is plotted as solid line in Fig. 12b. The functional form of hc(r) is inspired by the bound (8).
We then approximate the optimal base learning-rate (c) by computing the minimizer of he(r). The result of this step is plotted in red in Fig. 12a.
Assessing loss differences through scaling laws. In this section, we estimate with scaling laws how much more parameters or training data/steps would be needed to make up loss difference of 0.01 (see end of Section 5.1 for context). The Chinchilla law (Hoffmann et al., 2022) states that the loss L(N, D) for model with parameters N after training for D tokens can be estimated with
A No
B + D8
L(N,)) =
(12)
where E, A, B, B,a,8 3 are usually fitted from data. More concretely, assume we have trained a model of size N₂ for D1 tokens. To arrive at an improvement of 6 with new combination of the number of parameters and tokens to (N₂, D2), we obtain
8 = - L(N2, D2) = A
+ B
No
No
D
D2
Consequently, we can consider two cases:
Case 1: Fix N₂ N2. That is, we fix parameter size and look for the number of tokens by which we need to extend the training to improve the loss by 8. Solving the above equation then gives
8
D2
D
B
16

---

Learning Rate Scheduling for Large Model Training
Case 2: Fix D1 D2- This is the case where estimate the size that would achieve the desired loss improvement for the same training data. Similarly, this results in
8
N2
NT
In the settings of our experiments we have N1 E {124M, 210M] and Dj E 20.48B) Plugging in the constants by Besiroglu et al. (2024)) and using 0.01, yields'
Case I: Fix N1 N2. In this case, the scaling law results in D2 E (10.88B, 22.16B) for D₁ {10.24B, 20.48B), respectively. This means that we would need to train the models for 640M or 1.68B more tokens to match the adapted schedule.
Case 2: Fix D1 D2. In this case, obtain N₂ {129.0M, 220.1M) for N1 (124M, 210M). In other words, increasing the number of parameters by 5M or 10M would approximately result in the same loss after fixing the amount of tokens.
For both cases, the estimates from the scaling law match our general intuition and would require either noticeably training longer by 6 - 6% growing the model by 4 - 5%, in line with the argument at the end of Section 5.1. Also note that the (relative) additional cost implied by the Chinchilla law to obtain 0.01 loss improvement grows with the (extended) training length D1.
B.6. Miscellaneous Plots
0.10
0.20
wsd
T//7
1I+ 1/T-1
TT
0.08
Bound 0g (for M)
ou/VT)
0.15
0.06
0.10
0.04
0.05
0.02
0.00
0.00
2000
4000
6000
8000
0
1000
2000
3000
4000
Iteration
Iteration
Figure 20. (Leftt) The benefit of cooldown is reflected in the absence of logarithmic terms. Dark grey marks the bound of the constant schedule. (Right) Plotting the individual terms of the bound Ω. Ti/y TT2 with y > for the wsd schedule. The sudden drop the bound comes from the term 7/2.
'Batch size 50, two accumulation steps, two GPUs, sequence length 512, 100/200k steps. 'A 482.01, B 2085.43, E' 1.8172, 0.3478, 3 0.3658.
"Note that the Chinchilla scaling laws were obtained a different setup. In particular, we do not have access to the exact same dataset and tokenizer. which makes the scaling law not directly transferrable. However, our experiments are comparable in the general dataset composition (wehcrawl data extended with other sources) and training task (decoder-only language models). Moreaver, with the difference in vocabulary size (32k 50k),
scale the loss with the rough approximation of In(32 10") In(50 10*) 0.959 align the cross-entropy losses. This does not substantially change the results of this analysis.
17

---

Learning Rate Scheduling for Large Model Training
3.5
3.02
Cooldown Fraction 0.1,0.2,0..,00 0.9, 0.994
0.003
3.4
7
0.002
Final validation loss
3.00
0.001
Validation Loss
3.3
3.2
2.98
2.96
3.0
2.9
2.94
10000
20000
30000
40000
50000
0.2
0.4
0.6
0.8
1.0
Iteration
Cooldown fraction
Figure (Left) Analogous to Fig. 5 (right) with training curves. remove cooldown fraction 0.6 as its curve shows a spike and recovers only late. (Right) Analogous of Fig. (right) with real training data that shows parabola shape for fixed learning-rates.
1.0
(adapted)
(same LR)
[50k, 100k, 200k]
T [50k, 100k; 200k)
1.25
Decroasing factor p
0.8
1.00
Schodule fi
0.75
0.6
0.50
0.4
0.25
0.00
50
100
150
200
T/T
Iteration (unit 10a)
Figure 22. (Left) Decreasing factor computed with Corollary 3.3 for extended schedule up T (where Ti is length of the short run). See Section 5.1, (B2) for details. (Right): Extended schedule for the training runs in Fig. 10; note that this schedule is multiplied by 7 0.001.
18

---

Learning Rate Scheduling for Large Model Training
0.020
0.10
linear-docay
same base Ir+cooldown
0.015
0.08
adapted cooldown
0.06
0.010
0.04
0.005
0.02
0.000
5000
10000
15000
5000 Iteration
10000
15000
Iteration
Figure 23. (Left) Transferring the 1/sqrt schedule with linear cooldown from horizon Ti 4000 to Ta E [L.57, ).: (Right)
Adapting the cooldown length has only small benefits. Dashed horizontal lines mark bound for lincar-decay schedule with tuned y", See Section 5.1, (B1) for details.
3.3
wad, T [50A, 100%]
3.2
usd, T' E [50k, 100K usd (adapted schedule)
wed (adapted schedule) and (adapted cooldown)
3.1
wsd (adapted cooldown)
3.2
Validation loss
Validation loss
3.0
3.1
2.9
3.0
2.87
2.8
2.76
2.9
2.86
2.74
2.7
2.85
2.8
25
75
100
25
50
75
100
Iteration (unit 10y)
Iteration (unit 103)
Figure 24, Experiment with adapted cooldown length for 124M model (left) and 210M (right). See Section 5.1, (B1) for details.
19

---

Learning Rate Scheduling for Large Model Training
C. Additional Experiments
C.1. Additional Experiments on Imagenet
In order to corroborate our findings on additional data domains and architectures, we conduct additional experiments on training ResNet50 (He al., 2016) on Imagenet. We train all models with SGD with heavy-ball momentum.
Training is done using the timm library (Wightman, 2019). All runs using weight decay of 0.0001, momentum 0.9, batch size 4 x 256, and standard data augmentation techniques."
Figs. 25 and 26 confirm our previous findings on this additional training task. Note that the Imagenet training is in different regime as we are not training with single pass. Hence, we expect the validation set metrics to be confounded by generalization effects beyond training loss only.
0.4
cosine C JUZZR 0L04, 0LDGG QUEG
3.5
1.15
Salopiiniig bound
0.3
7. [
Validation loss
3.0
Validation lass
1.10
2.5
1.05
0.2
2.0
0.1
1.5
1.00-
1.0
0.95
200
400
600
800
1000
10
50
100
0.05
0.1
0.2
0.1
Iteration
Epoch
Bage learning
(a) Theoretical bound
(b) Validation loss Imagenet
(c) Leaming-rate sweep Imagenet
Figure 25. The theoretical bound (left) for range of base learning rates 7 (and setting D = G 1) qualitatively matches the empirical (validation) loss curves for training ResNet50 on Imagenet with SGD (middle). (We display a running average over five epoch in thick smoothen the plot, and the original data in thin.) (Right) We again find for the optimal base learning that it holds y (cosine) = 2-y" (wsd), as is predicted by the theory.
0.8
0.8
6.27
3,5
1122
0.4)
642
0.6
Insieedeaa, 0.40
Inseedeaa, C.01
0.7
Learsing cile
loes
3.0
Validution sap
0.4
0.6
2.5
0.2
2.0
0.5
1.20
BV00
0.0
Epoch
Epoch
Epoch
(a) Learning
(b) Train loss (epoch-wise averages)
Top-1 Validation accuracy
Figure 26. Training ResNet50 Imagenet with SGD: Plotting the three best base learning rates for both cosine and usd schodule, well as lincar-decay schedule with learning rate transfer following Section 5.2. Sudden drop of training loss (middle) and increase validation set accuracy (right) is clearly visible for wed schedule. Lincar-decay with zero-shot transfer of (multiplier exp(0.7) xs 2) improves over wad, which confirms the findings of Section 5.2. For validation set metrics, we display a running average over five epoch in thick smoothen the plot, and the original data in thin. Note that the final train loss of wsd appears slightly higher display the epoch-wise average of mini-hatch losses; due to the steep descent of the loss at the end of training this slightly distorts the plot.
C.2. Additional Experiments on OpenWebText2
We also train language models on a different dataset, namely OpenWebText2. Across three different model sizes we find (again) that the performance of wad (with 20% cooldown) matches the one of cosine when using base learning rate
'In timm we set configuration hflip= 0.5, vflip= 0, crop-pct= 1.0.
20

---

Learning Rate Scheduling for Large Model Training
half big (see Fig. 27). For wad, the sudden drop in the loss is clearly visible across all model sizes and training lengths. Training details identical to the OpenWebText2 experiments of Hägele et al. (2024).
(.M)
[JOUR, 25.04
[LD 26
50.06]
3.8
3.6
3.6
Validation Loss
P.M., 1R56, 1756
Validation Loss
[taue, SUM..1T..0 17.56, 15.00
Validation Lose
115
3.1
3.1
3.1
3.2
3.2
3.2
3.0
3.0
3.0
10000
15000
5000
2000
25000
20000
30000
Iteration
Iveration
Iseration
(a) 60M
(b) 93M
(c) 166M
Figure 27. Training three different model sizes on OpenWebText2. We observe the same characteristic drop of the loss for wad, as well as matching performance of wad and cosine. In each the base learning rate of cosine twice large as for wad.
C.3. Computing the Bound-minimizing Schedule
In this section, we showcase how the schedule that minimizes the bound Ω in Theorem 3.1 can be computed. However, we stress that this has more of a purpose of illustration than practical applicability: in general, the gradient norms will depend on the schedule, and thus create an interdependence of the schedule and gradient norms (see also Defazio et al. (2023)).
However, we will show that with constant gradient norm bounds (G, for all E N) - computing the schedule that minimizes Q2 converges to the linear-decay schedule. Consider 22 (M.-)), the right-hand side in Theorem E.2, function of the schedule Thr (71,... ........)p)))) We are interested in computing
arg min (r 91.7
subject to m..... > 0.
As we can not compute the above analytically, resort to using projected gradient descent (with momentum) in order to approximate the solution. We implement (7)rr) in Pytorch which allows us to use automatic differentiation. As starting point, we use constant schedule п = 1 for all = I,...,T" T. Fig. 28 (left) shows the value of the bound Ur over the optimization trajectory. We observe convergence of the optimized schedule to a linear-decay schedule (Fig. 28, right). Note that the shape of the schedule as well as its scale are optimized simultaneously; we plot the optimal bound linear-decay schedule and see that its value matches the bound at convergence. Though we did not observe any instabilities in the optimization procedure, this point we not verify whether the final point is global minimum.
100
1.0
PEPit decay
Bound for leneo-decay
Rescaled Schedule Wo mn
0.8
Bound Or
0.6
0.4
0.2
10-1
0.0
200
400
600
20
40
60
100
Optimization step
Figure Optimizing bound Qr with respect to schedule Th. Tr. (Left) Convergence the value the bound, matching the bound of a linear-decay schedule. (Right) Optimization trajectory from constant schedule (yellow) to linear-decay (purple).
21

---

Learning Rate Scheduling for Large Model Training
D. Auxiliary Lemmas
Lemma D.1 (Lemma 5 from Defazio et al. (2023)), Let (qt) be any sequence, and (we) be a positive sequence. Then, for any T N it holds
T
T-1
T
Wk
Σ
we(st - qm))
9T
We
+
Σ1=1 Wt t=1
1=1
21=k+1 Wt
-izh ut
t=k
Lemma D.2. Let T E N and! I < T. It holds
T+1-1
T+1-1
Σ s= (T +2 (T +1 1),
Σ
8² = (2T 3 21) + 1)(T + 1 1).
Proof. We refer to WolframAlpha: [link to first result]. [link to second result].
Lemma D.3. Let t E N. It holds
In(t) s In(t + +1] V S("
8+1 _ _ds S Σ H. S 1 + In())
E. Missing Proofs
The following lemma will be the basic inequality for subsequently proving Theorem E.2; it is standard result in the online learning and convex optimization literature (Zinkevich, 2003). Lemma E.1.
Let S k < T and let E IRd be measurable with respect to æg. It holds
1
Σ Blf(z) f(u)] - B u/12 +
(13)
Proof. From the update rule (2) and property (4) we obtain
|
u||2 || u1|2 -
- l/el
S || u1|2 2n 2n [f 8t)) f(u, + nh 19
Apply conditional expectation (conditioned on f) to obtain
El| ul|2 < || u||2 2mp [f(x)) s(w)] Ella
Apply total expectation (with respect to 1,... 7) and rearrange obtain
nh E 91
T[[[f(x)) - f(w)] s = w|12 BBllae+ w/12
2
Sum from - .,..., T to obtain the final result:
1 u1|2 +
Eff()) f(u)) s B
Σ
n
Elo
2
t=k
t=k
Theorem E.2 (Thm. 10 from Defazio et al. (2023)). Let the iterates (xt) be given by (2) with np > ffr 1,..., T. Let
J. E Rdd and D = |||| x. Then, it holds
D2
ΣT
2
9
- -
<
1
2
2
+=
T-1
T Σt=r
Σ E
2
Am
+==+1
tmk
22

---

Learning Rate Scheduling for Large Model Training
Remark E.3.
(i) Note that the result Theorem E.2 is an anytime result, in the sense that we can evaluate the right- hand side at any T without the knowledge of '7t for
(ii) A technical artifact of Theorem F.2 is that ay does not depend on yr by definition, however 1/r appears in the bound on the right-hand side. This is standard in the analysis of subgradient methods: in the proof, we bound - - - x. v 0. If one carries through this term to the end, then we obtain multiple terms the bound that depend on 'f-
Theorem 3.1 follows from applying Theorem E.2 with ij 7/p. We finally prove Theorem E.2.
Proof. First, apply Lemma E.1 with - and k 1 to obtain
Σ -Bif(E)- D2
1 +
Σ
E
9:
(14)
2
1=1
Define 9: = [lf();) - f(x))). Dividing by Et= le gives
D
ΣT= -t=l
19
Σ nhd S
(15)
E 1J t=1
2 2 7h
2 ΣE= 71
In order to apply Lemma D.I with W/ Ne need to bound the term
Σ ma 2k) 9x)= = (x))( f(xx))
(16)
Thus, apply Lemma E.1 with - 2y to obtain
Σ
nelac
9N]
n 11 Σ
EtEk
2 " T1 t=k
t=k
Now, combine Lemma D.1 with (15) and (16) get
2
Σ T ggl
B[f(xT) - f(xx)) = qr v
②Σ
2
t=1
1hb
1= 11
T 1
T
+
Σ
E
2
k=1
20=k+1 Th
Σt Tt t=k
23

---

Learning Rate Scheduling for Large Model Training
F. Mirror Descent Analysis
In this section, we extend the bound from Theorem E.2 to the stochastic mirror-descent method.
Notation. In this section only, we denote with || || an arbitrary norm (in contrast to the rest of the paper, where it denotes the standard Euclidean norm), and let || ||- denote its dual norm, defined by || == su (1t
Let wh Rd + R be continuously differentiable function that is ju-strongly convex with respect to || Define the Bregman divergence
(a(x,y) =- -xx) (by) -
y, Vu(y)).
follows (x, () 2 $ || y/1 from strong convexity of v. Further, we will need the following three- three-point-ie identiy (Beck & Teboulle, 2003, Lem. 4.1): for any z, y, E I2 it holds
B2(2, + (xr v) Bz v) = (Vi(()) - 2 xt).
(17)
Now, the iterates of (stochastic) mirror descent are given by: for n and ga E af(x,, st), compute
T++1 = arg min nu (9t, y σz) + (y,
(18)
yER"
We will now prove mirror descent version of Theorem E.2 (in fact, Theorem E.2 a special case with (2, y) 3|1 v) To do so, we first follow standard steps in miscer-descent analysis (Beck & Teboulle, 2003) obtain the basic this prove last-iterate bound in
inequality in Lemma F.1. In contrast the classical mirror-descent analysis, Theorem F.2.
Lemma F.1. Let the iterates (x+) be generated by (18). Let S T and let E Rdd be measurable with respect to rk- holds
Σ mB[[((1)) - [(4)] S ((B +
1
Σ
m E
gtl.
(19)
2/
Proof. For fixed y. we have V Bz VΦ(x) Ve(v). Thus, optimality conditions of (18) are
nıt + V+((++1) ((b(a))))V((())))..).)))))))))))
(20)
Then, we have
(F 84) - f(u, [()] nh 11 (1,
(u T1+1, (xb(Z)) (e((1++1 7h9tt) + (u Z1+1, Ve((441 Vab(Iz)) +
:=#1
(i) T141,
From (20), we have $1 0. From (17), we have $2 (u, (v x2+1) (2)41, From the (generalized) Cauchy-Schwar inequality combined with Young's inequality, have
y " | 12 +
$3 < 2
11
2μ
Using that (1) < - 12112 2=112, we obtain
22
m[f(xr,))) 8t) - s(), S (u, Bolu, x4441) +
|1 11. 2/
24

---

Learning Rate Scheduling for Large Model Training
Taking conditional expectation, have (a 86) f(u, 8e)] f(x) f(u). Finally, rearrange, take total expectation and sum from & : T. Using that B& (T++) 2 obtain
T Σ mL))(()) - f(u)] < (u,
T gell.
1
Σ
m
2μ t=
t=
Now, repeating the proof of Theorem E.2, but applying Lemma F.1 instead of Lemma E.I, we obtain the following bound.
Theorem F.2. Let the iterates (xt) be given by stochastic mirror descent (18) with nh > f ,....,T. T. Let x, R Then, it holds
E
El=1
mi
19
E ( -
Σ 21=1
2/1
Σi=1 7
T
T]k
re Blg 13.
2/4
Ltzk 7/t t=k
k=1
11=
25

---

Learning Rate Scheduling for Large Model Training
G. Analysis of wsd schedule
We first state Theorem 3.4 in its full version.
Theorem G.I. Let 1 S To < T. Assume that т = 1 for < To and = 1 Lo
T+I-1 for T2 2212 To- Further, assume
that (A3) holds with G. = G for some > 0 for all t E N. Then, for y = " from Corollary 3.3, we have
[(f( - 5
DG,
T42fa + HT+T--2 HT Hr-1-+1 3(T+Ta)
(T-1))(661)) (TT-
+ (T-To)² + Ho- T-
T + To
Note that for large 7, we have HT O(In(T - To - 1)) = o(T - To + 1). Thus, the last two terms can be summarized with o(1) as + cc.
In this section, we give further interpretation the bound Theorem G.1 in comparison to the constant and linear-decay schedules.
In Section 3.2, we derived that for large T and To 3T, the bound for wsd is approximately
DG
(| [x+))
1+28 B 31+3)
Ha-+sy-2 (y1-3)1+1]
VT
1 3
3(++))
To obtain concrete numbers, plugging in B 0.8 (that is, 20% cooldown) for wsd, and obtain
DG
([ - s(xx)) 2
2.2 - Ho.27+1)
VT
For example, if = 105, then 2.2(H1.TT22 Ho.2741) = 4.39. In comparison, we have:
constant schedule: the bound is DG vT
VI+ HT-
linear-decay schedule: the same bound from Corollary 3.3 results in (2 Hy 1-23 DG (Defazio al., 2023, Thm. T+1 VT
show the tighter bound DG (Defazio et al., vT 2023, Cor. 2), which is actually worst-case optimal for convex, Lipschitz problems (Zamani & Glineur, 2023).
13). However, with different (but less general) proof technique
Again for T 105, we have Hr-1 12.09 and (2 + HT-1-2/3) 2.0001. In conclusion, for this specific T the constant of the bound is roughly twice for wad compared to linear-decay, and 1/3 compare to constant schedule. Finally, we prove Theorem G.1.
7+1
Proof of Theorem G.I. From Corollary 3.3, we have
([f(xr) fxx)) s 2 Timm (m (I.
Thus, the rest of the proof will compute an upper bound of the right-hand side. First, for S T, we compute
T Σn
T n T41-To E( +1 (T++1-
T+1- Σ s
(T+2 2 () + 1
12To
t) T+1-To
2(T + To)
t=)
Here we made the change of variable T + 1 s and used Lemma D.2. Similarly, we get
T Σ Th
To-1
Σ 11t To -1 Σ = To - 11 (T + (T)
T
T
1 << To:
Σ Th +
2(T 31+1-T6)
1=1
1=1
t=
t=
[ + 2 20}.
26

---

Learning Rate Scheduling for Large Model Training
Note that this expression still correct if would plug in To- Next, compute the of squares. We start again with
T
7+1-1
1
(2T
12 To:
11
2
8
(T
+
1
To)
6(T -1 Tb)
t=1
s=1
And similarly
T Σ
To
T Σ
111 = To 1+ (2T + 2- To)
<< To:
2
+
6(T+1-7h)
t=1
t=]
(=1
(2T +3 2T) + 2 To)
To - 1+
6(T + - To)
27 +5 2To
To -
6
6(T + 1 - To)
1
2T + 4To +5 61 +
(T TT)
Here, used that
(T -3 2123:-2200----100--00---------1- 2 To) -7 (2T 2To) 2(T To) -27 5 (T +1.
Again, the expression we obtain for < To is correct we would plug 1 = To- Now we can try to compute the bound. We start with the easy ones: as G. G > for all E N, we obtain
1
D2
Ti D,T)
D2
2 1tt T To
T
G [2 + 4To + 5 6] + 677+1-2a
G =
(21)
D 1/t t=1
(T)
G
(T
21' + 47% - 1 + T+1-To
6(T + To)
3(T +T))
The last inequality is due to To T and thus 1 T+I-To
Next, for k I....,T T 1. we need to compute
GR
Lt Te
Et Tht t=k
Again, as G. = G>> 0 for all / E N, we omit G for now, and start with the case k 2 To:
T+1-k
(2T 20
11k
T+1-2
6(1+1-16 To)
2(2T + - 2k)
Σ Et=k+1 Ntt
Etzk 1/t
(T+2- -1-1))1++1-0
(T
3(T -1+-1-Tp) (T) k)
4[+1-T6)
Now, k < To:
1 42 +5 6k [7+1 To)
2
m1
>T
TT 2 + To 2h
2t=k+1 1 T]t
Lt 17t t=k
2[2T+4760 +5 + (1+1-1 15)
3T To - T + 2 2 To - 2k
27

---

Learning Rate Scheduling for Large Model Training
Now, compute
T
11k
nf
k++ 1]
Lt=k 7/t t=
2
To-1
2T + 4To + 5 6k + (7+1-Ta)
T
(2T + 3 - 2h)
=: +*))
3
[T T To 2kj TT + To
2
TT + - (o)(T k)
To
=
Then, holds [link to proof]
3
2T 2To + T- 33H -T
22
(T (T+1 - [T) -
T T +1
k=
To simplify this term a bit, we estimate
2 02 S
4 +
2
2HT TTo-
(T To) 2
- To
For Q1 we can bound the nominator by
2T 4To +
6k +
1+1-2 3(T + 2k) (T To 1)
T+1-B
S 3(T +2 + To -2k) - -(T - To),
where for the second term bound
T+1-T < 1 due to To < T. It holds [link to proof]
To
To - -1 (T To 2) + To)
T + To 2T + -2k]
Therefore, defining Ω3 = (T- +2](
we get
Tύ-1
3
Ω S
=3(H+47622 Hr-16+1) S2z
T To
2k
where we used T > To. Altogether, we have
2 22 + 3 (T - 2
2
+ 2Hr-To-1
(x) = (Ω4 - 92) < 2( -Hr-.+1)
+
T - To +
Multiplying this term with d² : we can plug the result in (21) to finally derive the bound
2 Ti Ti D, TyTT GLTTT) S
2DG
T427 3(T+1))
+ HT HT-To+1
(T- 77)
+ (T-
Hr- T-Tot1
T+T
3/7-76+211
28


Reviews:
[{"review_id": "beuEzBLqZI", "reviewer": "Reviewer_rpAe", "summary": "This work demonstrates that several empirical observations align with the last-iterate sub-optimality gap in convex optimization. Furthermore, the authors show that adjusting the learning rate in continual learning—an approach that theoretically improves this sub-optimality gap—also enhances real-world training performance. Lastly, the study establishes a convergence guarantee for the WSD schedule.\n\n\n## update after rebuttal\n\nThank you for your detailed response.\n\nIn light of the additional experiments, I will raise my score.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "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.\n\n* *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.).\nBesides the dataset choice (which we address in the rebuttal), are there any other reasons why the reviewer considers our empirical evaluation to be limited?\n* 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.\nThese 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.\n\n* *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).\n\n* 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. \n\n* 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:\n  - the results by Liu and Zhou do not hold for arbitrary schedules\n  - 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.\nExtending our analysis to the smooth case is definitely interesting and we are actually working actively in this direction.", "claims_and_evidence": "The claims linking the theoretical bounds to practical performance are supported by the experiments in the paper. However, the experiments are limited to a single synthetic problem and one real dataset. Consequently, the assertion that the theoretical bounds accurately reflect real-world behavior—allowing for extrapolation from the sub-optimality gap—is not strongly justified.", "methods_and_evaluation_criteria": "The methods and evaluation criteria are well-founded. However, experiments on additional datasets would further strengthen the analysis.", "theoretical_claims": "Given the last-iterate guarantee from previous work, which accommodates general step-size sequences, the convergence result is reasonable and unsurprising. It is well established that a linearly decaying step-size sequence leads to last-iterate convergence guarantees. Therefore, when the cooldown period is of the same order as the training duration, the observed behavior aligns with the theoretical expectations.", "experimental_designs_or_analyses": "In general, the experimental design and analysis appear valid. The behavior of WSD across the experiments aligns with findings from previous work.", "relation_to_broader_scientific_literature": "Understanding the relationship between theory and practice is an important topic for the machine learning community.", "essential_references_not_discussed": "The work of Liu and Zhou establishes a last-iterate convergence guarantee that applies to both convex Lipschitz and convex smooth problems, including mirror descent. These results can also be leveraged to obtain guarantees in the convex smooth case. Notably, (A) these guarantees rely on Lipschitz/variance bounds rather than the $G_t$ assumption used in this paper, and (B) the constants in this work may be slightly worse.\n\nZ. Liu and Z. Zhou. Revisiting the last-iterate convergence of stochastic gradient methods. In The Twelfth International Conference on Learning Representations, 2024.", "other_strengths_and_weaknesses": "**Strengths**\n- The work presents multiple observations that align well between theory and experiments.\n\n**Weaknesses**\n- The theoretical results are relatively straightforward. While the reviewer is not aware of prior results specifically for the WSD schedule, deriving such guarantees using existing general last-iterate convergence results does not require significant innovation.  \n- Additional experiments would substantially strengthen this work.", "other_comments_or_suggestions": "No additional comments.", "questions_for_authors": "- Given the last-iterate guarantee for convex smooth objectives established by Liu and Zhou (referenced above), should there be any notable differences in the predictions and alignment of this work, which assumes convex Lipschitz objectives? This question assumes that the number of steps is sufficiently large for the $1/T$ smooth term to be negligible.  \n- Can the authors provide justifications for Assumption A3 beyond a general Lipschitz assumption? While the reviewer is familiar with optimization theory, they are not aware of other instances where this assumption has been used, especially since it depends on the algorithm, as the authors mentioned.  \n\nOverall, while this work presents interesting observations, its limited empirical evaluation and lack of significant theoretical novelty make it a borderline case for acceptance.", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "FZoXOZpo5a", "reviewer": "Reviewer_hqnc", "summary": "This paper establishes a novel connection between empirical learning-rate schedules (e.g., cosine, WSD) used in LLM training and theoretical bounds on the loss at the final iterate of SGD in a non-smooth stochastic convex setting. Through empirical studies on Llama-style transformers, the paper demonstrates that the theoretical bound effectively predicts the actual loss curve based on the learning rate schedule. Building on this observation, the authors show how insights from these theoretical bounds can guide learning-rate tuning for continued training and enable learning-rate transfer across different schedules in practice.\n\n## update after rebuttal\n\nAfter carefully considering the authors’ responses in their rebuttal, my assessment remains unchanged, and I continue to recommend acceptance. The authors have adequately addressed my initial concerns, and the clarification provided further strengthens the contributions of the paper.", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their positive feedback, especially for the assessment that our experimental results are strongly supporting the claims/contributions.\n \n* 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.\n* 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.\n* 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.", "claims_and_evidence": "The main claim of the paper is that the theoretical performance bound on the final iterate of SGD in non-smooth stochastic convex optimization aligns well with the loss curve in LLM training. This claim is strongly supported by a series of experiments on language model training, as well as theoretical simulations.", "methods_and_evaluation_criteria": "The proposed application of the theoretical bound—schedule construction for continued training and learning-rate transfer across schedules—seems convincing and could be a valuable tool in practice.", "theoretical_claims": "The main theoretical framework is adopted from the existing work by Defazio et al. (2023). Building on these theoretical results, this paper derives the theoretical bound for different learning rate scheduling scenarios. These theoretical results appear to be correct.", "experimental_designs_or_analyses": "Overall, the experiments are well-designed, and the results look convincing.", "relation_to_broader_scientific_literature": "N/A", "essential_references_not_discussed": "The paper discusses all the essential references related to the main contribution.", "other_strengths_and_weaknesses": "Overall, this paper is well-written and makes a solid contribution by connecting non-smooth convex theory with language model training. This bridges a critical gap between optimization theory and deep learning practice, which I find to be very interesting.\n\nThe main strength lies in demonstrating that the theoretical performance bound has strong predictive power in practice.\n\nThere are a few minor weaknesses. First, while the non-smooth convex theory is derived from SGD, all the experiments in this paper use AdamW. This limitation is well acknowledged in the paper. The paper would be even stronger if it included experiments using SGD and compared the loss curves with the theoretical bound.", "other_comments_or_suggestions": "Here are some additional experiments I would suggest:\n1. Experiments using SGD on vision tasks could provide insights into the broader applicability of the theoretical bound.\n2. Testing a more diverse range of learning rate schedules beyond the popular cosine and WSD, and comparing them with the theoretical bound, would be interesting. For example, in Tissue et al. (2024), the authors tested various learning rate schemes—see Figures 3 and 4.\n\nHowever, the current paper remains valuable even without these additional experiments.\n\nTissue et al. (2024) - Scaling Law with Learning Rate https://arxiv.org/abs/2408.11029", "questions_for_authors": "1. Does the alignment between non-smooth convex optimization theory and learning rate scheduling apply only to “large” model training?\n2. Does the alignment hold only for AdamW, or can similar observations be made for other modern optimizers (e.g., Shampoo, Schedule-Free, Muon, MARS, etc.)?", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "y8Ws1tf9T5", "reviewer": "Reviewer_CJZu", "summary": "This paper studies learning-rate schedules in large model training, by bridging a new theoretical convergence analysis to empirical observations.\n\nIn particular, the first contribution (observation) is that for two popular schedules (cosine and wsd), the empirical loss curves of large model training, which is a non-convex optimization problem, has a similar shape with the theoretical bound predicted by the parameters in the convex case. \n\nThe second contribution is a new theoretical convergence analysis of wsd in the convex case, which roughly shows an $\\log T$ improvement over the best constant learning rate, providing an explanation to the success of wsd.\n\nThe final contribution is an application of the theory, for continue training and learning rate transfer.\n\n## update after rebuttal\n\nI have read other reviews and rebuttals. Though I agree with Reviewer rpAe that the results presented are not significant enough in either theoretical or empirical side, the new link does seem interesting and promising. I will keep my score.", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "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:\n\n* 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.).", "claims_and_evidence": "Yes.", "methods_and_evaluation_criteria": "Yes.", "theoretical_claims": "Yes. The only new theoretical result is Theorem 3.4.", "experimental_designs_or_analyses": "Yes.", "relation_to_broader_scientific_literature": "I think the results of the paper might be interesting to only only empirical ML researchers, but also optimization researchers.", "essential_references_not_discussed": "As far as I know, no.", "other_strengths_and_weaknesses": "The paper is well-written.\n\nThe theoretical result might not be a deep one from a technical perspective. However, it's very important to provide a rigorous explanation for a popular learning rate scheduler.", "other_comments_or_suggestions": "NAN", "questions_for_authors": "Is it possible to perform a rigorous convergence rate analysis for cosine?", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "dI28zOb0GI", "reviewer": "Reviewer_iNQu", "summary": "This paper proposes to find optimal settings of warmup-stable-decay (wsd) learning rate (LR) scheduler for training large language models (LLMs) from a theoretical inspiration. The wsd scheduler is a piecewise function by training iterations, the first part is a constant function (named based LR) and the second part is a linear function that decays the LR with a ratio proportional to iteration. Thus, the authors aim to investigate the configuration of two controllable hyper-parameters, i.e., base LR and a fraction of the decay period to the entire training iterations, by minimizing the convergence rate of the wsd scheduler.\n\nFirst, the convergence rate is derived from minimizing the last loss with gradient descent and the wsd scheduler. The optimal base LR and optimal convergence rate are derived by solving a convex optimization problem of convergence rate minimization with respect to base LR. Further, an iteration-free rate is derived with respect to the fraction. The authors utilize several simulation results to demonstrate the relevance of the proposed optimal rate to base LR and the fraction. The results show that the optimal base LR decreases with training iterations but increases with the fraction's value.\n\nThen, the authors try to apply the observations of simulation results in ideal convex optimization to real-world LLM training. Two simple strategies, i.e., decrease base LR and increase fraction rate, are proposed. Simulation results demonstrate the effectiveness of the synthetic convex optimization problem. Experimental results on 0.1B and 0.2B Llama-style transformers show the improvement of the proposed method. Moreover, the authors propose to fit a convergence rate function from several running results with the wsd scheduler, where the function is defined from the inspiration of the theoretical optimal rate, to demonstrate the alignment of the proposed method to the real-world LLM training.\n\n## update after rebuttal\nI keep my rating due to the following main concerns:\nThis paper is within the scope of the study of scaling law. Thus, the only way to demonstrate its effectiveness is by experimental results on heterogeneous LLMs, especially for large LLMs. However, the selected LLMs are only 124/210M Llama-style transformers, and there are no additional experiments on large LLMs in rebuttal.", "questions": "", "limitations": "", "rating": 2, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "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. \n\t\n* *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.\n* *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.\n* *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.)\n* *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?\n* *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? \n* *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.).\n* *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.\n\nOn Questions:\n1) In Equation (9), $\\gamma$ is arbitrary, and not necessarily equal to $\\gamma^\\star$.\n2) 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$.\n3) 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\n4) 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.", "claims_and_evidence": "Yes.", "methods_and_evaluation_criteria": "Not enough. \nFirst, training loss is not an important criterion for LLMs. Thus, the optimality of loss in training LLMs is not a convincing motivation for the so-called requirement of optimal learning rate.\nSecond, the selected LLMs are too small. The selected LLMs are less than 1B, which is not sufficient to demonstrate the effectiveness of the proposed method for LLMs.", "theoretical_claims": "Yes. The proofs are correct.", "experimental_designs_or_analyses": "The experimental designs are not enough to demonstrate the so-called effectiveness of LLMs, since the selected LLMs are extremely small, less than 1B.", "relation_to_broader_scientific_literature": "This paper aims to improve an existing learning rate scheduling method from theoretical inspirations of ideal convex optimization. However, the gap between strict assumptions and real-world neural network training is not well-addressed.", "essential_references_not_discussed": "No.", "other_strengths_and_weaknesses": "Strength:\n1. This paper conducts many simulations and experiments on two small LLMs.\n2. This paper proposes a simulation-based learning rate (LR) scheduling method. By sampling real-world running results from different LR configurations, a function inspired by theoretical optimal LR is proposed to fit the real-world running results. Then, the optimizer to the function yields the \"optimal\" LR configuration.\n\nWeakness: \nThe significance of this paper is limited. This paper tries to demonstrate that the empirical learning rate (LR) scheduling methods in training LLM align with its theoretical optimal solutions, which include three main components, i.e., theoretical proofs, LLM training, and alignment. This paper's soundness is poor in all three parts. \n1. The theoretical proofs are based on a strict convex optimization scenario, which is far from the high dimensional non-convex optimization scenario in neural network (NN) training. The assumptions are too strict that the loss function is extremely smooth, in which the function's gradient can be bounded by constants (A3 in the paper). Thus, it may be the reason for the results in Fig 4, where a larger decay rate leads to under-training (large loss).\n2. The selected LLMs are too small to demonstrate the effectiveness of the proposed method for LLMs. The simulation results are trivial under some specific loss functions and settings. First, most simulation results are based on extreme smooth loss ($G=1$) and a close initialization to optimum ($D=1$). Moreover, most observations (such as the specific settings of the decay fraction rate from Fig. 10, 11) are based on the 124M and 210M Llama-style transformers, which may not hold on large LLMs. \n3. The alignment is problematic. The connection of the proposed theorem to neural network training is achieved by approximating a linear function of the hyper-parameters in the theorem, which is not convincing since there are no standard solutions for the function approximation for all NN models. Further, most simulations and experiments are designed to evaluate the proposed method for optimal base LR the wsd LR decay method. However, the most crucial part is whether such a theoretical upper bound of loss is close enough, which is not evaluated in the paper. For example, from Fig. 12 (b), most losses are above the approximated function, which violates the upper bound requirement of the function.", "other_comments_or_suggestions": "Typo: \"The wsd schedule.\" paragraph, last sentence.", "questions_for_authors": "1. Does $\\gamma$ in Eq. (9) represent $\\gamma^{\\star}$?\n2. What is the $f$ in all simulation results? \n3. What's the meaning of Fig. 8？ How to distinguish the training before $T_1$ and after $T_1$? The illustrated result looks like the overrides of the former training processes.\n4. In Fig. 17, from the left part, the cosine scheduler converges faster than the wsd scheduler. However, the observation seems to be the opposite of the right part.", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Meta-review: This paper provides an analysis of the WSD schedule in stochastic convex optimization by leveraging some general theorems recently developed for understanding learning rate schedules in this setting. The authors also extract a particular anytime convergence result from prior work, and conduct an empirical investigation suggesting that this (convex) theory appears to have good predictive power for some (non-convex) transformer training tasks. Overall, the empirical results are perhaps of too small a scale to provide confidence in modern practical settings. However, even at this scale, optimization theory is notoriously poor at predicting empirical results, so this study may be of interest to the optimization theory community.
Decision: Accept (poster)