Stepping on the Edge: Curvature Aware Learning Rate Tuners

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.

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. The proposed method is compatible with any update direction, thus directly applicable to Adam or similar optimizers.

Weaknesses

The proposed CDAT tuner has several limitations: 1) It requires additional computation, such as vector-Hessian-vector products. How much does CDAT increase the runtime for one step? 2) 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$. 3) 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). Besides those questions whether CDAT can be effective for tuning the learning rate, I have the following doubt regarding its motivation: In 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.

Questions

* Why do you use the squared loss for the CIFAR dataset, instead of the cross-entropy loss? * 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) * 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))? * Fig. 3: how did you make sure that $\eta >0$ for the Quadratically Greedy Linesearch? Minor: * line 154: it should be "a direction", as there can be multiple eigenvectors for $\lambda_t$. * line 139: it would be easier to grasp if there is a short summary of the results of Vaswani et al.

Rating

5

Confidence

3

Soundness

2

Presentation

2

Contribution

2

Limitations

See weaknesses and questions. No code is provided.

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.

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. 2. **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. 3. **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: 1. **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. 2. 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. 3. 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.

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$)? 2. 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. These papers seem relevant for the discussion on EoS and its relation to the learning rate, and varying hyperparameters: [1] Universal Sharpness Dynamics in Neural Network Training: Fixed Point Analysis, Edge of Stability, and Route to Chaos (https://arxiv.org/abs/2311.02076) [2] Why do Learning Rates Transfer? Reconciling Optimization and Scaling Limits for Deep Learning (https://arxiv.org/abs/2402.17457) [3] Understanding Gradient Descent on the Edge of Stability in Deep Learning (https://arxiv.org/abs/2205.09745)

Rating

7

Confidence

4

Soundness

4

Presentation

4

Contribution

3

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.

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.

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.

Questions

Figure 16 seems to have the same figure accidentally repeated twice. I 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.

Rating

6

Confidence

4

Soundness

3

Presentation

3

Contribution

3

Limitations

n/a

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.

Strengths

1. The paper is well-organized and easy to understand. 2. The authors offered very detailed hyperparameter choices used in the paper.

Weaknesses

1. Theoretical ground is not solid - 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. - 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. - 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. - 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. [1] Lei Wu, Weijie J Su, https://proceedings.mlr.press/v202/wu23r.html [2] Jeremy M. Cohen, et al. https://arxiv.org/abs/2207.14484 [3] Jeremy M. Cohen, et al. https://arxiv.org/abs/2103.00065 [4] Aitor Lewkowycz, et al. https://arxiv.org/abs/2003.02218 [5] Dayal Singh Karla, et al. https://arxiv.org/abs/2311.02076 2. Experimental Results are not good: - It takes 500 epochs to train a ResNet-50 to $<80\%$ on CIFAR-10 - For most real-world settings, CDAT can not outperform existing simple cosine-annealing schedules.

Questions

Could the authors address my concerns in the weakness section? Please correct me if I misunderstood anything from the paper. I am willing to discuss this further.

Rating

6

Confidence

4

Soundness

3

Presentation

3

Contribution

3

Limitations

Yes

Thank you for rebuttal

Thank you for the detailled rebuttal, and for adressing all of my questions. * Regarding my concerns on practical use: from your rebuttal it seems that we agree that CDAT currently is not intended to be a fully practical optimizer yet, mainly because of several mentioned limitations like increased time per iteration, problems in the minibatch-setting, etc. While this is fine for me, it is indeed a limitation of the paper, because it will require a lot of further engineering and research effort to make it a practical method. * Regarding sensitivity on $\sigma$: indeed $\sigma=2.0$ seems to perform well, but still does not outperform (tuned) GD on 2 out of 4 tasks in Figure 5. My main concern is that a slightly different value of $\sigma=1.94$ can result in quite different loss curves (e.g. Fig. 5, Mixer), suggesting that the CDAT dynamics are sensititve to the choice of $\sigma$. Other replies: > thus the main motivation for the derivation of CDAT, namely to target EoS, seems not fully convincing in the first place. Thanks for clarifying, I now understood that the main focus is on the initial phase of training. In that case, I agree that line-search methods show a distinct behaviour than GD. I would kindly ask you to stress this in the final version, in order to avoid confusion. > the dynamics are not the same as what you write in (4) Redoing my calculations, I now obtain the correct model matching the one in (4). I recommend to explain this reparametrization in more detail (e.g. in appendix), as your notation does not match the one from Damian et al.. I am considering to raise my score, and will do so in the end of the discussion period, after spending more time to read the other reviews/rebuttals.

response

> Could the reviewer clarify what optimizer do they have in mind? Sorry, I was vague -- my point was that CDAT is not a state-of-the-art optimizer, which limits the significance studying it (especially in a way that leaves a lot of questions open).

Response

I thank the authors for addressing my concerns and questions. **On the performance of CDAT**: I would like to clarify that I do not find the lower performance of CDAT in the stochastic setting as an inherent weakness of the paper, but in a sense, it is an objective limitation that prevented me from giving an even higher score. However, I appreciate the authors' efforts to run these experiments beyond the setting where the theory applies. It would just have been even greater if CDAT (or a variation of it) was working in more practical settings. > [the fact that] the optimal scaling factor is mini-batch dependent [...] is not tested By mini-batch dependent, I referred to the dependence of every *single* mini-batch, and not the *size* of the batch, as the authors point out in lines 237-239. It would be interesting to see how the sharpness varies on a mini-batch level, i.e. ultimately how much the stochastic noise in the batch affects the curvature across training. I remain in solid favor of acceptance.