Straight to Zero: Why Linearly Decaying the Learning Rate to Zero Works Best for LLMs
Authors: Shane Bergsma
Nolan Simran Dey
Gurpreet Gosal
Gavia Gray
Daria Soboleva
Joel Hestness
Venue: ICLR 2025 Poster (2025)
Keywords: Learning rate schedules
Large language models (LLMs)
AdamW optimizer
Weight decay
Compute-optimal training
Primary area: foundation or frontier models, including LLMs
TL;DR: We perform a large-scale empirical study to establish that linear decay-to-zero is the optimal learning rate schedule for LLMs across a range of settings; some novel theoretical analysis helps explain why.

Abstract: LLMs are commonly trained with a learning rate (LR) warmup, followed by cosine decay to 10% of the maximum (10x decay). In a large-scale empirical study, we show that under an optimal peak LR, a simple linear decay-to-zero (D2Z) schedule consistently outperforms other schedules when training at compute-optimal dataset sizes. D2Z is superior across a range of model sizes, batch sizes, datasets, and vocabularies. Benefits increase as dataset size increases. Leveraging a novel interpretation of AdamW as an exponential moving average of weight updates, we show how linear D2Z optimally balances the demands of early training (moving away from initial conditions) and late training (averaging over more updates in order to mitigate gradient noise). In experiments, a 610M-parameter model trained for 80 tokens-per-parameter (TPP) using D2Z achieves lower loss than when trained for 200 TPP using 10x decay, corresponding to an astonishing 60% compute savings. Models such as Llama2-7B, trained for 286 TPP with 10x decay, could likely have saved a majority of compute by training with D2Z.

Published as a conference paper at ICLR 2025
STRAIGHT TO ZERO: WHY LINEARLY DECAYING THE LEARNING RATE TO ZERO WORKS BEST FOR LLMs
Shane Bergsma, Nolan Dey, Gurpreet Gosal, Gavia Gray, Daria Soboleva & Joel Hestness Cerebras Systems
(shane.bergsma joel @eerebras.................o..... net
ABSTRACT
LLMs are commonly trained with a learning rate (LR) warmup, followed by co- sine decay to 10% of the maximum (10x decay). In a large-scale empirical study, we show that under an optimal peak LR, simple linear decay-to-zero (D22) schedule consistently outperforms other schedules when training at compute- optimal dataset sizes. D2Z is superior across a range of model sizes, batch sizes, datasets, and vocabularies. Benefits increase dataset size increases. Leverag- ing a novel interpretation AdamW as an exponential moving average of weight updates, show how linear D)27. optimally balances the demands of early train- ing (moving away from initial conditions) and late training (averaging over more updates in order mitigate gradient noise). In experiments, (TPP) using D2Z achieves lower loss than when trained for 200 TPP using 10x decay, corresponding to an astonishing 60% compute savings. Models such as Llama2-7B, trained for 286 TPP with 10x decay, could likely have saved majority of compute by training with D2Z.
100M-parmmeter
model trained for 80 aakens-per-parameet
INTRODUCTION
Learning rate (LR) schedules play important role in training large language models. The origi- nal Transformers paper proposed a brief LR warmup followed by decay proportional to the inverse square root of the step number (Vaswani et al., 2017). This schedule has the advantage of not requir- ing prior specification of the total training steps. However, cooling down to a specific minimum LR acknowledged be "preferable when one knows the training duration in advance" (Zhai al., 2022) as it produces "slightly better results" (Raffel et al., 2020). In this paper, our main focus is finding LR schedules that achieve the minimum loss given a pre- preciiecified number of training tokens.
The "predominant choice" (Hu et al., 2024) in such training the "de-facto standard" (Hägele al., 2024)- - is warmup followed by cosine decay 10% of the maximum LR, an approach used in GPT3 (Brown et al., 2020), Gopher (Rae et al., 2022), Chinchilla (Hoffmann et al., 2022), BLOOM (Scao et al., 2023), Llama (Touvron et al., 2023a), Llama2 (Touvron et al., 2023b), Fal- con (Almazrouei et al., 2023), Pythia (Biderman al., 2023), etc. It is used "following Hoffmann et al." (Muennighoff et al., 2023), and is the default schedule in LLM codebases (Karpathy, 2024).
We present large-scale empirical study determine which schedules work best in which situations, and why. We focus both compute-eficiient and over-trained models. According to Chinchilla scaling laws (Hoffmann et al., 2022), the fewest FLOPs to achieve given loss is ob- tained when models are trained for around 20 tokens-per- parameter (TPP). It is also common train for more than 20 TPP because smaller, over-trained models are cheaper to serve (Touvron et al., 2023a). Our experiments (across var- ious model scales, vocabulary sizes, and dataset sources) reveal a consistent outcome: when all schedules use their optimal peak LR, linear decay-to-zero (D2Z) works best at compute-optimal TPP. Moreover, the relative benefit of D2Z over 10x (in terms of training, validation and down- stream loss) increases with TPP. Compute savings from D27, can be substantial for over-trained models (Figure 1).
3.2
10x
D2Z 80TPP
Smoothed training losz
10x 200TPP
D2Z 200TPP
60% fewer FL.PSs
2.6
J. Loss
2.4
Step:
25000
50000 75000 100000
Figure A 610M model trained for 80 TPP with Linear-D2Z has better train (and validation) loss than when trained for 200 TPP with 10xe

---

Published as a conference paper at ICLR 2025
Linear D22
Cosire-122
Weight upöare coeff. (cx)
1.0w05
Schelul (mn
4.0-03--
1.0-07
1.0e-09-
2.0-03
1.0-11
10ce03-
0-
10000
20000
30000
40000
30000
10000
20000
30000
40000
50000
Stop (f)
Weight update (+)
Figure 2: LR schedules and their update-combination duals: Each LR schedule, Toe (left) and weight decay, X, implies a weighted combination of weight updates, with combination coefficients Cta (right, log-scale) giving the contribution of ith update parameters 6, at step t (11IM scale, coefficients at final step). The more sudden the drop in LR, the less emphasis on valuable later up- dates, perhaps explaining why Step underperforms Cosine and Cosine underperforms Linear decay.
LLMs such Llama2-7B (Touvron et al., 2023b), trained for 286 TPP at 10x decay, could likely have saved most of their training compute by switching to D22..
To explain the success of Linear-D2Z, we build on recent work on the related topic of weight de- cay (Andriushchenko et al., 2023; Wang & Aitchison, 2024). First, decaying to zero helps because compute-efficient training includes long phase which gradient noise is the key factor slowing the loss reduction; with higher noise resulting from higher TPP smaller batches, vanishing LR works best. Second, we show that approaching zero linearly is beneficial via a novel interpretation of AdamW (Loshchilov & Hutter, 2017), the primary optimizer in L.MM pre-training. With AdamW, weights generated at each step are implicitly a weighted combination of weight updates. The com- bination coefficients depend Analyzing this dual of the LR schedule, we observe linear decay to produce a favorable combina- tion of prior weight updates (Figure 2). When LR drops abruptly, e.g., via step-decay or, lesser extent, cosine decay, later updates receive less emphasis. leading to worse model quality. The EMA dual perspective also reveals the implicit such as Sarmmup-Stabl-Deecae (WSD) (Hu et al., 2024; Bi et al., 2024; Hägele et al., 2024).
the learning rate schedule and weight decay settings (Section 3.1).
of recent "schedule-fred approaches
Our empirical study encompasses hundreds of models trained over a grid schedules and peak LRs, with scales ranging from 111M to 1.7B parameters, and datasets up to 137B tokens. We also compare LR schedules over the cross product of weight decay versus peak LR, and peak LR versus batch size. We confirm slight, consistent advantage of linear D2Z over cosine D2Z at 20 TPP. We also demonstrate that linear D2Z improves over continuous schedules such as WSD. Further, when dataset batch size changes, optimal peak LRs are much more stable when using D2Z compared to using lesser decay. This latter finding exposes LR decay ratio as an important confounder in prior work studying optimal hyperparameter transfer with μP (Yang & Hu, 2020; Yang et al., 2021).
2 BACKGROUND AND RELATED WORK
2.1 LEARNING RATE SCHEDULES
LR schedules have a long history in stochastic optimization, and are motivated by convergence bounds for stochastic gradient methods (Moulines & Bach, 2011; Bottou et al., 2018). For example, following Andriushchenko et al. (2023), consider SGD for a convex loss parameterized by d: with a constant LR n, the gap between the optimum and current loss step f can be bounded by:
B[C(0.) C(O.)] < (1 110 0. nja²
(I)
where Og are initial parameters, 0. is the loss minimizer, o² is a bound on the variance of gradient noise, and μ is measure of the objective's curvature (see also Bottou et al. (2018, Theorem 4.6)). A larger LR can decrease dependence on initial conditions (the bias term), but also increase the effect of gradient noise (the variance term). As training progresses, bias decreases exponentially in t, and the relative importance of variance increases. Bottou et al. (2018) note this motivates a strategy where L.R is high initially (to mitigate bias) and lowered later (to minimize variance).
2

---

Published as a conference paper at ICLR 2025
In practice, when and how to reduce the LR is rarely informed by ML theory. Many LLMs simply follow the 10x cosine schedule, which is noted to work slightly better than cosine D2Z in the in- fluential work of Hoffmann et al. (2022). It is also well established that using portion of longer (or extending a shorter) schedule is suboptimal compared to using a schedule that reaches its mini- mum only at the final training step (Li et al., 2019; Hoffmann et al., 2022; Hu et al., 2024; Hägele et al., 2024). Linear decay after warmup (Howard& & Ruder, 2018) has been used in LLMs, typically also to 10% of the peak (Henighan et al., 2020; Dey et al., 2023; Sengupta al., 2023). Dropping LR at specific milestones (step decay) is popular in vision models (He et al., 2016; Zagoruyko & Komodakis, 2016; Li et al., 2020), but has also been used in LLMs (Bi et al., 2024).
Kaplan et al. (2020) compared various decay functions and concluded the specific schedule was unimportant given a high enough average I.R, although decaying to zero "gives a fixed improvement close to the end of training." Few papers explicitly compare different LR schedules for large-scale training. and when comparisons are made (Shallue et al., 2019; Kaplan et al., 2020; Schmidt al., 2021; Hoffmann et al., 2022; Yang et al., 2021). they are not the primary focus. So. while some insights are gained, "comprehensive study" is usually regarded as "out of scope" (Aleph Alpha, 2024). One exception is Defazio et al. (2023), who found linear equals outperforms other com- mon schedules, including cosine, across a range of problems, including LLM training. In addition, Defazio al. develop convergence bounds that theoretically motivate linear as the optimal schedule. Unlike our work, they do not evaluate LL with different peak LRs or decay ratios.
Seeking to measure model quality at different training durations without having re-train sepa- rate models from scratch (Zhai et al., 2022; Hägele et al., 2024), researchers have adopted var- ious continuous schedules, such as constant, cyclic (Smith, 2017), etc. Following optimization theory (Moulines & Bach, 2011; Defazio et al., 2024) weight averaging provides an alternative to decay (Sandler al., 2023; Sanyal al., 2023; Busbridge et al., 2024), although is typically not as effective (Hägele et al.. 2024), and moreover may have hyperparameters that implicitly depend on (WSD) approaches have also been used in I.MM training (Hu et al., 2024; Shen et al., 2024a; Ibrahim et al., 2024; Hägele et al., 2024). These methods train at a constant LR, but decay from a checkpoint in a separate process when an intermediate model is needed. We show that the optimal constant LR of these methods implicitly depends on training duration, rendering them not truly schedule-free.
training duration (Defazio al., 2024). Warmup-Stable-DDer
2.2 MAXIMAL UPDATE PARAMETERIZATION (nP)
It is common to scale initial weights of neural networks such that activations have unit variance (Glo- rot & Bengio, 2010; He et al., 2015). Yet weights can become unstable after a few steps of updates, layer-wise LRs are imbalanced (Yang et al., 2021). In contrast, μP (Yang Hu, 2020; Dey al., 2024a) prescribes a re- parameterization of initial weight variances and LRs essentially, rules for scaling these values as model width (i.e., demodet) scales - so that activations and updates remain stable. uP also stabilizes embeddings, layer norms, and self-attention in Transformers. juP is seeing growing application in LLMs (Dey et al., 2023; Shen et al., 2024b; Hu et al., 2024), where acts to stabilize training and to enable transfer of optimal hyperparameters (HPs) across model scales.
With μP, base HPs can be tuned on a small proxy, or base, model, and transferred to larger models. Given the width of the proxy model, dy. and width of the target, de, μPP prescribes scaling factors to apply to HPs. The base LR î is scaled down to pīj. where p = dz/d,. In terms of LR schedules, the base LR ih is sealed every step to provide 7)- μP is convenient in our study the same base peak LRs, 77, at each model size, and observe trends that are scale-invariant.
can sweep
2.3 ADAMW WEIGHTS AS AN EXPONENTIALYY
MOVING AVERAGE (EMA)
An Adam W update a single training step, t, can be expressed as:
7 Vo, +
0. = (1 m)))0 1)
(2)
where 17 is the learning rate, my and 0n are (bias-corrected) running averages of the gradient and the squared gradient, respectively, and is a small constant added to prevent division by zero. The weight decay value, X, is typically set to 0.1 in LLM training (Brown et al., 2020; Hoffmann et al., 2022; Almazrouei et al., 2023; Aleph Alpha, 2024).
3

---

Published as a conference paper at ICLR 2025
While the running averages of and in Adam (Kingma & Ba. 2014) and AdamW aroonentiall-wweighee moving averages (EMAs), Wang Aitchison (2024) recently showed that the weights generated by AdamW can also be understood an EMA of the weight updates. That is, a standard EMA, yt, for time-varying quantity, Tt, can be written as:
Vi = (1 a a)3/-11 at,
(3)
where a is the smoothing parameter.
AdamW in Equation (2) can be seen as EMA by letting:
V/+ = 0,, a = 11,, and 24 1 m
(4)
Vo +
Wang & Aitchison note the quantity r=1/a, i... 1/G2), provides rough timescale over which up- dates are averaged. Weight decay can therefore be used to control the effective window over which updates are combined (smaller values of À increase r, increasing the role of earlier updates in 0z)- This perspective also motivates dynamic LR schedules. Initial should be small (high 1):)). to forget early updates, "while the final timescale around the total number of epochs [low m7. to ensure av- eraging over all datapoints." In fact, we will show that contribution of weight updates, 20, 21.... Er. to 0g at a particular step, t, cannot be determined by the instantaneous value of my. Rather, the contribution of any 2; to 0g requires looking at the full LR schedule holistically (Section 3.1).
Wang & Aitchison also motivate scaling rules for optimal À as model and dataset size vary. If dataset size increases by factor of K (i.e., Kx many steps), the EMA perspective recommends scaling A by l/K in order to expand 7 proportional to K. Moreover, with jaP, if model size increases and LR scaled by p (Section 2.2), the EMA perspective motivates scaling À by 1/p keep constant."
3 METHODS, CONCEPTUAL FOUNDATIONS, AND HYPOTHESES
We now present an extended EMA perspective on AdamW that accounts for time-varying LRs. We then introduce our conceptual model of LLM training, and connect it to the EMA perspective. Finally, we outline the specific testable hypotheses that follow from our conceptual model.
3.1 ADAMW AS CONVEX COMBINATION OF WEIGHT UPDATES, DRIVEN BY LR SCHEDULE
Consider a generic moving average, but now with time-varying smoothing, ag. i.c., yt (1 0t)yt-1 + 0,74. If we let (n = 1 (so that y1 = x1), we can express in terms of all inputs d::
y1 = 01
V2 (1
02) 02g22
Σ
II (1 11-00 -as)
i=1
Ø=++1
=
(5)
Where C.j (1 aç. It can be shown Vt, a, and hence, in standard EMA, each yt is a convex combination of inputs .i... 2g (all coefficients are non-negative and sum not restricted to appoentiallyydecceasiiii Ct i
to 1). However, with time-varying ag, decreases. Indeed, for any convex combination of inputs, there a corresponding smoothing sched- ule that generates the combination via the EMA. In terms of LR schedules for AdamW training. 0,, while a; = y becomes the smoothing parameter at step / (cf. Equations (3) and (4)).²
3.2 CONCEPTUAL MODEL: BIAS AND VARIANCE IN LLM TRAINING
Following Andriushchenko al. (2023), main conceptual premise is that in LLM training, there is an initial training phase focused on movement from initial conditions (bias reduction), followed by
While the EMA perspective does not account for the updates 2, themselves depending on parameters yt-1, we find it a useful part of our conceptual model of training, in that it helps predict experimental results.
When using a jP I.R scaling factor, p (Section 2.2), the smoothing parameter is nd přj..
4

---

Published as a conference paper at ICLR 2025
Bias reduction
Bias reduction
LR setting
LR sotting
maxLR
minLR
Variance reduction
Variance reduction
maxi R
mini R
Optimization steps
Optimization steps
(a) Constant, fewer steps (low TPP)
(b) Constant, many steps (high TPP)
reduction
reduction
LR setting
maxLR
LR setting
Variance reduction
maxLR
Variance reduction
minLR
Optimization steps
Optimization steps
minLR
(c) D22, fewer steps (low TPP)
(d) D2Z, many steps (high TPP)
Figure 3: Bias & variance in LLM pre-training: training duration increases (higher TPP), the importance of variance reduction - and having lower I.R - increases. With no decay (Constant, 3a, 3b), the optimal peak LR must drop significantly lower, neglecting bias reduction. With D27. (3c, 3d), periods of bias and variance reduction can both be enjoyed without large shifts in peak LR.
a later phase bottlenecked by gradient variance. Furthermore, analogous to prior work optimizing the convex loss gap via SGD (Equation (1)), argue the primary beneficial mechanism of LR decay in LLM training is to reduce gradient noise during later stages of training, while simultancously minimizing the impact bias reduction. The larger the dataset, the longer the proportion of training bottlenecked by gradient noise, and the greater the benefit from D2Z (Figure 3).
Variance reduction Per-step gradient variance is known to increase over the course of train- ing (McCandlish et al., 2018). Very recently, Zhang et al. (2024) showed that for an LLM of fixed size, training with larger datasets corresponds to larger critical batch sizes, which directly relates to larger (aggregate) gradient variance via the gradient noise scale (McCandlish et al., 2018).
In the EMA perspective, parameters 0, are combination of prior weight updates (indexed by i). The more variance at step 4, the more updates that should combined in order to reduce Here variance is reduced by combining updates across steps, rather than increasing batch size at a specific step. Now, with a constant LR, update coefficient Ct decreases exponentially in (t (t-1) i): the further back the update from the current step, the less it con- tributes. But if LR 7) decreases
Importance of decay, weight decay
Importance of high early LR
steps, each later coefficient Cr.m will be smaller due to a lower a'm. = A. Yet since co- efficients sum to one, earlier coefficients will correspondingly increase: coefficients are basically flattened and updates con- tribute more evenly across steps. The more we decay, the more updates we average over, and the more variance is reduced.
TPP
Figure 4: HP influence vs. TPP: Higher TPP means higher gradi- ent noise; LR decay & weight decay settings thus increase in importance with TPP.
Bias reduction Higher decay is preferable reducing peak LR because we must also minimize bias, i.e., the contribution of the initial random weights. In the EMM perspective, the contribution of these weights to 0, is C,2 2 az). For decaying IR, and where 0;- ΠiA < 1,
coefficient C, can be approximated (with equality when LR is constant, Appendix B.1):
1, (1-
(6)
where â is the average aj over the schedule. Exactly as in Equation (1), bias therefore decreases exponentially in the absolute number of steps, t, with a rate of decrease depending on the LR. Crucially, this means that as train for more total steps (i.e., higher pertn--per-prrrterr TPP), there is a decrease in the fraction of steps required for C,21 to become negligible. At higher TPP, bias reduction becomes relatively less important than variance reduction (Figure 4).
5

---

Published as a conference paper at ICLR 2025
3.3 EXPERIMENTAL HYPOTHESES
Hypothesis I: As TPP increases, the relative benefit of D227. over 10x decay will increase.
This hypothesis follows from the premise that gradient variance plays larger role at higher TPP, and greater LR decay (as in D2Z) allows for more updates to be averaged, and thus greater variance reduction. A related hypothesis that if we mcrease the batch size at each step, gradient variance will decrease, and so the benefits of D2Z over 10x decay should diminish. Note our conceptual framework does not say precisely at which TPP or batch size that D22, will first prevail; here we will rely on our empirical findings (Section 4) to fill in the theoretical gap.
Hypothesis 2: As TPP increases, the optimal peak LR will decrease, for all LR schedules.
Tuning peak LR is about trading off movement from initial conditions (requiring a high LR) and mitigating variance (requiring low LR). As TPP increases, and bias reduction plays smaller role, optimal peak I.R should decrease. We hypothesize the decrease will be greater with a Constant schedule, as Constant does not use decay to balance the conflicting demands of bias and variance (Figure 3). Moreover, optimal peak LR for other continuous LR schedules, like WSD and Cyclic, should also decrease with longer training durations. In this way, such "schedule-free" approaches are not truly schedule-free. This dependence is also obvious when plotting update coefficients for these schedules (appendix Figure 29): the higher the LR, the more emphasis on recent updates.
Hypothesis 3: Linear D2Z will improve Cosine D2Z.
While LR decay allows averaging final weight updates may not contribute to the EMA. From a loss surface perspective, as the LR approaches zero, take vanishingly small steps. Since Cosine reaches smaller steps faster than Linear (Figure 2, left), Cosine will make less progress toward the optimum loss. From the EMA perspective, this is equivalent Cosine having smaller Cr,; coefficients as approaches t (Figure 2, right). Note this problem is unique to Cosine D2Z and will not affect, e.g., Cosine 10x decay.
more weight updates, if the LR decreases too quickly, the
Hypothesis A high LR (not weight decay) reduces bias and achieves optimal loss.
Note that weight updates H, have a coefficient of 1/x in Equation (4). So, while 77 and X contribute equally to a, increasing Ả to reduce bias is counterproductive as weight updates will be scaled down proportionally, reducing movement from initial conditions. However, LR is in high enough range so that bias plays a minimal factor, both LR and WD should equally affect variance reduction.
4 EMPIRICAL ANALYSIS OF DECAY-TO-ZERO
4.1 EXPERIMENTAL SETUP
Experiments use a GPT-like LLM (Radford et al., 2019), with ALiBi embeddings (Press et al., 2022) and SwiGLU (Shazeer, 2020). Main paper models are trained on SlimPajama (Soboleva al., 2023) and evaluated over 1. held-out tokens (regardless of training TPP). Unless otherwise indicated, A=0.1 is used. Training runs use the same random seed, so all decay functions (Linear, Cosine, etc.)
and ratios (Constant, 10x, D2Z) have identical warmups, but note results are very consistent across seeds at this scale (appendix Figure 24). By default we use μP (standard parameterization results are in Appendices C.2 and C.10). μP HPs derived from smaller proxy model tuned using Linear- 10x schedule. Since we hypothesized that different LR schedules may enjoy their optimal peak LR different values, our experiments compare schedules when each is tuned its optimal peak LR; we sweep îj by factors of 2x around the μP proxy-tuned ij = 1.6e-02. Appendix A provides further details the model architecture (Table 2), dataset sizes (Table 3), and compute resources.
4.2 RESULTS
Finding I: Linear-D2Z is the optimal LR schedule across virtually all peak learning rates.
For (000mmmarmmter pP models trained to compute-optimal 20 TPP, the optimal Linear-D2Z setting achieves 0.77% lower loss than the optimal Linear-10x 10x setting (Table 1). At smaller, suboptimal

---

Published as a conference paper at ICLR 2025
Table 1: Comparison of schedules: validation loss for 610M μP models, 20 TPP. 7f=1.6c-02 the juP proxy-tuned peak (base) LR, is the LR after μul adjustment (p=0.1 125, see Section 2.2). D2Z is superior to 10x decay across most LRs. Linear-DD2 slightly outperforms Cosine-D22, in all cases.
Cosine
Linear 10x
Cosine
Linear D27
?
InvSqrt
Constant
10x
D27
6.5e-02
8.1e-03
2.789
3.035
NaN
2.667
2.611
2.605
3.2e-02
4.0e-03
2.710
2.850
2.604
2.606
2.574
2.571
1.6e-02
2.0e-03
2.671
2.768
2.590
2.591
2.578
2.573
8.1e-03
1.0e-03
2.665
2.722
2.598
2.600
2.595
2.590
4.0e-03
5.1e-04
2.691
2.711
2.634
2.635
2.637
2.633
2.0e-03
2.5e-04
2.762
2.739
2.707
2.710
2.717
2.714
2 TPP
20 TPP
200 TPP
Constant
3.8
Validation loss
4.5
D2Z
3.4
3.0
4.0
3.0
10"2
10
Max LR, 7
Max LR,
Max LR, /
Figure 5: Comparing decay schedules across TPP (111M scale): As TPP increases, Linear-D2 outperforms 10x, especially at the proxy-tuned peak LR (red lines).
peak LRs, 10x be better than D27. In appendix Figure 10, demonstrate very similar sults using the standard parameterization. Regarding the decay function, gains from Linear-D2Z over Cosine-D2Z are small, but perfectly consistent across peak LRs, exactly in line with recent work (Defazio et al., 2023; Lingle, 2024). Interestingly, Cosine-10x 10x is consistently slightly better than Linear-104 10x, showing that Linear itself is not always best, rather Linear plus D2Z is needed. Lacking cooldown phase, inverse square root (ImsSqrt) and Constant do not perform as well.
Appendix Figure 11 compares schedules as we sweep peak LR for 111M, 610M, and 1.7B models. For these and subsequent experiments, we use Linear decay.
Finding 2: As TPP increases, the relative improvement of D27 over 10x also increases.
For 111M models at only TPP, D2Z performs worse than 10x across all peak LRs (Figure 5, left). 610M results similar (appendix Figure 12). However, TPP increases, D2Z begins to outperform 10x, exceeding the best setting of 10x by 1.6% at 200 TPP, and performing 2.8% better the proxy-tuned fj=1. Ge-02 (marked in plots with red vertical line). While hypothesized D27. would surpass 10x at some TPP, our key finding is that, at both compute-efficiens (20 TPP) and over-trained sizes, decaying to zero is optimal. Even improves over Constant), but too much evidently hampers bias reduction. As TPP increases, optimal LRs for 10x shift lower as variance reduction gains in importance, but now it the lower peak LRs that hamper bias reduction, and 10x begins to lag D2Z. By 20+ TPP, D2Z is consistently superior.
TPP, some LR decay is beneficial (10x
Figure plots validation loss of 10x and D2Z at different TPP settings for different model sizes (all models are trained using the proxy-tuned peak LR). For the 610M model, D27. is also initially worse than 10x, but begins to surpass it around TPP, and by 200 TPP 2.6% better. As with training loss (Figure 1), an 80 TPP D2Z model can surpass 200 TPP 10x model in validation loss. Importantly, these trends also hold for downstream evaluation of the models (appendix Table 5).
Figure 6 also shows the advantage of using D2Z with 1.7B models trained 80 TPP; D2Z achieves roughly 3.05% lower loss than 10x at the proxy-tuned ij. At this scale, such gains provide real- proposal models in speculative decod-
world impact: we over-train models of this size for ing (Leviathan et al., 2023) and smaller models like Gemma-2B (Mesnard et al., 2024) and Phi- 1.5 (Li et al., 2023) are used in many production settings.

---

Published as a conference paper at ICLR 2025
3.8
10x
D2Z
3.6
111M-D27
B=63
B=
610M-10x
3.4
610M-D22
Validation loss
B=126
B-126
Valication lass
1.78- 10x
1.78. D2Z
B B-252
B 504
B B-252
B 504
3.0
2.8
2.6
10
10
Max LR, 7
Max LR, 7j
22
Figure 7: Comparison across batch sizes (iso-TPP) (610M): Comparing 10x (left) to D2Z (right). All mod- els trained to 20 TPP (smaller batches have more steps), As B decreases, the optimal LR (n) drifts significantly lower for 10x decay, less for D2Z.
TPP:
Figure 6: Loss gap between Linear- D2Z and 10x grows with TPP: for 111M, 610M, and 1.7B-scale models.
2.70
2.65
2.52
Cyclic
D2Z
2.60
10x
WSD
2.55
2.50
Validation
2.50
2.48
2.45
2.46
2.40
Constant, #11.66
D2Z,
4=
2.44
2.35
Constant, -=1.0e-3
D2Z, î 1.6e2
10 2
10x, fj 1.6e-2
Max LR, 7
Batch:
10000
20000
30000
40000
Figure 8: Comparison to continuous (610M, 80 TPP): D2Z surpasses Cyclic and WSD schedules.
Figure 9: Loss on training set batches (in original order) after models have been trained (610M, TPP)
Finding 3: Compared to Constant and 10x, optimal peak LRs are much more stable with D27.
Figure 5 reveals different levels of hyperparameter sensitivity when increasing TPP: optimal LR shifts substantially lower for Constant, somewhat lower for 10x, and hardly all for D2Z. The same trend holds at the 610M scale (appendix Figure 12). Moreover, with D2Z, loss is less sensitive to suboptimal LR (bowls are flatter). D2Z loss is also more stable (and lower) as we vary both weight decay (appendix Figures 17 and 18) and batch size, B. For example, the optimal LR already shifts significantly for 10x models at B=126, while D2Z. optima only begin to shift at B=63 (Figure 7). The superiority of D27. over 10x across B is more clearly observed in appendix Figure 13.
Crucially, if we keep the number of optimization steps constant (1 (11752), but only vary B (effectively increasing TPP proportional to B), we find the optimal learning rate does not decrease (i.e., with TPP) but rather increases (appendix Figure 14). Clearly, LR decay is primarily beneficial as a noise- reduction mechanism: with larger batches, we have less gradient noise, and can afford larger LR.
Finding 4: Linear-D2Z works better than WSD and Cyclic (continuous) LR schedules.
In Figure 8, compare D2Z to 10x, and to two approaches designed for continuous pre-training: Cyclic, which cycles LR up and down, and WSD (Section 2... For WSD, we simulate a model being retrieved after 80 TPP, and cool the LR to zero for the final 22.5% of steps (around the proportion recommended by Hägele et al. (2024), and equal to the cooldown duration in our 20 TPP models). Appendix Figure 29 provides full LR curves and update coefficients for all models in Figure 8.
8

---

Published as a conference paper at ICLR 2025
At its optimal LR, WSD works better than 10x, confirming results in Hägele et al. (2024). However, note the optimal peak LR shifts lower for both 10x and WSD at this TPP. Linear-D22 remains best here, around 0.84% better than the optimal WSD. Given the diminishing returns of high-TPP training (Figure 6), WSD would require significantly more training FL.OPs to reach the level of D27.
Finding 5: Constant and 10x, but not D27, strongly overfit to the end of the training data.
Figure 9 shows the loss of trained models (with frozen weights) the exactly- same ordered batches used during pre- vre-training. As predicted by the extended EMA perspective (Section 3), the higher the LR, the more models fit to late-stage training batches. It is striking that both Constant and 10x, but not D2Z, overfit to the very final portion. Extra adaptation of generative models to recent training sequences has long been observed (Graves, 2013). but to our knowledge this is the first evidence D2Z may help mitigate these effects. Since D2Z demonstrates the lowest loss on batches slightly before the final training phase, placing the highest-quality and most-recent data in the very final phase, while using D2Z (e.g., as in Dubey et al. (2024)), may be suboptimal. These findings also contradict Biderman et al. (2023), who found training order had little impact on memorization.
DISCUSSION
The confounding role of LR schedule We have shown that the more constant the LR schedule, the more sensitive the optimal peak n). These findings help explain LR sensitivity in prior work. For example, Shen et al. (2024b) were puzzled by their observation that "although the WSD scheduler could, in theory, continue in the stable phase forever ... the optimal LRs are different for different amounts of training tokens." " As noted under Hypothesis 2. LR schedules with long constant periods only appear "schedule-free" in the primal; from the dual (EMA) perspective, the higher the LR, the more emphasis is placed on recent updates (see appendix Figure 28a for Constant, Figure 29d for WSD). Filatov et al. (2024) and Yang et al. (2021, Figure 19) also observed major decreases in optimal LR with higher TPP; these results were both obtained with a constant LR. In contrast, Bjorck et al. (2024) observed less significant decrease of optimal 1) with TPP, but used 10x decay.
LR schedule may also affect the optimal 1j as batch size varies - particularly when using μP, Some prior μP studies (Yang al., 2021: Noci et al., 2024; Shen et al., 2024b) observed linear scaling of optimal LR with batch size, i.e., the so-called linear scaling rule (Krizhevsky, 2014; Chen et al., 2016; Smith et al., 2018). Others (Lingle, 2024) have observed square-root scaling, resonating with other prior studies (Hoffer et al., 2017: You et al., 2019b; Malladi et al., 2022). This discrepancy can be explained by noting the linear scaling results were all found with a Constant or WSD LR decay, while square-root was observed with Linear D2Z, again underscoring the greater stability of D2Z.
The confounding role of training duration While LR schedules have confounded studies varying TPP, TPP has analogously confounded studies evaluating LR schedules. Recall Kaplan et al. (2020) also saw a benefit from D2Z. In contrast to Chinchilla scaling (Hoffmann et al., 2022), in the Kaplan al. perspective, small models should be trained to high TPP (while larger models should be trained less). It is therefore not surprising Kaplan et al. saw benefits testing D2Z with small models; as we have shown, D2Z is especially effective at high TPP. contrast, in Figure 4 of Yang et al. (2021), Linear is worst of all schedules, and the gap between it and Constant and InvSqri grows with model width. But here, since training data small - and fixed, then TPP decreases as width increases. Thus training is in phase where bias reduction is paramount and D2Z is not effective. Notably, in their LLM training experiments at higher TPP, Yang et al. do report linear D2Z to work best.
With this context, we speculate why D2Z has not been adopted more widely in LLM pre-training:
First, it is common to evaluate hyperparameters on smaller training runs; unfortunately, with lim- ited training data (low TPP), D2Z misleadingly under-performs.
Secondly, coupling between LR schedule and optimal peak LR is problematic: with 10x decay, we may find a lower peak LR is optimal; if we then test D2Z with the same LR, it may be suboptimal for D2Z. Not seeing a benefit at this LR, we may conclude D2Z is inferior in general.
Finally, poorly controlled training dynamics may prevent D2Z models from training with their (higher) optimal LR. Indeed, when we initially compared D2Z and 10x using NanoGPT (Ap- pendix 1O)) 10x performed better at the default L.R. Raising the LR resulted in training diver-
9

---

Published as a conference paper at ICLR 2025
gence. Only after switching the numerical precision from 1f float 32 could D2Z succeed - and the model reach optimal loss.
The special benefit of D2Z Our results suggest a low a later in training can expand the timescale over which weight updates are combined. reducing noise in similar way increasing batch size. However, there is apparently a separate, independent benefit from a vanishing LR. Indeed, looking appendix Figure 29 for the 80 TPP comparison continuous schedules, are quite similar to the D22. curve, apart from missing the final drop. Moreover, they are flatter than the WSD EMA coefficients for the same peak LR, suggesting better integration of prior updates. Yet WSD performs better than 10x at all LR settings (Figure 8), In contrast, at 2 TPP (Figure 5, left), 10x performs better than D2Z at every LR setting. Prior work has shown large LRs allow exploration of the loss surface at height "above the valley floor" (Xing et al., 2018), while LR cooldown phases descend into a local minimum (Hu al., 2024; Hägele al., 2024). appears that descending into these minima is beneficial only after sufficient exploration of the loss surface.
see 10x coefficients
LIMITATIONS AND FURTHER EXPERIMENTS
While our findings strongly support linear D27. being optimal in our specific context, there are some limitations to keep in mind. First, LR schedules like D2Z require prior knowledge of the total number of training steps. But it is worth reiterating that even nominally schedule-free approaches such as Constant also require this knowledge in order to optimally set the LR value. In contrast, the extended EMA perspective of LR schedules enables derivation of a truly schedule-fre schedule, which we introduce in Appendix D.I. Second, our focus in this paper was specifically LLM training at compute-optimal and overtrained dataset sizes. For MI. problems with limited access to training data, D27. likely not the best strategy. Third, our work focuses on AdamW (the standard optimizer in LLM -re-traini While the extended EMA perspective of LR schedules will likely apply to other optimizers that use decoupled weight decay (as noted by Wang & Aitchison (2024)), it may not apply to approximate second order methods, such as Shampoo (Gupta et al., 201 8). Finally, for LLMs with unstable training dynamics that cannot tolerate high LRs, D2Z may not be beneficial. We experienced this first-hand when we initially trained NanoGPT (Appendix 10))
However, we also note the remarkable consistency of D2Z's success. While results the main paper used the SlimPajama dataset, consistent results were found with OpennWeTTeee (Appendix C. and a multilingual dataset (with a larger vocabulary) (Appendix C.I We also saw consistent findings with different parameterizations (Appendices C.2 and C.1 architectures (Appendix C.I weight sparsity settings (A (Appendix C.8), and training frameworks (Appendix C. C.10). Appendix C.11 describes scaling laws fit to models trained with D2Z and 10x, at model sizes up to 2.75B; results indicate a growing performance gap as scale increases. Finally, Appendix C.6 demonstrates that the benefits of weight decay are observed primarily when using LR D2Z, where raising weight decay can fine-tune the EMA update coefficients without affecting training stability.
7 CONCLUSION
Linear decay-to-zero is the optimal LR schedule for LLM training using AdamW. To be clear, less decay is beneficial for low pokens-per-parameete training, but there is no practical reason to perform such training with LLMs, since the same FLOPs could be used train smaller model, over more tokens, to a lower loss using D2Z. The superiority of D2Z was validated across a range of experimental conditions. Results suggest its relative benefit will increase as models increase in scale. Moreover, when using D27. and μP, the optimal peak LRR is less sensitive to changes in weight decay, dataset size, and batch size, i.e., there better hyperparameter transfer.
Our analysis was aided by our interpretation of AdamW's 's output as a convex combination of prior weight updates. D2Z overfits less the final training sequences, and is especially beneficial when gradient noise dominates training. As we enter a phase of applied ML where inference efficiency is a primary concern, there is strong motivation to study high-TPP training, where gradient noise is the bottleneck. While our results indicate that D2Z is key component of the solution here, further investigation is required, including into how and when to adjust hyperparameters such as weight decay, batch size, and learning rate, in the high-TPP context.
10

---

Published as a conference paper at ICLR 2025
REFERENCES
Aleph Alpha. Introducing I-LMM Transparent and compliant. Blog post, 2024. Accessed: 2024-09-07.
Ebtesam Almazrouei, Hamza Alobeidli, Abdulaziz Alshamsi, Alessandro Cappelli, Ruxandra Co- jocaru, Mérouane Debbah, Êtienne Goffinet, Daniel Hesslow, Julien Launay, Quentin Malartic, et al. The Falcon series of open language models. arXiv preprint
2023.
Maksym Andriushchenko, Francesco D' Angelo, Aditya Vare,, and Nicolas Flammarion. Why do we need weight decay in modern deep learning? arXiv preprint arXiv:3310.0441,, 2023.
Xiao Bi, Deli Chen, Guanting Chen, Shanhuang Chen, Damai Dai, Chengqi Deng, Honghui Ding, Kai Dong, Qiushi Du, Zhe Fu, al. DeepSeek LLM: Scaling open-source language models with longtermism. arXiv preprint arXiv:2401.02295. 2024.
Stella Biderman, Hailey Schoelkopf, Quentin Anthony, Herbie Bradley, Kyle O'Brien, Eric Hal- lahan, Mohammad Aflah Khan, Shivanshu Purohit, USVSN Sai Prashanth, Edward Raff, Aviya Skowron, Lintang Sutawika, and Oskar van der Wal. Pythia: A suite for analyzing large language models across training and scaling, 2023.
Johan Bjorck, Alon Benhaim, Vishrav Chaudhary, Furu Wei, and Xia Song. Scaling optimal LR across token horizons. arXiv preprint arXiv:2409.19913, 2024.
Léon Bottou, Frank Curtis, and Jorge Nocedal. Optimization methods for large-scale machine learning. SIAM review, 60(2):223-311, 2018.
Tom Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared D Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, et al. Language models are few-shot leamerr. Advances in Neural Information Processing Systems, 33:1877-1901, 2020.
Dan Busbridge, Jason Ramapuram, Pierre Ablin, Tatiana Likhomanenko. Eeshan Gunesh Dhekane, Xavier Suau Cuadros, and Russell Webb. How to scale your EMA. Advances in Neural Informa- fion Processing Systems, 36, 2024.
Jianmin Chen, Xinghao Pan, Rajat Monga, Samy Bengio, and Rafal Jozefowicz. Revisiting dis- tributed synchronous SGD. arXiv preprint (=Xiv:1604.0098) 2016.
Aaron Defazio, Ashok Cutkosky, Harsh Mehta, and Konstantin Mishchenko. When, why and how much? Adaptive learning rate scheduling by refinement. arXiv preprint arXiv:331000881,, 2023.
Aaron Defazio, Harsh Mehta, Konstantin Mishchenko, Ahmed Khaled, Ashok Cutkosky, et al. The road less scheduled. arXiv preprint urXiv:2405. /5682, 2024.
Nolan Dey, Gurpreet Gosal, Hemant Khachane, William Marshall, Ribhu Pathria, Marvin Tom, and Joel Hestness. Cerebras-GPT Open compute-optimal language models trained the Cerebras wafer-scale cluster. arXiv preprint #xXiv:2304.03208 2023.
Nolan Dey, Quentin Anthony, and Joel Hestness. The practitioner's guide to the maximal update parameterization. Blog post, 2024a.
Nolan Dey, Shane Bergsma, and Joel Hestness. Sparse maximal update parameterization: A holistic approach to sparse training dynamics. arXiv preprint (tt.:24005157433 2024b.
Abhimanyu Dubey, Abhinav Jauhri, Abhinav Pandey, Abhishek Kadian, Ahmad Al-Dahle, Aiesha Letman, Akhil Mathur, Alan Schelten, Amy Yang, Angela Fan, et The Llama 3 herd of models. arXiv preprint 2rXiv22407.21788 2024.
Oleg Filatov, Jan Ebert, Jiangtao Wang, and Stefan Kesselheim. Time transfer: On optimal learning rate and batch size in the infinite data limit. arXiv preprint arXiv:2410.05833 2024.
Leo Gao, Stella Biderman, Sid Black, Laurence Golding, Travis Hoppe, Charles Foster, Jason Phang. Horace He, Anish Thite, Noa Nabeshima. Shawn Presser, and Connor Leahy. The Pile: An 800GB dataset of diverse text for language modeling, 2020.
11

---

Published as a conference paper at ICLR 2025
Rong Ge, Sham M Kakade, Rahul Kidambi, and Praneeth Netrapalli. The step decay schedule: A near optimal, geometrically decaying learning rate procedure for least squares. Advances in neural information processing systems, 32, 2019.
Xavier Glorot and Yoshua Bengio. Understanding the difficulty of training deep feedforward neural networks. In Proceedings of the Thirleenth International Conference on Artificial Intelligence and Statistics (PMLR), 2010.
Aaron Gokaslan and Vanya Cohen. Open WebText corpus. Web page, 2019.
Priya Goyal, Piotr Dollár, Ross Girshick, Pieter Noordhuis, Lukasz Wesolowski, Aapo Kyrola, An- drew Tulloch, Yangqing Jia, and Kaiming He. Accurate, large minibatch SGD: Training ImageNet in 1 hour, 2018.
Alex Graves. Generating sequences with recurrent neural networks.
arXiv preprint
arXiv:1308.00500 2013.
Gavia Gray, Aman Tiwari, Shane Bergsma, and Joel Hestness. Normalization layer per- example gradients are sufficient to predict gradient noise scale in transformers. arXiv preprint arXiv:2411.00999 2024.
Vineet Gupta, Tomer Koren, and Yoram Singer. Shampoo: Preconditioned stochastic tensor opti- mization. In International Conference on Machine Learning, pp. 1842-1850. PMLR, 2018.
Alexander Hägele, Elie Bakouch, Atli Kosson, Loubna Ben Allal, Leandro Von Werra, and Martin Jaggi. Scaling laws and compute-optimal training beyond fixed training durations. arXiv preprint WrXiv:2405.18392 2024.
Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Delving deep into rectifiers: Surpassing human-level performance ImageNet classification. In Proceedings of the IEEE international conference on computer vision, pp. 1026-1034, 2015.
Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recog- nition. In Proceedings of the IEEE conference on computer vision and pattern recognition, pp- 770-778, 2016.
Tom Henighan, Jared Kaplan, Mor Katz, Mark Chen, Christopher Hesse, Jacob Jackson, Heewoo Jun, Tom B Brown, Prafulla Dhariwal, Scott Gray, et al. Scaling laws for autoregressive generative modeling. arXiv preprint arXiv:2010. 14701, 2020.
Torsten Hoefler, Dan Alistarh, Tal Ben-Nun, Nikoli Dryden, and Alexandra Peste. Sparsity in deep learning: Pruning and growth for efficient inference and training in neural networks. Journal of Machine Learning Research, 22(241):1-124, 2021.
Elad Hoffer, Itay Hubara, and Daniel Soudry. Train longer, generalize better: Closing the general- ization gap in large batch training of neural networks. Advances in Neural Information Processing Systems, 30, 2017.
Jordan Hoffmann, Sebastian Borgeaud, Arthur Mensch, Elena Buchatskaya, Trevor Cai, Eliza Rutherford, Diego de Las Casas, Lisa Anne Hendricks, Johannes Welbl, Aidan Clark, et al. An empirical analysis of compute-optimal large language model training. Advances in Neural Infor- mation Processing Systems, 35, 2022.
Jeremy Howard and Sebastian Ruder. Universal language model fine-tuning for text classification. arXiv preprint arXiv: 008.000..666666.. 2018
Shengding Hu, Yuge Tu, Xu Han, Chaoqun He, Ganqu Cui, Xiang Long, Zhi Zheng, Yewei Fang, Yuxiang Huang, Weilin Zhao, et MiniCPM: Unveiling the potential of small language models with scalable training strategies. arXiv preprint arXiv:2404.06390 2024.
Adam Ibrahim, Benjamin Thérien, Kshitij Gupta, Mats L Richter, Quentin Anthony, Timothée Lesort, Eugene Belilovsky, and Irina Rish. Simple and scalable strategies to continually pre-train large language models. arXiv preprint arXiv:2403.00 08766, 2024.
12

---

Published as a conference paper at ICLR 2025
Jared Kaplan, Sam McCandlish, Tom Henighan, Tom B. Brown, Benjamin Chess, Rewon Child, Scott Gray, Alec Radford, Jeffrey Wu, and Dario Amodei. Scaling laws for neural language models, 2020.
Andrej Karpathy, nanoGPT. gjtnb..cmmkaappthhhaaaGGG., 2024. GitHub repository.
Diederik P Kingma and Jimmy Ba. Adam: A method for stochastic optimization. arXiv preprint 1t/xiv:4//2.690000 2014.
Atli Kosson, Bettina Messmer, and Martin Jaggi. Analyzing & reducing the need for learning rate warmup in GPT training. In arXiv preprint arXiv:2410.23922, 2024.
Alex Krizhevsky. One weird trick for parallelizing convolutional neural networks. arXiv preprint arXiv:1404.5997, 2014.
Yaniv Leviathan, Matan Kalman, and Yossi Matias, Fast inference from transformers via speculative decoding. In International Conference on Machine Learning, pp. 19274-19286. PMLR, 2023,
Mengtian Li, Ersin Yumer, and Deva Ramanan. Budgeted training: Rethinking deep neural network training under resource constraints. arXiv preprint arXiv:1905.0475:. 2019.
Yuanzhi Li, Sébastien Bubeck, Ronen Eldan, Allie Del Giomn, Suriya Gunasekar, and Yin Tat Lee. Textbooks are all you need II: phi-1.5 technical report. arXiv preprint 034502300.05463,, 2023,
Zhiyuan Li, Kaifeng Lyu, and Sanjeev Arora. Reconciling modern deep learning with traditional optimization analyses: The intrinsic learning rate. Advances in Neural Information Processing Systems, 43:44454-14555, 2020.
Lucas Lingle. large-scale exploration of y-transfer. arXiv preprint arXiv:2404.05722 2024.
Liyunn Liu, Haoming Jiang, Pengcheng He, Weizhu Chen, Xiaodong Liu, Jianfeng Gao, and Jiawei Han. On the variance of the adaptive learning rate and beyond. arXiv preprint arXiv:1908. 1000::900..03655,, 2019.
Ilya Loshchilov and Frank Hutter. Decoupled weight decay regularization. In International Confer- ence on Learning Representations, 2017.
Sadhika Malladi, Kaifeng Lyu, Abhishek Panigrahi, and Sanjeev Arora. On the SDEs and scaling rules for adaptive gradient algorithms. Advances in Neural Information Processing Systems, 35: 7697-7711, 2022.
Sam McCandlish, Jared Kaplan, Dario Amodei, and OpenAl Dota Team. An empirical model of large-batch training, 2018.
Thomas Mesnard, Cassidy Hardin, Robert Dadashi, Surya Bhupatiraju, Shreya Pathak, Laurent Sifre, Morgane Rivière, Mihir Sanjay Kale, Juliette Love, al. Gemma: Open models based on gemini research and technology. arXiv preprint arXiv:2403.08295, 2024.
Eric Moulines and Francis Bach. Non-asymptotic analysis of stochastic approximation algorithms for machine learning. Advances in Neural Information Processing Systems, 24, 2011.
Niklas Muennighoff, Alexander Rush, Boaz Barak, Teven Le Scao, Nouamane Tazi, Aleksandra Piktus, Sampo Pyysalo, Thomas Wolf, and Colin A Raffel. Scaling cats-ronnt language models. Advances in Neural Information Processing Systems, 36, 2023.
Lorenzo Noci, Alexandru Meterez, Thomas Hofmann, and Antonio Orvieto. Why do learning rates transfer? Reconciling optimization and scaling limits for deep learning. arXiv preprint 17X44 2024.
Tomer Porian, Mitchell Wortsman, Jenia Jitsev, Ludwig Schmidt, and Yair Carmon. Resolving discrepancies in compute-optimal scaling of language models. arXiv preprint arXiv:2406. 221246699946 19146, 2024.
Ofir Press, Noah Smith, and Mike Lewis. Train short, test long: Attention with linear biases enables input length extrapolation. In International Conference Learning Representations, 2022.
13

---

Published as a conference paper at ICLR 2025
Alec Radford, Jeffrey Wu, Rewon Child, David Luan, Dario Amodei, and llya Sutskever. Language models are unsupervised multitask learners, 2019.
Jack W. Rae, Sebastian Borgeaud, Trevor Cai, Katie Millican, Jordan Hoffmann, Francis Song, John Aslanides, Sarah Henderson, Roman Ring, Susannah Young, et Scaling language models: Methods, analysis & insights from training Gopher, 2022.
Colin Raffel, Noam Shazeer, Adam Roberts, Katherine Lee, Sharan Narang, Michael Matena, Yanqi Zhou, Wei Li, and Peter J. Liu. Exploring the limits of transfer learning with a unified text-to-text transformer. Journal of Machine Learning Research. 2020.
Herbert Robbins and Sutton Monro. A stochastic approximation method. The Annals of Mathemat- ical Stafistics, pp. 400-407, 1951.
Mark Sandler, Andrey Zhmoginov, Max Vladymyrov, and Nolan Miller. Training trajectories, mini- batch losses and the curious role of the learning rate. arXiv preprint arXiv:2301.02312 2023.
Sunny Sanyal, Atula Neerkaje, Jean Kaddour, Abhishek Kumar, and Sujay Sanghavi. Early weight averaging meets high learning rates for LLM pre-training. arXiv preprint arXiv:2306.03241, 2023.
Teven Le Scao, Angela Fan, Christopher Akiki, Ellie Pavlick, Suzana Ilié, Daniel Hesslow, Roman Castagné, Alexandra Sasha Luccioni, François Yvon, Matthias Gallé, et al. BLOOM: 176b- parameter open-access multilingual language model, 2023.
Robin M Schmidt, Frank Schneider, and Philipp Hennig. Descending through crowded valley- benchmarking deep learning optimizers. In International Conference on Machine Learning, pp- 9367-9376. PMLR, 2021.
Neha Sengupta, Sunil Kumar Sahu, Bokang Jia, Satheesh Katipomu, Haonan Li, Fajri Koto, William Marshall, Gurpreet Gosal, Cynthia Liu, Zhiming Chen, et al. Jais and Jais-chat: Arabic- centric foundation and instruction-tuned open generative large language models. arXiv preprint arXiv:2308. 16149, 2023.
Christopher Shallue, Jaehoon Lee, Joseph Antognini, Jascha Sohl- Dick-Dicsttein, Roy Frostig, and George E Dahl. Measuring the effects of data parallelism on neural network training. Journal of Machine Learning Research, 20(1 12):1-49, 2019.
Noam Shazeer. GLU variants improve transformer, 2020.
Yikang Shen, Zhen Guo, Tianle Cai, and Zengyi Qin. JetMoE: Reaching Llama2 performance with 0.IM dollars. arXiv preprint arXiv:2404.07411 2024a.
Yikang Shen, Matthew Stallone, Mayank Mishra, Gaoyuan Zhang, Shawn Tan, Aditya Prasad, Adri- ana Meza Soria, David Cox, and Rameswar Panda. Power scheduler: batch size and token number agnostic learning rate scheduler. arXiv preprint arXiv:2408. 13359, 2024b.
Lestie N Smith. Cyclical learning rates for training neural networks. In 2017 IEEE Winter Confer- ence on Applications of Computer Vision (WACV), 464-472. IEEE, 2017.
Samuel L. Smith, Pieter-Jan Kindermans, and Quoe V. Le. Don't 't decay the learning rate, increase the batch size. International Conference on Learning Representations, 2018
Daria Soboleva, Faisal Al- 1-Kaateeb, Robert Myers, Jacob R Steeves. Joel Hestness, and Nolan Dey. SlimPajama: 627B token cleaned and deduplicated version of RedPajama. Blog post, 2023.
Howe Tissue, Venus Wang, and Lu Wang. Scaling law with learning rate annealing. arXiv preprint arXiv:2408. 1029, 2024.
Hugo Touvron, Thibaut Lavril, Gautier Izacard, Xavier Martinet, Marie-Anne Lachaux, Timothée Lacroix, Baptiste Rozière, Naman Goyal, Eric Hambro, Faisal Azhar, Aurelien Rodriguez, Ar- mand Joulin, Edouard Grave, and Guillaume Lample. LLaMA: Open and efficient foundation language models, 2023a.
14

---

Published as a conference paper at ICLR 2025
Hugo Touvron, Louis Martin, Kevin Stone, Peter Albert, Amjad Almahairi, Yasmine Babaei, Niko- lay Bashlykov, Soumya Batra, Prajjwal Bhargava, Shruti Bhosale, et al. LLaMA 2: Open foun- dation and fine-tuned chat models. arXiv preprint 1rXiv:2307.09280 2023b.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. In Advances in Neural Infor- mation Processing Systems, 2017.
Xi Wang and Laurence Aitchison. How to set AdamW's weight decay as you scale model and dataset size. arXiv preprini vrXiv:2405.13698 2024.
Mitchell Wortsman, Peter J Liu, Lechao Xiao, Katie Everett, Alex Alemi, Ben Adlam, John D Co- Reyes, Izzeddin Gur, Abhishek Kumar, Roman Novak, et al. Small-scale proxies for large-scale transformer training instabilities. arXiv preprint arXiv:2309. 14322, 2023.
Chen Xing, Devansh Arpit, Christos Tsirigotis. and Yoshua Bengio. A walk with SGD. arXiv preprint arXiv:1802.087771 2018.
Ruibin Xiong, Yunchang Yang, Di He, Kai Zheng, Shuxin Zheng, Chen Xing, Huishuai Zhang, Yanyan Lan, Liwei Wang, and Tieyan Liu. layer normalization in the transformer architecture. In International Conference on Machine Learning, pp. 10524-10533, PMLR, 2020.
Greg Yang and Edward J Hu. Feature learning in infinite-wite width neural networks. arXiv preprint arXiv:201 1. 14522, 2020.
Greg Yang. Edward Hu, Igor Babuschkin, Szymon Sidor, Xiaodong Liu, David Farhi, Nick Ryder, Jakub Pachocki, Weizhu Chen, and Jianfeng Gao. Tuning large neural networks via zero-shot hyperparameter transfer. In Advances in Neural Information Processing Systems, 2021.
Kaichao You, Mingsheng Long, Jianmin Wang, and Michael I Jordan. How does learning rate decay 2019a.
help modern neural networks? arXiv preprint 1900v
Yang You, Jing Li, Sashank Reddi, Jonathan Hseu, Sanjiv Kumar, Srinadh Bhojanapalli, Xiaodan Song, James Demmel, Kurt Keutzer, and Cho-Jui Hsieh. Large batch optimization for deep 2019b.
learning: Training BERT in 76 minutes. arXiv preprint
Sergey Zagoruyko and Nikos Komodakis.
Wide residual networks.
arXiv preprint
1605.07 1460077116600 2016.
Xiaohua Zhai, Alexander Kolesnikov, Neil Houlsby, and Lucas Beyer. Scaling vision transformers. In Proceedings of the IEEE/CVF conference computer vision and pattern recognition, pp. 12104-12113, 2022.
Hanlin Zhang, Depen Morwani, Nikhil Vyas, Jingfeng Wu, Difan Zou, Udaya Ghai, Dean Fos- ter, and Sham Kakade. How does critical batch size scale in pre-training? arXiv preprint arXiv:4410.1166,, 2024.
15

---

Published as a conference paper at ICLR 2025
Table 2: Model architecture and batch sizes for main experiments
Model
vocab. size
dmodel
Tlayers
dhead
dm
batch size
111M
50257
768
10
64
2048
192
610M
50257
2048
10
64
5461
504
1.7B
50257
2048
32
64
5461
504
Table 3: Training steps for main experiments
Model
TPP
Warmup
Steps
Tokens
111M
2
56
557
219M
111M
20
556
5568
2.19B
111M
200
5560
55680
21.9B
610M
2
118
1176
1,21B
610M
20
1175
11752
12.1B
610M
80
4700
47008
48.5B
610M
200
11750
117520
121B
1.7B
20
3322
33220
34.3B
1.7B
80
13288
132880
137B
A EXPERIMENTAL DETAILS
Table 2 provides details on model architecture and hyperparameters for the main experiments (i.e., results presented in the main paper). Table provides information on the training steps. All the models in our main experiments were trained on the SlimPajama dataset (Soboleva et al., 2023), cleaned and deduplicated version of the RedPajama dataset. We use the GPT-2 (Radford et al., 2019) vocabulary of size 50257, and a context length of 2048 tokens. Unless otherwise noted, the weight decay value, X, is by default set to 0.1, following standard practice in LLM pre-training (Brown et al., 2020; Hoffmann et al., 2022: Almazrouei et al., 2023; Aleph Alpha, 2024). Also following standard practice, do not apply weight decay or bias to LayerNorm layers. Validation loss is always computed over 1.1B held-out tokens, regardless of training TPP. By default we parameterize with μP (further details below).
For given TPP, all models have the exact same warmup phase: a linear warmup of the LR from O the peak value. In all runs, warmup was 10% of the total steps. LR warmup is standard practice in LLM pre-training."
All models in the main experiments were trained a Cerebras CS-3 system. 610M-parameter, 20 TPP models take roughly 6 hours each to train on single CS-3. If a training run did not complete due to numerical instabilities, the values are left off our plots or marked as NaN in Table 1.
Proxy model hyperparameter tuning To find the optimal μP hyperparameters (HPs), trained a 99M-maaaeeter proxy model using width dmodel=dp 256, with 24 layers and a head size of
While prior work has suggested L.R warmup is less valuable in modern Pre-LN Transformers (Xiong et al., 2020), various other studies have shown warmup leads to lower loss (Goyal et al., 2018; Liu al., 2019; Tissue al., 2024; Kosson al., 2024), and may reduce sensitivity peak I.R (Wortsman et al., 2023). In light of the similar benefits of D2Z, it would be interesting to investigate the value of warmup for models that are specifically trained using D2Z.
Table 4: Tuned hyperparameters for μP proxy model
Wasb T)
8.67e-02 1.62e-02
Oiiput 9.17
Qouupput
1.095
16

---

Published as a conference paper at ICLR 2025
Table 5: Downstream evaluations: for oloM-parameter models corresponding Figure A model trained for 80 tokens-per per-parameter with linear D2Z has equivalent downstream loss to the same model trained for 200 TPP with 10x decay.
Down-
MMLU (Avg.)
Reading Comp-
Truthfulness Bias
Commonsense Reasoning
stream (Avg.)
Open-
Lamb-
Lamb-
Truth-
Wino-
Hella-
CrowS- Puirs
grande
swag
Book
QA
OpenAl
ada
Stand.
SIQA PIQA
Are-e
RACE
QA
ful
10x 8STPP
23.6%
52.2%
43.9%
31.4%
46.7%
36.7%
32.8% 68.7%
47.6%
32.3%
39.8%
60.7%
43.05%
D
23.5%
53.4%
44.6%
31.6%
46.9%
37.3%
33.2%
68.8%
48.8%
33.4%
40.2%
60.8%
43.54%
10x @ 200TPP 23.3%
53.4%
46.6%
31.2%
46.2%
38.8%
32.2%
68.8%
47.9%
34.4%
38.4%
60.4%
43.46%
DIz@ 200TPP 24.7%
54.5%
48.2%
32.4%
50.0%
42.6%
32.9%
70.1%
50.4%
32.6%
38.9%
62.5%
45.00%
64. We trained this proxy model on 800M tokens with a batch size of 256 and context length 2048, using 10x decay, We randomly sampled 350 configurations of base learning rates, base initialization standard deviation, and embedding and output logits scaling factors, and used the top-performing values as our tuned HPs (Table 4).
B DERIVATIONS
B.1 DERIVATION OF EQUATION (6)
The initial coefficient is Ct,1 IT - az). Clearly, if a, = a is constant, we have Ct, =
(1 - a)(-1, Otherwise, given a, is small, we can use first-order Taylor expansion, e5%, 2 1 0zz and therefore:
II a
مد -a;
e- Ej=20
Assuming we only know the average value, ā [-1 Ej= aj we have:
1
4,1 R e a(t-1)
Given ā is also small, we reverse the earlier Taylor expansion, but now e- = ā, and:
Ct 6)1
That is, with a small time-varying smoothing parameter, the initial coefficient in an EMA, G, also decreases exponentially with the number of steps.
C ADDITIONAL EXPERIMENTAL RESULTS
In this section, we include some additional results to support the findings in the main paper. All validation losses reported in this section from models trained with Linear decay.
C.I DOWNSTREAM EVALUATIONS
Finding Trends in validation loss also hold for downstream evaluation.
17

---

Published as a conference paper at ICLR 2025
3.0
Constant
2.9
10x
Validation loss
D2Z
2.7
2.6
10-1
10m3
10
LR, F
Figure 10: Standard parameterization results: Validation loss for different LR and decay combi- nations, across peak LRs, for -0rammaeee models trained with the standard parameterization.
111M
610M
1.7B
Constant
3.6
10x
D2Z
Validation
3.4
2.8
2.5
3.2
2.4
2.3
3.0
10
10-22
10
Max LR. f
Max LR,
Max LR, A
Figure 11: Model size comparison: Validation losses 20 TPP. Across all model sizes, Linear- D2Z outperforms Linear- 10x. Note: Missing high-LR values in all plots correspond to failed train- ing runs due to NaN instabilities.
Table presents a variety of downstream evaluations of the four models presented in Figure 1. Differences between the models here are largely consistent with the differences in training and val- idation loss, showing that D2Z meaningful not just for the autoregressive training objective, but for real-world applications.
C.2 STANDARD PARAMETERIZATION
Finding Similar trends hold for the standard parameterization.
Figure 10 presents results for a 610M-parameter model trained with the standard parameterization. Here is therefore not ypP-corretted base LR, but rather LR that we swept directly for this model scale. Results are obviously quite similar to results using μP, suggesting the benefits of D2Z are not uP-specific. Further results using the standard parameterization. but for NanoGPT models, are in Appendix C010 below.
C.3 MODEL SIZES
Finding 8:
Improvement of D2Z over 10x is greater at 1.7B scale.
Figure presents results across 11 610M, and L.7B model sizes, all trained to 20 TPP. Note the absence of results for the highest LR setting the 1.7B-scale; at the very highest LR, numerical instabilities led to failed training runs. Otherwise, results are fairly similar across model sizes. At the proxy-tuned peak LR, the gap between D2Z and 10x is 0.81%, 16, 0.67%, and 1.56%, the 11 111M, 610M, and 1.7B scales, respectively. We further investigate the issue whether the gap between D27. and x0x varies with model size part of our scaling law experiments below (Appendix C.1
18

---

Published as a conference paper at ICLR 2025
TPP
TPP
TPP
3.4
3.0
Constant
10x
2.7
3.3
D2Z
Validation
3.2
2.6
2.5
2.6
3.0
- Max LR,
Max 17
Max LR, 8
Figure 12: Comparing decay schedules across TPP (610M scale): As TPP increases, Linear-D22 -D2 begins to outperform Linear-10x, especially at the proxy-tuned peak LR (red lines). The optimal LR also shifts significantly lower for Constant, somewhat lower for 10x, and hardly all for D2Z. Compare to Figure 5 for 111M models.
B=63
B=126
B=252
R=504
B=1008
2.8
Validation loss
10x
D2Z
2.7
2.6
[m 10-2
10-2
10 Max LR, !
10-2
10-2 Max LR,
Max LR, 8
Max LR, 9
Max LR, 9
Figure 13: Comparison across batch sizes (iso-TPP) (610M): Different batch sizes, B, but all trained to 20 TPP. As batch size decreases, the relative gain of D2Z over 10x increases. Same data as in Figure?
C.4 TPP
Finding Improvement of D2Z over 10x grows with TPP (610M scale).
As we vary TPP, consistently see increasing gains with D2Z. Here plot the results for the 610M-scale models in Figure 12, as a counterpart to main Figure 5.
C.5
BATCH SIZES
Batch size setup Default batch sizes 192 and 504 (Table 2) were selected for 111M 610M scales. These specific values were based
internal scaling law for optimal batch size function of compute FLOPs (similar to those in Bi et al. (2024) and Porian et al. (2024)). Batch size of 504 was then re-used for 1.7B training (later testing confirmed this be good setting at 1.7B for 20 TPP training). In our experiments varying batch size, we swept B by factors of 2 around these initial settings.
C.5.1 ISO-TPP TPP BATCH SIZE EXPERIMENTS
Finding 10: Improvement of D27. over 10x grows as batch size decreases (at constant TPP).
In main paper Figure 7, we presented results where all models train for the same number of total tokens (20 TPP), while only the batch size, B, varies. Figure 13 presents the same data, but with each B separated into a separate subplot; this lets us better observe how the differences between D2Z and 10x evolve as B changes. We clearly that for smaller B, gradient noise increases, the differences between D2Z and 10x also increase. This again demonstrates that LR decay is beneficial as noise reduction mechanism.
19

---

Published as a conference paper at ICLR 2025
B
B--26
8-252
504
Validation loss
3.0
10x D2Z
mmt 10-2
10-2
10-2 Max LR, 1
10-2 Max LR, 9
10-2 Max LR, 9
Max LR,
Max LR,
Figure 14: Comparison across batch sizes (iso-step) (610M): Different batch sizes, B, but all trained for 11752 steps (not iso-FLOP, smaller batches fewer TPP). In contrast, Figure 7 and appendix Figure 13 instead keep TPP constant. D2Z remains superior in the iso-Step context. More interestingly, here the optimal LRR actually increases batches and, correspondingly, the datasets increase in size.
We now discuss the related observation that the optimal LR decreases B decreases. Note the EMA perspective of Wang & Aitchison (2024) would predict linear scaling of i with B. This is because the total number of iterations scales with 1/B. Therefore, in order to keep the timescale constant (as a fraction of the total iterations), we would need to scale 7j (or X) proportional to any change in B. Although we do roughly linear scaling in the case of 10x models, the optimal is more stable with D27. our conceptual model, the purpose of the timescale to optimize variance reduction. But note that lowering 7 also has an effect it seems paying the cost of lower bias reduction is worth the benefits in variance reduction. Since using D2Z already enables better variance reduction, it seems, for D2Z. lowering i) has less value in the trade-off, so optimal f lowers to a lesser extent.
bias reduction. For 10x, as B changes,
C.5.2 ISO-STEP BATCH SIZE EXPERIMENTS
Finding 11: Optimal LR increases with batch size when training over the same number of batches/steps.
In the above batch size tests, we kept the dataset size constant (iso-TPP). Now consider the following two mechanisms for increasing the dataset size (TPP) of training run:
1. Fix the batch size, but increase the number of steps.
2. Fix the number of steps, but increase the batch size.
Approach (I) was taken in experiments increasing TPP. In that case, we found optimal LR to decrease, as a mechanism for coping with greater gradient noise, which grows with TPP. We now discuss Approach (2): fixing the number of steps, but increasing the batch size (iso-step).
In Figure 14, we keep the number of steps constant to 11752, so each model will see the same total number of batches, and increase by factors of two.* Since number of steps does not change, Wang & Aitchison (2024) do not prescribe any adjustments to the timescale, and hence adjustment is needed to LR (or weight decay). But clearly see that optimal LR increases as B increases (like a rolling wheel, LR bowls rotate clockwise as we move to the right). This is evidently because increasing B decreases gradient noise; with less noise, we can afford larger LR throughout training. More precisely, increasing î decreases the number of Adam W weight updates that are combined. Therefore, an increase î) can be viewed as way to reduce the effective number of batches that are combined, basically compensating for the fact per-step B is larger.
As B increases, gradient noise decreases, and we would expect benefits of LR decay also dimin- ish; differences between 10x and D2Z do seem to decrease. This aligns with prior work, e.g., You et al. (2019a) compared SGD to full-batch GD (with WideResNets for image classification), and showed that when gradient noise is eliminated by using full GD, LR decay is no longer beneficial.
'Note, for purposes of scale, results at B=504 are the same in Figure 13 and Figure 14; the latter just has larger range the y-axis.
20

---

Published as a conference paper at ICLR 2025
Constant
10x
D2Z
2.75
2-200--3
Validation loss
2.70
i -.0ee-3
8-8.1e-
2.65
1=1166-2
î 3.2c-2
6.55e-
2.60
2.55
10-5
10-3 3
10-1
10-5
10
10
10 3
10-1
11
1/1
nd
Figure 15: Weight decay results, grouped by a=nÀ varies (610M, 20 TPP): Validation losses for different settings of decay schedule (separated by subplot), maximum learning rate (7), marked by color), and weight decay (A, corresponding points the LR curves). Loss is plotted relative to a=nx. The optimal loss corresponds to high (but not too high) fj. As the smoothing a increases, Constant models suffer, 10x models see marginal gains, while D2Z models benefit significantly.
Constant
10x
D2Z
3.25
--2.0e-3
3.20
j=4.0e-3
Validation loss
7-1.e3
3.15
1=1.6-2
3.10
f=3.2e-2
6.5ee-2
3.05
3.00
10-5
10
10-1
10
10-3
10
10-5
10-3 n
10
Figure 16: Weight decay results, grouped by a=1iÀ varies (11 20 TPP): Validation losses for different settings of decay schedule (separated by subplot), maximum learning rate (7), marked by color), and weight decay (A. corresponding to points on the LR curves). Loss is plotted relative to a=nA. The optimal loss is obtained when ij is high, but not too high, typically around 1 1.6e-02 to 3.2e-02. As Figure 15, only D2Z models further improve increases.
C.6 WEIGHT DECAY
C.6.1
WEIGHT DECAY: RESULTS
Finding 12: Benefits of weight decay are observed primarily when using LR D27.
Figure analyzes the interaction between weight decay (WD) and LR for 610M models. Here we plot with the x-axis set to the (peak) smoothing parameter, a=1λ (on log scale), and y-axis set to validation loss. Note that for given decay ratio and fixed a, EMA update coefficients (Section 3) are identical (i.e., specific values of and À do not affect coefficients, only their product). Thus supports the extended EMA perspective to note that regardless peak LR, ň. models tend to reach lowest loss at around the same (Figure 15); that is. the same EMA timescale. 1/na, is optimal across different LRs. However, to reach optimal loss, models require î to be high enough, but not high at 022 and above, consistently encounter instabilities in training.
21

---

Published as a conference paper at ICLR 2025
Constant
10x
D2Z
2.75
À =1.0e-3
Validation loss
2.70
À -10ee-2
=1.0e-1
2.65
=1.ee0
2.60
2.55
10-5
10-3
10-1
10
10-3
10-1
10-6
10-3
10
n
n
1/
Figure 17: Weight decay results, grouped by À a=nd varies (610M. 20 TPP): Subset of the data in Figure 15, but now curves trace points with same weight decay À (in color) and L.R varies across each curve. Only D27. models significantly improve increase weight decay. D22. models also less sensitive to choice of X.
Constant
10x
D2Z
3.25
=1.0e-3
3.20
=10e-2
Validation loss
= 110e-1
3.15
=10ee0
3.10
3.05
3.00
10-3
10-3
10-1
10-8
10-3
10-1
10-5
10-3
10
a 7d
a n
a= = mn
Figure 18: Weight decay results, grouped by as a=πd varies (11IM, 20 TPP): Same data as in Figure 16, but now curves trace points with same weight decay À (in color) and LR 1 varies across each curve. Only D2Z models significantly improve as we increase weight decay. D2Z models are also less sensitive choice of X. See Figure 17 for 610M-mode results.
When increasing À with given peak LR (moving left-to- along curves), note behavior differs depending on L.R schedule: Constant models perform worse, Linear-10x sees marginal gains, and Linear-D27 benefits significantly. Results are similar for 111M models (Figure 16).
Figures 17 and 18 depict the same data, but with points now grouped by X.5 For higher or lower X, cannot reach the optimal timescale without taking out of its comfort zone. Therefore, with our 20 TPP models, >=0.1 emerges as good default setting because, in synergy with comfortable LR, A=0.1 corresponds
appropriate EMA timescale.
Similar to the findings Andriushchenko et al. (2023), we also confirm weight decay does not act like traditional regularizer here: beneficial settings of weight decay always improve both validation and training loss.
C.6.2 WEIGHT DECAY: FURTHER DISCUSSION
Learning rates reduce bias, weight decay best controls variance Recall that Hypothesis 4 made the observation that while LR 1) and weight decay À can both equally change the update coefficients,
NNote Figure 17 excludes some points from Figure 1.5: ically for -==l.6e-00 02. and we leave those off this figure order reduce clutter.
Figure 15 includes some additional À settings specif-
22

---

Published as a conference paper at ICLR 2025
only 77 is effective in moving the model away from initial conditions (due the 1/2 factor applied to the updates). This is confirmed by Figures 15 and 16: only when 1) is sufficiently high do models achieve optimal loss. Given such high-enough 7), the primary benefit of weight decay evidently to adjust o=riÀ (and hence the timescale l/x), to a setting that synergizes best with the decay schedule (for 610M models, around 4e-04). Weight decay is specifically useful in this context because training instabilities prevent reaching this optimal purely through increasing тj. However, given 1) and À both affect the timescale, a=nλ can be viewed as the effective or intrinsic LR (Li et al., 2020; Wang & Aitchison, 2024), but only later in training.
Since weight decay does not impact bias reduction, it follows that, when the timescale needs to be adjusted at high TPP due to the greater gradient noise, it should be more effective to lower WD than to lower LR, exactly as prescribed by Wang & Aitchison (2024). In ongoing concurrent work, we do observe lower loss from tuning WD compared to tuning LR. This stands in contrast to other recent work that recommends purely decreases in LR for high-TPP training (Shen et al., 2024b; Bjorek et al., 2024).
Why weight decay is only beneficial when using LR decay Our observation that weight decay only beneficial when using LR (decay) schedule aligns with both Loshchilov & Hutter (2017) and Andriushchenko et al. (2023) who also observed no benefit from WD when using constant LRs. In contrast, Aleph Alpha (2024) recently reported A=0.1 to improve over A=0.0 in LLM training; notably, they train with D2Z. We believe the reason for these observations can also be explained by the bias/variance trade-off, and by the fact, noted above, that weight decay primarily plays a role later in training. Since its LR is fixed, Constant negotiates the bias/varianc trade-off by using a L.R that is suboptimally low early and suboptimally high later on (refer back to Figure 3 in the main paper for visualization). Increasing À raises the already-too-high (late) even higher, further decreasing the timescale 1/a and hurting the model. In contrast, when using a LR schedule, weight decay can play its beneficial variance reduction role, as discussed above.
is also worth observing that if one is using D2Z with a LR that is suboptimally low (e.g., increasing it any further caused numerical stability issues, as observed with NanoGPT in Ap- pendix 1O),), then increasing weight decay can be very beneficial, For example, note the large improvements in loss as weight decay is increased for the rj=4.0e-3 and 1-220e-3 curves in Fig- ure 15.
Stability of D2Z loss with high A While Constant and 10x are much worse at higher values, D27. performs reasonably well even at x=1.0 (particularly at the 610M scale). Since increasing effectively decreases the timescale 1/a over which weight updates are integrated, it increases the variance over weight updates. Therefore, setting À too high hurts all models, but since D2Z has lower y later in training, it somewhat counterbalances the increase in X in the product a=nλ. Another view of this is that as fewer updates combined, noise increases; D2Z is evidently best at mitigating the negative effects of such noise through its maximal LR decay.
C.7 STEP DECAY
Finding 13: Step decay helps Constant, but does not provide benefits if already decaying the LR.
In Figure 19, we present results investigating the impact of Step decay on model training. Here, for 11 111M-parameter models, Step decay can improve the loss versus keeping the LR constant (left panel), but the resulting losses are still much worse than those obtained with D27. or 10x decay. We also tried applying Step decay to LR that had been following a 10x (middle panel) or D2Z trajectory (right panel); this approach always led to inferior results.
While it is likely possible to improve the quality of Step decay by tuning the positioning of the drop, we hypothesize that these efforts will not surpass D2Z, since dropping the LR will fundamentally always result in an over-emphasis of earlier updates, shown in Figure 2 and Figure 27. Moreover, introducing additional tunable hyperparameters (i.e., when and how much to decay) is a further drawback of the Step schedule.
23

---

Published as a conference paper at ICLR 2025
Constant
10x
D2Z
With Step
With Step
With Step Step
3.6
Validation loss
No Step
Step
3.4
3.2
3.0
10-2
10-2
10-2
Max LR, !
Max LR,
Max LR, fj
Figure 19: Impact of applying Step decay (111M, 20 TPP): Validation loss for different LR and decay combinations with and without Step decay. Step decay is applied after 90% of training, stepping to a value equal to 0.1% of the maximum LR. Dropping the LR in this manner helps Constant, although it is still below the level 10x and D2Z. Adding stepping to a 10x or D2Z schedule is not beneficial.
3.6
Constant
3x
Validation loss
10x
3.4
D2Z
3.2
10-3
10-2
LR, i
Figure 20: Comparison of decay schedules when using 93.75% weight sparsity (610M); Valida- tion losses for different LR and decay combinations, for a model trained with SuPar, with 93.75% unstructured weight sparsity. As the decay rate increases, loss improves, while the optimal peak LR is much more stable when using D2Z.
C.8 WEIGHT SPARSITY
Finding 14: Relative improvement of D2Z over 10x increases with weight sparsity.
Weight sparsity setup In this section, investigate the role of Linear-D27. in the context of models trained with unstructured weight sparsity, a promising direction for improving the effi- ciency of large neural networks (Hoefler al., 2021). We parameterize with μP's sparse extension, SyPar (Dey et al., 2024b). SyPar allows us to use the same μP hyperparameters as with dense mod- els, except we must now apply corrections due both model scaling (i.e., P. Section 2.2) and layer sparsity. For these experiments, we sparsified all non-embedding layers of our dense models by randomly fixing certain weights to zero for the duration training. models for 11752 total steps using batch size of 504, i.e., the same amount of training data used for training the corresponding A10M-paramens dense models to 20 TPP (Tables and 3).
trained
At 93.75% sparsity (1/16 density), the optimal D2Z model improves by 1.64% over the optimal 10x model, with a clear trend of optimal LRs shifting lower and loss becoming worse as we go from 10x to 3x Constant decay (Figure 20).
At the proxy-tuned peak LR of 1.6e-02, D2Z is 2.15% better than 10x. We also trained 93.75% sparse models with variety of other decay ratios between 10x and D2Z, and present these results in Figures and 22. Here we see a largely linear decrease loss with a linear increase in the decay ratio (i.e., linear decrease in the minimum LR). These are encouraging findings in the sense that
24

---

Published as a conference paper at ICLR 2025
3.08
3.07
3.08
3.00
3.07
Valdrien kow
3.05
3.04
Valdation
3
3.04
3.00
3.43
3.00
3.02
0.0000
0,0022 2.0004 0,0006 DUOOSS 0,0010 D.0012 0,0014 BU00L6
Mnnmum leorning
0,0000
0.0022 0,0001 0.0005 0.0008 0.0010 0.0122 0.0014 0.0016
Minimum ral6,
Figure 21: Decay ratio sweep, part (610M): Validation loss for models trained with SpPar, with 93. 75% unstructured weight sparsity, as we vary the LR decay ratio. The x-axis plots the minimum LR, i.e., the LR at the final training step. As the ratio increases from 10x (right-most point), where the minimum LR is 10% of the proxy-tuned rate, to 0% (D2Z, left-most point), validation loss decreases roughly linearly.
Figure 22: Decay ratio sweep, part 2 (610M): Validation loss for models trained with SuPar, with 93.75% unstructured weight sparsity, for LR decay ratio of 10x, 100x, 1000 and D2Z (using a different sparsity pattern than part 1). As in part 1, validation loss decreases roughly lin- early with the minimum LR.
3.4
D22
Validation lass
2.8
2.6
1/4
1/8
1/16
1/64 1/128
Weight density
Figure 23: Comparison of decay schedules when using varying weight sparsity (610M): Val- idation losses for models trained with SyPar, as unstructured weight sparsity increases (density decreases). All models trained with the same number of training tokens (11752 steps, equivalent to 20 TPP for the fully-dense models). As sparsity increases, the number of trainable parameters decreases, and thus tokens- pettrs----ooo--eeo-parrrraaa (and the benefit of D2Z) increases.
D2Z. can seemingly be used directly on a range of problems, without having to worry about tuning a cabblem-specifin LR decay ratio (e.g., 50x or 100x).
In Figure 23, we investigate the gap between D2Z and 10X at the proxy-tuned peak LR across different sparsity levels. Note that increasing sparsity effectively leads to corresponding decrease in the number of trainable parameters. Since we use a fixed number of training tokens in each case, as the number of parameters decreases, the number of Mikens-per-parameter (TPP) increases. In this sense, note the relative differences between D2Z and 10x are consistent with our results in Figure - as TPP (and gradient noise) increases, D2Z performs relatively better.
C.9 ERROR BARS
Finding 15: Run-to-run variation in results is low.
25

---

Published as a conference paper at ICLR 2025
Constant
3.8
10x
Validation loss
D2Z
3.4
32
3.0
10-2 Max LR, f
Figure 24: Stability across random seeds (11IM, 20 TPP): Validation loss variance for different maximum LRs at different decay ratios. Each point corresponds to the mean validation loss over 5 separate training runs with different random seeds; the enror bars give the standard deviation. Beyond the instabilities at high learning rates, run-to- variance is remarkably low at this scale.
Taken as a whole, our results are remarkably stable: empirical results for different model scales, training durations. batch sizes, weight decays, and weight sparsity settings largely behave pre- dicted by theory. Since all validation runs are performed on 1.IB tokens, any significant run-to-run variance must arise during training. To quantify this variance, repeated 111M-model 20 TPP training four additional times, resulting in total validation loss results for each original training run. Figure 24 confirms that run-to-run variance is remarkably low. The only significant variance arises in Constant results at the highest learning rate. Here, the training loss sometimes spikes at various points during training, and the final validation loss can be significantly higher for models that cannot recover sufficiently following the spike.
C. NANOGPT EXPERIMENTS
Finding 16: Improvement of D27. over 10x persists in NanoGPT models when numerical issues are solved.
We also compared 10x versus D2Z by training NanoGPT models using the NanoGPT codebase (Karpathy, 2024). NanoGPT uses the standard parameterization. We configured these models to be largely similar our 11rM-parreeter models, also using weight decay of 0.1, context length of 2048, and the GPT-2 vocab size of 50257. Key differences here are that we do not include bias weights, and we trained the WeenWebText dataset (Gokaslan & Cohen, 2019). Experiments run on Nvidia A10 GPUs.
As mentioned in Section 5, initially tested NanoGPT in balo precision. Here, at the default NanoGPT learning rate of 6c-4, 10x performed slightly better than D2Z. As we pushed the LR 50% higher (to 9c-4), both 10x and D2Z had higher loss, and by 1.2e-3, the loss from 10x doubled.
We suspected that numerical issues may be causing the instabilities, and repeated our experiments in fl 32. At this precision, we were able to successfully increase the LR, by factors of two, up times the default. At these levels, we see the familiar gains of D2Z over 10x (Figure 25). We note there is nothing fundamentally limiting about float1 6 precision itself indeed all our main experiments were done using this precision. Rather, faaat1 is simply problematic at a high learning rate in the NanoGPT codebase (sce Gray et al. (2024, Appendix C.2) for potential root cause of this instability).
These experiments demonstrate that comparison between D2Z and 10x may serve as a kind of diagnostic of whether a model is being trained at optimal peak LRs: if a 100M+ model trained to 20 TPP does not see roughly 1% gains from using D2Z, it is likely the LR is not high enough. In order to raise the LR further, efforts to stabilize the model, perhaps including μP or other tech- niques (Wortsman et al., 2023), may be warranted.
26

---

Published as a conference paper at ICLR 2025
3:276
Validation loss
3.250
3.225
3.200
3.175
10x
3.150
D2Z
10-3
10-2
LR, 9
Figure 25: NanoGPT results (111M, 20 TPP): Validation loss for different LR and decay combi- nations, for a NanoGPT model. With floatl6 precision, we were not able to train above 6e-04 without instabilities (first point in eurves). Moreover, at 7j 6e-04, 10x performed better than D2Z. After switching to float 32, we were able to train at higher r) values, where D2Z demonstrates its familiar superiority over 10x.
Table Model architecture and batch sizes for scaling law experiments.
Model
vocab. size
dmodel
Talayors
dhead
din
batch size
272M
84992
1024
14
64
2912
128
653M
84992
1536
18
128
4096
160
1.39B
84992
2048
24
128
5472
208
2.75B
84992
2560
32
128
6832
256
C.11 SCALING LAW EXPERIMENTS
Improvement of D2Z over 10x grows across model scales, for different model and training setup.
Finding 17:
Encouraged by the results of D2Z at smaller scales, we began testing D2Z in some of our frontier model efforts. Frontier models do not provide scope for hyperparameter tuning at scale. Thus becomes important to derive scaling laws forecast loss at larger scales, based on the loss with sequence smaller models.
For this set of experiments, tested Llama-style (Touvron al., 2023a) architectures, except using LayerNorm instead of RMSNorm, and multi-head attention instead of group-query attention. We μP here well, and batch size scaling law to determine optimal batch size for each model scale. These models also use ALiBi embeddings (Press et al., 2022) and SwiGLU (Shazeer, 2020). Here, the context length is 8192 tokens, and use tied embeddings. We also re-tuned pP proxy-model hyperparameters.
We used bilingual data mix of English, Arabic and source code samples mixed in 2:1:0.4 mix ratio. English data is from the Pile (Gao et al., 2020), Arabic uses proprietary dataset (Sengupta al., 2023), and the source code comes from the GitHub portion the Pile.
To derive scaling laws to compare D2Z and 10x models, we used the power law functional form y cx.", where is the pre- pre-training FLOPs, y the loss the Pile validation set, and c and are parameters to be fit. To fit these parameters, we transformed our Loss-to-Fl -o-FF-FPPPP equation, y/ cr" logarithmic form, log(y) = log(c) + mlog(x)) (x), and the slope and intercept parameters of this line using least-squares linear regression. We trained models four sizes to compute-optimal 20 TPP (Table 6), and computed total FLOPs spent as well as validation loss on the Pile. We then fit the power law free parameters to obtain our scaling laws.
Encouragingly, here we find the scaling law slope of D2Z roughly 2.5% better than 10x decay (Figure 26). This translates similar to our earlier results at 1.7B scale. Projecting our scaling law to 70B model trained to a compute optimal 20 TPP. D2Z would achieve roughly 2% loss improvement over 10x decay.
improvement of roughly 1% at 1.3B and 2.7B scales, broadly
27

---

Published as a conference paper at ICLR 2025
2 10"
272M
10c L(C) 36, 30-MS
2.75B
1.2%
D2Z: L(C) 38.6C
2.5 10"
Relative loss improvement (%)
Validation Loss, L(C)
1.0%
0.8%
233 10"
0.6%
653M
1.396
x 10
272M
1:398
0.4%
x.1x 10"
2,758
0.2%
2x 10"
0.0%
10 Compute FLOPS, C
10
10
10"1
Compute FLOPS, C'
Figure 26: Scaling law comparison of decay schedules: Loss-to-FLOPS scaling law fits for models trained to 20 TPP (8192 context lengths, Llama architecture, Pile validation data). The power law fit for D2Z models has steeper slope than the scaling law for 10x models, implying a growing performance gap for larger model sizes.
D ADAMW AS CONVEX COMBINATION OF WEIGHT UPDATES, DRIVEN BY LR SCHEDULE: FURTHER DETAILS
We have begun using plots of the weight update coefficients internally to proactively analyze LR schedules. These plots give insight into the timescale over which training data is integrated into the model parameters, and can be obtained without actually performing any training. Moreover, can also use the extended EMA perspective in order design novel LR schedules that have desirable blends of weight updates. Optimizing these coefficients thus provides dual to optimizing the LR decay function, which is why refer these coefficients the dual coefficients the LR schedule. In this section, we first describe the design of one such LR schedule, and then we follow with example plots for this and other schedules.
D.1 TRULY SCHEDULE-FREE LR SCHEDULES
Recall that, for moving average with time-varying smoothing, (1 a1) + OgTt, we use Ct2 (II' (II) as)) a to denote the dual coefficients, i.e., the contribution of input æ; output Yi at time t. Constant schedules, or schedules with a long constant phase such WSD, are not truly "schedule-free" because their peak LR setting affects the (1 aş) terms in the dual. Different LRs will correspond different effective timescales over updates, and thus different LRs will be optimal for different training durations."
In contrast, we now derive a schedule such that coefficients are always weighted equally, at every training step. First, it can be shown that:
Cri+1
0:-1
Cτ.
(1 0:+1)a;
For the coefficients to be uniform at any step t, we require this ratio to Unity of this ratio
implies that smoothing evolves 0;+1 [T+m).) Assuming a fixed weight decay (so as m))),
coefficients will be equal if the LR evolves:
(8)
(1 + 7;
TThe different shapes of the coefficients for different IRs can he seen in Figure 28a for Constant and Figure 29d for WSD.
28

---

Published as a conference paper at ICLR 2025
10-4
10
10000
10000
10"4
100
Output step (3)
20000
Ourput stop (4
20000
(4p0s Boj P3
10
100-
30000
(gezs
30000
10
10
40000
60000
50000
10000
20000
30000
20000
50000
10000
20000
30000
400000
50000
Weight update
Weight update (1)
(a) Linear
(b) Cosine
10
10-0
10
20000
10-7
20000
Outpot step
Eopoo 3o)
Output stop
30000
10-
30000
10
10
40000
40000
10-1a
50000
10-13
50000
10-1a
50000
40000)
50000
Weight update (+)
Wiight update (5)
(c) Step
(d) Constant
Figure 27: EMA perspective: weight update contribution across steps: Convex combination of weight updates, with color indicating value of combination coefficients Cg each Ct gives the contribution of the ith update (across x-axis) to model weights 0, across steps t (y-axis). Note that LR schedules and coefficients corresponding the final step only were presented earlier in Figure (except for Constant). Coefficients correspond to settings for 11 111M-paraeetrr models trained to 200 TPP: 155660, 7f=1.6e-2,, p=1/3, X=0.1.
We can initialize 77o to some value and iterate Equation (8) to generate the full L.R schedule. We call this the Rational schedule since it is both rational expression of f) and very reasonable approach: regardless of how long we train, all weight updates contribute equally. At each step we effectively decrease all prior coefficients such that they now equal the coefficient of the current update, azg-
With >=1 and no warmup, this schedule evolves like 1/n. However, in order to distance ourselves from the initial conditions, we can warmup the LR the usual manner the desired peak LR, then switch on rational decay, ensuring equal coefficients going forward. The full schedule is depicted Figure 28d. However, such a schedule will almost surely not work as effectively as Linear D27. for fixed-duration training, since it lacks the cooldown phase where gradient noise is minimized. Indeed, 1/2 has performed relatively poorly in prior studies (Ge al., 2019; Defazio et al., 2023). However, introduce it here promising schedule for continuous pre-training."
D.2 LR CURVES AND DUAL COEFFICIENTS
In this section, we provide some extra figures that were referenced in the main paper. Figure 27 shows the dual coefficients every step of training, using color to indicate the coefficient value (log-scale). Every horizontal row/step of Figure 27 reflects the coefficients at that step, essentially providing version of Figure but each step. Figure 28 provides the LR schedules and dual coefficients for some of the schedules discussed in the main paper, as well as our proposed Rational
In fact, Vo decay has long been regarded optimal for strongly-convex loss (Robbins & Monro, 1951), and optimal in other contexts when combined with averaging (Defazio et al., 2023). Since averaging is an alternative cooldown (Hägele al., 2024), Rational plus averaging is compelling schedule-free direction.
29

---

Published as a conference paper at ICLR 2025
1.0e+00
1-2.0e00
-4.00
(ou))
LR Schddll (t)
6.0w03
-8.1e-33
1:0e-00
-322-00
1.0g
1.e-066
0-
1.0-08
4000
6000
10000
12000
4000
8000
10000
12000
Step (f)
Weight update (4)
Constant
0.0++00
1-2.0x 03
-4.0x-03
6.0e-03
-8.1e 03
LR Schodln
=3 Re-02
8.0e 03
6.5e
1.0e-04
1.0e:06
1.0w088
4000
12000
4000
8000
12000
Step (+)
Weight update (1
(b) Linear D2Z
2.0e033
Warmu-100 100
1.0++00
Warmup 300
LB Schaliis (m))
1.5e-03
Wammup-
Warmup=30 3000
1.0e-00
1.0e00
Weight updatte
1.0c-04
5.0m04
1.0w06
1.0e00
2000
1000
6000
5000
10000
12000
2000
4000
8000
8000
100000
12000
Step
Weight
(c) InvSqrt
1.0c+00
8.0w03
03
-4.0--m 05
6.0c
-8.1
1.0c-02-
6--02
LR Sclwbik
=3.2
-6.5c
1.0w04
2.0w03
Weight
1.0e-06
0-
1:0c000
2000
1000
6000
8000
10030
12000
2000
4000
0000
8000
10000
12000
Step (1)
Weight update (1)
(d) Rational
Figure 28: LR curves and dual coefficients for various common LR schedules well as the proposed Rational approach (Appendix D.I): Dual coefficients shown at final training step (right side) for 610M-parameted 20 TPP training (t=11752, p=l/s, >=0.1). For InvSqrt, vary the warmup and fix 76- for all curves. Note coefficient values are plotted on log-scale
30

---

Published as a conference paper at ICLR 2025
4,00
1. 1.0e-01
3.0m03
--1 Se-02
1.0m-3
Weight coel::
LR Schddul
1.0c 05
2.0m03
1.0m03
1.0e-09
0
1.0e11
10000
33000
30003
40000
20000
30000
40000
Step (1)
Weight update (+)
(a) Linear D27.
4.0w00
- De
3.0w00
9 -1.0
1.0e-03
7--3 2.2022
LR Schokln
1.0u-05
2.0e
1.0c-07
1.0c
1.0a-11
200000
40000
30000
Step (1)
Weigt update (+)
(b) Linear 10x
1.0c 01
- B.2e-03
LB Schalis (m))
3.0e-03
3.2e-
6e-22
1.0c-03-
1.0m-05
2.0e-03 03
1.0c 07
1.0c-03
1.0c-11
10000
20000
30000
40000
10000
20000
30000
40000
Weight updote (+)
(c) Cyclic
5.0+00
1.0c-C1
1e-28
-1 de
-3 3.0
LR Sclebik
1.0+-05
1.0x-77
1.0+03
1.0c-09
0-
1.0s:11
10000
20000
30003
40000
10000
2000
30000
40000
Stop (1)
Weight update (5)
(d) WSD
Figure 29: LR curves and dual coefficients for standard vs. continuous schedules: Comparison of standard Linear schedules versus Cyclic and WSD. On the left are the exact LR schedules used in Figure 8 evaluations, i.e., 610M-parameter, 80 TPP training (p=l/8, >= Dual coefficients shown at at final training step 1=47008 (right side). Note coefficient values are plotted on log-scale.
31

---

Published as a conference paper at ICLR 2025
schedule, which combines all the prior weight updates (after warmup) equally at every step. Finally, Figure 29 gives the LR schedules and dual coefficients for the comparison to WSD and Cyclic in Figure 8.
is worth re-iterating that the dual coefficients can be computed separately from any actual training. They are mathematically equivalent to the LR schedule itself and simply provide a perspective how the weight updates combine to form parameters, when using the Adam optimizer. Further- more. is also worth noting the vertical bar at step in the dual coefficient plots; this bar reflects the coefficient on the initial, random weights. To some extent, this CI value can serve as an indicator of how far the model has moved from initial conditions, i.e. the extent to which the model has reduced the bias. In general, if C1 is too high (e.g., when it outweighs the sum of the other coefficients), then bias likely significantly hindering learning. Two effective ways to reduce C1 and thus reduce bias are to (1) raise the peak LR, and (2) train for more TPP. In contrast to these two methods, raising weight decay, X, can also decrease C1 but counterproducttive for reducing bias because also reduces the scale of weight updates, as noted in Hypothesis 4. Likewise, using smaller batches also reduces C1 but is likewise counterproductive the batches become too small, to the extent that gradient noise increases. However, D2Z is more robust such noise than 10x Constant decay. Further investigating the interplay of x, batch size, and peak learning rate important future work.
32


Reviews:
[{"review_id": "qnVQ87wwy9", "reviewer": "Reviewer_sx5Y", "summary": "This paper analyses the role the learning rate scheduler for LLM pre-training. The authors mainly focus on two questions:\n- can we replace the widely used cosine learning rate scheduler by a simpler linear decay?\n- can we do better than the commonly used rule to take the final lr as 1/10 of the peak lr?\nThe authors answer positively those two questions, demonstrating that a linear lr decaying to 0 works best for small models (most experiments are on 600m models, a few on 1.7B models) on the slimpajama dataset.\nThe authors conduct many ablations on the subject, including training length, sensitivity to peak lr, and role of weight decay.\nThe authors also propose a theoretical framework to try to understand this, based on unrolling EMAs with varying step-sizes.", "questions": "- what base lr is used to train the 1.7B models?\n- It is doubtful that d2z is optimal: with d2z, the last iterations have an almost zero learning rate, yielding no progress and effectively wasting these last steps. This effect is worsened for the cosine schedule, which spends more time at low lr's. Having a non-zero value for the last lr would be better: how about a decay to, e.g., 100x? I understand that this leads to another hyperparameter to tune, which is cumbersome, but this could also be discussed.", "limitations": "", "rating": 6, "confidence": 4, "soundness": 2, "presentation": 3, "contribution": 2, "strengths": "- this paper studies a very important problem, that of the training dynamics of LLMs\n- it is well-written and easy to read\n- it makes a compelling story\n- the paper convincingly demonstrates the importance of completely annealing the learning rate for small models on one dataset", "weaknesses": "- The main weakness of this work is the small scale of the experiments. For instance, the base learning rates for the small-scale models used in this paper are very high (cf table 2 and fig.10, best lr is around 3e-3), so decaying it to 10x still yields a high LR at the end of training. For bigger models, the base lr is much lower (typically 3e-4 or 1e-4 for 7B models, even lower for larger models). Therefore, the final lr when using the 10x technique is low in this case. Hence, it is very likely that the impact of d2z would be far less in these cases (for instance, in Fig. 10, there is barely any difference between d2z and 10x for these learning rates).\nIn order to address this concern properly, this paper needs to show that the benefits of d2z do not completely vanish at large scales. A similar base lr sweep to that conducted for the 600M models for the 1.7B model would benefit the paper greatly.\n\n\n- The methods section seems very vague:\n    - There are no formal statement, everything is hand-wavy\n    - The main problem with the discussion regarding the EMAs, and summarizing everything with the variables $c_{t, i}$, is that it makes it seem that the variables $x_t$ in the parameter’s update EMA (eq.3) are external variables to the problem, while clearly they depend on the parameters $y_t$ themselves. This makes all the discussion in section 3 flawed.\n\n- The impact of the dataset is another thing that is not discussed, while it will change the learning curves. This warrants at least a discussion, clarifying that the phenomenon observed here might not happen for datasets of different sources/quality.\n\n[Edit after discussion, see replies in the forum: score 3-> 6]", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "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": "XUURngYwQD", "reviewer": "Reviewer_3io8", "summary": "This manuscript empirically demonstrates that, with a fixed maximum learning rate, a simple linear decay-to-zero (D2Z) schedule often outperforms the widely-used cosine decay to 10% of the maximum, as well as other popular learning-rate schedules. It also shows that linear D2Z is robust to variations in maximum learning rate, weight decay, dataset size, and batch size.", "questions": "Subsection 3.3 has a weak connection to the overall manuscript.  Can you further clarify its relevance to the main findings about D2Z.", "limitations": "", "rating": 5, "confidence": 4, "soundness": 2, "presentation": 3, "contribution": 2, "strengths": "The manuscript presents extensive experiments that validate the effectiveness of D2Z for training GPT-like models with 111M, 617M, and 1.7B parameters. It effectively illustrates D2Z's robustness across different maximum learning rates, weight decay values, dataset sizes, and batch sizes.", "weaknesses": "- Linear D2Z appears to function more as an engineering trick than a comprehensive algorithm. While it demonstrates empirical effectiveness, the theoretical justification is somewhat unconvincing. The authors attempt to explain its success through theoretical convergence analysis (Eq. (1)) and update rules involving weight decay (Eq. (3), Eq. (4)), but a unified framework to interpret the results clearly is lacking. It seems the authors selectively apply certain perspectives to support their experimental findings. Furthermore, the convergence bound in Eq. (1) applies only to convex optimization problems with SGD, which are significantly different from the training of large language models (LLMs) using Adam. In summary, a more rigorous mathematical framework that connects the empirical results to theoretical principles is needed.\n\n- Although the manuscript presents simple scaling law experiments in Appendix B.10, the experimental details are unclear. The authors are encouraged to provide more comprehensive information on scaling law experiments for D2Z to demonstrate its applicability to larger models.\n\n- D2Z requires prior knowledge of the total number of iterations, which is not ideal for continuous pretraining scenarios. In contrast, WSD and cyclic schedules  are  more effective for continuous pretraining.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "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": "z6ulkjAU4E", "reviewer": "Reviewer_f1Pq", "summary": "The submission studies learning rate scheduling for LLM pre-training, in particular, the paper notices that, in high TPP setting, i.e. when the model is trained much longer than what scaling laws recommend, a linear decay to zero (D2Z) scheduling gives optimal performance and converges faster than the commonly used cosine decay to 0.1 max_lr. To be more specific, the method utilizes a recently proposed interpretation AdamW, which suggests that the model weights at a certain time step is an exponential moving average of past updates, with the smoothing factor determined by learning_rate * weight_decay_coefficient. Based on this observation, the authors suggest that linear decaying to zero allows for more updates to be averaged, therefore reducing the variance, which is a crucial component in the later phase of the optimization (which is the pattern for high TPP setting). \n\nThe paper then empirically verifies the claims on an **extensive** set of experiment settings, all confirming the benefit of linear D2Z: Including better validation loss in high TPP setting, less sensitivity to the choice of batch size, learning rate setting. Additional experiment results are also presented in the appendix.\n\nOverall, I find the paper well written, the presentation is very clear, and the empirical evidence convincing.", "questions": "- I really like the author's argument on why, given a fixed value of eta * lambda, modifying eta does not have the same effect as modifying lambda, since having a lambda too large would affect the bias reduction. I wonder if the authors have considered mitigating this issue by modifying the weight initialization scale: The EMA perspective suggests that the weight will end up having a magnitude ~ 1 / lambda, and if lambda is large, the relative difference (bias) between the final weights and initialization of fixed magnitude will indeed be very small, however, one can consider additionally scale the initialization by 1 / lambda, such that the bias could potentially be reduced at larger lambda. Such a strategy is also suggested by Wang & Aitchison, (2024) in Appendix A, Eq. 21 (but under fully scale invariant network assumption).\n\n- I wonder what does the author think of the reason behind Finding 4, i.e. weight decay is beneficial only used together with certain learning rate scheduling. Such behavior is also observed when training ResNet on CIFAR-10, Loshchilov & Hutter, 2017 noticed that weight decay can improve test accuracy only combined with learning rate scheduling.\n\n- Notice that the dataset scaling rule Wang & Aitchison (2024), also implied that the learning rate should increase as batch size increase, in particular, the authors suggest that the key hyperparameter is M * eta * lambda, where M = number_of_training_point / batch_size, therefore as batch_size increase, eta * lambda should also increase. But Wang & Aitchison (2024) mainly considered high-epoch settings with the dataset being repeated multiple times (e.g. standard CIFAR-10 training pipeline) so their claims may not be transferrable to LLM pre-training where the same training data point almost never show up twice.\n\n- In a very extreme and unrealistic setting, where we can afford full batch training, in which case I wonder if weight decay can be entirely removed since there will be no gradient noise at all?\n\n- I believe such interpretation can also be applied to nonadaptive optimizers, e.g. standard SGD or Sign-GD with momentum, would the authors think the claims still hold true?\n\n- (Minor) Why is the batch size chosen as multiples of 63? This seems to be a weird number.", "limitations": "", "rating": 8, "confidence": 3, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "- The paper studies an important problem: How to schedule the learning rate, a critical but under-studied issue, until recently by Defazio et al. (2023). \n\n- The paper also studies the role of weight decay (as well as its interaction with learning rate and learning rate scheduling), which is also under studied, until recently by Andriushchenko et al. (2023) and  Wang & Aitchison, (2024)\n\n- The argument, despite not being supported by strict proof, is built upon classic optimization theory (bias-variance decomposition) and simple linear algebra (EMA perspective of AdamW), so I find the claims fairly rigorous.\n\n- The authors explore a wide range of confounding factors that could affect the conclusion on \"best learning rate schedule\", such as peak learning rate, batch size, weight decay, training length, and observe similar pattern across all settings, therefore I find the conclusion very solid and would consider starting using linear D2Z in my own LLM experiments.", "weaknesses": "- I think there is room for improvement in the writing: Section 3.1 and 3.2: The paper's main argument is about the optimality of linear D2Z, but these two sections do not talk about, e.g. why cosine decay is suboptimal (which was only briefly mentioned in the caption of FIg.2) or why decay to zero is better than decay to 0.1 max_lr. Although I do understand the connection between the suboptimality of other scheduling and the reducing variance perspective, more explicit arguments can be made to improve the readability.\n\n- The cyclic learning rate scheduling (or is it just cosine learning rate decay?) is used in experiments as a baseline but not explained or discussed in, e.g. Figure. 2.\n\n- It would be nice if more details about the computational resources, model configuration (e.g. number of self attention layer, head, hidden dimension size) can be revealed.\n\n- It would be nice if the experiment code can be released.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "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:

Decision: Accept (Poster)