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
Primary area: foundation or frontier models, including LLMs
Keywords: Learning rate schedules, Large language models (LLMs), AdamW optimizer, Weight decay, Compute-optimal training

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.

Decision: accept (poster)
Reviewer scores: 6, 5, 8

Reviews and discussion:
<thread>
creation date:
1730393696533

thread_type:
official_review

<thread_object>
creator: reviewer_f1pq

<content>
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). 

The 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.

Overall, I find the paper well written, the presentation is very clear, and the empirical evidence convincing.

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). 

- 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)

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

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

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

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

- It would be nice if the experiment code can be released.

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).

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

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

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

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

- (Minor) Why is the batch size chosen as multiples of 63? This seems to be a weird number.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730423518460

thread_type:
official_review

<thread_object>
creator: reviewer_3io8

<content>
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.

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.

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

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

questions:
Subsection 3.3 has a weak connection to the overall manuscript.  Can you further clarify its relevance to the main findings about D2Z.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730765847090

thread_type:
official_review

<thread_object>
creator: reviewer_sx5y

<content>
summary:
This paper analyses the role the learning rate scheduler for LLM pre-training. The authors mainly focus on two questions:
- can we replace the widely used cosine learning rate scheduler by a simpler linear decay?
- can we do better than the commonly used rule to take the final lr as 1/10 of the peak lr?
The 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.
The authors conduct many ablations on the subject, including training length, sensitivity to peak lr, and role of weight decay.
The authors also propose a theoretical framework to try to understand this, based on unrolling EMAs with varying step-sizes.

strengths:
- this paper studies a very important problem, that of the training dynamics of LLMs
- it is well-written and easy to read
- it makes a compelling story
- 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).
In 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.


- The methods section seems very vague:
    - There are no formal statement, everything is hand-wavy
    - 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.

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

[Edit after discussion, see replies in the forum: score 3-> 6]

questions:
- what base lr is used to train the 1.7B models?
- 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.
</content>
</thread_object>

</thread>
<thread>
creation date:
1732030298002

thread_type:
public_comment

<thread_object>
creator: ~laurence_aitchison1

<content>
title:
Request to chat.

comment:
This is awesome work!  And thanks for the shout-out to the EMA interpretation of AdamW.  I was wondering if you'd like to chat?  You can look up my details in Wang and Aitchison :-).

(I am not a reviewer, AC etc. on this manuscript).
</content>
</thread_object>

</thread>
<thread>
creation date:
1732993370929

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
Improvements to Section 3 (1/2)

comment:
We thank all of the reviewers for their very helpful comments and their support for the paper's core contribution. The paper presents a large-scale, hypothesis-driven study of learning rate schedules. Reviewers all noted the experiments were extensive, effective, and convincing, and they found the paper well-written. Reviewer sx5Y initially had some concerns whether the results would generalize to larger models and datasets, but concerns were alleviated via discussion and additional large-scale experiments.

The remaining concerns center on the methods and writing in Section 3. Reviewer f1Pq urges "more explicit arguments" to improve the readability, while Reviewer 3io8 would like to see stronger theoretical foundations for the success of linear decay-to-zero. Reviewer sx5Y has concerns regarding the formal correctness of the EMA perspective from Wang & Aitchison.

These concerns ultimately all arise from suboptimal organization and writing in Section 3. We failed to make clear that the bias/variance view and extended EMA perspective serve as our *conceptual model* of LLM training, rather than as a formal mathematical theory. The difference becomes subtle because our conceptual model does use mathematics to help predict training dynamics, but it does not let us, for example, formally derive the optimal LR schedule, as has been done in related work under various assumptions. We rely on experiment to fill in theoretical gaps. The value of any conceptual model is whether it makes useful predictions, but we failed to connect our conceptual model to the specific hypotheses that we tested.

We provide here a revised version of Section 3 that addresses these issues. We greatly thank all the reviewers for helping to strengthen the paper.

___
___

# 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 extended EMA perspective. Finally, we outline the specific testable hypotheses that follow from our model.

## 3.1: AdamW as convex combination of weight updates, driven by LR schedule

*largely as in submitted paper, minus second paragraph*

## 3.2: Conceptual model: bias and variance in LLM training

Following [Andriushchenko et al. (2023, Section 4.1)](https://arxiv.org/abs/2310.04415), our main conceptual premise is that in LLM training, there is an initial training phase focused on movement from initial conditions (bias reduction), followed by a later phase bottlenecked by gradient variance. Furthermore, analogous to prior work optimizing the convex loss gap via SGD (Eq. 1), we argue the primary beneficial mechanism of LR decay in LLM training is to reduce gradient variance during later stages of training.

Per-step gradient variance is known to increase over the course of training ([McCandlish et al., 2018](https://arxiv.org/abs/1812.06162)). Very recently, [Zhang et al. (2024)](https://arxiv.org/abs/2410.21676) showed that for an LLM of a 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](https://arxiv.org/abs/1812.06162)).

In the EMA perspective, parameters at step $t$ are a convex combination of prior weight updates. The greater the update variance at each step, the greater the number of updates that should be combined in order to reduce overall variance. While the principle is the same, note here variance is reduced by combining updates *across* steps, rather than by increasing batch size at a *specific* step. Now, with a constant LR, update coefficient $c_{t, i}$ decreases exponentially in $t - i$. However, with a decaying LR, scaling down later coefficients $c_{t, i}$ via a lower $\alpha_i = \eta_i\lambda$ corresponds to scaling *up* all $c_{t, k}$ where $k < i$ (cf. Eq. 5). This has the effect of flattening the coefficients, effectively averaging over more outputs (see appendix Figure 26 for a visual contrast between *Constant* and D2Z coefficients, on a log scale). The more we decay, the more updates we average over, and the more variance is reduced.
</content>
</thread_object>

</thread>


Metareview:
metareview:
The paper studies the impact of learning rate scheduling on LLM training, focusing on a linear decay to zero schedule. It empirically demonstrates that this schedule outperforms cosine decay schedules in terms of convergence and robustness across varying hyperparameters. The work also introduces a perspective of AdamW as an exponential moving Average (EMA) perspective.


### **Strengths**
- Tackles an important practical problem
- Extensive experiments across various models and datasets demonstrate the interest in the specified learning rate schedule

### **Weaknesses**
- No theory or strong justification for the schedule
- The empirical gains are quite small (see e.g. Table 1)

Overall, the paper reports interesting empirical findings and I therefore lean toward acceptance. I encourage the authors to enhance the discussion regarding the shortcomings of the paper.