Where Do Large Learning Rates Lead Us?

We thank the Reviewer for their thoughtful review of our work! We now address the raised concerns and questions.

The experiments in the study are limited to specific datasets (CIFARor synthetic) and neural network architectures (Resnet).

We have conducted additional experiments showing that our findings transfer to other practical scenarios. Please see the general comment for more detail.

The conclusion that “only a narrow range of initial LRs slightly above the convergence threshold lead to optimal results” is hard to apply in practice, as determining the convergence threshold beforehand is challenging and problem-specific.

We fully agree with this remark, however, in our general comment we also point out that precisely determining this range (subregime 2A) is often not necessary to achieve a good final solution as long as a more advanced LR schedule is used than a simple LR drop.

Answering Q1
The convergence threshold is understood as a value in the LR range that divides the LRs allowing the loss to converge (if trained with these fixed LR values) and the remaining LRs. In theory, finding this threshold could be done relatively efficiently using, e.g., binary search, however, as we answered above, it is not needed in practice.
Please refer to our general comment for a more detailed discussion on this topic.

Answering Q2
Takeaway 1 essentially provides a recipe to obtain optimal solutions after fine-tuning with a constant small LR or weight averaging: start training with a moderately high LR from subregime 2A, a relatively narrow range above the convergence threshold. However, as we show in Appendix D, training even with significantly larger initial LRs from subregime 2B can get similar results if one chooses just a bit more complex LR schedule. Hence, we argue that for practical schedules with gradual LR decay it is often sufficient to take some reasonably large initial LR, which do not allow for convergence, to achieve good final quality.
Please see the general comment for further discussion.

Answering Q3
A totally correct note, we thank the reviewer for the found inaccuracy that we will fix in the following text revision!
What we should have written was: "In the third regime, both the angular distances and the error barriers approach their upper limits (angular distance of $\pi/2$, which is a typical angular distance between two independent points in high-dimensional space, and random-guess error), completing the trend established in subregime 2B".

We thank the Reviewer for their valuable feedback and overall positive assessment of our work! We respond to the comments and questions as follows.

Some of the paper's results have already been shown by prior work, e.g., large LRs can lead to sparse features [1].

The results of [1] are closely related to ours, and we discuss them in Section 1.1. However, there are still important differences between this study and ours.
The referenced work examines feature sparsity in terms of the fraction of distinct, non-zero activations at some hidden layer of a neural network over the dataset. In contrast, we study the importance of features in the input data for prediction. This approach allowed us to identify feature sparsity as the model preference for the most task-relevant features in the data when training with optimal initial LR values, which is a novel result. In short, our works consider different definitions of “features”: whether they are internal representations of data within a model or patterns in the input data itself.
Furthermore, we came to a different conclusion because we found that feature sparsity (as we understand it) behaves non-monotonically w.r.t. LR, while [1] suggests a monotonic trend.

how to choose LR to reach the best regime (2A) seems to remain an issue in practice

This indeed could be the case if one wanted to locate the optimal initial LRs for weight averaging or fine-tuning with a decreased LR at the end of training. However, more advanced LR schedules in practice allow for substantially higher initial LRs with similar final performance.
Please see our general comment for further discussion.

It would be great if the authors could discuss whether different training regimes are also identifiable from the training loss alone.

This is a very good point and we regret that we did not reflect it clearly in the text.
The distinctive features of different regimes were first established in [2] based on various metrics including training loss and gradient norm (please refer to Fig. 1 in [2]). For instance, regimes 1 and 2 can be easily distinguished by the behavior of the training loss/error: whether it reaches low values (convergence) or hovers at some non-zero level. Therefore, no oracle access to the test accuracy is necessary to identify regimes.
We will improve our wording to avoid further misinterpretations.

the overall evaluation of the paper is still a bit limited

We have conducted additional experiments supporting our findings, please see the general comment for more detail.

Answering Q1
As we answered above, based on our findings, in practice one simply needs to choose some large enough LR value from regime 2, i.e., not allowing for convergence, and use a gradually decaying LR schedule to obtain near optimal results.

Answering Q2
That is a very intriguing and non-trivial question that requires further investigation. We attempted to provide some initial intuition based on the loss landscape and feature learning analysis. As we briefly discuss in Section 7, we conjecture that feature sparsity and mode proximity/linear connectivity observed in subregime 2A are closely related. We hypothesize that pre-training in this regime helps the model filter out unnecessary input features and emphasize the most relevant ones, which is beneficial for further fine-tuning/SWA. This process is dually represented in the training dynamics as finding a stable linearly connected basin of good solutions in the loss landscape. By contrast, training in regime 1 converges to the nearest minimum admissible by the chosen LR value, not allowing enough time for feature consolidation. Again, this interpretation is currently largely speculative, but deserves further study.

Answering Q3
As our additional experiments (in the general comment) suggest, all key findings of our work remain valid in other practical settings, including training on Tiny ImageNet and using the ViT model. Also, for example, our synthetic example, which substantially differs from the CIFAR image classification with convolutional networks, possesses all the properties of training in different regimes (Appendix G). Overall, we expect that our results can be transferred to any overparameterized training setting that allows convergence, so that regime 1 is reachable.

[1] Andriushchenko Maksym et al. SGD with large step sizes learns sparse features. In International Conference on Machine Learning, 2023.
[2] Kodryan Maxim et al.. Training scale-invariant neural networks on the sphere can happen in three regimes. Advances in Neural Information Processing Systems, 2022.

We are grateful to the Reviewer for their positive and constructive feedback on our work! We address the concerns and questions below one by one.

Scale-invariant setting
As the Reviewer rightly noted, our main experiments are conducted in a specific scale-invariant setting to provide effective control over the learning rate value. Since our study is more fundamental than practical in nature, we decided to follow this setup in line with prior work examining the effect of learning rate on training dynamics. This allowed us to isolate the impact of scale invariance on the effective learning rate of the model and then extend our main findings to more practical scenarios. In our general comment, we present additional experimental results that further confirm that our claims apply to conventional training settings as well.

Finding the optimal LR range
It could indeed be computationally burdensome to find the exact convergence threshold for a given training setup. However, although the optimal LRs for fine-tuning with a constant small LR or weight averaging usually lie in a relatively narrow range just above this threshold, in practice even taking substantially larger initial LRs can give similar results if one chooses a slightly more complex LR schedule (Appendix D). Hence, we may conclude that for regular LR schedules it is not so important to precisely determine the convergence threshold but to choose some reasonably large initial LR above it.
Please see our general comment for further discussion.

Answering Q1
Thanks for the comment, this is indeed a poor formulation, which we will correct in the next text revision!
What we wanted to say is the following. LRs from regime 1 allow training to converge to some minima, however these minima a) have non-optimal generalization (compared to the fine-tuned/SWA solutions of subregime 2A) and b) are “unstable” in the sense that increasing LR (within the same regime 1) after convergence can knock the model out of the current minimum to a new minimum, which is perhaps better but still belongs to the minima of the first regime and therefore is not optimal.

Answering Q2
Indeed, local geometry, from the point of view of training/test error barriers, cannot be used to separate regimes 1 and 2A, as in fact we don’t have barriers in either case. In the first case, because all the points (pre-trained, fine-tuned, and SWA) simply coincide, which is suggested by the angular distance on the plots. In the second case, because the localized basin contains close high-quality solutions, which are linearly connected. Error barriers are, however, useful for separating subregimes 2A and 2B, as in the latter case fine-tuned/SWA solutions lose linear connectivity. At the same time, regimes 1 and 2 can be easily distinguished by the behavior of the training loss/error: whether it reaches low values (convergence) or hovers at some non-zero level (see, e.g., Figure 1 in [1]).
We will clarify this more in the text.

Limitations

The architecture and datasets tested are restrictive.

We have conducted additional experiments supporting our findings, please see the general comment for more detail.

[1] Kodryan Maxim et al. Training scale-invariant neural networks on the sphere can happen in three regimes. Advances in Neural Information Processing Systems, 2022.