Why Warmup the Learning Rate? Underlying Mechanisms and Improvements

Authors: dayal  singh kalra
maissam  barkeshli
Venue: NeurIPS_cc_2024_Conference
Published: 2024-05-14 11:59:54

Abstract: In modern deep learning, it is common 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 with 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 target learning rates in turn makes hyperparameter tuning more robust while improving the final performance of the network. We uncover different regimes of operation during the warmup period, depending on whether the network 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.

Main paper: 
 introduction: 
 One of the most important choices to make in gradient-based optimization is the learning rate (step size) η. If η is too small, then learning may take place too slowly or the model might get stuck in unfavorable regions of the loss landscape. If η is too large, training will typically diverge. In practice, it is common to pick a dynamical learning rate schedule η t [2, 4, 40, 26] . Modern learning rate schedules for deep learning typically consist of a warmup period where η t is increased linearly from zero to a target value η trgt over a warmup time T wrm [13, 34] . After the warmup period, it is common to eventually decay the learning rate, for example via a cosine decay schedule [34, 26, 40] .Given that warmup is standard in the practitioner's toolkit, it is important to understand it deeply and identify improvements. In modern settings, perhaps the earliest work to use warmup was [14] , which used a small constant learning rate for the first few epochs of training and then switched to a larger learning rate. A linear warmup schedule was later introduced in [13] . The intuition given was that to scale the minibatch size in SGD by a factor of k, it is natural to also scale the learning rate by a factor of k, provided the model is not changing too rapidly and successive gradients are roughly aligned. However at the beginning of training, the model is changing rapidly, so it is natural to start with a lower learning rate and gradually increase it to the target value after the network has stabilized.Other explanations suggest that since the network is initialized randomly, the gradient steps at the beginning of training are not meaningful, and thus it would be harmful to take large steps in such directions [40] , so it makes sense to take smaller steps early in training. The analysis by [12] suggests that warmup primarily limits the magnitude of weight updates in the deeper layers, preventing large instabilities. It has also been suggested that the key benefit of warmup arises for adaptive optimizers, such as Adam: [23] argues that the variance of the adaptive learning rate is large during early training because the network has seen too few training samples; it is asserted that this large variance is harmful, and that warmup acts as a variance reduction method by allowing the network to collect accurate statistics of the gradient moments before using larger learning rates. Alternatively, it is also sometimes stated that the initialization may start the model off at places in parameter space that are unstable, difficult to optimize, and easily lead to divergence, and that warmup can help alleviate this [40] .The above explanations are varied and do not clearly demonstrate why and to what extent warmup is necessary. A loss landscape perspective was given in [10] (and summarized in [26] Ch. 8), which argued that an important effect of warmup is to gradually reduce the sharpness (the top eigenvalue of the Hessian of the loss), thus causing the model to leave poorly conditioned areas of the loss landscape and move towards flatter regions which can tolerate larger learning rates. They argue that the mechanism for this is similar to the dynamical stability (catapult) mechanisms studied in [35, 22] .Our contributions. Here we perform extensive studies on the effect of learning rate warmup across architectures (FCNs, ResNets, and Transformers), initializations and parameterizations, datasets (CIFAR-10, CIFAR-100, TinyImageNet, WikiText), and for both SGD and Adam.We demonstrate through systematic experiments that by far the primary benefit of learning rate warmup is to allow the network to tolerate larger learning rates than it otherwise would have. This builds on the observations of [10] by showing that any other benefits are marginal, disentangling the effect of warmup duration and target learning rate, and by extending the empirical evidence to include adaptive optimizers and Transformers.For SGD, the maximal allowable learning rate is determined by the sharpness (the top eigenvalue of the Hessian of the loss). As we discuss in Section 4, we find that there are several qualitatively distinct regimes and mechanisms at play. These depend on whether the network starts off in a sharpness reduction or progressive sharpening phase [18, 19, 5] , which in turn depends on the initialization and parameterization. We further find that the performance of the network is largely determined by the target learning rate. For a fixed target learning rate, increasing the warmup time provides only marginal benefit, which arises by keeping the network further away from the divergence (failure) boundary. The ability of the network to withstand a larger target learning rate in turn makes hyperparameter tuning of the target learning rate more robust, since the network responds well to a larger window of target learning rates, possibly explaining the popularity of warmup.We then investigate Adam in detail, and show that the underlying mechanisms of warmup are similar to the SGD case, but with sharpness replaced by a preconditioned sharpness (the top eigenvalue of the pre-conditioned Hessian, defined below) . Our results disagree somewhat with prior results [23] on the underlying reason for warmup's benefits: We find that the key issue is not observing too few training samples, but rather that the pre-conditioned sharpness typically starts off at high values (even in the large batch case), causing considerable instabilities at high learning rates. Such instabilities, which may be retained in Adam's memory, can result in performance degradation and even training failures. Warmup mitigates such instabilities by gradually pushing down the preconditioned sharpness, enhancing performance, and preventing training failures. We propose a simple alternative initialization for Adam, which we refer to as GI-Adam, which provides benefits similar to warmup and consistently improves over standard Adam by inducing lower preconditioned sharpness at initialization, thus pushing the training failure boundary to higher target learning rates. This also demonstrates a different way to remove the bias correction of RMSProp with momentum.Our analysis shows how much of the time spent during the warmup period is wasted. We show that this wasted time can be saved by making use of the catapult mechanism [22] to effectively estimate the initial sharpness scale by line search, providing a more principled choice of η init . Our experiments show that, depending on the target learning rate and initial sharpness, one can dramatically reduce the warmup time, and in some cases remove it altogether.2 Notations and Preliminaries SGD(-M): Given gradients g t := ∇ θ L(θ t ) at step t, Stochastic Gradient Descent with momentum updates the parameters θ t using learning rate η t and momentum m t with coefficient β. The update equations are: m t+1 = g t + βm t and θ t+1 = θ tη t m t+1 . β = 0 corresponds to SGD.Adam: Adam [20] updates the parameters θ t according to the equations: Figure 1 : Training loss and sharpness trajectories of FCNs trained on a 5k subset of CIFAR-10 with MSE loss using GD. In the sharpness plot, the dashed lines represent the 2 /ηt curves, and when λ H t is above these curves, training exceeds the instability threshold (η > η c ). (top) µP with η trgt = 1 /λ H 0 , (bottom) SP with η trgt = 32 /λ H 0 . Similar mechanisms are observed across different architectures, loss functions, and mini-batch sizes, as shown in Appendix E. and vt = vt 1-β t 2 are the bias-corrected moments, ϵ is a small scalar used for numerical stability andm t+1 = β 1 m t +(1-β 1 )g t , v t+1 = β 2 v t + (1 -β 2 )gP t = (1 -β t 1 ) diag √ vt + ϵI is the preconditioner.Linear Warmup: This is defined by the schedule η t = η init + (η trgtη init ) t /Twrm. The warmup rate is α := (ηtrgt-ηinit) /Twrm. T wrm = 1 corresponds to constant learning rate. Unless otherwise specified, we set η init = 0 when referring to linear warmup. We propose strategies for selecting η init in Section 6.The sharpness is defined as the maximum eigenvalue of the Hessian of the loss λ H t := λ max (∇ 2 θ L), with subscript t indexing the training step. Adaptive optimizers, such as Adam, effectively perform gradient descent in a transformed space determined by ϕ := P 1/2 θ (for details, see Appendix B.1). Hence, their stability is determined by the largest eigenvalue of the pre-conditioned Hessian, denoted by λ P -1 H := λ max (P -1 ∇ 2 θ L), rather than the sharpness itself. Parameterizations in Neural Networks: The mechanism of warmup and its effectiveness is heavily influenced by the network parameterization (see Sections 4 and 5) . Standard Parameterization (SP) [33] is a staple in common libraries [28, 3] . Another notable parameterization is the Neural Tangent Parameterization (NTP) [17] , which along with SP resides in the kernel learning class at infinite width. Ref. [37] proposed Maximal Update Parameterization (µP) which exhibits feature learning at infinite width. Neural network parameterizations significantly impact training dynamics [19] . 
overview of training instabilities and the self-stabilization mechanism: 
 One important underlying mechanism of warmup is intimately tied to training instabilities. These training instabilities, often referred to as 'catapults' [22, 5] , arise when the learning rate η exceeds a critical threshold η c , where both η and η c generally change with time. The critical learning rate η c is influenced by a variety of factors, including the choice of optimizer [5, 6] , mini-batch size [35, 6] , and model properties such as depth, width, parameterization, and initialization [18, 19] . For a detailed overview of instability thresholds, see Appendix B.When the instability threshold is exceeded (η > η c ), two cases arise: (i) if the learning rate is higher than the instability threshold but smaller than a maximum stable learning rate (which varies with time), i.e., η c < η < η max , training stabilizes through a self-stabilization process and training continues, (ii) if the learning rate exceeds this maximum stable learning rate η > η max , training experiences severe instabilities. For SGD, these can result in training divergence, characterized by the loss increasing to infinity. For Adam, training may cease, resulting in a training failure, where the loss fails to improve significantly over its initial value, as we demonstrate in Section 5.The self-stabilization mechanism of GD can be understood through both empirical observations (Figure 1 ) and a theoretical model. We first describe a model derived by Ref. [7] that effectively captures this phenomenon. The model assumes that the top eigenvector u changes slowly through training and can be treated as constant and considers a cubic approximation of the GD dynamics around a reference point θ * . The dynamics along the projection x t := u T (θ tθ * ) is given by two coupled non-linear equations:EQUATIONwhere α := -∇λ H •∇L quantifies the instantaneous change in sharpness and β := ∥∇λ H ∥ 2 controls the non-linear change in sharpness. In this model, an instability arises when η t > η c = 2 /λ H t . Ref. [7] considered a constant learning rate η and assumed progressive sharpening (α > 0). In contrast, we consider a time-dependent learning rate and allow α to attain both positive and negative values in order to analyze different warmup mechanisms.The self-stabilization mechanism manifests as a four-step process [22, 7] . Below, we describe the four steps of the self-stabilization mechanism using the above model and the T wrm = 64 trajectories illustrated in Figure 1(c, d ):(1) Approaching instability: Due to increasing learning rate and/or progressive sharpening, training approaches the instability threshold η t = η c = 2 /λ H t . In Figure 1 (d), this occurs within the first 10 steps due to increasing learning rate.(2) Blow up: On exceeding the instability threshold (η > η c ), Equation (1) predicts exponential growth in x t , empirically manifesting as a sharp increase in loss, as observed in Figure 1(c ).(3) Sharpness reduction: For small enough learning rates, |x t | (and the loss) continues to increase until the higher-order term in the sharpness update equation causes a decrease in sharpness (x t > α /β). This is observed as an abrupt decrease in sharpness in Figure 1(d) . If the sharpness fails to decrease over extended steps, it may result in training divergence (e.g., see T wrm = 1 case in the same figure).(4) Return to stability: Once the sharpness has decreased appreciably so that η t λ H t < 2, stability is restored and the loss eventually decreases.While the self-stabilization process for more complex optimizers remains poorly understood, a qualitatively similar mechanism is observed in practice, as we will see in the later sections. 
warmup mechanisms of gradient and adaptive methods: 
 This section analyzes the underlying mechanism of warmup through the lens of sharpness dynamics. A key finding is that warmup decreases sharpness in two ways: by allowing a natural sharpness reduction effect at early times and/or by forcing sharpness reduction through training instability at later times.Learning rate warmup is intrinsically tied to sharpness dynamics, as sharpness determines the instability threshold η c . As the learning rate is increased during warmup, training instabilities can be triggered. Assuming the warmup rate is not too high, these instabilities induce a temporary increase in the loss and a decrease in the sharpness to restore stability through the self-stabilization mechanism. Ultimately this allows the model to adapt to the increased learning rate. In other words, a primary goal of warmup is to gradually reduce sharpness, guiding training towards flatter regions that can accommodate training at higher learning rates [10] .However, digging deeper, we find that training has a 'natural' preference for sharpness evolution throughout the training course [19] . Before exceeding the instability threshold (η < η c ), training naturally experiences either a progressive increase or decrease in sharpness, as observed in Figure 1 , which is unrelated to warmup. For instance, consider the sharpness trajectories with T wrm = 1024 in the above figure. In Figure 1 (b), sharpness has a natural preference for increasing, whereas in Figure 1 (d), it tends to decrease on its own. This natural sharpness evolution can defined as the sharpness evolution under gradient flow, corresponding to α in Equation (1) . The interplay between this natural sharpness evolution and the deliberate intervention of warmup to reduce sharpness can result in completely distinct dynamics. Below, we use the model described by Equation (1) and the experiments in Figure 1 to describe these distinct dynamics.(C1) Natural Progressive Sharpening (α > 0; top row of Figure 1 ): The combined effect of naturally increasing sharpness while the learning rate is also being increased results in a "head-on collision" at which the instability threshold is exceeded (η t > η c ). This causes the loss to increase, leading to a decrease in sharpness. Once the sharpness has decreased appreciably, the stability is restored (η t < η c ). As training proceeds, both sharpness and learning rate continue to increase, again surpassing the instability threshold. This results in a persistent catapult cycle, characterized by η t ≈ 2 /λ H t ≈ η c , for the remainder of the warmup period, as seen in Figure 1 (b). (C2) Natural Sharpness Reduction (α < 0; bottom row of Figure 1 ): The network is naturally already reducing its sharpness during early training. However, if the learning rate is increased sufficiently quickly, eventually the instability threshold will be reached (akin to a "rear-end collision"), causing the loss to increase. For small enough learning rates, the increased loss induces a dramatically more pronounced decrease in sharpness than would naturally occur, ultimately restoring stability (η t < η c ). To exceed the instability threshold again, the learning rate must significantly increase to account for the decreased sharpness, potentially requiring considerable training steps. Consequently, training experiences one or more separated catapults during the warmup phase, as seen in Figure 1(c, d ). This contrasts with the progressive sharpening case, where training enters a continuous catapult cycle after reaching the instability threshold for the first time. Notably, training may eventually reach a very flat region of the landscape during warmup, with gradients pointing towards increasing sharpness (e.g., T wrm = 64 in Figure 1(d) ). Upon reaching such a region, the dynamics aligns with the natural progressive sharpening scenario.When natural sharpness reduction is significant (large negative α), warmup may not need to actively reduce sharpness. Instead, it may "piggy-back" on the inherent sharpness decrease, resulting in a completely different warmup mechanism, which does not rely on instabilities to facilitate training at higher learning rates.The above two scenarios can be interpreted as cooperative or competitive dynamics between warmup and the natural evolution of sharpness. When training inherently undergoes sharpness reduction, it cooperates with warmup in decreasing sharpness. Conversely, if the natural trajectory of training is towards increasing sharpness, it opposes the warmup's effort, leading to a persistent cycle of catapults.The Effect of Warmup Duration: Given a fixed target learning rate η trgt , increasing the warmup duration T wrm delays the point at which training exceeds the instability threshold η c , allowing the sharpness to evolve freely before reaching this point. In the sharpness reduction case, sharpness can significantly decrease by the time this threshold is reached, lowering the need for warmup to decrease sharpness actively. Consequently, increasing T wrm results in catapults that are both delayed and smaller in magnitude, as seen in Figure 1(d) . As the catapults become less intense on increasing the warmup duration, the model can train at higher target learning rates without diverging. For extended warmup durations, warmup may not actively reduce sharpness in these sharpness reduction cases and instead it leverages the inherent sharpness decrease.In the progressive sharpening case, increasing T wrm allows the sharpness to naturally increase. As a result, training exceeds the instability threshold for the first time at a relatively lower learning rate compared to the constant learning rate case. Although warmup has to now undertake more work in decreasing sharpness, it does so in a more gradual manner since increasing the warmup duration amounts to a lower warmup rate ηtrgt /Twrm. As a result, the fluctuations observed on exceeding the instability threshold are much smaller in magnitude, as seen in Figure 1(a, b ).Small vs. Large Initializations: So far, we have outlined different warmup mechanisms without describing specific conditions that typically exhibit them. Small initializations, such as those using maximal update parameterization (µP) [37] in the large width limit or appropriately using normalizing layers (e.g. standard Transformer architectures, see Figure 17 in Appendix E.5), are characterized by a small initial network output. Such initializations start in flat regions where gradients point toward increasing sharpness [19] , placing them in the progressive sharpening category (C1). As we will see in Section 5, such initializations may not significantly benefit from warmup as they already start in a flat region. In contrast, large initializations, such as FCNS, CNNs, ResNets with Standard Parameterization (SP) initialized at criticality [29, 31] or Transformers with the last layer-norm removed, undergo an early sharpness reduction, categorizing them into sharpness reduction category (C2). As the primary effect of warmup is to reduce sharpness, we expect such large initializations to considerably benefit from warmup. Notably, large initializations can eventually undergo progressive sharpening at later training stages [18, 19] and adhere to the second mechanism, especially for prolonged warmups. Natural sharpness change provides an intuitive way of determining whether an initialization is 'small' or 'large': if training from a given initialization exhibits sharpness reduction, it suggests the existence of naturally flatter initializations that could be chosen instead. This observation is particularly helpful for analyzing Adam in the later sections and motivates modifications to improve it.The warmup mechanism of SGD-M, while at its core is similar to that of vanilla SGD, has a few subtleties. Here we summarize the major differences, leaving details to Appendix E.2.During early training, the loss may decrease non-monotonically on incorporating momentum, even at small learning rates. Such oscillations are also observed when quadratic loss functions are optimized using GD with momentum [11] . These oscillations make it challenging to differentiate between warmup-induced catapults and fluctuations in loss due to the intrinsic effects of momentum. Nevertheless, we can still observe loss spikes correlated with an abrupt decrease in sharpness at large learning rates, as detailed in Appendix E.2.Additionally, the instability threshold η c itself evolves differently during training. It changes from 2 /λ H 0 at initialization to (2+2β) /λ H t later in training. Moreover, the late-time instability threshold is significantly influenced by the batch size, exhibiting a much smaller value than SGD for the same batch size. These properties make it more challenging to analyze the training dynamics of SGD with momentum. Nonetheless, the fundamental warmup mechanisms closely mirror the vanilla SGD case. We leave a more detailed analysis of the early training dynamics of SGD-M for future studies.Adaptive optimizers effectively perform gradient descent in a transformed space given by ϕ = P 1/2 θ, as we show in Appendix B.1. This analysis suggests that their local stability should be determined by the largest eigenvalue of ∇ 2 ϕ L = P -1 ∇ 2 θ L, which we refer to as the pre-conditioned Hessian. Indeed, this is what we observe in Figure 2 , which shows the training loss, pre-conditioned sharpness λ P -1 H , and sharpness trajectories for full batch Adam. In these figures, sharpness is significantly smaller than its instability threshold (2+2β1) /ηt ≈ 4000, indicating that sharpness does not determine stability. Instead, loss catapults are associated with λ P -1 H exceeding its instability threshold.The pre-conditioned sharpness starts high for both progressive sharpening (simple-µP) and sharpness reduction (SP) scenarios considered in the previous section. For simplicity, we considered a simpler version of µP, detailed in Appendix D.2.1. In particular, for µP models, λ P -1 H 0 ∼ 10 5 despite being initialized in a flat region as measured by sharpness, while for SP models, λ P -1 H 0 ∼ 10 6 . These large initial values arise because P 1 = (1β 1 ) diag(g 2 0 ) + ϵI and g 0 has components that are near zero.(a) The large λ P -1 H 0 can lead to training failures if the learning rate does not start sufficiently small. We put forward modifications to improve Adam in Section 6; here we continue characterizing the warmup mechanisms of Adam.Given that the pre-conditioned sharpness consistently starts high and decreases during early training, this behavior can be viewed as an extreme example of the natural sharpness reduction scenario (C2) described in the previous section. Training Adam at high initial learning rates without warmup can cause large catapults, as seen in Figure 2 (d), potentially leading to training failures. Increasing the warmup duration allows the pre-conditioned sharpness to naturally decrease. This prevents the loss from spiking during early training and avoids training failures. In the later stages of training, the pre-conditioned sharpness may continue reducing or exhibit progressive sharpening. From here on, the dynamics follows the warmup mechanisms discussed in the previous sections, with sharpness replaced with pre-conditioned sharpness. Similar to the momentum case, Adam's stability threshold at late training times significantly decreases for smaller batch sizes [6] , also shown in Appendix E.4. 
impact of warmup on training and generalization: 
 Here we investigate the impact of warmup on training efficacy and generalization by disentangling the role of η trgt and T wrm . Our key findings are that generalization capability is primarily determined by η trgt and that Adam is particularly sensitive to large catapults. The role of increasing T wrm is to (i) allow the network to tolerate larger η trgt , and (ii) move training further away from the divergence (failure) boundary, leading to a marginal improvement in generalization.Experimental Setup: We consider WideResNets (WRNs) and Transformers (LM) parameterized in either SP or µP. WRNs are trained on CIFAR-10, CIFAR-100, and Tiny-ImageNet, employing data augmentation. Transformers are trained on the next token prediction task using the WikiText dataset. These models are trained with MSE or cross-entropy (xent) loss functions using SGD or Adam optimizers for a fixed training budget of T = 10 5 steps unless otherwise specified. Training begins with a linear warmup phase from η init = 0 to η trgt over T wrm steps. After warmup, training continues at η trgt for the remaining training budget. In some cases, following the warmup period, we decrease the learning rate using cosine decay [24] . Further details are provided in Appendix D. Figure 3 presents heatmaps that show the best test accuracy achieved during training, plotted in the η trgt -T wrm plane for different parameterizations and loss functions. These phase diagrams of warmup also show the convergence-divergence boundary, with empty cells indicating training divergences, illustrating the interplay between warmup duration and the maximum trainable η trgt . Below, we discuss the crucial insights these results provide into warmup's role in training dynamics.Longer Warmup Facilitates Training at Higher Learning Rates: These phase diagrams reveal that an extended warmup duration facilitates training at higher target learning rates. This benefit is particularly noticeable for large initializations (like SP) and MSE loss. In contrast, the advantage is less pronounced when using cross-entropy loss and smaller initializations (like µP). The diminished benefit for µP is likely due to its initialization in a relatively flat region of the loss landscape, which can already facilitate training at higher learning rates at initialization. This consistent increase in maximum η trgt with warmup durations can be understood through the lens of warmup mechanisms described in the previous section. As observed in Figure 1 , when the warmup duration is increased, loss catapults occurring on surpassing the instability thresholds become milder. This effectively pushes the divergent boundary to higher learning rates.Final Performance Primarily Depends on the Target Learning Rate: A closer look into these phase diagrams reveals that, slightly away from the divergent boundary, the test accuracy primarily depends on the target learning rate and nominally on the warmup duration. Based on the model performance, we can categorize these phase diagrams into two distinct cases: (i) models that fail to achieve optimal performance when trained with a constant learning rate (e.g., Figure 3 (c)), and (ii) models that attain optimal performance without warmup (e.g., Figure 3(b) ). The first scenario corresponds to models with large initializations. Increasing the warmup duration improves performance by facilitating training at higher learning rates. Yet, similar performance is observed for different warmup durations, suggesting that the primary gain comes from the target learning rate, rather than the duration itself. The second case arises for flat initializations, which can already train at large learning rates, and resultantly the optimal performance is already achieved without warmup. While increasing warmup duration facilitates training at even higher learning rates, it does not enhance performance. Nevertheless, it does broaden the range of optimal learning rates, reducing the need for precise tuning of the target learning rate, and making training more practical and robust. We conclude that warmup can serve two key purposes: (i) it can significantly improve model performance in large initialization cases, and (ii) extend the range of optimal target learning rates for small initializations, making it easier to tune the target learning rate. In Appendix F.2, we demonstrate that these results hold on incorporating momentum and employing cosine learning rate decay.Figure 4 (a) shows the warmup phase diagram of Adam. Increasing the warmup duration enables training at higher learning rates by allowing the pre-conditioned sharpness to decrease naturally, thereby reducing the severity of catapults. These large catapults, which may persist in Adam's memory, can lead to performance degradation and training failures. Thus, in addition to facilitating training at higher rates similar to SGD, warmup further improves Adam's performance by addressing its vulnerability to large catapults, justifying its widespread use with Adam. Below, we discuss the distinct properties of Adam phase diagrams in detail.Training Failures of Adam: Remarkably, we find that models trained with Adam always exhibit training failures rather than divergences where the loss grows without bound, as further demonstrated in Appendix G. In cases of training failure, we often observed that certain layers or residual blocks output zero, leading to vanishing gradients. This implies that the model gets stuck at a critical point and is unable to train further. Understanding this unexpected phenomenon requires further study.Performance Degradation prior to Failure Boundary: Test accuracy in these phase diagrams declines well before the failure boundary, in contrast to SGD where optimal learning rates are observed near the divergence boundary. This discrepancy stems from Adam's property of retaining a memory of gradient magnitudes. At large learning rates, along with the loss, the gradients spike during early training, as seen in Figure 28 in Appendix G. While the gradients decrease after a few steps, the second moment v remains large for an extended period, leading to a small effective learning rate ηP -1 . As a result, training struggles to escape high-loss regions. Therefore, a longer warmup is more beneficial for Adam compared to SGD, as it is crucial to stay away from the failure boundary. 
improved hyperparameter initialization schemes for optimizers: 
 Initial Learning Rate Selection for Warmup: Setting the initial learning rate to η init = 0 is common practice in warmup [27, 8] . Our analysis reveals that the primary effect of warmup is to facilitate training at higher learning rates by annealing sharpness (or pre-conditioned sharpness for Adam). From this perspective, starting with η init = 0 appears suboptimal, as it can significantly delay the learning rate from exceeding the instability threshold, thus delaying the primary effect of warmup.An effective strategy involves setting η init = η c to induce loss increase and thereby sharpness decrease right from initialization. We introduce a straightforward search method that only uses forward passes to estimate the initial critical learning rate η c . The method consists of two stages: (i) an exponential search, starting from an initial guess η 0 , iteratively multiplies η 0 by a factor k > 1 until the loss increases. This identifies an interval [η lwr , η uppr ] containing η c , (ii) a binary search further narrows down [η lwr , η uppr ] by evaluating the loss at the midpoint η mid = (ηlwr+ηuppr) /2. If the loss increases, η uppr is updated to η mid ; otherwise, η lwr is set to η mid . This process is repeated until the loss in the next step L(θ 1 ) satisfies the condition L(θ 1 ) < L(θ 0 )(1 + δ), for some hyperparameter δ > 0. It is worth noting that there can be instances where there is no loss catapult at initialization. In our experiments, this only occurs for Transformers trained using SGD. In such scenarios, the prescribed approach is not applicable and one can resort to heuristics, such as setting the initial learning rate to a fraction of the maximum stable learning rate, such as η init = ηtrgt /10.In Section 4.3, we observed that the pre-conditioned sharpness for Adam starts at a high value, even for low sharpness initializations like µP, and can lead to training failures at large learning rates. We propose Gradient Initialized Adam (GI-Adam), which initializes the second moment using the gradient squared, v 0 = g 2 0 . In Appendix I.2, we show that a bias correction is not required when the second moment is initialized using the gradients. As a result, GI-Adam can be viewed as standard Adam with an automated warmup given by η t = η trgt 1β t 2 . This simple trick reduces the initial pre-conditioned sharpness by around two orders of magnitude (more precisely by a factor of √ 1β 2 ) at initialization, preventing large catapults, as illustrated in Figure 6 (c.f. Figure 2(d-f )). Moreover, it consistently shows improvement over standard Adam across datasets and prevents training failures by pushing the training failure boundary to higher η trgt , as shown in Figure 4 (b). We provide additional results for different datasets in Appendix F.3. In Appendix F.4, we show that GI-Adam consistently performs on par or better than RAdam [23] while offering a simple modification to Adam.To further assess that the primary cause of instability during early training is the large pre-conditioned sharpness, we randomly initialize v 0 but with the same norm as the gradients at initialization. Like GI-Adam, this also results in improved performance as shown in Appendix I.3.We further reduce the pre-conditioner size by removing bias correction for the first moment, referred to as Flat-Adam. As demonstrated in Appendix I.4, this modification eliminates the initial decrease in pre-conditioned sharpness. We leave a comprehensive evaluation of Flat-Adam for future work. 
discussion: 
 Our analysis provides new insights into the role of warmup across optimizers and parameterizations. We found compelling evidence that the primary effect of warmup is to facilitate training at higher learning rates and stabilizing the training dynamics by keeping it away from the failure (divergence) boundary. Looking under the hood, we found a variety of underlying mechanisms, which also suggested several improvements for hyperparameter initialization. In Appendix A, we provide practical guidance for practitioners on choosing the warmup duration.Our analysis also motivates a potential parameter-free warmup strategy, which we refer to as persistent catapult warmup. The central idea behind this strategy is to repeatedly induce catapults aimed to progressively reduce sharpness, thereby facilitating training at higher learning rates. We present encouraging preliminary results in Appendix C and defer further development to future work.The maximum learning rate can be written as η max = cmax /λ P -1 H . Here we showed how warmup effectively decreases λ P -1 H , which is a local measure of sharpness. There is also another possible effect of warmup, that it can cause an increase in c max , which can be viewed as a more non-local measure of sharpness. Further analysis is required to understand how warmup helps increase c max .Limitations: Our experiments were conducted on relatively small-scale datasets and models, and further investigations are needed to understand the generalizability of our findings to larger-scale settings. For Adam, we did not explore the dependence on hyperparameters β 1 , β 2 , ϵ.How to Select the Warmup Duration? Given a target learning rate η trgt , if the training loss during the warmup period exhibits large instabilities (loss spikes), the warmup duration T wrm should be increased until such instabilities are sufficiently small. This effectively moves training away from the divergent / failure boundary, as illustrated in Figure 3 . This is particularly crucial for Adam, as large instabilities can be detrimental and lead to considerable performance degradation without divergence, as discussed in Section 5.2.How to Select the Target Learning Rate? As the primary effect of warmup is to anneal sharpness by increasing the learning rate beyond the instability threshold, it suggests that the target learning rate should be at least greater than the instability threshold at initialization.When to Decay the Learning Rate? Figure 21 suggests that employing learning rate decay at small learning rates can result in performance degradation for a fixed training budget. Therefore, the learning rate should be decayed at large target learning rates only. The underlying intuition is that we use large target learning rates to train in a flat region of the landscape. However, these large learning rates restrict training to go into sharper regions of the basin and learning rate decay helps.Leveraging µP for Effecient Training: Our analysis suggests that the primary role of warmup facilitates training at higher learning rates by gradually reducing sharpness. Given this perspective, beginning training with flat initializations, such as µP, is advantageous. These initializations might allow for achieving optimal performance without the need for warmup, as observed in Figure 3 .For adaptive optimizers, such as Adam, a pre-conditioner P t+1 multiplies the gradients in the update equationEQUATIONNext, we define ϕ t := P 1/2 t+1 θ t , then the gradient and Hessian are∇ ϕ L = P -1/2 ∇ θ L and ∇ 2 ϕ L := P -1 ∇ 2 θ L.As a result, the update equations for ϕ is given byEQUATIONWe observe that the update equation of an adaptive optimizer mirrors that of SGD under a change of variables. Using the stability arguments of SGD, we conclude that λ max (P -1 t+1 ∇ 2 θ L(θ t )) determines the stability.Lewkowycz et al. [22] showed that for wide networks in NTP/SP trained with MSE loss and SGD, this critical learning rate is 2 /λ H 0 early in training. Further investigation by Kalra and Barkeshli [18] demonstrated that sharpness reduction during early training causes η c to increase with depth and 1 /width. In such scenarios, η c can be as large as 40 /λ H 0 . Cohen et al. [5] demonstrated that sharpness at late training times for GD with momentum coefficient β oscillates above (2+2β) /η, suggesting η c ≳ (2+2β) /λ H t at late training times. Expanding on this, Cohen et al. [6] analyzed adaptive optimizers and found that for Adam, the pre-conditioned sharpness λ P -1 H oscillates around (2+2β1) /η(1-β1) at late training times. The instability threshold also depends on the mini-batch size [35] and is often observed to be smaller than their full batch counterparts [5, 6] .This section describes the method for estimating the instability threshold η c at initialization (or generically, any point θ t ) using only forward passes. The method consists of two stages: Exponential Search: An exponential search, starting from an initial guess η 0 , iteratively multiplies η 0 by a factor k = 2 until the loss increases. This identifies an interval [η lwr , η uppr ] containing η c . The detailed algorithm is described in Algorithm 1. Unless specified, we use η 0 = 10 -4 as our initial guess. If the loss already increases at the initial guess η 0 , we set η init = η 0 .Algorithm 1 Exponential Search 1: Input: (Initial weights: θ 0 , Initial guess: η 0 ) 2: Output:Interval [η lwr , η uppr ] containing η c 3: Evaluate initial loss L(θ 0 ) 4: Initialize η ← η 0 5: θ 1 ← Optimizer(η, θ 0 ) 6: Evaluate L(θ 1 ) 7: while L(θ 1 ) < L(θ 0 ) do 8: if η ≥ η trgt then 9: η ← η trgt 10: break 11: end if 12: η ← 2η 13: θ 1 ← Optimizer(η, θ 0 ) 14:Evaluate L(θ 1 ) 15: end while 16: η uppr ← η 17: η lwr ← η/2 18: return [η lwr , η uppr ] Binary search: A binary search further narrows down [η lwr , η uppr ] by evaluating the loss at the midpoint η mid = (ηlwr+ηuppr) /2. If the loss increases, η uppr is updated to η mid ; otherwise, η lwr is set to η mid . This process is repeated until the loss in the next step L(θ 1 ) satisfies the condition L(θ 1 ) < L(θ 0 )(1 + δ), for some δ > 0. The algorithm is detailed in Algorithm 2. In all our experiments, we set δ = 0.1. end if 16: end while 17: return η c ← η uppr While both η 0 and δ are additional hyperparameters, the method does not heavily depend on these choices. A poor initial guess of η 0 would only take a few more iterations to find an interval [η lwr , η uppr ]. Meanwhile, any small value of δ ∈ (0, 1] is effective in finding η c , as small initial loss spikes have minimal impact on the overall dynamics. Note that for Adam, a small δ (∼ 0.01) has to be selected to ensure that we do not observe large catapults.Our analysis also motivates a potential parameter-free warmup strategy, which we refer to as persistent catapult warmup. The central idea behind this strategy is to repeatedly induce catapults aimed to progressively reduce sharpness (or pre-conditioned sharpness), thereby facilitating training at higher learning rates. Given a target learning rate η trgt , the strategy consists of the following steps:1. Start with a 'stable' reference point θ * , defined as a point where the loss decreases in the next step and estimate the interval [η lwr , η uppr ] containing η c , as described in Appendix B.3.2. Induce a catapult by increasing the learning rate to η = η uppr .3. Continue training and wait until the loss falls below the reference point, i.e., L(θ t ) < L(θ * ). This new point now becomes the stable reference point.4. Repeat the above steps until the target learning is achieved, i.e., η = η trgt .Here, the initial stable reference point is the model's initialization. The detailed algorithm is described in Algorithm 3 in Appendix C.Figure 7 compares persistent catapult warmup (shown in black) with linear warmup. The persistent catapult warmup facilitates training at higher learning rates without the need to specify warmup duration. Since η c serves as an indicator of sharpness, persistent catapult warmup utilizes the local sharpness information to automatically determine the warmup rate, resulting in an adaptive non-linear warmup. This adaptive approach eliminates the need for manual tuning of the warmup duration, allowing for a more efficient and effective warmup.Although persistent catapult warmup is a promising approach to warmup, it requires specifying how large a catapult should be induced, which introduces another hyperparameter. Nevertheless, persistent catapult warmup motivates the development of parameter-free warmup strategies that could simplify the training process. We leave further development of parameter-free warmup to future work.This section provides additional experimental details. All models were implemented using the JAX [3] , and Flax libraries [15] . The key results can be reproduced using the GitHub repo: https: //github.com/dayal-kalra/why-warmup. We consider standard image classification datasets such as CIFAR-10, CIFAR-100 [21] , and Tiny-ImageNet [1] . The images are normalized to have zero mean and unit variance. For MSE loss, we use one-hot encoding for the labels.Data augmentation: For various image classification tasks, we employ data augmentation techniques, applied in the following order: random horizontal flips, random cropping, and mixup [41] .We consider the next token prediction task on the Wikitext-2 and Wikitext-103 datasets [25] . The Wikitext-2 dataset consists of ∼ 2M tokens, whereas the Wikitext-103 dataset has ∼ 0.1B tokens. We use Byte Pair Encoding (BPE) tokenizer [32] with a Whitespace pre-tokenizer. For Wikitext-103 experiments, we considered a standard vocabulary size of 50, 257, whereas our initial Wikitext-2 experiments employed a smaller vocabulary size of 4096.This section describes the models considered, including their parameterization and initialization details. We adopt parameterizations outlined in Table 9 of Ref. [38] . Unless otherwise specified, we employ ReLU non-linearities and initialize the weights with a truncated normal distribution 1 , with a variance σ 2 w = 2.0 in appropriate parameterizations (details below), except for the last layer, which has a weight variance of σ 2 w = 1.0. All biases are initialized to zeros.Standard Parameterization (SP): For SP, the weights are initialized with truncated Gaussian distribution N (0, σ 2 w/fan in) and the biases are initialized to zero.Maximal Update Parameterization (µP): For µP, different schemes are employed for the intermediate and last layers. The intermediate layers are initialized using N (0, σ 2 w/fan out ) and the layer outputs are scaled by the factor fanout /fanin. In comparison, the layer weights are initialized with N (0, σ 2 w/fan in ), and the final output is rescaled by the factor 1 /fanin. Conveniently, for SGD, the learning rate does not scale with width in the above µP formulation. In comparison, for Adam, the learning rate corresponding to input, intermediate, and output layers are rescaled by the factors 1 / √ fanout, 1 / √ fanin and 1 /fanin. Since we are utilizing µP only to obtain flat initializations, we omit the additional scaling of the learning rate for Adam in some experiments (e.g., Figure 2 ). As a result, the instability threshold is only dependent on the target learning rate η trgt during late training, rather than on the largest learning rate across layers. We refer to this parameterization as 'simple-µP' for Adam.Fully Connected Networks (FCNs): We consider fully connected networks with a constant width of n and a depth of d layers. These networks are denoted by FCN-d-n. Unless specified, we considered d = 4 layer FCNs with width n = 512.We consider WideResNets [39] with d layers, S stages, and a widening factor of k, denoted by WRN-d-k. The number of channels in each stage s ∈ [0, S) is given by 2 s × 16 × k, with the input layer having 16 channels. For example, WRN-16-4 consists of S = 3 stages, each with [2, 2, 2] layers, and the corresponding number of channels in each stage is [64, 128, 256] . In all our experiments, we use LayerNorm instead of BatchNorm.Transformers: We consider Transformers with GPT-2 style architecture [30] . These models use sinusoidal positional embeddings [34] and are implemented in the Standard Parameterization (SP) with GELU activation [16] . We initialize all layers using the σ 2 w/fan in scheme, except for the embedding layers, as they do not involve matrix multiplication [9] . We consider both Pre-LN [36] and Post-LN [34] EQUATIONEQUATIONHere, β = 0 corresponds to SGD. In all experiments incorporating momentum, the default value of the coefficient is set to β = 0.9.Adam: Given gradients g t at step t, Adam [20] updates the parameters θ t using learning rate η t and the first two moments of the gradient m t and v t with their coefficients β 1 and β 2 , respectively. The equations governing the updates are:EQUATIONEQUATIONEQUATIONwhere mt = mt1-β t 1 and vt = vt 1-β t 2are the bias-corrected moments, and ϵ is a small scalar used for numerical stability. The pre-conditioner for Adam is given by:EQUATIONIn all experiments, the default values are set to β 1 = 0.9, β 2 = 0.999, and ϵ = 10 -8 , unless otherwise specified.GI-Adam: Given gradients g t at step t, GI-Adam updates the parameters like Adam but with the second moment initialized using the gradients at initialization, i.e., v 0 = g 2 0 . The equations governing the updates are:EQUATIONEQUATIONEQUATIONwhere vt = vt1-β t 2is the bias-correction second moment and ϵ is a small scalar used for numerical stability.Flat-Adam: Given gradients g t at step t, Flat-Adam updates the parameters θ t using learning rate η t and the first two moments of the gradients m t and v t with their coefficients β 1 and β 2 , respectively. In contrast to Adam, we do not apply bias correction for the first moment and initialize the second moment at initialization using the gradients v 0 = g 2 0 . The equations governing the updates are:EQUATIONEQUATIONEQUATIONwhere vt = vt1-β t 2is the bias-correction second moment and ϵ is a small scalar used for numerical stability. The pre-conditioner for Flat-Adam is given by:EQUATIONWarmup linearly increases the learning rate from an initial value η init to a target value η trgt over T wrm training steps. The learning rate η t at step t is given by:EQUATIONHere, α := (ηtrgt-ηinit) Twrm is referred to as the rate of warmup. Under the above definition, constant learning rate training corresponds to T wrm = 1. T wrm = 1 corresponds to constant learning rate. Unless otherwise specified, we set η init = 0 when referring to linear warmup.In several experiments, we employ learning rate decay following the warmup phase. Specifically, we use cosine learning rate decay, which is detailed below.Cosine Decay: Towards the end of training, it is typical to reduce the learning rate to a small value. Cosine decay is a commonly used method for decaying the learning rate from an initial value of η trgt down to a value η min over T cos steps, according to the rule:EQUATIONwhere ρ governs the rate of decay, with ρ = 1 being the standard. Note that with ρ = 0, the learning rate is not decayed and instead maintained at η trgt . In the above expression, t counts the steps from the initiation of cosine decay and not the current training step. As per standard practice, we consider ρ = 1 and decay the learning rate to η min = ηtrgt /10.For SGD, target learning rates η trgt are exponentially sampled using the initial sharpness λ H 0 . Starting with η trgt = 1 /λ H 0 , subsequent rates are sampled until divergence as 2 x /λ H 0 for values of x increased in integer steps starting from zero. For WRNs trained with Adam, we sample target learning rates exponentially as η trgt = 2 x × 10 -5 , where x is incremented in integer steps starting from zero until training failure. For Transformers, we sample the learning rate in a similar fashion but starting from 10 -4 and increment x in steps of 0.5.We measured sharpness / pre-conditioned sharpness using the JAX implementation of the LOBPCG sparse eigenvalue solver with the tolerance set to 10 -9 and maximum number of iterations to n iter = 1000. In most cases, the solver converges within 40 iterations. We performed these computations in float64, as the solver would not converge with float32 in some cases.In certain instances, the pre-conditioned sharpness computation did not converge within 1000 solver iterations. Moreover, we observed that the solver converges on restarting it with a new initial guess of the eigenvector within 40 iterations. To address these edge cases, we employed the following method: if the solver did not converge within 100 iterations, we restarted it with a new initial guess for the eigenvector. We allowed for at most 10 restarts with the maximum number of iterations set to n iter = 1000 in the last attempt. In all reported cases, the solver converges using this method. Figure DetailsThe phase diagram experiments typically required about an hour on per run on an A100 GPU. Consequently, each phase diagram consumed approximately 100 A100 hours of computational time. With a total of 16 phase diagrams, this equates to 1600 A100 hours dedicated solely to phase diagram computations. Additionally, the warmup mechanism experiments, which were conducted over 2000 steps, required sharpness estimation. The FCN experiments required approximately 1200 A100 hours, while the WRN mechanism experiments consumed 1600 A100 hours. The experiments concerning the initial learning rate took about 20 A100 hours. This brings the total computational time amounted to approximately 4500 A100 hours. Preliminary experiments took about 1000 A100 hours. Hence, we estimate the total computational cost to be around 5500 A100 hours.This section presents additional trajectories for warmup mechanisms discussed in Section 4 covering various architectures, loss functions, and optimizers.Figure 8 shows that the warmup mechanisms for full batch GD are also observed in the SGD with a batch size B = 512. The results for other optimizers in the mini-batch setting are discussed in their respective sections.While the warmup mechanisms of SGD with momentum are fundamentally similar to those of vanilla SGD, three key differences arise, as discussed below.First, the training loss can decrease loss in an oscillatory fashion during training [11] . To illustrate this, consider the full-batch GD with momentum. The middle row of Figure 9 demonstrates the sharpness reduction case with the learning rates well below the stability threshold (η << η c ). Despite being far below these thresholds, the loss does not decrease monotonically but converges in an oscillatory fashion. This makes it challenging to differentiate between warmup-induced catapults and fluctuations in loss due to the intrinsic effects of momentum. Nevertheless, we can still observe loss spikes correlated with an abrupt decrease in sharpness at large learning rates, as seen in the bottom row of the same figure. Similar to the SGD case, we observe these catapults are delayed and become smaller in magnitude on increasing the warmup duration.Next, the stability threshold η c for SGD with momentum evolves during training. For simplicity of explanations, we again consider the full batch GD case. The stability threshold η c for SGD with (f) Figure 9 : Training loss and sharpness trajectories of FCNs trained on 5k subset of CIFAR-10 using MSE loss and full batch GD with momentum β = 0.9: (top) µP with learning rate 1 /λ H 0 (middle) SP with learning rate 1 /λ H 0 , and (bottom) SP with learning rate 32 /λ H 0 . The dotted lines in the sharpness figures correspond to the (2+2β) /ηt curves, while dashed lines show the 2 /ηt for reference. momentum changes from 2 /λ H 0 at initialization to (2+2β) /λt late in training. At initialization, the momentum vector is set to zero m 0 = 0, and the stability is given by vanilla GD threshold 2 /λ H 0 . As training progresses, the momentum m t increases in magnitude, and the instability threshold at late training time becomes (2+2β) /λt. The bottom row of Figure 9 show the sharpness trajectories with both 2 /ηt and (2+2β) /ηt curves. For T wrm = 64, the learning rate curve 2 /ηt causes an abrupt decrease in sharpness, which is coupled with a loss spike. For longer warmup durations, the sharpness decreases before training exceeds the 2 /λ H t . Finally, the instability threshold η c for SGD with momentum significantly decreases for smaller batch sizes. Figure 10 shows the training trajectories under the same setup as in Figure 8 , but with momentum coefficient β = 0.9. The late-time sharpness trajectories oscillate well below the (2+2β) /ηt (and even below 2 /λ H t ), whereas the vanilla SGD counterpart oscillates on the 2 /ηt curve. This indicates a strong dependence of the instability threshold on batch size.Besides these three differences, we note that the warmup mechanisms of SGD with momentum are similar to the vanilla SGD case. We leave a thorough analysis of the early sharpness dynamics of SGD with momentum for future works.The warmup mechanisms for models trained with cross-entropy loss exhibit trends similar to those observed with MSE loss with one crucial difference. Near convergence, sharpness first increases and then abruptly decreases. The decrease in sharpness towards the end of training is observed in previous studies analyzing SGD with fixed learning rate [5] . Additionally, we observe higher fluctuations compared to the MSE loss case. Figure 11 shows trajectories of FCNs under different parameterizations trained on CIFAR-10 with cross-entropy loss using vanilla SGD. Meanwhile, Figure 12 shows the loss and sharpness trajectories of FCNs in SP trained on CIFAR-10 with cross-entropy loss using full batch GD with and without momentum. As discussed in Section 4.3, the instability threshold for Adam is determined by the pre-conditioned sharpness λ P -1 H and not by the sharpness itself. Moreover, training dynamics falls under the sharpness reduction case as the pre-conditioned sharpness starts off large and reduces considerably during the first few training.Figure 13 shows the training trajectories of FCNs trained with Adam in the same setting as in Figure 2 but with a batch size of B = 512. Similar to the SGD with momentum case, the late time sharpness oscillates far below the instability threshold ( (2+2β1) /ηt(1-β1)), suggesting that the instability threshold heavily decreases with a smaller batch size. We note similar findings by Ref. [6] .Next, Figure 14 show the warmup mechanism of FCNs trained with cross-entropy loss using Adam under the full-batch setting. Similar to the SGD case, the pre-conditioned sharpness decreases towards the end of training. In the previous sections, we confined our analysis to FCNs to thoroughly explore the effects of different optimizers and loss functions. This section expands on those results by demonstrating that the observed warmup mechanisms apply to ResNets and Transformers as well. The Resnet experiments also employ data augmentation as detailed in Appendix D.1.Figures 15 and 16 show the training trajectories of WideResNets (WRNs) trained on CIFAR-10 with MSE and cross-entropy loss using SGD. These trajectories generally reflect the warmup mechanisms discussed in Section 4. However, certain additional features obscure the clarity of these mechanisms. Notably, we observed a significant sharpness spike on the first training step when using longer warmup durations, which automatically resolves in the subsequent step. The magnitude of this spike increases with longer warmup periods. Further analysis revealed that this phenomenon is associated with an initial increase in the first LayerNorm parameters, which also resolves automatically by the second step. Beyond this observation, the training trajectories align with the warmup mechanisms described in the main text. These experiments demonstrate that Transformers trained on language modeling tasks exhibit warmup mechanisms consistent with those discussed in the main text.This section presents further results related to the phase diagrams of warmup shown in Section 5. Figure 19 shows the test accuracy heatmaps of WRN-16-4 trained on CIFAR-100 and Tiny-ImageNet. These models are trained using cross-entropy loss using SGD with a batch size of B = 128. Additional phase diagrams for Adam are presented in Appendix F. Figure 21 shows that incorporating momentum and cosine decay (for details, see Appendix D.3.3) minimally affects the warmup phase diagrams. While the conclusions regarding warmup presented in the main text remain unaffected, we note a few interesting observations. First, the divergent boundary shifts leftward on incorporating momentum, indicating that momentum permits smaller target learning rates without warmup, and warmup helps SGD-M more. Meanwhile, cosine decay has a minimal effect on the divergent boundary.Additionally, we observe a performance enhancement by incorporating momentum, especially at small learning rates. In contrast, a decaying learning rate beyond warmup degrades performance at small learning rates while improving at higher ones. Finally, incorporating both momentum and cosine decay leads to further enhancement, indicating a synergistic interaction between the two. Figures 26 and 27 and table 1 compare the performance of GI-Adam with RAdam. We observe that GI-Adam performs on par or better than RAdam while offering a simpler modification to Adam. These large loss catapults cause the gradients g to spike during early training, leading to a substantial increase in its second moment v. While the gradients return to a lower value after a few training steps, the second moment remains large in magnitude for a prolonged period. These large values of v result in a small effective learning rate, which hinders training to escape these high-loss regions. Consequently, the models remain stuck in a suboptimal state rather than converging. We refer to this as a training failure.Upon closer examination of the individual layers during training failures, we found that certain layers or residual blocks output zero. This results in vanishing gradients except for the last layer bias and training halts. We defer the detailed analysis of Adam's failures to future work. This section provides additional results for the initial learning rate selection. Figure 29 shows the number of steps T reach required to reach the target learning rate and the effective number of steps saved T save for WRNs in µP. Note that in such cases, η c ≈ η max . We observe that T reach = 1 for a wide range of learning rates, saving almost the entire warmup duration. Figures 30 and 31 show similar phase diagrams for WRNs in SP and µP trained with Adam. This section presents additional results for GI-Adam. We provide further insights into the mechanisms and interpretations of GI-Adam.Figure 32 shows the training trajectories of FCNs with different parameterizations trained with GI-Adam. Notably, the pre-conditioned sharpness starts at significantly lower values than standard Adam. Specifically, for the µP model, the initial pre-conditioned sharpness λ P -1 H is around 2000 instead of the value 10 5 observed for Adam (c.f. Figure 2 ). Remarkably, this almost eliminates initial reduction. Similarly, the pre-conditioned sharpness for the SP model starts around 10 4 instead of 10 6 . Notably, in the SP scenario, there is no initial spike in the T wrm = 1 (c.f. Figure 2 ), demonstrating that this simple modification effectively reduces instabilities during the early training.In this section, we show that a bias correction is not required when the second moment is initialized with the gradients at initialization in GI-Adam. Therefore, employing a bias correction as in the original Adam algorithm in this case serves as an automated warmup given by η t = η trgt 1β t 2 . The moving average of the second moment is given by:EQUATIONwhere v 0 = g 2 0 . Following standard assumptions, we assume that the second moment of the gradient is constant during early training E[g 2 t ] = σ 2 . Taking the expectation of the above equation over the gradient distribution yieldsEQUATIONSimplifying the above equation, we haveEQUATION)This result demonstrates that when the second moment is initialized with the gradients at initialization, it does not require bias correction, as the expected value of the second moment is equal to the constant σ 2 . If we apply the usual bias correction on top of initializing the second moment with the gradients, we effectively downscale the second moment by a factor 1β t 2 . Assuming small enough ϵ, this can be viewed as a multiplicative factor to the learning rate. As a result, GI-Adam is equivalent to having a natural warmup given by η t = η trgt 1β t 2 . To further assess if the primary cause of instability during early training is the large λ P -1 H , we randomly initialize v 0 but with the same norm as the gradients at initialization. We refer to this as Randomly Initialized Adam (RI-Adam). Like GI-Adam, this also results in improved performance as shown in Figure 33 .Figure 34 shows the training trajectories of FCNs with different parameterizations trained with Flat-Adam. We observe that further removing bias correction for the first moment completely removes the initial reduction in the pre-conditioned sharpness. • The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions). • The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.) • The assumptions made should be given (e.g., Normally distributed errors). • It should be clear whether the error bar is the standard deviation or the standard error of the mean. • It is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have a 96% CI, if the hypothesis of Normality errors is not verified. • For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are out of range (e.g. negative error rates). • If error bars are reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables in the text. 
experiments compute resources: 
 Question: For each experiment, does the paper provide sufficient information on the computer resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?Answer: [Yes]Justification: We provide the details in Appendix D.6.Guidelines:• The answer NA means that the paper does not include experiments.• The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage. • The paper should provide the amount of compute required for each of the individual experimental runs as well as estimate the total compute. • The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't make it into the paper). 
code of ethics: 
 Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics https://neurips.cc/public/EthicsGuidelines?Answer: [Yes] Justification: We adhere to the NeurIPS Code of Ethics.Guidelines:• The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.• If the authors answer No, they should explain the special circumstances that require a deviation from the Code of Ethics. • The authors should make sure to preserve anonymity (e.g., if there is a special consideration due to laws or regulations in their jurisdiction). 
broader impacts: 
 Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?Answer: [NA] Justification: This paper analyzes the underlying mechanisms of warmup. We do not think there are any societal impacts of the work that are worth mentioning here.Guidelines:• The answer NA means that there is no societal impact of the work performed.• If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does not address societal impact. • Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations. • The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is a direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster. • The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology. • If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML). 
safeguards: 
 Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?Answer: [NA] Justification: We do not release data or models.Guidelines:• The answer NA means that the paper poses no such risks.• Released models that have a high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters. • Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images. • We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make a best faith effort.12. Licenses for existing assets Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?Answer: [Yes]Justification: We have used standard image datasets and architectures and cited the original work.Guidelines:• The answer NA means that the paper does not use existing assets.• The authors should cite the original paper that produced the code package or dataset.• The authors should state which version of the asset is used and, if possible, include a URL. • The name of the license (e.g., CC-BY 4.0) should be included for each asset.• For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided.• We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurIPS Code of Ethics and the guidelines for their institution. • For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.Department of Physics and Joint Quantum Institute, University of Maryland, College ParkInstitute for Physical Science and Technology, University of Maryland, College Park38th Conference on Neural Information Processing Systems (NeurIPS 2024).for details, see https://jax.readthedocs.io/en/latest/_autosummary/jax.nn.initializers. truncated_normal.html 
 
 Appendix:  


Decision: This paper provides a detailed empirical analysis of the warmup mechanism in neural network training, focusing on its role in improving training stability and enabling the use of larger learning rates. The authors also introduce the GI-Adam optimization method, which aims to enhance the initialization of the Adam optimizer. Strengths of the paper include the solid experimental validation of the warmup process, particularly its influence in different sharpness evolution regimes (as noted by 6XDG and kKLj), and the extension of these findings to Transformer models, which addressed concerns about the generality of the results (raised by LAU1 and YVuK). Additionally, the paper’s insights into the interplay between warmup duration and learning rate are appreciated for their potential practical applications in training neural networks (mentioned by kKLj and YVuK).

However, there are some weaknesses. ZUpa expressed concerns about the novelty of the paper, specifically the lack of a clear distinction from previous research, which was not fully alleviated by the authors’ response. Additionally, 6XDG pointed out that while the proposed GI-Adam method is reasonable, it does not offer significant improvements over existing alternatives like RAdam, which limits its practical impact.

While concerns remain regarding the novelty and the practical effectiveness of GI-Adam when compared to existing methods, I believe these issues can be adequately addressed if the manuscript clearly discusses these points in the main text. The numerical analysis of the warmup mechanism, focusing on different regimes during training, provides significant insights that are valuable to the NeurIPS community. Therefore, I recommend accepting this paper.

Reviews: review: paper_summary: paper_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.
strength_weakness: strength_weakness: 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: questions: I don't have questions.
limitations: limitations: Please refer to the Weaknesses
paper_summary: paper_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.
strength_weakness: strength_weakness: 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: 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.
limitations: limitations: The authors adequately discuss the limitations of the paper in the conclusion.
paper_summary: paper_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.
strength_weakness: strength_weakness: 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: questions: See Weaknesses.
limitations: limitations: N/A
paper_summary: paper_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.
strength_weakness: strength_weakness: - 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: 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)
limitations: limitations: Please refer to ``Weaknesses''.
paper_summary: paper_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.
strength_weakness: strength_weakness: 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: 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?
limitations: limitations: Yes

Comments: title: Responses to minor questions by Reviewer 6XDG
comment: Thanks for the response. I have no further questions.
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.
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.
> 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.
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.
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.
Thanks for the authors' response. I have no further questions and raised my score.
Many thanks to the authors for the detailed response. I feel that this paper provides great insights into Warmup. I have raised my score.
Thank you for the reply. I think that these additions to the paper will make it even stronger.
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.