Stepping on the Edge: Curvature Aware Learning Rate Tuners
Authors: Vincent Roulet
Atish Agarwala
Jean-Bastien Grill
Grzegorz Michal Swirszcz
Mathieu Blondel
Fabian Pedregosa
Venue: NeurIPS 2024 poster (2024)
Keywords: deep learning optimization
learning rate tuner
progressive sharpening
edge of stability
Primary area: optimization_for_deep_networks
TL;DR: We investigate the interplay between curvature dynamics and learning rate tuners through a new learning rate tuning method.

Abstract: Curvature information -- particularly, the largest eigenvalue of the loss
Hessian, known as the sharpness -- often forms the basis for learning rate
tuners. However, recent work has shown that the curvature information undergoes
complex dynamics during training, going from a phase of increasing sharpness to
eventual stabilization. We analyze the closed-loop feedback effect between
learning rate tuning and curvature. We find that classical learning rate tuners
may yield greater one-step loss reduction, yet they ultimately underperform in
the long term when compared to constant learning rates in the full batch regime.
These models break the stabilization of the sharpness, which we explain using a
simplified model of the joint dynamics of the learning rate and the curvature.
To further investigate these effects, we introduce a new learning rate tuning
method, Curvature Dynamics Aware Tuning (CDAT), which prioritizes long term
curvature stabilization over instantaneous progress on the objective. In the
full batch regime, CDAT shows behavior akin to prefixed warm-up schedules on deep
learning objectives, outperforming tuned constant learning rates. In the mini
batch regime, we observe that stochasticity introduces confounding effects that
explain the previous success of some learning rate tuners at appropriate batch
sizes. Our findings highlight the critical role of understanding the joint
dynamics of the learning rate and curvature, beyond greedy minimization, to
diagnose failures and design effective adaptive learning rate tuners.

Stepping on the Edge: Curvature Aware Learning Rate Tuners
Vincent Roulet" Google DeepMind vroulet@google..om com
Atish Agarwala' Google DeepMind thetish@google.com com
Jean-Bastien Grill Google DeepMind jbgrill0google.co com
Grzegorz Swirszez Google DeepMind swirszcz0google.com
Mathieu Blondel Google DeepMind mbl obblodde@google.... com
Fabian Pedregosa
Google DeepMind pedregosa@googgee.com com
Abstract
Curvature information particularly, the largest eigenvalue of the loss Hessian, known the sharpness often forms the basis for learning rate tuners. However, recent work has shown that the curvature information undergoes complex dynamics during training, going from phase of increasing sharpness to eventual stabilization. We analyze the closed-looy feedback effect between learning rate tuning and curvature. We find that classical learning rate tuners may yield greater one-step loss reduction, yet they ultimately underperform in the long term when compared to constant learning rates in the full batch regime. These models break the stabilization of the sharpness, which we explain using a simplified model of the joint dynamics of the learning rate and the curvature. To further investigate these effects, we introduce a new learning rate tuning method, Curvature Dynamics Aware Tuning (CDAT), which prioritizes long term curvature stabilization over instantaneous progress on the objective. In the full batch regime, CDAT shows behavior akin prefixed warm-up schedules on deep learning objectives, outperforming tuned constant learning rates. In the mini batch regime, we observe that stochasticity introduces confounding effects that explain the previous success of some learning rate tuners at appropriate batch sizes. Our findings highlight the critical role of understanding the joint dynamics of the learning rate and curvature, beyond greedy mininization, to diagnose failures and design effective adaptive learning rate tuners.
Introduction
The learning rate, a.k.a. stepsize, the main hyperparameter controlling the efficiency and stability of gradient-based training of deep neural networks. The learning rate is typically adjusted through predetermined schedule often consisting of a warm-up phase, where the learning rate is gradually increased to a peak, followed by an annealing phase, where it is decreased to zero (Goyal et al., 201 Loshchilov and Hutter, 2016). Tuning the shape of the schedule (warm-up time, peak learning rate, decay scale and shape) is essential for good performance. Despite recent efforts to understand their effectiveness, the optimal shape of these schedules remains an area of active research (Liu al., 2020; Shi et al., 2023). The cost of tuning these schedules has led to interest in automatic selection of these hyperparameters with learning rate tuners methods which aim automatically adjust the learning rate through training.
Equual contribution.
38th Conference on Neural Information Processing Systems (NcurIPS 2024).

---

These methods have roots in traditional optimization theory, including inexact linesearch with Armijo- Goldstein criterion (Armijo, 1966; Nocedal and Wright, 1999) and Polyak stepsizes (Polyak, 1964), which select the learning rate via estimates of the gap to optimality of the objective. The Armijo- Goldstein criterion a crucial component of popular full-batch convex oetinizetrii optimizers, such as L--FGS (Liu and Nocedal, 1989). Recent efforts have adapted linesearches to stochastic optimization, with some partial empirical successes and with some approaches offering convergence guarantees (Galli et al., 2023; Mutschler and Zell, 2020; Vaswani et al., 2019). Similar efforts have been made for Polyak stepsizes (Berrada et al., 2020; Loizou et al., 2021), in addition to new methods which combine distance to optimality with online learning convergence bounds (Cutkosky et al., 2023; Defazio and Mishchenko, 2023; Ivgi et al., 2023; Mishchenko and Defazio, 2023).
Classicall-innsiree methods, however, have generally struggled to gain traction in deep learning. This is partly due to their design, which prioritizes convex, Lipschitz-conninuuus and/or smooth (Lipschitz-continuon) gradients) objectives. In contrast, the loss landscape of deep networks is known to be non-convex (Li et al., 2018), and non-Lipschitz continuous (Hochreiter et al., 2001). Moreover, non-linear models, especially neural networks, will commonly undergo dramatic changes in geometry during training (Arora et al., 2022; Jastrzębski et al., 2019; Jastrzebski et al., 2020; Kalra et al., 2023; Kopitkov and Indelman, 2020; Lewkowycz et al., 2020; Wu et al., 2018). In particular, most models undergo a phase of progressive sharpening where the sharpness, the largest eigenvalue of the Hessian, increases during training (Cohen et al., 2021). These potentially detrimental effects are mitigated by non- con-linear stabilization arising from the discreteness of the dynamics namely, the edge of stability (EOS) phenomenon (Cohen et al., 2021). This causes large Hessian eigenvalues to stabilize at the critical value for a given learning rate in an equivalent smooth setting (for example, max Hessian eigenvalue stabilizes Amax = 2/1 for learning rate 1)) (Cohen et al., 2023, 2021). The early training time behavior corresponds to the regime where there is the most feature learning (Cohen et al., 2023, 2021), and is the main focus of this work; at late times, the large eigenvalues of the Hessian usually drop below the edge of stability. Gilmer et al. (2022) considered EOS stabilization as a leading candidate for the necessity of the warm-up procedure; as the learning rate 1) increases, Amax is effectively annealed.
This raises some natural questions. How do these sharpness dynamics affect the performance of learning rate tuners? What insights can we gain to design better tuners for deep learning? Our work takes first step answering these questions, starting with study of some classical learning rate tuners: a linesearch ensuring sufficient decrease and an approximately greedy method that minimizes quadratic approximation of the objective. Specifically, we find the following.
We empirically observe that classical learning rate tuners qualitatively underperform their constant learning rate counterparts across several deep learning benchmarks, in the full batch regime, for which these methods were originally designed.
Our empirical analysis of curvature dynamics reveals that classical learning rate tuners generally undershoot the edge of stability. This undershooting creates a snowball effect of ever-increasing sharpness and ever-decreasing learning rates.
We propose a theoretical model that effectively captures these empirically observed failures.
Our analysis suggests that stabilizing the sharpness may be a more important goal for the long-term success of training, compared to greedily optimizing the objective. To explore this idea, we propose the Curvature Dynamics Aware Tuning (CDAT) method, which dynamically drives the learning rate to the EOS. our exploration, we find the following.
We observe empirically that the proposed learning rate tuner can outperform fine-tuned constant learning rate counterparts full batch regime.
We analyze the sharpness dynamics induced by CDAT in these examples and observe that the progressive sharpening is mitigated by the tuner, increasing learning rates at early times before stabilizing, akin to an automatic warm-up schedule.
We propose a theoretical model that clarifies the dynamical mechanisms by which CDAT maintains proximity to the EOS, while highlighting the limitations of existing models of curvature dynamics.
Our work suggests that the design of learning rate tuners benefits from exploiting curvature sta- bilization rather than focusing on loss decrease. The introduction of simple learning rate tuners can also refine our understanding of sharpness dynamics through feedback loop effects. Additional experiments and experimental details are presented in Appendix B and Appendix C respectively.
2

---

PlesNeB34 with loss
MieeTTeB with loss subeet Cifar1
croes-entropy Shukesspere
with stoss-entrppy Imageriel
subset Cilar10
Tran Lo85
Tean Lows
Train Loss
Truin
3000
10300
15000
2000
4000
00:00
5000
6912
10:00
15:00
2030
600
1900
1600
2000
Epoch
Eppcch
Epoch
Epoch
GD
Linesearch
GD
Quad Greedy GD
RMSProp -i- Linesearch RMSProp - Quad. Groody RMSProp
Figure 1: Simple learning rates tuners qualitatively underperform their constant learning rate counterparts. Gradient descent or RMSProp with a tuned constant learning rate versus self-tuned gradient descent by linesearch method (1), quadratically greedy rule (3) various datasets, architectures and losses in a full batch regime. The linesearch may perform better at early times but stalls in the long term.
2 The Interplay Between Learning Rate Tuners and Curvature Dynamics
A leitmotif in the design of learning rate tuners has been to select the learning rate to ensure a maximal or sufficient decrease of the objective each iteration. We focus here on two canonical examples. Polyak stepsizes and typer-gradien descent are also briefly examined in Appendix B, Fig. 13.
2.1 Canonical learning rate tuners failures in deep learning
The first classical approach we consider linesearch (Is) method that selects the learning rate n such that the objective f satisfies certain decrease criterion (Armijo, 1966; Nocedal and Wright, 1999). Formally, given current parameters wy and an update direction uy the learning rate of chosen such that
[(w + u.) S (wwr) + enl u V ([())..
This rule assumes that uy is a descent direction (w un < 0), which ensures the existence of learning rate satisfying (1). This holds true for simple Gradient Descent (GD) preconditioned variants like RMSProp (Hinton et al., 2012). In the criterion (1), is usually small constant set to 10-4 or O. A valid learning rate is searched with usual backtracking linesearch (Appendix C).
The second method we consider involves selecting the learning rate at each iteration to minimize quadratic approximation of the objective. Formally, the objective I at parameters wy can be approximated along an update direction u, by a quadratic approximation L as
1 n
[w,, + quu) Ro of (TE We, u.) = f(wx) + 1/V H (wx) u. Tu.
ue (w:))u
(2)
Provided that this quadratic approximation is strongly convex in 1 (u, V fw,(u, > 0), the minimum of the quadratic approximation gy(r; Wi, we) is reached for the quadratically greedy (qg) learning rate 1788 given by
Linear model with squared loss on subset of Cifar10
0.500
0.475
Linesearch GD
Train Loss
0.450
Quad. Greedy GD
-V [((w))
98 V 21 42 f(w))
0.425
0.100
0.375
Setting the learning rate by minimizing the quadratic ap- proximations (3) is simple intuitive idea studied for ex- ample by Schaul et al. (2013), Martens and Grosse (2015, Section 6.4). This approach as well as linesearches are effective on simple linear problems (Fig. 2). While their rationale originates in non-stochastic optimization, they have been analyzed in the context of stochastic optimiza- tion for deep learning (Schaul et al., 2013; Vaswani al.,
0.350
2000
4000
6000
8000
Epoch
Figure 2: Classical learning rate tuners can be effective on linear models.
2019).
3

---

ResNet34 with les on subset Citar10
Leatring Nate X Sharptees
Messien Sharpnees
Learing Raate
122
Gradiert Norm
10-2
12-
0050
10000
12000
2000
10000
15000
2000
10000
19000
0000
10000
10000
Epoch
Epoch
Epoch
Epoch
GD
Linesearch
GD
Quad. Greedy GD
Figure 3: Classical learning rate tuners undershoot the edge of stability. Learning rate, sharpness, their product, and the gradient norm evolution of a constant learning rate and learning rate tuners, full batch gradient descent. Learning rate decreases by 3 orders of magnitude for tuners (1st panel) while sharpness increases (2 panel). Their product remains relatively steady, just below the edge of stability (3th panel). The gradient norm increases by less than a factor of 10, consistent with slow training at late times (4th panel).
2.2 Analyzing learning rate tuners through curvature dynamics
Full batch regime. We revisit the performance of the learning tuners presented in Section 2.1 in the full batch regime on deep learning problems in Fig. 1. As demonstrated in Fig. 1, a linesearch (1) or the quadratically greedy rule (3) qualitatively underperform their constant learning rate counterpart in the deep learning benchmarks considered. Notably, all these results are obtained despite being in a full batch regime, for which these methods are originally designed. To understand the failures of these upproaches, we consider several measures presented in Fig. 3 (see also Fig. 12).
First, we observe consistent decrease in the chosen learning rate time, spanning several orders of magnitude (1st panel of Fig. 3). This is surprising, as none of these approaches explicitly encode a decreasing learning rate mechanism. Specifically, the linesearch always initiates its search with guess larger than the previously selected learning rate (see Appendix C for implementation details). Decreasing learning rates are theoretically optimal for non-smooth objectives (Nesterov et al., 2018), such as the ones induced by using the ReLU activation; however in our example, the gradient norm does not increase beyond one order of magnitude (4th panel of Fig. 3). This suggests both that an increase in gradient norm is not the primary cause of learning rate decrease, and also explains why the learning rate decrease is correlated with slower progress on the training loss.
Following the work of Cohen et al. (2021), ) we analyze the dynamics of the sharpness, that is the largest eigenvalue of the Hessian, Amax (V2 (wx)). In the 2nd panel of Fig. 3, observe that while sharpness stabilizes for gradient descent, it does not exhibit the same behavior for the considered learning rate tuners. By plotting the product of the learning rate 7)- and the sharpness (3rd panel of Fig. 3), we find that this product can exceed the stability threshold of 2, eventually stabilizing below this threshold for constant learning rate gradient descent. In contrast, for the learning rate tuners, this product neither surpasses the stability threshold nor stabilizes around 2 in the long run. Therefore, these classical learning rate tuners do not operate the edge of stability.
From a theoretical perspective, objectives are typically classified as either smooth or non moo-sh Smooth objectives have gradients that are Lipschitz-continuous, at least locally around any point. Non-smooth objectives, on the other hand, may contain points with kinks (non-differentiable points). However, this taxonomy might not fully capture the curvature dynamics observed by Cohen et al. (2023, 2021) for constant learning rates, and in Fig. for the classical learning rate tuners. In particular, the concept of smoothness might not be entirely relevant in the context of deep learning, where its local estimate (the spectral norm of the Hessian, also known as sharpness) can continue to increase throughout training. To push the limits of classical smoothness assumptions, we consider in Section 3 a learning rate tuner that propels the optimizer at the edge of stability or above, a regime that usual smoothness assumptions would theoretically prohibit.
Mini-batch regime. The results presented in Fig. in the full batch regime do not contradict the success of linesearches at medium batch size observed by Vaswani et al. (2019) in the stochastic regime. This observation is illustrated in Fig. 14, and was previously reported by Roulet et al. (2023). We simply point out that the success of linesearches observed by Vaswani al. (2019) may not be entirely attributable to the method's original rationale.

---

1.0
0.50
p:: 18/A
0.25
0.8
16-
0.6
10
00.25
-0.50
0.
75
Steps
Steps
Steps
The poor performance of classical learning rate tuners, understood in a simplified model. The dynamics of learning rate 17, sharpness Xmax. and normalized centered sharpness y = r//maxx - are examined in the simplified model (5). With a constant 7), Amax stabilizes and oscillates around 0 (blue). Classical learning rate tuners often quickly equilibrate around y/- -6, which we model using 1 = 1.9 (orange). This equilibration of y away from zero prevents stabilization in Amax, leading to an increase in Amax. and corresponding decrease in 17.
Figure 4:
The actual success of linesearches in stochastic regime may instead be explained by the attenuated progressive sharpening observed such a regime (Agarwala and Pennington, 2024; Cohen et al., 2021; Jastrzębski et al., 2017). Moreover, linesearches applied mini-batches tend to select larger learning rates than they would in a full-batch regime (Mutschler and Zell, 2020) potentially allowing them to avoid undershooting the full objective's 's edge of stability.
2.3 Theoretical analysis
The sharpening effects can be understood theoretically. Previous work has shown that the stabilization provided by EOS is due to non-linear interaction between the component of the gradient in the largest eigendirection, and the dynamics of the largest eigenvalues themselves (Agarwala et al., 2023; Damian et al., 2023). We can use these analyses to understand why there is no stabilization for some classical learning rate tuners.
We start with the model from Damian et al. (2023), which focuses on the dynamics in the largest eigendirection of the Hessian. We consider unique eigenvector for simplicity: we don't 't observe degeneracy in the eigenspace of the largest eigenvalue in any practical models. Given an objective parameterized by parameters wi, let A be the largest eigenvalue of the Hessian V f i.e.. d = A(w,)= = Amax (wz)). Let v be its normalized eigenvector; the model assumes slow eigenvector change, so treated as a fixed direction. The joint dynamies of X₂ and the projection Jt == TTw, can then be written as
I++1 = - 7/ De)ae =
= = (a ((ca-bai bx,)) + X
Here, = -V\w) Vv(ww) corresponds to the instantaneous change of À along the negative gradient (the update direction), and == ||>X(=)] encodes the non-linear negative feedback between 2y and Aç. Both a and b are considered constant along iterations. These equations are derived by Damian al. (2023) using a Taylor expansion of the iterates combined with a coupling argument. We provide intuition for the model in Appendix A.
In the original model, the learning rate 1); is also fixed 17. This leads to the following dynamics: while 17 < 2, the magnitude of I, decreases. This, in turn, leads to an increase in λe. Eventually, m) > 2 and (arel increases. This eventually leads the bug term becoming large, which decreases X₁- There a range of learning rates over which this dynamic leads to quasi-stable oscillations of X around the edge of stability value 2/1 (Fig. 4. blue curves).
When using a learning rate tuner, 7: is also a dynamical variable. This introduces the additional
complication of shifting edge of stability. Therefore, it is advantageous to analyze the dynamical
system using normalized variables (Agarwala et al., 2023). We define y/2 == Ae 2, where y = 0 corresponds to the EOS, and Pt x² This gives the dynamical equations (Appendix A.2)
P++1 + pr. U M [n,. (a bp:)]
31
+
2
M
(5)
1/t
7
We must then supply rule for 7 In Fig. 3, we observed that in the full batch setting, the learning rate multiplied by the sharpness appears to quickly approach a threshold of (corresponding to y -)), and then varies slowly below the EOS threshold.
5

---

squared
Citar10
NancLM with entropy subeet of Tiny Shakespeare
10
0.3
Learnnng Rate
1ir
Tain Lose
Learing
10-1
0.-
Trei
101
2000 Epoch
1002
3000
Epoch
Epoch
CDAT Scale
RMSPr00
RMSProp CDAT 2.0
GD CDAT Seate 1.0 Scale
GD CDAT Scale 2.05
RMSP mo CDAT Scale 1.0
RMSProp CDAT Scale 2,06
AMSProp
RMSProp CDAT
MikerT/B with squared loss on subset of Gifar10
VITTI/32 with cross entrapy loss on subeet Imagenst with edge EMA
10
10"
Leatring Raee
Loaning Palln
Tran Loos
X2-
10-
22-
Train
10-
12-
12-
2000
4000
2000
2000 Epoch
Epoch
Epoch
GD Mom
Mom CDAT
Adam
CDAT Scale 2.0
GD Mam CDAT Soale 1.0
GD Mom CDAT Scale 2.05
Adam CDAT Scale 1.0
Adam CDAT Scale 2.06
Mom CDAT Scale
Adam CDAT
COAT
Figure 5: Enforcing optimizers to stay on edge (a 2.0) improves performance over greedy approximation (o 1.0). Train loss and learning rate behaviors for fine-tuned optimizers vs self-tuned counterparts with CDAT on various datasets, architectures, losses in a full batch regime. Tuning the learning rate "on edge" (a 2) improves performance over greedy tuning (a as well as constant learning rate.
We model the varying learning rate
(6) Notably, this schedule explicitly proposed by Cohen et al. (2021) (see
Tp 2
This maintains y4 also Fig. 15). In this regime, Pt decreases monotonically, aligning with the original goal of these methods to decrease the loss (Fig. 10). However, this eliminates feedback for controlling the increase in A₁ resulting in significant progressive sharpening (Fig. 4, orange curve).
Consequently, when attempting to enforce monotonicity, learning rate tuners may inadvertently disrupt the non-linear stabilization that makes gradient descent robust and effective for training deep neural networks. Continually undershooting the EOS triggers a snowball effect of decreasing learning rate and increasing sharpness. If there is no corresponding increase in gradient norms, this causes optimization to slow down.
The poor performance of the classical learning rate tuners in Fig. therefore appear strongly correlated with their tendency to undershoot the edge of stability in the normalized sharpness coordinate In the following, we focus on understanding tuners that prioritize training at or near the edge of stability.
3 Optimizing on the Edge of Stability
Based on our observations in Section 2, we design learning rate tuners that position the underlying optimizer on the edge of stability (y 0). We analyze a tuner capable of operating both slightly below and slightly above the EOS in order to exploit nonlinear stabilization.
Formally, we investigate generalization of the quadratically greedy rule from Section 2, which sought 7/, to minimize the quadratic approximation y in (2). We instead choose the learning rate to be on edge by seeking the largest value of 17 such that qy is smaller or equal to the original value of f.
v((w)) u = max 11 2 0: gr (7:: We, uz) S f(w)) = 2 u. V vuf(w:)u
where the last formula holds provided that u, (us(w))u > (convex quadratic) and V ([(.) Tuy < O (ue is descent direction).

---

RosNic144 with squared subset of Citarto
Learning Srrepneee
Hossiar on Sharpesss
Gradient Norm
5000
10000
"5900
5000
19000 Epoch
19303
5000
10000
15000
Epoch
Epocc
GD COAT Scale 1.94
GD CDAT Scale 2,06
Scale
GD CDAT Scale
Figure 6: Optimizing on edge induces different curvature dynamics. Sharpness, product between learning rate and sharpness, and gradient norm evolutions for gradient descent with CDAT. By putting the learning rate on edge (ar = 2), the sharpness does not ever increase and actually decreases slightly over time. GD with CDAT operates slightly above the edge constantly during training. Its gradient norm evolution is akin to a fine-tuned constant learning rate baseline.
For it ((wx), and if V (we) ) is aligned with the eigendirection Umax associated with the largest eigenvalue Amax of H, we recover the familiar of = =2/ 22/Xmax- Note however, that contrarily to using directly 7) 2/ Amax. the on-edge rule can naturally take into account the alignment with Umax (see Fig. 15). We note that we recover the edge of stability even when the updates are given by the gradient multiplied by a preconditioner, e.g. U -- for matrix P. In this case, we have -u Hu/ u g - 9 p- HP- g/g p-1pp for H V2 f(,)) g = V f(w)). This is maximized when p-1/2 lies in the largest cigendirection of the PSD matrix H P- pp112pp-11 which for H. This is exactly the edge of stability for adaptive methods (Cohen et al., 2023).
gives us the learning rate тソ- = 2/ //maxx where Amax is the maximum eigenvalue of
We note that the only difference between this and the quadratically greedy rule is a factor of 2 in the numerator. Inspired by this observation, and with eye towards robustness, define our Curvature Dynamics Aware Tuning (CDAT) rule by:
cdat 1/1
for nı, max - (w:)) we, 0), de |u, 2 (wx) + 6.
0 d,"
The scaling factor lets us interpolate between greedy (a 1) and on-edge (a 2). We are most interested in the behavior near a = 2. (also studied in Rosca et al. (2023)). In (8), the max function takes care of the case where u, an ascent direction (V f(,,) T u, > 0). the absolute value takes care of cases where the objective has negative curvature in the update directions (see Appendix C for additional justification). and we simply set & O as we always observed non-negligible positive curvature. The definitions of the numerator and the denominator de allow for the possibility of exponential moving averages (EMA) of each quantity such as u.+1 = (1 Bedar) + Sedañi for Sedat referred to the CDAT EMA parameter thereafter. We observed that smoothing the estimates of n, and de by an EMA is particularly relevant when the updates are themselves defined through an exponential moving average as in Adam, or when using the proposed rule in a stochastic setting.
CDAT has two major advantages: it sensitive information from all eigenvalues of V 2 fw)), and it depends on updates u, coming from any base optimizer. We will take advantage of these properties to explore the behavior of "on edge" optimization variety of settings.
3.1 On edge optimizers in practice
Full batch regime. Fig. 5 presents results for training with CDAT across various optimizers, architectures, datasets, and losses. Overall, selecting the learning rate to be on edge (a 2) is on with or better than a fine-tuned constant learning rate and is always better than a quadratically greedy approach (a = 1). This observation holds even though the quadratically greedy rule ensures larger instantaneous decrease (Fig. 16). One notes that targeting slightly above the edge (a 2.0625) provides even better performance than the on edge rule (a = 2) on all examples except the MLP Mixer on CIFAR10. 10. However, targeting higher above the edge (a 2.5) generally gives diverging results in the short or long terms. To integrate the proposed rule with the Adam optimizer, we also observed that the estimation of the curvatures through ne, d, in (8) was necessary.
7

---

ResNet55 squared
Cifar10
CDAT EMA :.00
CDAT EMA :0.9
CDAT EMA : 00.90
28-
18403
24-
26-
2.4
Scaling Pusto
22
22
Scaling Fade
22
1.8
18
1.6
18
16-02
1e4C
1.0
08
1e-00
08
1e405
120
1004
22.6
1004
Batch
Figure 7: Stochasticity shifts the optimal sealing. Normalized performance of gradient descent with momentum equipped with CDAT in a stochastic regime with varying batch sizes. In a mini-batch regime, the optimal scale decreases as the batch size decreases. Using an exponential moving average smooths out the performance of the CDAT rule over batch sizes.
Remarkably, all choices around the edge (1.9375, 2.0, 20625) show a progressive increase of the learning rate that results generally in a better performance than the constant learning rate counterparts, except for RMSProp on the NanoLM experiment. The increasing learning rate behavior is akin to the warm-up up phase generally hard-coded by a scheduler. Fig. 19, we observe that the CDAT rule displays similar behavior as warm-up schedules, yet it may not fully capture the benefits of prefixed schedules.
In Fig. 6, we analyze the dynamics of the curvature when optimizing on edge. We observe that the sharpness be pushed reduce the iterations st panel of Fig. 6). The CDAT rule may operate constantly slightly above the edge (2nd panel of Fig. 6). By reducing the sharpness, the algorithm may be able take larger stepsizes and converge faster. Sensitivity to architecure's width and depth, as well as weight decay, are also analyzed in Fig. 18.
Mini batch regime. The CDAT rule be used in stochastic regime by replacing f in (8) by its stochastic counterpart (mme) on mini-batch my. However, two difficulties may arise.
First, the edge rule is motivated by the sharpening effects of the overall objective, which can be overestimated or underestimated by single mini-batch. Previous work shows that the trace of the Hessian may best capture the sharpening and stabilization effects in a stochastic regime (Agarwala and Pennington, 2024; Wu and Su, 2023); it is unclear what function of the Hessian spectrum, the CDAT rule captures in the stochastic regime. As result the optimal scaling factor may vary with the
In Fig. 7, we observe that the optimal scaling of the on-edge rule is proportional to the batch size up to some size. particular, at specific batch sizes, we observe that the greedy rule (a = 1) outperforms the on-edge rule. This result is consistent with the good performance of linesearches or greedy rules in a mini-batch regime previously mentioned and observed in Fig. 14. We also observe in Fig. 7, that integrating an EMA into the estimation of the edge in (8) smooth out the selection of the optimal scaling factor.
Finally, the sharpening effects are known to be generally mitigated in the stochastic regime (Agarwala and Pennington, 2024; Cohen al., 2021 1; Jastrzębski al., 2017). The benefits of the on edge rule appear also subdued in this regime (Fig. 8, Fig. 20, Fig. 21).
3.2 Modeling CDAT dynamics
The classical optimization framework is insufficient to fully explain the benefits of CDAT. For example, on a convex quadratic objective, a = is the optimal choice, and a > results (in the worst case) in divergent algorithm. However, we can use simplified model to begin understanding the joint dynamics of the learning rate and sharpness under CDAT.
We approximate the gradients around stationary point w., where (f(w,) 0, as Vf(w)) = Hi; for W. WA W., and H being a symmetric matrix. In this scenario, the learning rate given by CDAT is eddtt aw H (w, H' Consider the case where H has two eigenvalues À and v, with
In this case the CDAT learning rate can be written as
X²p: v²gh X2p/ + v2
T
odar
0 33p: + våge 13pt/9 + v3
(9)
8

---

ResNet50 with squared Cifar00
Batch size: 256
0.4
Learning Rlate
Train Lase
Test Error
02-
72-
0.1
4.4-
Badio
12"
100
20c
230
400
201
120
Pin
302
280
550
100
200
R60
402
Epoch
Epoch
Epoch
Schadulad GD Mom
GD Mom
VITS/1 with cross-entrope loss on Imagenet
Batch size: 256
Leariig Puste
Train Loos
06-
Epoch
Epoch
Epoch
Adam
Schaduled Adam
Adam On Edga
Figure 8: The performance of CDAT is subdued in the stochastic regime. Fine-tuned constant, scheduled, and self-tuned with CDAT learning rates in stochastic regime. In stochastic regime, CDAT can also exhibit a form of learning rate warm-up (top figure). However, the interplay between sharpening and learning rate are known be mitigated in stochastic regime which may explain the underperformance of CDAT in this regime (bottom figure).
Here Pt. 9t are the projections pr (w, o (w, v ) respectively onto the eigendirections v and U. associated with X, v. Therefore, v/ interpolates between its minimum value o/x to the larger value a a/v, depending the alignment ratio pc/ge St. For 2v/A learning rates both above and below the EOS.
2, this rule achieve
We can gain additional insight by modeling dynamical λ extending the model of Section 2.3. While model (5) captures the dynamics the largest eigendirection c, here we aim to model the dynamics in the orthogonal subspace. To simplify, we consider the eigendirections V,VL. and small eigenvalue y fixed. We then model the gradients as ff(w) zs HpW, with H, = Av + ve_vl_ update in the direction t1 using gradient descent vgi, obtain the following dynamical system describing the CDAT learning rate tuner:
A.P: + +22ge XpP: vgg.' 91+1 (1-111)) 1
7/1+1
P++1 (1 + 2
(10)
Combining this with the update rule for y given in (5) completes the model.
There are two important regimes of behavior in this model. First, ye 0, Pt will increase and eventually y/ will decrease as in the normal EOS case. If yt < 0, the key threshold is y/- < - In this case, the ratio pt/ decreases leading to an increase in т) according to the on edge rule. hp, 0 (as it is if Pt: has become small due to < O), then we see from (5) that this leads to an increase in yt- This suggests that CDAT has tendency to push yt closer to the EOS sending towards 0 if the learning rate is driven by the eigendirections corresponding to smaller cigenvalues.
Numerical simulations on this model (Fig, 9) suggest that this effect can indeed cause remarkably small values of (3rd panel of Fig. 9). We emphasize that this is due to the joint dynamics of π) (induced by the learning rate tuner), and d Pt, and 9r (induced by GD). There are also important limitations in this model's 's ability to fully explain CDAT's behavior. For example, the model predicts runaway sharpening for range of stable and useful settings for scale centered around 2. This modeling limitation likely stems as well as higher-order terms, which empirically tend
(2nd panel of Fig. 9), and divergence for a > In practice, we saw a
from neglecting the dynamics orthogonal to stabilize EOS dynamics (Agarwala et al., 2023).

---

0.000
35x10
0,5
--13
37x10
- 2.0
3.005
36x19
0.3-
36x
3.000
0.1-
33x
0.065
-0.1
31x10
400
Steps
Steps
Steps
Figure 9: A simple model partially captures the benefits induced by the proposed CDAT rule. Dynamics of theoretical model of CDAT (10). For o 2, feedback stabilizes close to the EOS (y = 0), which stabilizes Amax (orange). For a 2 - and small e (blue, 0.1), model predicts that Amax slowly grows (middle), but predicts that stabilizes to a value -e < y/ < 0 (right).
Conclusion and Future Directions
Summary. Our empirical results showed that simple linesearches and approximate greedy learning rate tuners underperform constant learning rate approaches in the full batch regime - despite being better on individual steps. The idea that "locally greedy" methods perform poorly long time scales has been shown in other settings as well, including evolutionary dynamics Agarwala and Fisher (2019). Our experiments and theoretical work suggest the failure of these classical tuners due to the fact that they suppress the feedback which stabilizes sharpness in the fixed learning rate setting, As the sharpness increases, tuners are forced to take smaller steps, which ends up leading to slower learning.
We find, in contrast, that prioritizing stability of the sharpness yields tangible benefits. Our CDAT method pushes the network towards the edge of stability via a dynamically driven process. also naturally displays some form of progressive increase of the learning rate akin to prefixed warm-up schedules. CDAT also sheds light on the more complicated dynamics in small mini batch regime, where estimation of a locally greedy rule may actually place the optimizer the edge of stability of the full batch objective.
Limitations and future directions. We explored some limitations of the current modeling frame- work in Section 2.3 in particular, the failure to capture stabilization due to higher order terms. Developing improved models (either analytically or numerically) would allow for powerful tools from other disciplines to aid algorithm design particularly. methods from control theory. For example, state feedback schemes can be designed through the analysis of nonlinear dynamical systems to ensure asymptotic stabilization (Isidori, 1995, Chapter 7). We believe cross disciplinary approach will be useful for designing the next generation of learning rate tuners.
The proposed CDAT rule may also help to understand and refine the design learning rate schedules through scaling ladders (Wortsman al., 2024). Recent work has shown that transfer of learning rates over different scales is related to consistency of curvatue dynamics (Noci et al., 2024); this suggests that approaches like ours may be useful increase predictability of optimal learning rates across scale.
We thank James Martens and Mihaela Rosca for fruiftul discussions on related ideas. We also thank the reviewers for their insightful comments that helped us refine the manuscript.
Acknowledgements.
10

---

References
Agarwala, A. and Fisher, D. S. (2019), 'Adaptive walks on high-dimensional fitness landscapes and seascapes with distance-dependent statistics', Theoretical population biology 130, 13 49.
Agarwala, A., Pedregosa, F. and Pennington, J. (2023), Second-order regression models exhibit progressive sharpening to the edge of stability, in 'International Conference Machine Learning', PMLR, pp. 169-195.
Agarwala, A. and Pennington. J. (2024), 'High dimensional analysis reveals conservative sharpening and a stochastic edge of stability,, arXiv preprint arXiv:2404. 19261.
Almeida, L. B., Langlois, T., Amaral, J. D. and Plakhov, A. (1999), Parameter adaptation in stochastic optimization, 'On-line learning in neural networks', pp. 111-134.
Armijo, L. (1966), 'Minimization of functions having Lipschitz continuous first partial derivatives', Pacific Journal of mathematics 16(1), 1-3.
Arora, S., Li. Z. and Panigrahi, A. (2022), Understanding gradient descent on the edge of stability in deep learning, in "International Conference Machine Learning', PMLR, pp. 948-1024.
Ba, J. L., Kiros, J. R. and Hinton, G. E. (2016), 'Layer normalization', arXiv preprint arXiv: 1607.06450..
Baydin, A. G., Comnish, R., Rubio, D. M., Schmidt, M. and Wood, (2018), Online learning rate adaptation with hypergradient descent, in 'International Conference on Learning Representations"
Berrada, L., Zisserman, A. and Kumar, M. P. (2020), Training neural networks for and by interpolation, in International conference machine learning', PMLR.
Blondel, M. and Roulet, V. (2024), 'The Elements of Differentiable Programming', arXiv preprint arXiv:2403. 14606
Cohen, J., Ghorbani, B., Krishnan, S., Agarwal, N., Medapati, S., Badura, M., Suo, D., Cardoze, D., Nado, Z., Dahl, G. E. and Gilmer, J. (2023), Adaptive gradient methods at the edge of stability, in 'NeurIPS 2023 Workshop Heavy Tails in Machine Learning'.
Cohen, J., Kaur, S., Li, Y., Kolter, J. Z. and Talwalkar, A. (2021), Gradient descent on neural networks typically occurs at the edge of stability, in 'International Conference on Learning Representations?
Cutkosky, A., Defazio, A. and Mehta, H. (2023), Mechanic: A learning rate tuner, in 'Advances in Neural Information Processing Systems'.
Dagréou, M., Ablin, P., Vaiter, S. and Moreau, T. (2024), How to compute hessian vector products?, in 'ICLR Blogposts 2024'. http::::///rrrrbopppee
Damian, A., Nichani, E. and Lee, J. D. (2023), Self-stabiliizatin:: The implicit bias of gradient descent at the edge of stability, in 'International Conference Learning Representations
DeepMind, Babuschkin, I., Baumli, K., Bell, A., Bhupatiraju, S., Bruce, J., Buchlovsky, P., Budden,
D.. Cai, T., Clark, A., Danihelka, I., Dedieu, A., Fantacci, C., Godwin, J., Jones, C., Hemsley, R., Hennigan, T., Hessel, M., Hou, S., Kapturowski, S., Keck, T., Kemaev, I., King, M., Kunesch, M., Martens, L., Merzie, H., Mikulik, V., Norman, T.. Papamakarios, G., Quan, J., Ring, R., Ruiz, F., Sanchez, A., Sartran, L., Schneider, R., Sezener, E., Spencer, S., Srinivasan, S., Stanojević, M., Stokowiee, W., Wang, L.. Zhou, G. and Viola, F. (2020), 'The DeepMind JAX Ecosystem'.
URL: http://githuh.c http://gihuu..com/gggggg/// www.googge-deemmiiiiii.
Defazio, A. and Mishchenko, K. (2023), Learning rateinn---ttteeee rate-free learning by d-adaptation, in 'International Conference on Machine Learning', PMLR, pp. 7449-7479.
Dehghani, M., Gritsenko, A., Arnab, A., Minderer, M. and Tay, Y. (2022), Scenic: A jax library for computer vision research and beyond, in 'Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR)', pp. 21393-21398.

---

Deng, J., Dong, W., Socher, R., Li, L.-J., K. and Fei-Fei, L. (2009), Imagenet: A large-scale scale hier- archical image database, in '2009 IEEE Conference Computer Vision and Pattern Recognition', IEEE, pp. 248-2288 255.
Dosovitskiy, A., Beyer, L., Kolesnikov, A., Weissenborn, D., Zhai, X., Unterthiner, T., Dehghani, M., Minderer, M., Heigold, G., Gelly, S., Uszkoreit, J. and Houlsby, N. (2021), An image is worth 16x16 words: Transformers for image recognition at scale, in 'Internationa Conference on Learning Representations'
Galli, L., Rauhut, H. and Schmidt, M. (2023), Don't be monotone: Relaxing stochastic line search in wver-parameterized models, in 'Advances in Neural Information Processing Systems'.
Ghorbani, B., Krishnan, S. and Xiao, Y. (2019), An investigation into neural net optimization via Hessian eigenvalue density, 'International Conference on Machine Learning'. PMLR, pp. 2232- 2241.
Gilmer, J., Ghorbani, B., Garg, A., Kudugunta, S., Neyshabur, B., Cardoze, D., Dahl, G. E., Nado, Z. and Firat, O. (2022), loss curvature perspective on training instabilities of deep learning models, in 'International Conference Learning Representations'.
Goyal, P., Dollár, P., Girshick, R., Noordhuis, P., Wesolowski, L., Kyrola, A., Tulloch, A., Jia, Y. and He, K. (2017), 'Accurate, large minibatch SGD: Training imagenet in I hour', arXiv preprini DoXiv:/70.00000)
He, K., Zhang, X., Ren, S. and Sun, J. (2016), Deep residual learning for image recognition, in 'Proceedings of the IEEE conference on computer vision and pattern recognition', pp. 770-778.
Hinton, G., Nitish, S. and Swersky, (2012), 'Divide the gradient by running average of its recent magnitude', COURSERA: Neural Networks for Machine Learning.
Hochreiter, S., Bengio, Y., Frasconi, P. and Schmidhuber, J. (2001), 'Gradient flow in recurrent nets: the difficulty of learning long-term dependencies', A Field Guide to Dynamical Recurrent Neural Networks pp. 237-244.
Howard, J. (2019), 'imagenette'.
URL: https://gtttt..c.mm///o wwt.//sssaaimmmmgeneteeee
loffe, S. and Szegedy, C. (2015), Batch normalization:
Accelerating deep network training by reducing internal covariate shift, in 'International conference machine learning', pmir, pp. 448-456.
Isidori, A. (1995), Nonlinear control systems, third edn, Springer.
Ivgi, M., Hinder, O. and Carmon, Y. (2023), DoG is SGD's best friend: A parameter-free dynamic step size schedule, "International Conference on Machine Learning.. . PMLR, pp. 14465-14499.
Jastrzebski, S., Kenton, Z., Arpit, D., Ballas, N., Fischer, A., Bengio, Y. and Storkey, (2017), "Three factors influencing minima in SGD', arXiv preprint arXiv: 1711:462.......
Jastrzebski, S., Kenton, 7., Ballas, N., Fischer, A., Bengio, Y. and Storkey, A. (2019), On the relation between the sharpest directions of DNN loss and the SGD step length, in 'International Conference on Learning Representations'
Jastrzebski, S., Szymezak, M., Fort, S., Arpit, D., Tabor, J., Cho, K. and Geras, K. (2020), The break-even point on optimization trajectories of deep neural networks, in 'International Conference on Learning Representations'
Kalra, D. S., He, T. and Barkeshli, M. (2023), 'Universal sharpness dynamics in neural network train- ing: Fixed point analysis, edge stability, and route to chaos", arXiv preprint urXiv:2311.02076
Karpathy, A. (2015), 'The unreasonable effectiveness of recurrent neural networks', http:// karpathy.. github...... 1o///////0 2t5/////21/mmmefffetiieeee
Kopitkov, D. and Indelman, V. (2020), Neural spectrum alignment: Empirical study, in 'International Conference on Artificial Neural Networks', pp. 168-179.
12

---

Krizhevsky, A., Hinton, G. et al. (2009), Learning multiple layers of features from tiny images, Technical report, University of Toronto, ON, Canada.
LeCun, Y., Cortes, C. and Burges, C. (2010), 'Mnist handwritten digit database', ATT Labs [Online]. Available: http://ww..cc...oo/////////ineee 2.
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 urXiv:2003.02211
Li, H., Xu, Z., Taylor, G., Studer, C. and Goldstein, T. (2018), "Visualizing the loss landscape of neural nets", Advances in neural information processing systems 31.
Liu, D. C. and Nocedal, J. (1989), 'On the limited memory bfgs method for large scale optimization', Mathematical programming.
Liu, H., Li, Z., Hall, D. W., Liang, and Ma, (2024), Sophia: A scalable stochastic second-order optimizer for language model pre-training, in 'International Conference on Learning Representa- tions'.
Liu. L.. Jiang, H., He, P.. Chen, W., Liu, X., Gao, and Han, J. (2020), On the variance of the adaptive learning rate and beyond, in 'International Conference on Learning Representations'
Loizou, N., Vaswani, S., Laradji, I. H. and Lacoste-Julien, S. (2021), Stochastic Polyak step-size for SGD: An adaptive learning rate for fast convergence, in 'International Conference on Artificial Intelligence and Statistics', PMLR.
Loshchilov, I. and Hutter, F. (2016), SGDR: Stochastic gradient descent with warm restarts, in 'International Conference Learning Representations'
Martens, J. and Grosse, R. (2015), Optimizing neural networks with Krotecker-factored approximate curvature, in 'International conference on machine learning', PMLR, pp. 2408-2417.
Mishchenko, K. and Defazio, A. (2023), 'Prodigy: An expeditiously adaptive parameter-free learner', arXiv preprini arXiv2300.00101/
Mutschler, M. and Zell, A. (2020), 'Parabolic approximation line search for DNNs', Advances in Neural Information Processing Systems 33, 5405-5416.
Nesterov, Y. et al. (2018), Lectures on convex optimization, Vol. 137, Springer.
Nocedal, J. and Wright, S. J. (1999), Numerical optimization, Springer.
Noci, L., Meterez, A., Hofmann, T. and Orvieto, A. (2024), "Why do learning rates transfer? reconciling optimization and scaling limits for deep learning', arXiv preprint 17444
Polyak, B. Т. (1964), 'Some methods of speeding up the convergence of iteration methods', USSR computational mathematics and mathematical physics 4(5), 1-17.
Riedmiller, M. and Braun, H. (1992), Rprop: a fast adaptive learning algorithm, 'Proc, of the Int. Symposium Computer and Information Science VII'.
Rosca, M., Wu, Y.. Qin, C. and Dherin, B. (2023), 'On continuous time model of gradient descent dynamics and instability in deep learning', Transactions on Machine Learning Research.
Roulet, V., Agarwala, and Pedregosa, F. (2023), On the interplay between stepsize tuning and progressive sharpening, in 'OPT 2023: Optimization for Machine Learning'.
Schaul, T., Zhang, S. and LeCun, Y. (2013), No more pesky learning rates, in 'International conference on machine learning', PMLR, pp. 343-351.
Shi, B., Su, W. and Jordan, M. (2023), *On learning rates and Schrödinger operators', Journal of Machine Learning Research 24(379), 1-53.
Tolstikhin, I. O., Houlsby, N., Kolesnikov, A., Beyer, L.,, Zhai, X., Unterthiner, T., Yung, J., Steiner, A., Keysers, D., Uszkoreit, et (2021), 'MI.P-mixe:: An all-MLP architecture for vision', Advances in neural information processing systems 34, 24261-24272.
13

---

Vaswani, S., Mishkin, A., Laradji. I., Schmidt, M., Gidel, G. and Lacoste-Julien, S. (2019), Pain- less stochastic gradient: Interpolation, line-search, and convergence rates', Advances in neural information processing systems.
Wortsman, M., Liu, J., Xiao, L., Everett, K. E., Alemi, A. A., Adlam, B., Co-Reyes, J. D., Gur, I., Kumar, A., Novak, R., Pennington, J., Sohl-Dickstein, J., Xu, K., Lee, J.. Gilmer, and Korblith, S. (2024), Small-scale proxies for large-scale transformer training instabilities, in 'International Conference on Learning Representation'.
Wu, L., Ma, C. et al. (2018), 'How SGD selects the global minima in over-parameterizee learning: A dynamical stability perspective', Advances in Neural Information Processing Systems 31.
Wu, L. and Su, W. (2023), The implicit regularization of dynamical stability in stochastic gradient descent, in 'International Conference on Machine Learning', PMLR, pp. 17656-37668
14

---

A Theoretical Model
A.1 Intuition for model of curvature dynamics
In Section 2.3, we extend the model from Damian et al. (2023) in order to understand the curvature dynamics. Here we provide some intuition for the basic structure of the model.
The analysis in Damian et al. (2023) is based on a 3rd order expansion of the loss function. A third order approximation of f around some wy reads
1 f(w) ~ f(wo) + vf(wo) (w - wo) + 2
(ww-wo) -
T f(w))(( - wo)
v f(wa)lw wu, w wa, w wul.
Assuming without loss of generality that wi 0. s(wo) = O, we can write:
1 f(w) R ((wo) + zu V f(wo) + 8 (wo))l, w. w]
(11)
Consider the gradient descent dynamics with learning rate yy under this model. We can write:
W++++ - me ~ - (V (V + (w))wx))
(12)
f(wo)) where H(w,) = V f(wo)) + f wr... -] the Hessian at the current parameter Wy,-
Let v be the direction of the largest eigenvalue of H (wo).. For small enough w; wo, this is good approximation of the largest eigendirection of H (wx). Consider the dynamics of the projection Iy u 2/, We make the additional assumption that v ff(ww) 0 (which can be achieved with coordinate transformation). We then get
T++1 (1 (1-10 m))))
(13)
where A == A(wx) = Amax (wy) is the largest eigenvalue of the current Hessian H (wx). The dynamics of d are generally slower than the dynamics of we (and x1), and are governed approximately by
A+11 d = (V\(w))) (w141 - we) =-n((X(we)) Vf(w))
(14)
This gradient is given by
VA(1;) Vi H(w))u) V (woo(lo v.
(15)
For very small Fe. it has been observed (see e.g. Cohen et al. (2021)) that A X, is increasing during most of training. Early in this regime H₁ is often very close 0 (that is, the gradient has small component in the direction). Therefore we are interested in the contribution of the gradient orthogonal v, [f(w:)]
the dynamics. We define a, -(VA(w)) ( as the instantaneous change in the top eigenvalue, due to gradient contributions orthogonal v. We assume that this contribution is positive and, simplify, independent of t, that is, >> 0.
(1 v fwe),
From (13), while A 2/n, xy. is decreasing and from (14) A₁ is increasing.
Once d₂ > 2/1,, the dynamics of x, changes from convergence towards increasing values with alternating sign. We can write down the contribution to the gradient from I; using third order expansion of the loss. have:
(wo)) v. Vwo + x.v) (((w)) + x, (wx + 2
(16)
We can now understand how these terms contribute to the dynamics of A The first term contributes to sharpening (increasing A₂) via the constant a. The second term has oscillating sign and its long term contribution to X; is small. From Equation 13 we know that
Tt+1 + Tz. RJ - 1)Ax)a
(17)
Near the edge of stability, uj,, close to 2, and the sum of successive x, are small; therefore we neglect the linear 2, term in Equation 16.
15

---

10
10
106
4
R 10:4
10-2
10
1,9
Ferk
25
50
75
100
200
Steps
Steps
Figure 10: Dynamics of the largest eigenvalue projection p. For constant learning rate, cy- cles between stability and instability and Amax stabilizes (blue). If learning rate tuner sets "e 2(1 /) Amax,t- p decays to 0 and there is no negative feedback preventing sharpening.
Dynamics p for toy model of CDAT. For 2, p is stable, which induces sta- bilization sharpness (orange). For a 2 - e for small (bluc, 0.1), p decreases but the ratio of to the orthogonal projection remains constant which stabilizes y near O.
Figure 11:
The sign of the x² term is constant and there are no such issues. Therefore, in the non-linear regime the dynamics of XS can be approximated by:
? f(w)lu, v, [e A4+1 A na
(18)
2
we define V (w)) u, u, fuog v. (guaranteed to be positive), then we arrive at the dynamical equations presented in Equation 4:
It+1 - St = (2 - - de+1
A -1(a bx,)
(19)
This system leads quasi-stable oscillation of X around 2/1;; when X < 2/n, at goes to and eventually d increases; when A 2/1,, increases in magnitude, eventually leading to decrease Ae.
This heuristic argument is made rigorous in the appropriate regimes in Damian et al. (2023); there they compute the difference between the full dynamics and the dynamics with the gradient projected away from the largest eigendirection using coupling argument. The key point is that higher order contributions the gradient are guaranteed be anti-aligned with the gradient of the large Hessian eigenvalues. This leads to restoring force which counteracts eigenvalue growth (progressive sharpening) the largest eigenmode becomes unstable.
A.2 Dynamics of rescaled variables
In Section 2.3, we used Equation to derive equations in == x22 and == 1/1 - 2. We arrived at
Pt+1 = (1 + 2t) Ue+1 = 7++1 [17 (a bpe)] +
9/t+1
Ye + 2
7t++
=]
(20)
171
m7
The dynamical equation for p obtained by squaring the equation for an in Equation 4. To derive the equation for yt, we first derive the dynamics of n Ae+1 as
meA++1 = =] [n bp:)l + no
(21)
We then have
Mc+1 mmde+1 - + -2|+2
Ut+1
(22)
Th
Evaluating completes Equation
A.3 Dynamics of the projection on the largest eigenvector
We present results on the dynamics of p in the various models; these were omitted from the main text in order to simplify the presentation. For constant learning rate, p initially decreases until EOS crossed, after which it enters into cycle of increase and decrease (Figure 10, blue). For our model
16

---

of linesearches, where 7) = 2(1 - -)/ decays to quickly and there is no mediation of sharpening (orange, 0.1).
For our model of CDAT presented in Section 3.2, stabilizes for a = (Figure 11, orange). For 2, the model still predicts decay of p, but the ratio of pt the orthgonal component ge remains constant (Figure 11, blue). This fixed ratio stabilizes to a value near 0.
In practice, the higher order terms in the dynamics provide additional stability, in the on-edge model, which allows p to stabilize as well, see Fig. 12, Fig. 17. The key is that these terms can operate when y is close to for long periods of time. These results suggest that additional model development is required to understand the behavior of learning rate tuners which target the EOS.
B
Additional Experiments
B.1 Further analyzes of base learning rate tuners
Fig. 12 completes Fig. 3 with measures of sharpening and learning rates the settings considered in Fig. 1. For RMSProp we considered the preconditioned Hessian following the observations done by Cohen et al. (2023) that for adaptive gradient methods such RMSProp Adam, the sharpness of the preconditioned Hessian, rather than the sharpness of the Hessian, defines the edge of stability. Namely, recall that RMSProp takes updates of the form
W++1 wi - ff forP: = diag( tiag(V//= 1/ E
for
1/2 32)9 f24-11 V-1 0, 9t Vf(w)).
with 32 exponential moving average parameter. The preconditioned Hessian takes then the form
H. = 1/2 ((w
and we report \max (H.).
We observe similar behaviors in these regimes as in Fig. 3. Namely, the sharpness or preconditioned sharpness ever increase (2nd panels of Fig. 12), while the learning rates ever decrease (1xt panels of Fig. 12). The constant learning counterpart can operate above the edge of stability while the self-tuned methods generally avoid crossing the edge of stability (3 panels of Fig. 12).
B.2 Analyzing additional learning rate tuners
We consider the performance of two additional classical learning rate tuners, Polyak stepsize (Berrada al., 2020; Loizou al., 2021; Polyak, 1964) and
descent (Almeida et al., 1999; Baydin et al., 2018) akin to the resilient backpropagation scheme (Riedmiller and Braun, 1992).
Briefly, Polyak stepsizes consider setting the learning rate
7) min { f(w)) f [V](w 7max
(23)
where / min,, (w) the minimum of the objective set to by assuming that neural network overfit the data, and Tmax is maximal stepsize selected as 1 or 100 in experiments (we take the best instance).
Hyper-gradient descent considers updating the stepsize towards maximal decrease of the objective. Namely, defining the objective obtained after one step he(r) = f + muu), the algorithm updates Me by gradient step ht resulting priori in 7h+1 l}: a [w, + m)4)) u, for given hyper-learning rate a. Almeida et al. (1999); Baydin et al. (2018 8) argued for using multiplicative updates of the form
M nt (1-
Vf(w, + u) ue | f(w, + 12ll)llll
(24)
Intuitively, the learning rate increases the update aligned with the negative gradient direction and decreases otherwise. Resilient backpropagation (Riedmiller and Braun, 1992) adopts similar logic
17

---

componentwise. In our experiments we vary 3 and select the best instance, see Appendix C.5 for more details.
We observe that Polyak stepsizes (top figure of Fig. 13) generally select larger learning rates than the constant learning rate counterpart. The efficiency of Polyak stepsizes is not reached by the CDAT rule with = 2 but with a slightly larger scale o 2.06. The efficiency of the Polyak stepsize method, in particular compared
simple linesearch, in a full batch regime, has also been reported by Roulet et al. (2023). The proposed CDAT rule may capture the benefits of aggressive learning rates taken by Polyak stepsizes in smoother way by allowing various scales.
On the other hand, the learning rate counterpart (bottom figure of Fig. 13). We also observe slow, yet steady, progressive sharpening when using the hyper-gradient descent. As with the linesearch method or the quadratically greedy rule, the hyper-gradient descent focuses selecting learning rate that decreases the loss, which appears, across those tuners, to potentially suppress effective stabilization effects naturally appearing with constant learning rate.
descent performs just on par with the fine-tuned constant
B.3 Base learning rate tuners in a stochastic regime
In Fig. 14, we report the performance of classical learning rate tuners (linesearch or quadratically greedy method) in a stochastic regime for varying batch-sizes. As observed previously by Vaswani al. (2019) or Roulet al. (2023), a linesearch for example can perform well in stochastic regime. Note that the two approaches (linesearch and quadratically greedy method) display similar behaviors (just as they displayed similar behaviors in the full batch regime). This hints that, rather than playing with the numerous hyperparameters of linesearch we may focus simply on an additional scaling factor for the quadratically greedy rule, which motivated the proposed CDAT rule.
B.4 Targeting the edge of stability using the exact sharpness
Cohen et al. (2021, Appendix F) reported bad performance of adaptive learning rate tuners selecting the stepsize
7 = 2/ Ammaamm' f(w))
which may fix the learning rate just at the edge of stability. Note that such a definition does not take into account the additional alignment of the update with the largest eigenvector. Our proposed diagnostic rule CDAT rather considers the edge of stability given by a local approximation of the objective along the update to take into account the alignment of the update with the largest eigenvector the Hessian. We ran experiments with a rule
(25) that lets the scaling factor vary just as done with CDAT. The only difference is in the estimation of the base estimate of the edge of stability (CDAT does with the help of a quadratic approximation of the objective, while the rule (25) uses exact computation of the sharpness). In Fig. 15, we observe that setting the scale a 2 2 leads to poor performance as previously observed by Cohen et al. (2021). Note however that by setting the scaling much above (like 3) such rule may outperform constant learning rate. This hints that the rule (25) misses the alignment of the update with the largest eigenvector, which motivated the CDAT rule.
= Amaxx ( V f(w)).
B.5 Analyzing instantaneous gains versus long-term gains
In Fig. 16, we investigate the difference of instantaneous decrease using the quadratically greedy rule (CDAT with a = 1) compared to the on edge rule (CDAT with o = 2). Throughout training, the quadratically rule ensures a larger instantaneous decrease as intended through its definition as a learning rate that minimizes the loss. Yet, in the long term, the quadratically greedy rule underperforms the on edge rule (Fig. 5).
B.6 Additional metries for the CDAT rule
In Fig. 17, we additionally measure the alignment of the updates with the largest eigenvector and the angle between successive updates. We observe that the CDAT rule for z 2 behaves similarly as the constant learning rate counterpart. In particular, the updates tend to quickly be in opposed directions. The quadratically greedy rule does not demonstrate such behavior.
18

---

B.7 Sensibility analysis to architecture hyperparameters
In Fig. 18, we study CDAT for simple MLPs in a full batch regime on the MNIST dataset. Our goal is to understand the benefits of the proposed CDAT rule for varying hyperparameters. First, we analyze the sensibility to width and depth of MLP in a similar fashion Cohen et al. (2021, Appendix D) did to analyze progressive sharpening.
We observe that accrued gains be obtained with the CDAT rule for larger widths (top left panel of Fig. 18). Note that Cohen et al. (2021, Appendix D) found less sharpening at higher widths. In terms of depth (top right panel of Fig. 18), the CDAT rule works best with larger depths while note slight shift of optimal scaling factors from 2 to slightly below 2.
The CDAT rule appears to work best with small while its benefits fade with larger weight decay (no difference between greedy a = 1 and on edge 2). Finally, while the method naturally finds smaller train losses with larger subsets of data, we do not observe a sigmificant shift of relative performance between scales as the size of the data increases (bottom right panel of Fig. 18).
weight decay (bottom left panel of Fig. 18)
B.8 CDAT rule versus prefixed schedule in full batch regime
In Fig. 19, we compare the proposed CDAT rule with prefixed schedules in a full batch regime. We observe that while placing the optimizer on edge could improve constant learning rate counterparts, prefixed schedules can outperform the CDAT rule. This points out that the feedback loop exploited by CDAT may miss some additional nonlinear effeets that could further enhance self-tuning rules.
B.9 Detailed performances of CDAT in stochastic regime
In Fig. 20 and Fig. 21, we detail the performances the CDAT rule the stochastic regime for varying batch sizes. In the stochastic regime, recall that the heatmap of the performance of CDAT in terms of batch-size heavily depended on the appropriate scaling factor Fig. (in comparison scaling factor of approximately
2 appeared generally good in the full batch regime). In both Fig. 20 and Fig. 21, we observe that the method may generally work better at larger batch sizes. Understanding better the right statisties estimate well appropriate estimators of the edge of stability in a stochastic regime a future direction.
19

---

ReaNet34 squared loss on subeet Cita110
Rate Sharoniss
12P
Hoss Sharppesz
Leaming Rata
10P
Hession
Laaring
101
02-
50.00
10000
10000
5000
10020
15000
5200
10000
0.0
5000
10000
15600
Epach
Epoch
Epoch
Epoch
Linesearch
Quad. Gready
MixerTi8 with squared loss subset of Ci
123
Learnig Sharpness
1.0
Hossian Sharpness
Learing Pate
0.8
ton
10-4
9592
4000
190n
2006
5000
6000
9002
4000
2008 Epoch
3060
4000
Epach
Epoch
Epach
Linesearch
Greedy
NanoLM with cross entropy loss on subset Tiny Shalesspeare
Procand Snarpness
Loarning Parte Starpnoss
Precend Hessian Algn
Learing Ritts
10"
06
100
04-
10
9005
2000 Epoch
Epech
Epoch
RMIS
Linesearch RMSProp
Groady RMSProp
VITTV32 with cross entropy loss on subset of Imagenet
Precond Hossian Sharpeiess
Loaming Rate x Sharpness
Precond Hessian Algn Eig
Learnig Paata
107
ca
12*
2000
4000
3000
3002
4050
2000
3200
RMISProp
Linesearch RMSProp
Groady RMSProp
Figure 12: Learning rates and sharpness dynamics of baseline learning rate tuners. Learning rate, Hessian or preconditioned Hessian sharpness, their product, and the alignment between the update and the largest eigenvector of the Hessian the preconditioned Hessian. As in Fig. 12, we observe that linesearch (1) quadratically greedy (3) learning rate tuner display decreasing learning rates along training. The sharpness of the Hessian (for GD) preconditioned Hessian (for RMSProp) keep increasing for self-tuned baselines while they stabilize for the constant learning rate counterparts. The self-tuned methods perform generally below the edge of stability or at least much less above than the constant learning rate counterpart.
20

---

ResNet34 squared on Ciar100
10-4
Leaminng Rate X Sharpeeco
Leariigg Patto
Hessian Sharprens
3+
Train Loss
t0
5000
10393
1526e
10000 Epoch
15/00
19900
15000
1903
10508
1900
Epach
Epoch
Epoch
Polysk
ResNet54 with squared loss on subeet CFar10
102
Lssennng Rata X Sharpnoze
Hessien Sharpress
sstrring Flats
Train Laba
19
Epoch
Epoch
Epoch
Epech
GD
Hyper GD
GD CDAT Scale 2.0
GD CDAT Scale 2.08
Figure 13: Analyzing additional learning rate tuners. Train loss, learning rate, sharpness, and their product along training with gradient descent using a constant or self-tuned learning rate with various tuners. Polyak stepizes
(23) are effective in full batch regime (top figure) outperforming CDAT on edge (a = 2). The effectiveness the Polyak stepsizes are partially captured by an aggressive CDAT rule placing the optimizer edge 2.06. On the other hand, hyper-gradien descent (24) performs just on par with the constant learning rate counterpart in this regime. It also displays an ever- iver-increasing sharpening akin to the observed for linesearch or the quadratically greedy rule Fig. 3.
21

---

FicsNats0 with squared loss on Citar10
Batch size: 64
0.30
Irain Lo35
Test Enor
2:30
Learning Rate
0.16
4.40
0.06
450
200
480
Epoch
Epoch
Batch size: 256
Learning Patie
Tels Ertor
10
450
200
400
Epoch
Epoch
Epoch
Batch size: 10024
860
3.025
Leamiing Rate
Tran 6055
0.30
0.87
00
0.20
0.685
200
420
200
480
200
400
Epoch
Epoch
Epoch
SGD
Linosearch SGD
Quad. Groody SGD
Figure 14: Classical learning rate tuners work well in stochastic regime. In stochastic regime, we observe that, e.g., linesearch can perform on or even better than the constant learning counterparts. However, this good performance is not explained by the common belief that linescarches work well in a full batch regime (Fig. 1). The linesearch and quadratically greedy rule perform similarly in this setting.
MLP(256, 256] with squared subsat
0.4
Lscanng Rate X Sharpners
10-
Tain Loas
Leariing Rata
0.1
1000
500e
Epech
Epoch
QD
GD Exact CDAT Scale 2.5 GD Exact CDAT Scale 2.06
GD Exact CDAT Scale 2.0
GD Exect CDAT Scale
CDAT Scale 3.0
Figure 15: On edge with exact sharpness. We consider rule similar to CDAT (8) but using the exact sharpness, that the largest eigenvalue of the Hessian, base learning rate while varying an additional scale factor (see (25)). By using the exact sharpness, scaling factor of a = 2 leads now poor performance, while scaling factor of a = performant. By using the exact sharpness (25) we do not take into account the actual alignment of the update with the largest eigenvector which may explain the shift of optimal scaling factors in this case.
22

---

Training with CDAT Scale 1,0
Training with CDAT Scale 2.0
Dift in lese after one updale
bhv CDAT Scale 2.0 and 1,3
Dift loss one update
1.0
GD CDAT Scale 1.0
GD CDAT Scale
Train loss
01
bhw
100
300
200
300
400
100
220
300
Steps
Stepe
Steps
Figure 16: The quadratically greedy rule ensures larger instantancous decrease of the loss. Train losses and difference in loss between one step using the on-edge rule (CDAT, scale=2) and one step using quadratically greedy rule (CDAT, scale=1). Results averaged over 10 initializations and disjoint sample subsets of CIFAR100. MLP architecture: single hidden layer of size 1024, ReLU activations, trained with GD and cross-entropy loss in full batch setting. We plot f(w: - u)) f(w; + 8 u,) along training performed either with the quadratically greedy rule (a = 1) or the on-edge rule (a 2). In both cases, this difference is positive meaning that the quadratically greedy rule ensures larger instantaneous decrease the loss. Yet the quadratically greedy rule underperforms in the long term (see Fig. 5, also holds in this full batch setting).
ResNer34 squared
Heaslan Algn Elg
Angle Updatus
0.5
6.c
-0.-
60
10000
10003
Epoch
Epoch
CDAT
GD COAT Scale 1.0
GD CDAT Scale 2.0
Figure 17: Further analysis of the CDAT rule This is the same setting in Fig. 6 except that we plot the alignment between the largest eigenvector of the Hessian and the update, and the angle between successive updates. We observe that the CDAT rule for 2 behaves similarly as the constant learning rate counterpart. In particular, the updates tend to quickly be in opposed directions. The quadratically greedy rule does not demonstrate such behavior.
23

---

Square Edge
SGD Edge
258
3.0
1e:00
Staking Factor
0.00
4.100
1+01
mndr
2.7 09709/25
2.11-55
1+:00
1/007%
175
1:5
1e44
Subset of Mhiit MUP Square Lass SGD On Edge
Subset Mrist MLP Square Loos SGD On Edge
Hiddon sizes (256, 256
1e:00
Final train losalnit bain 000
Hidden sibes [256, 256, 256]
le40
Boaling Facdor
1+d1
2/0
1+:00
1,20
fec4
10
0.
1e-65
0.0011
6d01
2016
4000
6/92
Waight Decay
Data Size
Figure 18: Improvements of CDAT rule edge for varying hyperparameters. Varying width, depth, weight decay and size of the subset considered when using the CDAT rule with varying scaling factors. We observe that accrued gains can be obtained with the CDAT rule for larger widths (top left panel). In terms of depth (top right panel), the CDAT rule works best with larger depths, while we note there slight shift of optimal scaling factors from 2 to slightly below 2. The CDAT rule appears to work best with small or no weight decay (bottom left panel) while its benefits fade with larger weight decay (no difference between greedy a 1 and on edge o 2). Finally, while the method naturally finds smaller train losses with larger subsets of data, we do not observe a significant shift of relative performance between scales as the size of the data increases (bottom right panel).
ResNe134 squared Citar10
NanoLM with cross entropy loss subeet of Tiny Shakespeare
10*
Train Loos
0.-
10
Leaming Rate
10-
0.
Learring
Trai
10+
104
10
2050
Epoch
Epach
Epoch
Scheduled
PMSP
Scheduled RMSProp
GD CDAT Scate 1.0
-
GD CDAT Scale 2.0
RMS Prop CDAT Scale 1.0
HMSProp CDAT Scale
MikerTiBB with squared loss subeet Cier110
VITTi332 with cross entropy loss subeet Imagenee with EMA
10"
102-
Learning Rate
Train Love
10-1
A2
Lees
Learing
12-
10-
12-
10-4
12-
10-4
1000
3000
2000 <000
1003
2000 Epoch
9000
4000
1000
2000
3000
2250
1060
2000
5000
4000
Epoch
Epoch
Epoch
GD Mom CDAT Scale 2.0
Adam
Scheduled Adam
GD Mom CDAT Scale
Scheduled GD
Adam CDAT Scale
Adam CDAT Scale
Figure 19: CDAT rule may not fully capture the benefits of pre-fixed schedules. Train loss and learning rate behaviors for fine-tuned optimizers with or without schedules vs self-tuned counterparts with CDAT on various architecture, datasets, losses in full batch regime. While the CDAT rule displays a behavior to warmup schedules. it does not completely catch the benefits of pre-fixed schedules. Note for the top-left part that the performance of the pre-fixed schedule akin to the performance reported with CDAT 2.5 at early times suggesting that varying scaling factor, or taking higher order dynamics may be important to fully capture the benefits of warm-up schedules.

---

ReeNe150 with squared loss on Cifar10
Batch size:
244
0.8
035
Laaming Paln
835
220
0.15-
6.10-
Epoch
Epach
0.40
0.0
0.5
9.2
0.7
Laarning Paate
3.0-
Text
0.
0.3
0.4
03-
800
552
800
40
0.2
100
350
s/ce
160
Epoch
Epoch
Epocn
Batch
8.50
Leaniigg Piale
225-
320
8:15
0.66
0.00
100
300
100
300
400
500
100
300 Epcch
460
Epoch
Epoch
Batch
3.5
3.5-
Learing Paate
925-
820
0.10
0.6-
0.00
032
300
400
100
250
430
506
135
200
200
Epooh
Epoch
Escon
Batch size: 512
335
330
Learting Ratla
8.50-
815-
210
335
104
8:0
100
20
300
420
500
100
290
400
500
180
200
300
400
Epoch
Epoch
Epoch
Batch sizee: 1024
8.40
12
M0
339
830
Leaming Ratt
330
818
3.0-
855-
300
200
Edo
100
Epoch
Mom
Schoduled
CDAT
Figure 20: Performance of fine-tuned algorithms in stochastic regime with SGD Momentum. In the stochastic regime with SGD momentum, we observe that the CDAT rule may outperform the constant learning rate counterpart (particularly for large batch sizes) while performing on par or underperforming the scheduled learning rate counterparts. Interestingly, a warm-up phase appears naturally induced by the CDAT rule (right plots).
25

---

VITS/16 with cross cross- entropy loss on Imagenat
Batch size: 64
10-
Leaening Rate
10-4
Train
104
06-
05
64-
088
Epoch
Batch sizer 256
Learning Patte
Trai Loos
10-4
05-
04-
8.
100
Epoch
Epoch
Epoch
Bettch size: 1024
1.0
G-
00
Learming Tutle
Train Losz
75 Epoch
102
00
126
25
100
120
Epoch
Epoch
Adam
Scheduled Adam
Adam CDAT
Figure 21: Performance of fine-tuned algorithms stochastic regime with Adam. In the stochas- tic regime with Adam we observe varying performance of the CDAT rule, generally on par under- performing the baselines. Several factors can explain the underperformance. The scaling factor is not well understood, and a finer grid search could improve the scheme. Similarly, a better estimate or better understanding of the edge of stability stochastic regime could improve the approach. Tackling curvature dynamics by analyzing feedback effects as in the CDAT rule may help such design.
26

---

Experimental Details
C.1 Datasets
MNIST. MNIST is an image classification dataset of handwritten digits (LeCun al., 2010). We normalized the pictures so that each pixel is between 0 and 1. We did not standardize the data. We only used this dataset to test varying MLP architectures in Fig. 18. See Appendix C.5 for any additional relevant details.
CIFARIO. CIFARIO is an image classification dataset of colored images of size 32 x 32 with 10 classes (Krizhevsky et al., 2009). We normalized the pictures so that each pixel is between 0 and 1. We did not standardize the data. In the full batch regime, we considered a subset 4096 samples. In the mini batch regime, we considered the full dataset of 50, 000 samples for training (dropping out the remainder batch) and tested on the 10, 000 validation samples.
Tiny Shakespeare. This consists in 40, 000 lines of Shakespeare from variety of Shakespeare's plays (Karpathy, 2015). The task consists in character-level prediction task among 64 characters. In the full batch regime, we considered a subset of 2048 samples consisting of blocks of 64 characters.
Imagenet. Imagenet is image classification dataset of various images from the web (Deng et al., 2009). The images have various sizes. The original Imagenet- K dataset contains 1000 classes. For the full batch experiments, we consider the Imagenette (Howard, 2019) subset that consists in only 10 classes and took 1024 samples out of it. We consider the usual prepreocessing for Imagenet as detailed in the Scenic library (Dehghani et al., 2022). Namely, for training we consider random cropping and random flip at training. For testing, we center crop the images. Each time the cropping reduces the colored images 224 224 size. In the mini-batch regime we consider the complete training dataset of 1.2 million images (Imagenet- IK), dropped the remainder batch, and reported test error on the 50, 000 validation images.
C.2 Architectures
Residual Network (ResNet). We considered the standard ResNet architectures (ResNet34, ResNet50) of He et al. (2016) as implemented in the Scenic library (Dehghani et al., 2022). For the examples with ResNet34 we removed the batch normalization layers (loffe and Szegedy, 2015). For the examples with ResNet50 we replace the batch normalization layers with layer normalization layers (Ba al., 2016).
Multi-Layer Perceptron (MLP) Mixer. We consider the standard MLP Mixer architec- tures (Tolstikhin et al.. 2021) as implemented in the Scenic library (Dehghani et al., 2022). By Mixer Ti/8, we mean the tiny model of Mixer provided in the Scenic li- http///oooooaa///nniiic.ooo
brary (see sttpps://gt //github.b com/ wwwm/ooo///eeeeeccoo baselines/ configsfii.i////// imaagmnt//////////m/meo/ooo.aaee weeanggnneiimmggneeeeeeee //t/ggaaannnn..e/eeeeeeeeae.aeee mugreg_mixeee confir py) with patches of size 8 x 8. We removed dropout (both layer and depth wise).
Nano Language Model (NanoLM). We consider simpel srquence-to-sequeee Transformer model implemented in Optax (https://gg g/gbttttu..... com com/google- pooogle-deemmndd/// wo--eeeind/opptt//lllllo /ttp////lt/o//aaaaaaaaaao examples//o stamples/nanolm.. ianolm..nnp.. ipynb). The model consists of stacked transformer blocks, each of which contains multi-head attention layer followed by feed-forward layer. Layer normalization is used used within the transformer blocks to improve training stability. Finally, a dense layer maps the model's output to the vocabulary size, producing probabilities for each character the next potential character. The only difference with respect to the previous code is that we removed dropout (for deterministic training) and, as mentioned in the previous subsection, reduced the size of the dataset to be able to estimate the sharpness.
Vision Transformer (VIT). We consider the standard Vision Transformer architecture (Doso- vitskiy et al., 2021) as implemented in the Scenie library (Dehghani et al., 2022). By ViT Ti/32, we mean the tiny model of Vision Transformer provided in the Scenic li- brary (see https: //gtt//////t/b///e/ com/ ctt.r///oooeeeeeeerrrrreeeeeetto wwwaaaii//eeencpoojececcc////
27

---

wttp:///mmmmeeeeiinggeee.tttttee 32. We removed dropout (both layer and depth wise).
py) with patches of size 32 x
Weight decay. In all examples, except mentioned otherwise, we consider a fixed weight decay of 10-5. See Fig. 18 (bottom left panel) for analysis of sensitivity of the proposed CDAT rule with varying weight decay.
C.3 Algorithms
In all experiments, when selecting the "best" tuner considered the average train loss the last 5 iterates.
Fine-tuning base optimizers. The implementation of all optimizers taken from the Optax library (DeepMind et al., 2020). In all experiments, we fix all hyperparameters of the base optimizer (S)GD, (S)GD with Momentum, RMSProp, Adam) their default values: 0.9 for the momentum of (S)GD with Momentum, 0.999 for the EMA parameter of the second moment of RMSProp and Adam, 0.9 for the EMA parameter of the first moment of Adam. We fine-tune the learning rate logarithmic base of 2 around base learning rate such as 10-3 or 10 (depending on the algorithm, the architecture and the mini-batch size in the stochastic regime as detailed below in Appendix C.5), while making sure that the grid is sufficiently large such that the best learning rate is found inside the grid and not as the smallest the largest.
For the scheduled versions of the base optimizers, we consider three shapes: linear warm-up followed by constant learning rate, linear warm-up followed by linear decay, linear warm-up followed by cosine decay. The number iterations for the warm-up period is chosen as a fraction of the overall number of steps detailed in Appendix C.5. We also varied the horizon for the decaying schedules, again Appendix C.5.
Implementation of the linesearch procedure. in Section 2, we consider the following criterion
To implement the linesearch procedure described
f(we + TH. uz) s (1+ (),(w,) + cTj," C u..
Compared to (I), added relative decrease hyperparmeter å as we observed that the linesearch can sometimes stay stuck at vanishing learning rates otherwise.
To find a valid criterion we consider a usual backtracking linesearch that starts from guess T1.00 min(c, 1/- 1}. Choosing c+ = +o0 means that we start with an initial guess of at each iteration. The learning rate then decreased by factor e_ until the criterion is satisfied. Formally, the selected stepsize is then
vh max(Vi.. = = f(w + m, S (1 + 01)())))))))))))))))))())))))0 + CTj. (w.) u..).
We the search until the criterion is satisfied in the full batch regime and for maximum of 30 iterations the mini-batch regime. In the experiments, we consider the following variations.
[0, 10 10.100..00 0.55).
c +oo].
[0.8,
E 10 3] for = 0 and = 0, for c E [10- ..0.5..
Implementation of quadratically greedy tuner and CDAT. To implement the quadratically greedy tuner or CDAT, we compute the denominator u V2 fw as the second partial derivative of f along that is,
u 22 f()j 8² [w)] = ((Os()[[))))l
where [u] amounts to a Jacobian vector product (jvp) computed with forward mode auto-diff in differentiable programming languages such as JAX (DeepMind et al., 2020).
Computing the denominator in the CDAT rule by forward mode automatic differentiation enables a much lower memory consumption than using Hessian vector products (see, e.g., (Blondel and Roulet, 2024, Chapter 8), (Dagréou et al., 2024) for more details). The computation of the denominator by applying twice forward mode automatic differentiation still incurs approximately three times the
28

---

memory necessary to compute the objective (Blondel and Roulet, 2024, Chapter 8). The computa- tional cost of computing the denominator is also approximately three times the computational cost of computing the objective. The above approximations are done with the following reasoning. The second partial derivative requires to follow the graph of computation but with three variables, one for the parameters, one for copy of the update direction, one for another copy. At each node in the computation graph, the program computes the original computation, computes its first derivative along the first copy of the update direction, and computes the second derivative along the second copy.
In practice we observed for, e.g., the experiment on the full Imagenet dataset in mini-batch that the proposed CDAT rule required twice the wall time of the constant or scheduled learning rates counterparts. For this project, we considered CDAT as diagnostic tool to understand the interplay between curvature and learning rate tuners. For future work, the cost of computing the approximate edge may be circumvented or amortized by using, e.g., parabolic approximations as done by Mutschler and Zell (2020), or by computing it at given intervals as done by Liu al. (2024).
Further justification for the CDAT formula. In (8) we took the absolute value of the denominator to deal with concave approximations. Eigenvalue modifications in a Newton method are discussed by Nocedal and Wright (1999, Section 3.4). Taking the absolute value is one possible option. In practice, we observed positive curvatures along the update direction such that this choice did not matter.
C.4 Metrics implementation
Sharpness estimation. We estimated the sharpness by power iteration method run for 1000 iterations with an early stopping criterion defined by less than 10 relative accuracy. We accessed the Hessian by Hessian vector products, which limited the size of the full batch datasets considered on TPUs. The power iteration priori returns the largest eigenvalue in magnitude [X] MmImax and not necessarily the largest positive eigenvalue Amax- But in practice the largest eigenvalue in magnitude is the largest eigenvalue, see, e.g., (Ghorbani et al., 2019) for an in-depth study of the spectrum of the Hessian along the iterations of deep learning.
C.5 Additional experimental details per figure
We detail here any additional detail per figure not detailed in the summary above.
Fig. 1. The ResNet34 has batch normalization layers. For the GD baseline ResNet and the MLP Mixer the constant learning rate was tuned on a grid {10-3 3 2', é E (-l..... 7)}. For the RMSPRop baseline on the NanoLM and VIT. the constant learning rate tuned on grid {10 . 2', (-1.... 7}} and {10-b 2', E {O.... 7}} respectively.
Fig. 2. We consider a linear classification (in other words using an MLP without hidden layers) on the subset of CIFARIO detailed above. We search the constant stepsize of gradient descent in {10-3 2,i E (0, 2)]}. The grid centered around 1/ ( (wo))ll, that the optimal stepsize in convex smooth setting. As for above experiments, the largest learning rate in the grid led to divergence.
Fig. 3. This is the same setting as in the first panel of Fig. 1.
Fig. 4. For constant learning rate, model settings were given by a 3- 10-2,66 3. 10-1, For fixed = -0.1 training, a = 1, 6 = 0.5, 7o = 1.0.
Fig. 5. We considered the same settings as in Fig. 1. The grid search for GD on ResNet34 and RMSProp ManoLM the same in Fig. 1. For the MLPMixer, we fine-tuned GD with momentum on a grid {10-2 2 : E (1,...,8) 8]}. For the VIT, we fine-tuned Adam on a grid (10-5 2 : (1.....8)).
Fig. 6. This is exactly the same setting in the first panel of Fig. 5.
29

---

Fig. 7. We considered the full dataset of CIFARIO with ResNet50 with layer normalization in place of batch normalization.
Fig. 8.
See the details given for Fig. 20 and Fig. 21.
Fig.
For both values of a, a = 5 10-110-220000 10-2,b= 10-1, , -= 0.1, Ag 18, 1/ 0.05, 9o = Po
Fig.
10.
Same settings Figure Fig. 4.
Fig.
11.
Same settings as Figure Fig. 9.
Fig. 12.
This is the same setting as in Fig.
Fig. 13. We consider the same setting as in Fig. 3. For the Polyak stepsizes (23), we let 7max vary between 1 and 100 and select the best. For the hypergradient descent, let the hyper learning rate vary in 3 E {10', {-3....,00) 0}}.
Fig. 14. We considered again the ResNet50 with layer normalization instead of batch normalization. For the constant learning rate baseline we searched over grid of {mmm- 2" E {-1,..., 5H} for Tm 10 mm 4096 for the batch size.
Fig. 15. We considered a simple MLP with hidden sizes (256, 256, 256), ReLU activations. We tuned the constant learning rate baseline {10-3 33 2', {)1,
Fig. 16.
Details are provided in the legend.
Fig. 17. This is the same setting as in Fig. 6.
Fig. 18. In this figure, the MLPs considered use ReLu activations. If not detailed, the weight decay set to 10-5 and the subset considered of size 8192.
Fig. 19. The settings are the same as in Fig. 5, For the constant learning rate baselines we searched on a gird (Trmee 22' i 9}}. (-3.....
The base learning rate Thuse was chosen to be 10-3 for ResNet, 10-2 for the Mixer, 10 10-4 for the NanoLM and VIT. For the schedules' shapes, we searched over linear warm-up. linear warm-up with linear decay, linear warm-up with cosine decay. The initial and end learning rate were set to O. The horizons for the schedules were chosen in [N N/2, N/4] for N 8192 for the NanoLM, steps was searched in (0.05, 0.1,
16384 for the VIT, Mixer and ResNet. The fraction warm-up
Fig. 20. For the constant learning rate baseline, we consider searching the best constant learning rate on a grid {mmm 2', E {-1,...,77 7}} for 17 = 10-4 v /m/ 4009 where m denotes the varying batch size.
For the scheduled baseline, we consider the variants presented above (linear warm-up followed by constant, linear warm-up followed by cosine decay, linear warm-up followed by linear decay) with varying fraction of warm-up steps (0.05, 1, 0.2) and an initial learning rate of 0, final learning rate of for a fixed horizon of 512 epochs, and peak learning rate searched over {mm -4', {4, ....... 9}}.
The scaling factor of CDAT was searched parameter Beas in the computation of the numerators and denominators of the edge in {0, 0.9, 0.99). The best parameters found for CDAT can be inferred from Fig. 7. Namely, we found that non-zero EMA parameter for the estimation of the edge decay was essential for good performance and that the best scaling factor varied with the batch size. For example, at batch size 256 the best scaling factor 1.8 with Beda = 0.9.
grid (0.4, 0.6, ... 2.8). and we also tuned the EMA
Fig. 21. For the constant learning rate baseline, we consider searching the best constant learning rate on a grid {mm 20, {-1,...; for 7/m 10-1 / //m/124 1024 where m denotes the varying batch size.
30

---

For the scheduled baseline, we consider linear warm-up followed by cosine decay, with a fraction of warm-up steps of 0.1 and an initial learning rate of 0, final learning rate of 0 for a fixed horizon of 128 epochs, and peak learning rate searched over (7/m -4, i E .,]..
The scaling factor of CDAT was searched grid (0.4, 0.6,...- 2.6), and we also tuned the EMA parameter Bedat in the computation of the numerators and denominators of the edge in {0, 0.9, 0.99}.
C.6 Assets license and computing ressources
Assets. All experiments are done the open-source JAX ecosystem (DeepMind et al., 2020): archi- tectures are taken from Scenic (Dehghani et al., 2022), datasets from TensorFlow Dataset, algorithms from Optax. The datasets are MNIST (LeCun et al., 2010), (Creative Commons Attribution -hare Alike 3.0 license) CIFARIO (Krizhevsky et al., 2009) (no available license), Imagenet (Deng et al., 2009) (ImageNet explicitly permits the use of the dataset for non-commercial research purposes, however there is single license since the images are scrapped from different sources with different licenses), TinyShakespeare (Karpathy. 2015) (Apache 2.0 license in TensorFlow dataset, though the works of William Shakespeare in the public domain).
Computing resources. Experiments have mostly been run on Tensor Processing Units (TPUs) v2 (180 Tera Floating- Point Operations per Second (TFLOPS), 64 GB High Bandwidth Memory (HBM)). Experiments on MLP Mixers required TPUs v3 (420 TFLOPS 128 GB HBM). Very small scale experiments on MNIST with MLPs were run on CPUs. In terms of wall time, as discussed in Appendix C.3, we observed that the CDAT rule can be twice slower than the constant or scheduled learning rate counterparts. We consider CDAT as a diagnostic tool and leave as future work efficient implementations. Preliminary experiments and additional attempts to further adapt the momentum parameter on edge are not reported.
Authors contributions.
Vincent Roulet conducted the experimental work from the failures of linesearches to the analysis of the CDAT rule in various settings.
Atish Agarwala developed the theoretical model, with associated figures, interpretations and comments. He also helped to guide the experimental study with the insights gathered by the model. Jean Bastien Grill did an initial empirical study that gathered first intuitions on the method. He participated in the discussions and contributed to the writing. Grzegorz Swirszez participated in the discussions.
Mathieu Blondel participated in the discussions, contributed to the writing and proposed an alternative rule using a Gauss-Newton approximation of the objective.
Fabian Pedregosa initiated the project, performed larger scale empirical study of the CDAT rule on the MI.Commons benchmark, participated in the discussions, and contributed to the writing.
NeurIPS Paper Checklist
1. Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
Justification: This is an exploratory work that poses questions and attempts at answers from experiments and simplified models.
2. Limitations
Question:
Does the paper discuss the limitations of the work performed by the authors?
Answer: [Ye]]
We clearly point out the subdued effects of the proposed CDAT rule in the stochastic framework. In addition, we discuss the limitations of the current simplified models of the sharpness dynamics to explain the performance of CDAT in full batch regime.
Justification:
3. Theory Assumptions and Proofs
31

---

Question: For each theoretical result, does the paper provide the full set of assumptions and complete (and correct) proof?
Answer: [Yes]
Justification: Paper includes numerical experiments on models, derived from previous work (Damian et al., 2023). Modeling assumptions are clearly labelled, and accompanied with simulations. We also discuss the limitations the modeling assumptions.
4. Experimental Result Reproducibility
Question: Does the paper fully disclose all the information needed to reproduce the main ex- perimental results the paper to the extent that affects the main claims and/or conclusions of the paper (regardless whether the code and data are provided or not)?
Answer: [Yss]
Justification: We present detailed experimental setups in Appendix C.
5. Open access to data and code
Question: Does the paper provide open access to the data and code, with sufficient instruc- tions to faithfully reproduce the main experimental results, as described in supplemental material?
Answer: [No]
Justification: We are not able to open source the code yet. We hope to release it with the arxiv version of the manuscript later. Note that the data is in open access, and the algorithms are based on open source libraries such as JAX (DeepMind et al., 2020), Optax, and Scenic (Dehghani et al., 2022).
6. Experimental Setting/Details
Question: Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Answer: [Yes]
Justification: We present detailed experimental setups in Appendix C.
7. Experiment Statistical Significance
Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Answer: [No]
Justification: We privileged varying settings (datasets, architectures, losses, base optimizers) over varying the seeds to provide global overview of the CDAT rule diagnostic tool. We also provided additional experiments with varying mini-batch sizes (Fig. 20, Fig. 21), varying depths, widths, weight decays in MLP setting (Fig. 18). Our goal has been extract qualitative conclusions rather than quantitative conclusions. We also show in all transparency that the proposed rule does not outperform prefixed schedules (Fig. 19), and highlight limitations in the stochastic regime (Fig. 8).
8. Experiments Compute Resources
Question: For each experiment, does the paper provide sufficient information on the com- puter resources (type of compute workers, memory, time of execution) needed reproduce the experiments?
Answer: [Yes]
Justification: This is detailed in Appendix C.6.
9. Code Of Ethics
Question: Does the research conducted in the paper conform, in every respect, with the NeurlPS Code of Ethics https: https://teeee httpa/::///e/..//eeeee httpc////clectttisstiiieeeeeeeee
Answer: [Yes]
10. Broader Impacts
32

---

Question:
Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [NA]
Justification: No societal impact of the work is foreseen.
11. Safeguards
Does the paper describe safeguards that have been put in place for responsible release of data models that have high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Question:
Answer: [NA]
Justification: The paper poses no such risks
12. Licenses for existing assets
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?
Question:
Answer: [Yes]
Justification: This is detailed in Appendix C.6.
New Assets
Question: Are assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [NA]
Justification: The paper does not release new assets.
14. Crowdsourcing and Research with Human Subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well details about compensation (if any)?
Answer: [NA]
Justification: The paper does not involve crowdsourcing nor research with human subjects 15. Institutional Review Board (IRB) Approvals or Equivalent for Research with Human Subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?
Answer: [NA]
Justification: The paper does not involve crowdsourcing nor research with human subjects
33


Reviews:
[{"review_id": "ObQhCstaPg", "reviewer": "Reviewer_7YAk", "summary": "The paper proposes a novel learning rate tuning method, CDAT, that leverages the largest Hessian eigenvalue information during training. To illustrate the feedback loop between learning rate selection and sharpness dynamics, and to emphasize the importance of stepping on the edge of stability, the authors introduce a toy model. In the full batch setting, CDAT is shown to outperform constant learning rate schedules. Furthermore, under the mini-batch regime, the authors demonstrate that CDAT can naturally discover a warmup schedule in specific scenarios.", "questions": "Could the authors address my concerns in the weakness section?\nPlease correct me if I misunderstood anything from the paper. I am willing to discuss this further.", "limitations": "Yes", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. The paper is well-organized and easy to understand.\n2. The authors offered very detailed hyperparameter choices used in the paper.", "weaknesses": "1. Theoretical ground is not solid\n- For stochastic gradient descent, it was known that the correct measure of sharpness was the trace of Hessian instead of the largest eigenvalue [1], so using the largest eigenvalue of Hessian for SGD cases is not reasonable. \n- For adaptive optimizers like Adam, the correct quantity to measure is the pre-conditioned Hessian [2], the largest eigenvalue of which is usually very large at initialization. In Fig 8 Adam on the edge, the initial learning rate is very large, which is most likely because the authors used Hessian instead of pre-conditioned Hessian.\n- For cross-entropy loss, [3] has shown that the sharpness will reach $2/\\eta$ and then come down later. This is not covered anywhere in the paper.\n- At early stage of training, using $\\eta \\sim 2 / \\lambda_{max}$ does not mean it will remains on the EoS. As the [4, 5] have shown, there could be a catapult phase for $\\eta > 2 / \\lambda_{max}$, or depending on the initialization, the sharpness might decrease or increase at an early stage.\n\n[1] Lei Wu, Weijie J Su, https://proceedings.mlr.press/v202/wu23r.html\n\n[2] Jeremy M. Cohen, et al. https://arxiv.org/abs/2207.14484\n\n[3] Jeremy M. Cohen, et al. https://arxiv.org/abs/2103.00065\n\n[4] Aitor Lewkowycz, et al. https://arxiv.org/abs/2003.02218\n\n[5] Dayal Singh Karla, et al. https://arxiv.org/abs/2311.02076\n\n2. Experimental Results are not good:\n- It takes 500 epochs to train a ResNet-50 to $<80\\%$ on CIFAR-10\n- For most real-world settings, CDAT can not outperform existing simple cosine-annealing schedules.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their detailed feedback, and address their comments here.\n\n> _\"Theoretical ground is not solid\"_  \n- 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). \n\n> _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.''_  \n- 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. \n- 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.\n\n> _\"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.''_\n- 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\n$$ \n\\frac{u^\\top H u}{-u^\\top g}  = \\frac{g^\\top P^{-1} H P^{-1} g}{g^\\top P^{-1} g}.\n$$\nIf we were to maximize the above ratio, we would get\n$$\n\\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}),\n$$ \nand 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). \n- 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. \n\n> _\"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.''_ \n\n- 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. \n\n> _\"Experimental results are not good.''_\n- 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. \n\nPlease let us know if we can clarify our thinking further; we look forward to a fruitful dialogue with you.\n\n*References*:\n- (1) The Road Less Scheduled, https://arxiv.org/abs/2405.15682  \n- (2) Second-order regression models exhibit progressive sharpening to the edge of stability, https://arxiv.org/abs/2210.04860  \n- (3) High dimensional analysis reveals conservative sharpening and a stochastic edge of stability, https://arxiv.org/abs/2404.19261  \n- (4) SAM operates far from home: eigenvalue regularization as a dynamical phenomenon, https://arxiv.org/abs/2302.08692", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "2iKBBxZGw6", "reviewer": "Reviewer_cxBm", "summary": "This papers studies the behavior of several automatic learning rate schedulers in deep learning.  First, the paper studies two classical learning rate tuners - line search, and quadratically greedy (which chooses the learning rate that minimizes a quadratic Taylor approximation in the negative gradient direction).  The paper shows that while these two methods work well on a linear model, they severely outperform fixed-step-size GD on deep learning problems, in the full-batch setting.  The paper explains these results by noting that fixed-step-size GD automatically controls the sharpness along its trajectory, whereas these two classical learning rate tuners do not do this; hence, as training goes on, the sharpness increases and the estimated step sizes decrease.  The paper shows that this intuition can be reproduced in a simplified model of sharpness dynamics that was proposed in a prior work.   The paper then suggests a new automatic learning rate rule called \"CDAT\" which uses some fixed scaling $\\sigma$ of the quadratically greedy schedule; $\\sigma=2$ corresponds to a rule which always chooses as a step size that yields no predicted change in loss, under a local quadratic model.  The paper shows that in the full-batch regime, CDAT usually does as well as, or better than, GD with a fixed step size, though not as well as GD with an optimally tuned learning rate schedule.   The paper finds that in the _minibatch_ regime, the classical line search methods work well and CDAT performs worse; the paper explains this by noting that progressive sharpening is attenuated in the minibatch regime.  Finally, the paper shows that the benefits of CDAT can be partially captured by a simplified model of sharpness dynamics.", "questions": "Figure 16 seems to have the same figure accidentally repeated twice.\n\nI appreciate the ablation on MNIST, but I am not sure if the findings will transfer to other settings.  For example, I am skeptical about the finding that the efficacy of CDAT goes away with large weight decay.", "limitations": "n/a", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "I think the paper will be valuable in shedding light on the use (or lack thereof) of classical learning rate schedulers in deep learning.  The paper is high quality, original, and clear.", "weaknesses": "The main weakness of the paper is that it studies optimizers which are not state-of-the-art in the first place, and then does not fully or satisfactorily explain the behavior of these optimizers.  However, these weaknesses are arguably to be expected given that we are still in the very early stages of the theoretical study of optimization in deep learning.  In particular, no other papers contain stronger analyses than the analyses here.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for reading our paper and providing valuable feedback. \n\nWe answer below their comments.\n\n> _\"it studies optimizers which are not state-of-the-art in the first place\"_\n\n- 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. \n\n> _\"Figure 16 seems to have the same figure accidentally repeated twice.\"_\n\n- Thanks for catching the mistake on Fig. 16, we corrected this. The rightmost figure offered a similar conclusion as detailed in the caption.\n\n> _\"I am skeptical about the finding that the efficacy of CDAT goes away with large weight decay.\"_\n\n- 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.   \n\n_References_:  \n(1) Prodigy: An expeditiously adaptive parameter-free learner, https://arxiv.org/abs/2306.06101    \n(2) Dog is sgd’s best friend: A parameter-free dynamic step size schedule, https://arxiv.org/abs/2405.15682  \n(3) Why do Learning Rates Transfer? Reconciling Optimization and Scaling Limits for Deep Learning, https://arxiv.org/abs/2402.17457", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "4AzQ15i1do", "reviewer": "Reviewer_J393", "summary": "The paper investigates the consequences of the sharpness dynamics on step-size tuners' design. In particular, given a learning rate $\\eta$, the sharpness exhibits a progressive sharpening phase towards the Edge of Stability (EoS) threshold of $2/\\eta$ , where it stays for a large part of training time. First, the authors analyze the sharpness dynamics for two popular step-size tuners (linesearch and quadratically greedy tuners). In both cases, it is observed a significant decrease in learning rate over time, which coincides with a sharpening phase to ever-increasing thresholds. This leads to a suboptimal schedule of ever-decreasing learning rates that might perform better in a single step, but fails to deliver good performances at longer time scales. To fix this vicious cycle, the paper proposes Curvature Dynamics Aware Tuning (CDAT), which takes into consideration the alignment of the gradient with the Hessian, and it is designed to operate \"at the edge\" through a scalar multiplier $\\sigma$. The intuitions are corroborated with simple theoretical analyses on the interplay between sharpness dynamics and learning rate schedules, which (compared to when the learning rate is fixed) have both a time dependence.", "questions": "1. The authors provide an experiment where the tuner uses the exact sharpness instead of CDAT, which takes into account how the update is aligned with the largest eigenvector. What conclusions can be drawn on the learning rate and Hessian interplay? For instance, that at the beginning of training, the updates are not aligned with the leading eigenvector, which allows you to take larger steps (i.e. by increasing $\\sigma$)? \n\n2. Varying width and depth. The authors have an experiment (Fig. 18) where either the width or the depth is increased. There is a line of work that studies how hyperparameters (such as the learning rate) transfer from small to large scale (Yang et al, 2022 https://arxiv.org/abs/2203.03466). What are the implications of this paper's results in this context? It is my understanding that [2] can be related/discussed.\n\nThese papers seem relevant for the discussion on EoS and its relation to the learning rate, and varying hyperparameters:\n\n[1] Universal Sharpness Dynamics in Neural Network Training: Fixed Point Analysis, Edge of Stability, and Route to Chaos (https://arxiv.org/abs/2311.02076) \n\n[2] Why do Learning Rates Transfer? Reconciling Optimization and Scaling Limits for Deep Learning (https://arxiv.org/abs/2402.17457)\n\n[3] Understanding Gradient Descent on the Edge of Stability in Deep Learning (https://arxiv.org/abs/2205.09745)", "limitations": "The authors provide extensive experimental details that make the results fully reproducible. Also, the Appendix provides a lot of interesting ablations, such as varying width and depth, additional base learning rate tuners, and comparisons with commonly used prefixed schedules. In general, I find the suite of experiments very extensive. Finally, I also appreciate that the authors provide an extensive limitation section, setting the boundaries to which aspects of the interplay between sharpness and learning rate dynamics can be captured by CDAT.", "rating": 7, "confidence": 4, "soundness": 4, "presentation": 4, "contribution": 3, "strengths": "1. **Novelty of the idea**. Compared to classical analysis (e.g. linear models) where the sharpness $\\lambda_t$ is fixed at initialization, in neural networks the sharpness dynamics exhibit the consistent phenomenon of progressive sharpening towards the value of $2/\\eta$ (idea crystallized in Cohen et al. 2021). Given that the sharpness also provides a bound on the maximum step size allowed based on the local landscape, there is an interesting interplay between sharpness dynamics and learning rate, especially when the learning rate is varied either with a fixed schedule (e.g. warmup and decay) or with a step-size tuner. Thus, by analyzing step-size tuners, the paper is a step toward understanding this delicate interplay. To my knowledge, this is the first work toward this direction and it's the main strength of the paper.\n\n2. **Clarity**. The paper is exceptionally well-structured and flows nicely. The structure goes back and forth from empirical evidence to the theoretical model which provides intuitive justification. Overall, the paper is also very well-written and self-contained, and it provides all the necessary (and sufficient) context. The experiments are also portrayed schematically and intuitively.  \n\n3. **CDAT tuner**. The idea of the new proposed scheduler is simple and effective, and captures many interesting properties of the interplay between sharpness and learning rate, outperforming the base tuners in the full-batch case. Also, the authors put in a little extra engineering to ensure the stability of the optimizers, which makes intuitive sense.", "weaknesses": "I am overall strongly in favor of acceptance. However, my score is not higher for the following reasons:\n\n1. **Practical Limitations of CDAT**: I understand that the main purpose of the paper is to diagnose optimization properties at the edge of stability, and to study the interplay between learning rate and sharpness. However, there are a few practical limitations of the proposed tuner CDAT. First, it does not outperform the baseline of a constant learning rate in the practical deep learning use-case of mini-batch optimization, even after tuning the $\\sigma$ parameter which is supposed to take into account stochastic batch noise. Furthermore, it has an additional hyperparameter $\\sigma$ that has to be tuned. Thus, it loses the advantages of a learning rate tuner in the first place.  \n\n2. The authors attribute the lower performances of CDAT for the stochastic regime to (1) the optimal scaling factor is mini-batch dependent, and (2) that the sharpening effect is mitigated in the stochastic regime. However, (1) is not tested. Also, the fact that even after tuning the scaling factor $\\sigma$, CDAT underperforms is a partial indication that something else (beyond the fact the stochastic batches lower the sharpness threshold below $2/\\eta$) is responsible for the drop in performance. Again, it could be because you need a different $\\sigma$ per batch, but this is not experimentally validated. Also, this would further limit the applicability of the tuner.\n\n3. I would appreciate it if the theoretical model from Damial et al (2022)(together with its underlying assumptions) is summarized, at least in the Appendix.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We sincerely thank the reviewer for the positive feedback and to appreciate the exploratory nature of this work. \n\nBelow we answer their comments.\n\n> _\"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$\"_  \n- 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.  \n- 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.\nWe 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.  \n\n> _\"[the fact that] the optimal scaling factor is mini-batch dependent [...] is not tested\"_  \n- 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.  \n\n> _\"I would appreciate it if the theoretical model from Damian et al (2022)(together with its underlying assumptions) is summarized\"_  \n- 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. \n \n> _\"What conclusions can be drawn on the learning rate and Hessian interplay? [...] the updates are not aligned with the leading eigenvector [...] ?\"_  \n- 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.  \n- 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. \n\n> _\"What are the implications of this paper's results in [the] context [of [2]]?\"_  \n- 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.\n    \n*References*:  \n(1) The Road Less Scheduled, https://arxiv.org/abs/2405.15682", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "nk3pAKRDxR", "reviewer": "Reviewer_jCtW", "summary": "This paper intents to design new learning rate tuners based on insights from the Edge of Stability phenomenon. Based on the conjecture that classsical linesearch methods undershoot the edge of stability, and that this causes poor performance, the paper proposes a new learning rate tuner, that targets ceratin sharpness levels.", "questions": "* Why do you use the squared loss for the CIFAR dataset, instead of the cross-entropy loss?\n* In the Damian et al. paper (ref. [34]), the dynamics are not the same as what you write in (4): in particular, they define quantity $y_t$ where you write $\\lambda_t$, and those are not the same in my understanding. Can you clarify? (The quantity $y_t$ in [34] is roughly sharpness - 2/lr)\n* 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))?\n* Fig. 3: how did you make sure that $\\eta >0$ for the Quadratically Greedy Linesearch?\n\nMinor:\n\n* line 154: it should be \"a direction\", as there can be multiple eigenvectors for $\\lambda_t$.\n* line 139: it would be easier to grasp if there is a short summary of the results of Vaswani et al.", "limitations": "See weaknesses and questions. No code is provided.", "rating": 5, "confidence": 3, "soundness": 2, "presentation": 2, "contribution": 2, "strengths": "Learning rate tuning is tedious and time/compute consuming, and classical linesearch techniques often underperform. Therefore, this paper adresses an important problem for machine learning practicioners. The connection between Edge of Stability dynamics, and LR scheduling looks novel to me.\nThe proposed method is compatible with any update direction, thus directly applicable to Adam or similar optimizers.", "weaknesses": "The proposed CDAT tuner has several limitations:\n\n1) It requires additional computation, such as  vector-Hessian-vector products. How much does CDAT increase the runtime for one step?\n2) The performance looks highly sensitive to the choice of $\\sigma$ (see Fig. 5). In particular, choosing $\\sigma=2.0$ does not always beat tuned GD, even for rather small and easy datasets (subset of CIFAR10). It seems that CDAT replaces the necessity to tune the learning rate with tuning $\\sigma$.\n3) Only few experiments are conducted in a stochastic regime, which is the one that is interesting for practical use, and the results are not very convincing (e.g. Fig. 8 bottom).\n\nBesides those questions whether CDAT can be effective for tuning the learning rate, I have the following doubt regarding its motivation:\n\nIn Figure 3, Gradient Descent (GD) is also below the EoS for much of the training (third plot, from epoch 10000 onwards, and similar in Fig 12 for MixerTi/8). The paper conjectures that the poor performance of LR tuners comes from the fact that they stay below EoS systematically (Fig. 4 and end of section 2). However, from the observation above, it seems to me that the same reasoning would apply to GD with constant learning rate for some experiments. Thus, the main motivation for the derivation of CDAT, namely to target EoS, seems not fully convincing in the first place. In case I missed something important here, I am happy to discuss further.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for reading this manuscript and providing valuable feedback. \n\nWe answer their comments below.\n\n> _\"The proposed CDAT tuner has several limitations\"_\n- 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.\n\n> _\"How much does CDAT increase the runtime for one step?\"_\n- 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.\n\n> _\"CDAT replaces the necessity to tune the learning rate with tuning sigma.\"_\n- 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$.\n\n> _\"few experiments are conducted in a stochastic regime […] the results are not very convincing\"_\n- 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. \n- 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.\n\n> _\"[...] thus the main motivation for the derivation of CDAT, namely to target EoS, seems not fully convincing in the first place.\"_\n- 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]).\n- 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.\n\n> _\"In the Damian et al. paper (ref. [34]), the dynamics are not the same as what you write in (4)\"_\n- 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.\n\n> _\"Why do you use the squared loss for the CIFAR dataset, instead of the cross-entropy loss?\"_\n- 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).\n\n> _\"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))?\"_\n- The dynamics does indeed reflect model (5); thank you for catching that error, we will correct the reference.\n\n> _\"Fig. 3: how did you make sure that $\\eta>0$ for the Quadratically Greedy Linesearch?\"_\n- 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.\n\nThanks for catching the typos that we will correct in our manuscript. We look forward to a fruitful discussion that can further elucidate any concerns.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:
We sincerely thank the reviewers for reading our paper carefully, and providing numerous insightful comments. 

We answer each reviewer's comments separately and provide here some answers to the main common comment.

> "Experimental results are not good." (Reviewer 7YAk)
> "it does not outperform the baseline" (Reviewer J393)
> "the results are not very convincing" (Reviewer jCtW)

- The purpose of this paper is to underscore the interplay between sharpness dynamics and learning rate tuners through an extensive set of experiments and a simple rule (CDAT) playing the role of a probe. Sharpness dynamics are best understood in the full batch setting following the earlier work of Cohen [24, 25]. Such a setting reveals unexpected behaviors of classical learning rate tuners (Fig. 1, 3), questioning preconceived beliefs [9, 11]. The proposed CDAT rule is used as a **diagnostic tool** (as reviewer J393 points out) to examine whether the sharpening/edge of stability dynamics (usually studied for constant learning rate) can help designing more efficient learning rate tuners. The experiments in the full batch setting reveal that, across *optimizers* and tasks, a simple scaling of $\sigma=2$ ("on edge") provides remarkable performance, at least much better than usual learning rate tuners. 
- Experiments in the stochastic regime reveal, in full transparency, additional challenges. We preferred to present extensive experiments (21 figures) showing both opportunities and pitfalls of a simple idea (placing the optimizer "on edge'') rather than focusing on pure performance. The pure performance viewpoint in a stochastic regime may have led to some misunderstandings about the efficiency of e.g. linesearch methods in deep learning [9, 11] that did not carefully diagnose the separate effects of loss landscape and stochasticity.
- Recent learning rate tuners (Prodigy (1), schedule-free (2), DoG (3), etc…) have also focused solely on the issue of stochasticity motivated by the classical online learning framework. For all these recent optimizers, a warmup is necessary, which the theory cannot explain.  Our work takes an orthogonal approach: focus on a non-stochastic regime from an empirical viewpoint to understand the role of loss landscape. We may then understand how warm-up can be induced by closed-loop feedback effects on sharpness and learning rate dynamics. 
- Finally, we are unaware of other work putting the recent findings in sharpness dynamics of optimizers into practice (though the reference (4) provided by reviewer J393 can have practical impacts on scaling laws). The CDAT rule sheds light on the practical benefits of the study of edge of stability/sharpening that drew much attention in past years. It also highlights the importance of studying such dynamics in closed loop feedback scenarios such as the ones induced by learning rate tuners.

_References_:  
(1) Prodigy: An Expeditiously Adaptive Parameter-Free Learner, https://arxiv.org/abs/2306.06101  
(2) The Road Less Scheduled, https://arxiv.org/abs/2405.15682  
(3) DoG is SGD's Best Friend: A Parameter-Free Dynamic Step Size Schedule, https://arxiv.org/abs/2302.12022  
(4) Why do Learning Rates Transfer? Reconciling Optimization and Scaling Limits for Deep Learning, https://arxiv.org/abs/2402.17457

Meta-review: The paper examinates a simple idea: choosing an adaptive learning  to be always `on the edge' instead of in the stable regime as is usually done. The paper makes some interesting observations, and proposes a clean and simple explanation for them. All reviews are in agreement for acceptance after the discussion period.
Decision: Accept (poster)