Stepping on the Edge: Curvature Aware Learning Rate Tuners

We thank the reviewer for reading this manuscript and providing valuable feedback.

We answer their comments below.

"The proposed CDAT tuner has several limitations"

• Our goal with this study is to underscore the interplay between sharpness dynamics and stepsize tuners. The CDAT rule serves as a probe to question both the design of learning rate tuners and the study of sharpness dynamics beyond constant learning rates settings. CDAT does stabilize the sharpness dynamics that cause classical tuners to fail (Fig. 1) and seems to induce an automatic warmup (Fig. 5). It also demonstrates the limitations of the current theory on the sharpening and edge of stability phenomena. We agree that CDAT is not yet a mature optimizer but we believe the concepts it introduces will help move the field forward.

"How much does CDAT increase the runtime for one step?"

• As detailed in Appendix C.3, CDAT required twice the wall-clock time of the constant or scheduled learning rates counterparts. The computation of $g_t^\top H_t g_t$ is done by two forward mode autodiff avoiding the memory cost of Hessian vector products. We leave for future work an optimized implementation of the rule by e.g. computing $g_t^\top H_t g_t$ at fixed intervals.

"CDAT replaces the necessity to tune the learning rate with tuning sigma."

• The usual selection of the learning rate may largely vary with the optimizer, the architecture and the dataset considered. For example a learning rate of 10^{-3} is generally used for Adam while the learning rate of GD/SGD can vary much more. Though the CDAT rule does not always outperform a constant tuned learning rate, it surprisingly performs well across tasks and optimizers for $\sigma=2$.

"few experiments are conducted in a stochastic regime […] the results are not very convincing"

• Linesearches and other stepsize tuners have indeed generally been tested in a stochastic regime (see also Fig. 14). The potential drawbacks of e.g. linesearches have generally been attributed to small batch sizes [11, Fig. 7]. The analysis of stepsize tuners in a full batch regime shows a different picture: linesearches or greedy methods fail because of phenomena like progressive sharpening. One of our contributions is to question general beliefs about the design of learning rate tuners which may have missed important properties by jumping into the performance in a stochastic regime.
• Rather than adding yet another optimizer to the long list of existing ones, we preferred to provide numerous experiments to diagnose the benefits and/or challenges (such as adapting the scaling to the batch size, Fig. 7) that arise when adapting the learning rate to the sharpness dynamics.

"[...] thus the main motivation for the derivation of CDAT, namely to target EoS, seems not fully convincing in the first place."

• We apologize for the confusion. For fixed step size, the typical dynamics of EOS is as follows: at early to intermediate times, training reaches the edge of stability and stabilizes there. For intermediate to late times, the sharpness drops below the edge of stability and training converges. The early time behavior corresponds to the regime where there is the most feature learning/the Hessian changes the fastest; late time corresponds more to the final convergence (see e.g. [24, 25]).
• The key issue with classical stepsize tuners is that they stay below the edge of stability in the early-intermediate dynamics, where the geometry of the loss landscape is developing the most quickly. This is what leads to the runaway process by which the sharpness keeps increasing, which can sometimes prevent training from even getting to the intermediate/late time dynamics where the Hessian changes slowly, and training converges (see Fig. 3). Therefore stabilizing the early dynamics is key to training success. It remains an open question as to how long it is good to stay at the EOS. We will include a more detailed discussion of this point.

"In the Damian et al. paper (ref. [34]), the dynamics are not the same as what you write in (4)"

• We wrote down the equations with $\lambda$ instead of $y$ (after converting appropriately) in order to link to the previous discussion of $\lambda$. We switched back to the variable $y$ when analyzing the more detailed model.

"Why do you use the squared loss for the CIFAR dataset, instead of the cross-entropy loss?"

• We followed Cohen [24, 25] that first analyzed the sharpening and edge of stability phenomena on squared losses. We nevertheless demonstrate the behavior of algorithms in various other settings and losses (see Fig. 1, 5, 8, 12, 19, 21).

"Why is Fig. 4 using model (4), and not model (5), if the goal is to illustrate learning rate tuners, where the LR depends also of time (model (5))?"

• The dynamics does indeed reflect model (5); thank you for catching that error, we will correct the reference.

"Fig. 3: how did you make sure that $\eta>0$ for the Quadratically Greedy Linesearch?"

• In Sec. 1, we considered the quadratically greedy rule only with optimizers that did not incorporate a momentum term. In that case, the numerator $-g^\top u$ of the quadratically rule is always positive since the updates $u$ are either the negative gradients $-g$ (for GD) or an element-wise positive scaling of the negative gradients (for RMSPROP). The denominator $u^\top H u$ could be negative only (i) if the Hessian is not positive definite at the updates, (ii) the updates do not belong to the eigenspace of positive eigenvalues. The eigenvalues of the Hessian are generally overwhelmingly positive, and, in practice, we have never observed a negative learning rate. The CDAT rule with $\sigma=1$, that clips the learning rate to positive values can also be seen as a "safe'' quadratically greedy rule.

Thanks for catching the typos that we will correct in our manuscript. We look forward to a fruitful discussion that can further elucidate any concerns.

We sincerely thank the reviewer for the positive feedback and to appreciate the exploratory nature of this work.

Below we answer their comments.

"it does not outperform the baseline of a constant learning rate in the practical deep learning use-case of mini-batch optimization [...] it has an additional hyperparameter $\sigma$"

• Yes, the mini-batch case suggests that a rule placing the optimizer "on the edge'' needs to take into account stochasticity. We have been first and foremost interested in the full-batch case because the performance of e.g. linesearches reported in the stochastic setting left out the intriguing poor performance in large or full batch cases (see also [31]). While the CDAT rule still requires additional mechanisms to handle stochasticity (the recent schedule-free optimizer (1) may be combined with the proposed rule for example), we preferred providing numerous ablation studies to understand the interplay between sharpness dynamics and learning rate tuners in diverse scenarios.
• Note that at larger batch sizes, a scaling factor of 2 still works well (Fig. 7 and 20). On the other hand, the behavior of usual tuners such as linesearches suffer from unexpected changes in behavior as the batch size changes: at both small and large batch sizes they can perform worse than constant learning rate counterparts (see e.g. [31]). We hope that the insights provided in the full-batch setting gathered in the provided experiments will help refine the design of learning rate tuners in the stochastic case. We agree that in the stochastic case the parameter $\sigma$ still requires to be tuned. On the other hand, we note that in the full batch setting the CDAT rule with a scaling of $\sigma=2$ performs well across different optimizers. This is particularly surprising given that an optimizer such as Adam generally requires a learning rate of $10^{-3}$ while the learning rates of e.g. SGD with momentum vary greatly with the problem.

"[the fact that] the optimal scaling factor is mini-batch dependent [...] is not tested"

• We would like to point the reviewer to Fig. 7, which shows that the optimal scaling factor varies with the mini-batch size. Kindly let us know if that doesn't answer the question.

"I would appreciate it if the theoretical model from Damian et al (2022)(together with its underlying assumptions) is summarized"

• Thank you for the suggestion. We will add a complete introduction to the model of Damian et al (2022) for the final version. The gist of the model is summarized in page 6 line 161.

"What conclusions can be drawn on the learning rate and Hessian interplay? [...] the updates are not aligned with the leading eigenvector [...] ?"

• One interesting feature of CDAT is that it allows for learning rates to be instantaneously larger than the edge of stability. If the updates are not aligned with the large eigenmodes (as is the case during early iterations), the tuner chooses large step sizes - since the instability won’t matter over a few steps. More generally, the dynamical nature of the alignment allows CDAT to be more flexible and respond to changing curvature in a way that overall leads to stabilization at the EOS without sacrificing optimization.
• We would also like to highlight Fig. 5 top left panel: placing the optimizer well above the edge ($\sigma=2.5$) does not necessarily lead to divergence in early optimization stages. The optimizers may in fact benefit from a mechanism that ensures no growth in sharpness at the start, beyond thinking whether the optimizer is then unstable according to some quadratic approximation. It is still unclear whether choosing a learning rate rule can actually extract some third order derivative information that drives the sharpness into lower values.

"What are the implications of this paper's results in [the] context [of [2]]?"

• Thank you very much for the interesting reference! Our goal is somewhat orthogonal to this work, as we aim to harness the sharpening effects to derive an "automatic warmup'', as well as an automatic selection of the learning rate once a relevant warmup phase has cooled down. Yet, the provided reference suggests that a fixed warmup may also transfer across scales. The question is then what shape the warmup must take and what is it controlling. The CDAT rule may give a post-hoc explanation to this question. Our hope is to let a learning rate tuner such as CDAT capture the right warmup shape and peak learning rate in a stochastic regime with small to medium models and then use scaling laws to transfer the found warmup schedule to larger models.

References:
(1) The Road Less Scheduled, https://arxiv.org/abs/2405.15682

We thank the reviewer for reading our paper and providing valuable feedback.

We answer below their comments.

"it studies optimizers which are not state-of-the-art in the first place"

• Could the reviewer clarify what optimizer do they have in mind? We already added additional learning rate tuners (Polyak stepsizes and hypergradient descent) in Fig. 13 and plan to add also recent tuners such as Prodigy (1) or DoG (2). In practice Adam and its variants are generally still considered as the main optimizers used in practice, hence we provided experiments with those.

"Figure 16 seems to have the same figure accidentally repeated twice."

• Thanks for catching the mistake on Fig. 16, we corrected this. The rightmost figure offered a similar conclusion as detailed in the caption.

"I am skeptical about the finding that the efficacy of CDAT goes away with large weight decay."

• As pointed out by J393, a future step will be to understand the interplay between stepsize tuners and sharpness dynamics in terms of scaling laws (using the findings of (3)). In particular, we agree that ablation studies as done in Fig. 18 will benefit from placing ourselves in a rich feature learning limit.

References:
(1) Prodigy: An expeditiously adaptive parameter-free learner, https://arxiv.org/abs/2306.06101
(2) Dog is sgd’s best friend: A parameter-free dynamic step size schedule, https://arxiv.org/abs/2405.15682
(3) Why do Learning Rates Transfer? Reconciling Optimization and Scaling Limits for Deep Learning, https://arxiv.org/abs/2402.17457

We thank the reviewer for their detailed feedback, and address their comments here.

"Theoretical ground is not solid"

• The experiments and models focus on the full batch regime just as Cohen [24, 25] did to uncover the edge of stability phenomenon. The theory is grounded in the full batch regime following previous work [34]. We clearly point out the limitations of the model in a stochastic regime referring to [24, 32, 33] (we thank the reviewer for the additional reference that we will add).

For [SGD], [...] the correct measure of sharpness [is] the trace of Hessian instead of the largest eigenvalue [1], so using the largest eigenvalue [...] for SGD [...] is not reasonable.''

• We would like to start by clarifying what we believe is a critical misunderstanding. The CDAT rule does not use the largest eigenvalue of the Hessian. In fact as demonstrated by [24, Appendix F], and further explored in Fig. 15, using the exact sharpness does not provide gains unless the scaling is much larger than 2 (3 in that case). The CDAT rule incorporates the alignment of the gradient with the Hessian as reviewer J393 also identified. That's what the model in Sec. 3.2 points out too.
• The trace of the Hessian indeed controls the EoS in highly stochastic settings; the largest eigenvalue controls stability due to first moments, while the trace controls second moments (3). The trace becomes more important at later times, e.g. during learning rate decay. This means that in SGD settings, there are ranges of batch sizes/stages of training where the largest eigenvalue can still control stability/EOS behavior ([24, Fig. 24]; (4, Fig. 5)). The CDAT tuner seems to have the most beneficial effect in the early stage of training (to induce some warm-up). In this first work we designed CDAT around these first moment effects and focused on empirical studies, but we hope to incorporate late time stochastic effects into future work.

"For adaptive optimizers [...], the correct quantity [...] is the pre-conditioned Hessian [...]. In Fig 8 [...], the initial learning rate is very large, which is [...] because the authors used Hessian instead of pre-conditioned Hessian.''

• The CDAT rule takes care of this implicit preconditioning. Both the numerator and denominator depend on the update $u$. For pre-conditioned optimizers, $u = -P^{-1} g$, where $g$ is the gradient, and $P$ is the diagonal preconditioner. The reciprocal of the CDAT rule then reads

$$\frac{u^\top H u}{-u^\top g} = \frac{g^\top P^{-1} H P^{-1} g}{g^\top P^{-1} g}.$$

If we were to maximize the above ratio, we would get

$$\max_g \frac{g^\top P^{-1} H P^{-1} g}{g^\top P^{-1}g} = \max_v \frac{v^\top P^{-1/2} H P^{-1/2} v}{\||v||^2} = \lambda_{\max}(P^{-1/2}H P^{-1/2}),$$

and the CDAT rule "on edge" would then be $2/\lambda_{\max}(P^{-1/2}H P^{-1/2})$, the edge of stability of adaptive gradient methods without momentum [25] (again the CDAT rule does not $\lambda_{\max}$ but incorporates the alignment).

• The initial large learning rate in Fig. 8 is likely due to the fact that the preconditioner itself is highly variable in the very first iterations.

"For cross-entropy loss, [...] the sharpness will reach $2/\eta$ and then come down later […] using $\eta\sim 2/\lambda$ does not mean it will remain on the EoS [...] depending on the initialization, the sharpness might decrease or increase at an early stage.''

• The curvature dynamics for constant learning rates is indeed complex. However, the complexities pointed out by the reviewer (that will be added) need to be revisited in the context of variable learning rates defined through a learning rate tuner. Learning rate tuners introduce closed loop feedback effects, and so even more complex behavior. To revisit these results, a first set of experimental results need to be laid down and that's what we present. The complexities pointed out by the reviewer do not seem crucial to observe the failure of classical learning rate tuners. On the other hand, the stabilization mechanisms at EoS (captured by the models of [34], and observed in many nonlinear models trained at large batch size during some appreciable fraction of training) can provide first insights on the closed loop feedback effects. That's what the CDAT rule puts to the test. CDAT provides a closed loop feedback that encourages and is compatible with EoS dynamics - but does not put a hard-constraint towards EoS. Our experiments with CDAT show some novel behaviors. For example, in Fig. 6 $\lambda_{\max} \eta$ stabilizes below 2 for GD, while for CDAT $\lambda_{\max} \eta$ remains slightly above 2. We also observe that using the CDAT rule, the sharpness may even decrease at later times (Fig. 6). These empirical findings are new and crucial to put the EoS studies into practical use.

"Experimental results are not good.''

• See main rebuttal. We agree that CDAT is not yet a mature solver. We preferred to focus on an extensive set of experiments analyzing the idea rather than adding yet a new optimizer whose performance would be left to be diagnosed by peers. None of the references provided by the reviewer provided a new state of the art solver, neither did they propose new methodologies to design new solvers.

Please let us know if we can clarify our thinking further; we look forward to a fruitful dialogue with you.

References:

• (1) The Road Less Scheduled, https://arxiv.org/abs/2405.15682
• (2) Second-order regression models exhibit progressive sharpening to the edge of stability, https://arxiv.org/abs/2210.04860
• (3) High dimensional analysis reveals conservative sharpening and a stochastic edge of stability, https://arxiv.org/abs/2404.19261
• (4) SAM operates far from home: eigenvalue regularization as a dynamical phenomenon, https://arxiv.org/abs/2302.08692