It is common in deep learning to warm up the learning rate $\eta$, often by a linear schedule between $\eta_{\text{init}} = 0$ and a predetermined target $\eta_{\text{trgt}}$. In this paper, we show through systematic experiments using SGD and Adam that the overwhelming benefit of warmup arises from allowing the network to tolerate larger $\eta_{\text{trgt}}$ {by forcing the network to more well-conditioned areas of the loss landscape}. The ability to handle larger $\eta_{\text{trgt}}$ makes hyperparameter tuning more robust while improving the final performance. We uncover different regimes of operation during the warmup period, depending on whether training starts off in a progressive sharpening or sharpness reduction phase, which in turn depends on the initialization and parameterization. Using these insights, we show how $\eta_{\text{init}}$ can be properly chosen by utilizing the loss catapult mechanism, which saves on the number of warmup steps, in some cases completely eliminating the need for warmup. We also suggest an initialization for the variance in Adam which provides benefits similar to warmup.
Paper
Similar papers
Peer review
Summary
The paper analyses different aspects of warmup in the gradient based training focusing on SGD, ADAM and their variants under two types of parameterization. It shows through mostly empirical analysis that warmup facilitates training at higher learning rates and stabilizes the training dynamics by keeping it away from (what they call) divergence boundary (that results in failure). Furthermore, it suggests several improvements for hyperparameter initialization that shorten the training process and improves generalization.
Strengths
Choosing the learning rate is critical for training large models. The paper proposes a nice analysis of the warmup procedure that agrees with some of the previous observations. And it suggests useful tips for practitioners.
Weaknesses
The figures are not quite clear, especially because the captions do not describe the figures well enough. The observations are made mostly from empirical study.
Questions
The study is done on convnets and resnets applied to images. Training LLMs is a more complex task. Would these ideas apply to attention based architectures? Or more widely, do the mechanisms described in the paper depend on the architecture?
Rating
7
Confidence
4
Soundness
3
Presentation
3
Contribution
3
Limitations
Yes
Summary
This paper studies the mechanisms of the warmup technique. The authors experimentally demonstrate that the primary benefit of warmup is its ability to enable the network to handle larger learning rates.
Strengths
Warmup is an essential trick for training modern deep neural networks, and understanding its role is a critical and open issue. This paper makes a significant contribution to this exploration. From the perspective of training stability, the author highlights that the use of warmup enables network training to utilize larger learning rate. Additionally, the author notes that different initializations correspond to different stability regimes at the beginning of training: sharpness reduction and progressive sharpening. Warmup is particularly important for maintaining training stability in the sharpness reduction regime.
Weaknesses
- The experiments are primarily conducted on ResNet models on CIFAR, where warmup is not essential. However, for Transformer-based models, warmup appears to be indispensable, particularly in applications such as language model pretraining and Vision Transformer (ViT) training. - The observations in this paper are all based on experimental results. The findings would be more convincing if they could be theoretically validated in some settings (even in toy settings).
Questions
- It is natural and insightful that extending warmup's time can tolerate the use of larger LR, but can the authors further clearly and quantitatively indicate this relationship? - Does the measure of sharpness, $\lambda_{\max}(H)$, apply to randomized algorithms? For example, for GD, the stability condition is typically $\lambda_{\max}(H)\leq 2/\eta$ ; however, for SGD, $\lambda_{\max}(H)$ may no longer be suitable sharpness measure about training stability, and it might be $||H||_{\rm F}$ [1]. As the authors show in Figure 1(d) for GD, the results fully support the author's claim; however, in Figure 11(d) for SGD, the results do not completely align with the author's claim. - A closely related work [2] should be discussed. In Section 5 and Fig 3 in [2], the authors also discussed how initialization, warmup, and SGD noise influence progressive sharpening or sharpness reduction. [1] Wu et al. The alignment property of SGD noise and how it helps select flat minima: A stability analysis. (NeurIPS 2022) [2] Ziyin et al. Loss Symmetry and Noise Equilibrium of Stochastic Gradient Descent. (2024)
Rating
7
Confidence
4
Soundness
3
Presentation
3
Contribution
3
Limitations
Please refer to ``Weaknesses''.
Summary
The authors explain the mechanisms of the warmup technique showing that with warm up the loss of NN will go to a flatter space than direct optimization. Further, based on analysis, the authors propose a new optimization algorithm called GI-Adam.
Strengths
1. The authors explain why the warm-up technique can help networks converge better. 2. With the analysis, the authors show that the initialization of Adam is not "correct". Thus, the authors proposed a new initialization of Adam called GI-Adam.
Weaknesses
1. The conclusions are from FCNs and WideResnet, which can be trained well without warm-up. Does the conclusion still hold for some "hard" models and datasets (e.g., Transformer)? 2. GI-Adam is used in [1]. [1] Zhang, Yushun, Congliang Chen, Naichen Shi, Ruoyu Sun, and Zhi-Quan Luo. "Adam can converge without any modification on update rules." Advances in neural information processing systems 35 (2022): 28386-28399.
Questions
See Weaknesses.
Rating
6
Confidence
4
Soundness
3
Presentation
3
Contribution
3
Limitations
N/A
Summary
The paper examines the learning rate warmup technique from the perspective of its influence on the evolution of loss sharpness for different optimizers (GD, SGD(-M), Adam) and network parametrizations (Maximal Update Parameterization – $\mu$P and Standard Parametrization – SP). The authors demonstrate that the warmup allows the network to tolerate larger learning rates by gradually reducing sharpness through a series of loss catapults and self-stabilizations. Experiments with different network parameterizations reveal how warmup influences training in progressive sharpening and sharpness reduction regimes and the lower importance of warmup for training $\mu$P parameterized networks. Based on the empirical analysis, the paper proposes two practical training heuristics: (1) initializing the learning rate at an estimated critical value to eliminate unnecessary warmup steps and (2) introducing the GI-Adam optimizer, which initializes Adam's second moment with a squared gradient.
Strengths
1. The paper provides solid experiments demonstrating that the gradual self-stabilization mechanism induced by the warmup is observed for different network parameterizations and optimizers. 2. The paper discusses the specifics of the warmup effect on networks in different parameterizations and confirms that networks in $\mu$P parametrization benefit less from it. 3. The paper explains why warmup may be unstable for Adam and proposes a simple heuristic on how to deal with this instability. This heuristic potentially may be useful in practice. 4. The paper includes an extensive study of warmup hyperparameters (warmup length and maximal learning rate) and suggests how to choose their optimal values. I specifically like the Persistent Catapult Warmup idea from the appendix and think it is promising. 5. The paper is clearly written and easy to follow.
Weaknesses
My main concerns are related to the level of novelty of the empirical analysis of the warmup and the effectiveness of the proposed practical modifications: 1. The novelty and significance of the empirical analysis in the first part of the paper seem limited. The study heavily relies on two previous works. Gilmer et al., 2021 (https://arxiv.org/pdf/2110.04369) investigate training instabilities from the perspective of sharpness, including a very similar analysis on how warmup decreases the sharpness and a discussion on how starting training from a flat initialization makes warmup much less important. Karla et al., 2023 (https://arxiv.org/pdf/2311.02076) demonstrate that $\mu$P and SP exhibit different natural evolutions of sharpness. The paper combines these two perspectives, however, I am not sure if this combination leads to new insights. The warmup works similarly in both progressive sharpening and sharpness reduction regimes. The paper points out the instability in Adam warmup and the lower importance of warmup for $\mu$P parametrized models (similar to the discussion in Gilmer et al., 2021), but both of these insights are not related to the different sharpness evolution regimes. 2. The authors claim that starting warmup from the critical learning rate $\eta_c$ is an effective strategy. However, this claim is not obvious, and adding empirical evaluation for it would improve the paper. For example, the experiment varying starting warmup learning rate can be added to demonstrate that starting from $\eta_c$ value results in a shorter or more stable warmup than for lower and higher values. 3. Test accuracy heatmaps are not provided for the experiments on the initial learning rate selection for warmup. Hence, it is not clear whether the proposed strategy results in high-quality solutions. Moreover, high values of $T_{\text{save}}$ are observed for the small $\eta_{\text{trgt}}$ and large $T_{\text{wrm}}$ (Figure 5b), but this configuration is clearly suboptimal since much shorter warmups work well for the low learning rates. For the most practically interesting hyperparameter regions associated with the shortest effective warmup for each learning rate, $T_{\text{save}}$ is negligible. 4. The maximum test accuracy achieved by GI-Adam appears indistinguishable from that of baseline Adam in most cases, so it is difficult to say whether the GI-Adam improves the training. Moreover, adding standard deviations to this comparison seems important due to the noisy behavior of the test accuracy. Also, GI-Adam is not the first Adam modification that increases stability at the beginning of training (see, e.g., RAdam from https://arxiv.org/pdf/1908.03265v4). A more accurate discussion of such methods and comparison with them would benefit the paper. 5. The paper lacks experiments with Transformer architecture, while it is a primary case of using Adam optimizer with warmup. Minor comments 1. Line 64: the idea that warmup is unnecessary if training is stable with the chosen learning rate is obvious and widely used in practice. 2. Lines 185-186: I would not say that $\mu$P does not benefit from warmup at all. As shown in Figure 3a, a longer warmup in the case of MSE training extends the range of converging learning rates.
Questions
I would kindly ask the authors to address the main concerns from the Weaknesses section and focus on the following questions: 1. Could you please summarize the main novel insights of the empirical analysis part of the paper compared to Gilmer et al., 2021 and Karla et al., 2023, and explain why analyzing the warmup behavior in different sharpness evolution regimes is important? 2. Could you please provide any experiments demonstrating that starting warmup from the critical learning rate $\eta_c$ is an effective strategy? 3. What is the test accuracy with optimal warmup hyperparameters for baseline Adam, Adam with $\eta_{\text{init}}=\eta_{c}$ and GI-Adam? Is there any statistically significant difference between them? 4. Is the same self-stabilization mechanism observed when training Transformers? Additional minor questions: 1. Why do you use different target learning rates for $\mu$P and SP in Figure 1? It is a bit confusing since the target learning rate and parameterization are changed between the two experiments at the same time, and it is not clear which change results in which effect. At the same time, Figure 2 uses an identical learning rate for both initializations. 2. How do you measure the loss value for initial learning rate selection and the squared gradient for GI-Adam in the stochastic variants of the algorithms? Do you use a single batch or estimate these values over multiple batches? Using a single batch may result in higher variance in the estimates, which could be undesirable. On the other hand, estimating over several batches would incur additional computational costs. There also exist several related works which the authors may find interesting: * Lobacheva et al., 2021 (https://arxiv.org/pdf/2106.15739) report an effect similar to warmup self-stabilization when training scale-invariant networks with weight decay and a constant learning rate. The decreasing weight norm increases the effective learning rate, which eventually leads to training instability and catapults. This periodic behavior allows the network to achieve flatter optima with higher test accuracy after several cycles. * A different cyclical behavior, the Slingshot effect, is observed in adaptive optimizers like Adam, as shown by Thilak et al., 2024 (https://openreview.net/pdf?id=OZbn8ULouY). This effect occurs in the terminal phase of training and involves a rapid growth of the last layer norm before catapulting, followed by an improvement in test performance. * In your experiments, optimal test performance is achieved with large learning rates close to the convergence boundary. However, Kodryan et al., 2022 (https://arxiv.org/pdf/2209.03695) show that training networks with weight decay and learning rates larger than optimal usually lead not to divergence but to a noisy stabilization of test error. Moreover, Andriushchenko et al., 2022 (https://arxiv.org/pdf/2210.05337) demonstrate that further reducing the learning rate from these stabilized solutions results in the model learning sparser features and achieving better final test performance.
Rating
6
Confidence
4
Soundness
3
Presentation
3
Contribution
3
Limitations
The authors adequately discuss the limitations of the paper in the conclusion.
Responses to minor questions by Reviewer 6XDG
Here, we provide replies to minor questions by Reviewer 6XDG. > Why do you use different target learning rates for ... The different target learning rates for $\mu$P and SP in Figure 1 are chosen deliberately because we need to satisfy two requirements: (i) $\eta_{trgt}$ cannot be so large that training fails, and (ii) $\eta_{trgt}$ needs to be large enough so that warmup has a non-trivial effect. (i) and (ii) together give distinct viable values of $\eta_{trgt}$ for SP vs. $\mu$P. In comparison, we found that Adam's learning rates are relatively stable across parameterizations, which is why we used the same learning rate for both initializations in Figure 2 of the submission. > How do you measure the loss value... For both initial learning rate selection and squared gradient for GI-Adam we use a single batch for computations. Nevertheless, we did not observe any performance degradation for commonly used batch sizes as shown in Figure 6 of the PDF attached to the global response. For the initial learning rate estimation, small errors in estimating $\eta_c$ are not expected to impact training, as small initial loss spikes have minimal impact on the overall dynamics. For most models used in practice, $\eta_{\text{max}}$ is at least $4-8$ times larger than $\eta_c$ [1] and hence small errors in estimating $\eta_c$ would still be in the catapult phase $\eta_c < \eta < \eta_{\text{max}}$. However, we do agree with the reviewer that for small enough batch sizes, this estimate can be error-prone and or incur additional computational costs. We will mention it in the limitations section. [1] Lewkowycz, A., Bahri, Y., Dyer, E., Sohl-Dickstein, J. and Gur-Ari, G., 2020. The large learning rate phase of deep learning: the catapult mechanism. arXiv preprint arXiv:2003.02218 > There also exist several related works ... We thank the reviewer for bringing these prior works to our attention. We will discuss them in the related works section.
Thank you for the detailed response and additional experimental results! I am still a bit confused about the analysis of the two warmup regimes and the importance of the difference between them. Based on the paper and rebuttal, I would summarize the main results as follows: * warmup influences training dynamics differently depending on the behavior of the sharpness at the beginning of training (progressive sharpening or sharpness reduction), * warmup is more important for the sharpness reduction regime. Could you please provide some clarifications on the following concerns regarding these results: 1. I fail to see significant differences in the warmup influence between the two regimes. In both of them, training with a high initial learning rate may lead to strong catapults in training, and warmup allows the network to experience smaller sequential catapults instead. More catapults can be observed in the progressive sharpening regime, but it is unclear to me why it is important. 2. The importance of the warmup seems to be much more related to the difference between the sharpness of the initialization and the critical sharpness for the target learning rate than to the decreasing/increasing sharpness. The warmup may be crucial for training networks with progressive sharpening if we want to use learning rates higher than the critical threshold. At the same time, warmup is unnecessary for training networks with sharpness reduction with small enough learning rates. Why does increasing or decreasing sharpness define the importance of warmup and not just the difference between the sharpness at initialization and the critical sharpness for the target learning rate? 3. In most experiments, the initial sharpness reduction quickly transitions to progressive sharpness during warmup steps, and all warmup catapults take place in the progressive sharpening regime (see Fig. 2 in the paper and Fig. 1,2 in the rebuttal). This observation makes the claim that warmup is more important for the sharpness reduction regime even more confusing.
We thank the reviewer for their questions and comments. We hope other concerns regarding the initial learning rate, GI-Adam, and Transformers have been resolved. > The importance of the warmup seems to be much more related to the difference between the sharpness of the initialization and the critical sharpness for the target learning rate than to the decreasing/increasing sharpness. The warmup may be crucial for training networks with progressive sharpening if we want to use learning rates higher than the critical threshold. At the same time, warmup is unnecessary for training networks with sharpness reduction with small enough learning rates. Why does increasing or decreasing sharpness define the importance of warmup and not just the difference between the sharpness at initialization and the critical sharpness for the target learning rate? We agree with the reviewer that the necessity of warmup is based on the initial sharpness relative to the target learning rate. In our paper, when we made the statement that sharpness reduction regimes necessitate more warmup, it was for the case where the target learning rate was large and close to the optimal value. For such fixed optimal target learning rates, we found that networks that start off in the sharpness reduction phase necessitate warmup more than networks that start off in the progressive sharpening regime. The reason for this is that empirically there is a direct correlation between the initial sharpness and whether one observes sharpness reduction or progressive sharpening. Sharpness reduction in early training implies that the initial sharpness is "large," which then implies that warmup will be more important. We will modify the wording of our paper to make this point more clear and avoid this confusion. We would like to further point out that the main utility of understanding which regime we are in is that it can provide a way of defining whether the sharpness is "large" or "small." This in turn can give a clear indication of whether there is a better choice of initialization. For example, if one starts with a given sharpness and observes sharpness reduction phenomena, then it indicates that there is naturally a flatter initialization that one can pick. Similarly, if one starts with that same sharpness but sees progressive sharpness phenomena, then it is unclear, perhaps even unlikely, that a flatter initialization can be found. This was precisely what led us to discover GI-Adam. If we had found that the usual Adam initialization has high preconditioned sharpness but starts off in a progressive sharpening phase, then it would have been less clear that there might be a better possibility. > In most experiments, the initial sharpness reduction quickly transitions to progressive sharpness during warmup steps, and all warmup catapults take place in the progressive sharpening regime (see Fig. 2 in the paper and Fig. 1,2 in the rebuttal). This observation makes the claim that warmup is more important for the sharpness reduction regime even more confusing. This argument does not take into account the convergence/failure boundary. If we start off in progressive sharpening, the network can tolerate relatively high warmup rates (i.e. small $T_\text{wrm}$) according to Fig. 4(a, b). But if we start off in the sharpness reduction regime, even if we crossover to progressive sharpening after $10-20$ steps, those first $10-20$ steps will severely limit the warmup rate: High warmup rates will cause the training to diverge/fail. Therefore this early-time training dynamics is crucial for setting a maximum speed limit on warmup. We note that it is possible that the deeper understanding developed here might lead to more sophisticated warmup schedules; one can imagine starting with a low warmup rate while the network is in the sharpness reduction phase and transitioning to a high warmup rate after some time.
> I fail to see significant differences in the warmup influence between the two regimes. In both of them, training with a high initial learning rate may lead to strong catapults in training, and warmup allows the network to experience smaller sequential catapults instead. More catapults can be observed in the progressive sharpening regime, but it is unclear to me why it is important. We believe the importance of understanding the two qualitatively distinct regimes can be summarized as follows: 1. Having as deep an understanding of the underlying dynamics as possible is intrinsically valuable. Such developments in understanding may be followed by unanticipated innovations and can inform future decisions about the design of algorithms, initializations, and architectures. We think these unanticipated developments are likely to be most important. Our understanding of the different regimes is already important for the practical innovations proposed in our work, as we explain below. 2. As we mentioned in our previous reply, our proposal for GI-Adam was motivated entirely by understanding that the usual Adam has an unnecessarily large preconditioned sharpness, leaving it deep in the sharpness reduction phase and that there is a simple tweak to the initialization that can reduce the need for warmup. We note that this analysis also led us to the understanding that the explanations of the RAdam paper were incorrect and that the RAdam algorithm was unnecessarily complicated. 3. In our paper we have suggested the persistent catapult warmup schedule, for which we have developed some preliminary analysis and left a complete development for future work. The hyperparameters of this schedule, $\delta$ in Algorithm 3 in Appendix C, which specifies the amount of increase in loss that can be tolerated, depend crucially on whether the training is in the sharpness reduction regime or the progressive sharpening regime and the strength of the individual loss catapults. 4. Our suggestion for picking $\eta_{init} = \eta_c$ was originally motivated by analyzing the sharpness curves in the progressive sharpening case, and realizing how much time was being wasted. If we had not analyzed the early-time training dynamics carefully, this realization would have eluded us. 5. Our analysis shows clearly that the sharpness reduction regime is suboptimal because it requires a longer warmup duration for optimal target learning rates. This is an important observation because it demonstrates there are ways to achieve the same test accuracy while requiring fewer warmup steps. 6. The sharpness reduction regime shows that catapult / self-stabilization effects are not always the only mechanism by which warmup works. Another mechanism is via the natural sharpness reduction effect, where sharpness naturally reduces on its own even without a catapult. Therefore warmup can allow the network to tolerate larger learning rates without ever inducing catapults. This is a novel point that was not mentioned in prior work to our knowledge.
Thank you for the explanation, it is much more clear than the one in the initial paper submission. After reading other reviews and authors' responses, I am increasing my score to 6. The concerns from my initial review were related to the three main points: the importance of the analysis of two warmup regimes (weakness 1), the importance of the proposed practical tricks (weaknesses 3 and 4), and some technical questions about experiments (weaknesses 2-5). Below, I summarize how my evaluation of each of them changed after the rebuttal: 1. Two warmup regimes. It was my main concern, and the authors reasonably addressed it in the last response. Specifically, I find the discussion of the correlation between the sharpness level and its increasing/decreasing behavior at the initialization and point 6 on the difference between the two regimes very helpful. I would ask the authors to incorporate these ideas in the final version of the paper. I think adding a discussion on the correlation between the sharpness level and its behavior is especially important. Without it, the conclusion about the connection between sharpness behavior and the importance of the warmup seems vague and unsupported. 2. Practical tricks. I still believe that shortening the warmup does not improve effectiveness significantly. The results of Gi-Adam are not better than alternatives (RAdam), but after the rebuttal, I think the method is simple, well-motivated, and works better than the Adam baseline. Overall, I think the proposed practical modifications are reasonable, even if they are not the strongest contributions of the paper. 3. Technical questions. The authors fully addressed all my technical questions and provided additional experiments.
We thank the reviewer for the insightful discussion. We agree that adding a discussion on the sharpness level and its training behavior will significantly improve the paper. We will incorporate this discussion along with other insightful suggestions in the final manuscript.
Summary
This paper finds that the learning rate warmup allow the network to tolerate larger learning rates. It gradually reduces the sharpness and forces the model to leave poorly conditioned areas of the lossscape and move toward flatter regions which can tolerate larger learning rates.
Strengths
1. This paper analyzes the mechanisms of warmup. 2. It also proposes a GI-Adam strategy, which is better than Adam.
Weaknesses
1. Though the author claims to find that warmup allows for larger learning rates, it has been found in existing work such as "Gilmer, J., Ghorbani, B., Garg, A., Kudugunta, S., Neyshabur, B., Cardoze, D., Dahl, G.E., Nado, Z. and Firat, O., 2022, March. A loss curvature perspective on training instabilities of deep learning models. In International Conference on Learning Representations." Further elaboration on the difference and novelty will make the paper more convincing. 2. They claimed to find "wasted time can be saved by making use of the catapult mechanism", but it seems this has been revealed in "Lewkowycz, A., Bahri, Y., Dyer, E., Sohl-Dickstein, J. and Gur-Ari, G., 2020. The large learning rate phase of deep learning: the catapult mechanism. arXiv preprint arXiv:2003.02218." Not only the above two points, I can't well understand the novelty of this paper. I suggest further summarize the contribution part. 3. Could you give more explainations on the accuracy maps, such as Figures 3 and 4?
Questions
I don't have questions.
Rating
5
Confidence
3
Soundness
2
Presentation
3
Contribution
2
Limitations
Please refer to the Weaknesses
Thanks for the response. I have no further questions.
Thank you for the reply. I think that these additions to the paper will make it even stronger.
Many thanks to the authors for the detailed response. I feel that this paper provides great insights into Warmup. I have raised my score.
Thanks for the authors' response. I have no further questions and raised my score.
Decision
Accept (poster)