The large learning rate phase of deep learning
Authors: Aitor Lewkowycz
Yasaman Bahri
Ethan Dyer
Jascha Sohl-Dickstein
Guy Gur-Ari
Venue:  (2021)
Keywords: deep learning
wide networks
training dynamics

Abstract: The choice of initial learning rate can have a profound effect on the performance of deep networks. We present empirical evidence that networks exhibit sharply distinct behaviors at small and large learning rates. In the small learning rate phase, training can be understood using the existing theory of infinitely wide neural networks. At large learning rates, we find that networks exhibit qualitatively distinct phenomena that cannot be explained by existing theory: The loss grows during the early part of training, and optimization eventually converges to a flatter minimum. Furthermore, we find that the optimal performance is often found in the large learning rate phase. To better understand this behavior we analyze the dynamics of a two-layer linear network and prove that it exhibits these different phases. We find good agreement between our analysis and the training dynamics observed in realistic deep learning settings. 



Under review as a conference paper at ICLR 2021
THE LARGE LEARNING RATE PHASE OF DEEP LEARNING
Anonymous authors Paper under double-bling review
ABSTRACT
The choice of initial learning rate can have a profound effect on the performance of deep networks. We present empirical evidence that networks exhibit sharply distinct behaviors at small and large learning rates. In the small learning rate phase, training can be understood using the existing theory of infinitely wide neural networks. At large learning rates, we find that networks exhibit qualitatively distinct phenomena that cannot explained by existing theory: The loss grows during the early part of training. and optimization eventually converges to a flatter minimum. Furthermore, we find that the optimal performance is often found the large learning rate phase. To better understand this behavior analyze the dynamics of a two-layer linear network and prove that it exhibits these different phases. We find good agreement between our analysis and the training dynamics observed in realistic deep learning settings.
INTRODUCTION
Deep learning has shown remarkable success across a variety of tasks. At the same time, our theoretical understanding of deep learning methods remains limited. In particular, the interplay between training dynamics, properties of the learned network, and generalization remains a largely open problem.
In tackling this problem, much progress has been made by studying deep neural networks whose hidden layers are wide. In the limit of infinite width, connections between stochastic gradient descent (SGD) dynamics of neural networks, compositional kernels, and linear models have been made. These connections hold when the learning rate is sufficiently small. However, a theory of the dynamics of deep networks that operate outside this regime remains largely open.
In this work, we present evidence that SGD dynamics change significantly when the learning rate above critical value, 'eritt determined by the local curvature of the loss landscape at initialization. These dynamics are stable above the critical lcarning rate, up to a maximum learning rate 7mma Training these large learning rates results in different signatures than observed for learning rates 7) < Meeii: the loss initially increases and peaks before decreasing again, and the local curvature drops significantly early in training. We typically find that the best performance is obtained when training above the critical learning rate. Empirically, we find these two learning rate regimes are robust, holding across variety of architectural and data settings.
Figure highlights our key findings. We now describe the main contributions of this work.
1.1 TRAINING WITH A LARGE LEARNING RATE LEADS TO A CATAPULT EFFECT
Consider deep network defined by the network function f0,, where are the model parameters and the input. We define the curvature X, training step & be the max eigenvalue of the Fisher Information Matrix, F. == Eg oo (0,, x") Amari et al. (2000); Karakida et al. (2018). Equivalently, x, is the max eigenvalue of the Neural Tangent Kernel Jacot al. (2018).
Figure shows the results of training several deep networks with mean squared error (MSE) loss using SGD with range of learning rates. The loss and curvature are measured at every step during

---

Under review as a conference paper at ICLR 2021
1.00
0.95
0.85
Curvature
0.75
0.70
campuut boundary
0.65
catapuit divergent boundary
finil
Wide Blennt 25-10
0.60
(b)
10'2
10
Learning rate
Weight correlation
Figure I: Large learning rates lead large weight movement and better performance. (a) A visualization of gradient descent dynamics derived in our analytic model. A 2D slice of parameter space is shown, where lighter color indicates higher loss and dots represents points visited during optimization. Initially, the loss grows rapidly while local curvature decreases. Once curvature is sufficiently low, gradient descent converges to flat minimum. call this the catapult effect. See Figures S2 and SI for more details. (b) Confirmation of our predictions a practical deep learning setting. Line shows the test accuracy of a Wide ResNet trained on CIFAR-10 as a function of learning rate, each trained for a fixed number of steps. Dashed lines show our predictions for the boundaries of the large learning rate regime (the catapult phase), where expect optimal performance to occur. Maximal performance is achieved between the dashed lines, confirming our predictions. See Section for details.
the early part of training. We notice the following effects, which occur when the learning rate is above the critical value Merit = 2/>o, where Ag is the curvature initialization."
1. During the first few steps of training, the loss grows significantly compared its initial value before it begins decreasing. We call this the catapult effect.
2. Over the same time frame, the curvature decreases until it is below 2/n.
We can build intuition for these effects using loss landscape considerations. Consider a linear model where the curvature of the loss landscape is given by Ag Here, curvature means the largest eigenvalue of the linear model kernel. The model can be trained using gradient descent as long as the learning rate 17 obeys y) < 2/ When 1) 2/ Ao. the loss diverges and optimization fails.
Next, consider a deep network. If we train the model with learning rate 17 > Meriit, we may again expect the loss grow initially, assuming the curvature is approximately constant the neighborhood of the initial point in parameter space. This is the effect observed in Figure 2. However, unlike the linear case, optimization may still succeed if gradient descent able to navigate to an area of the landscape that has lower curvature X, such that 1) < 2/X. This is indeed what we observe in practice.
In Figure we show that optimal performance typically occurs when a network is trained in the large learning rate regime. As discussed further in Section 2, this is true even when the compute budget for smaller learning rates increased to account for the smaller step size. This consistent with previous observations in the literature, which showed correlation between performance and the flatness of the minimum (Keskar et al., 2016).
1.2 AT LARGE WIDTHI A SHARP DISTINCTION BETWEEN LEARNING RATES REGIMES
The large width limit of deep networks has been shown to lead to simplified training dynamics that are amenable to theoretical study, as in the case of the Neural Tangent Kernel (Jacot et al., 2018). In this work we show that the distinction between small and large learning rates becomes sharply defined at large width. This can be seen in Figures 2c, 2f, which show the curvature of sufficiently
wide networks after the initial part of training, as a function of learning rate. When 77 < Mertit the curvature is approximately independent of the learning rate, while for 1 Marit. the curvature is lower than 2/n--
'The critical learning rate depends the scale of initialization through Ag-
2

---

Under review as a conference paper at ICLR 2021
4
2/Au
0.3
2/A
5-0.0
Trainin loss
12/Aa
11:122
0.2
0.1
10"
100
40
41
(a)
(b)
(c)
0
0.2
40
13/Aa
Touning bas
2/1
0.3
(d)
Figure 2: Early time catapult dynamics, (a,b.c) A hidden layer fully-connected network with ReLU non- lincarity with width 2048 trained MNIST (Beiit 6.25). (d.c.f) Wide ResNet 28- trained on CIFAR-ID (Meris 0.18). Both networks are trained with vanilla SGD; for more experimental details see Appendix A. (a,d) Early time dynamics of the training loss for learming rates in the linear and catapult phases. (b) Maximum value of the loss a function of the learning rate. (e) Early time dynamics of the curvature for learning rates in the n 250 (for FC) and t r 30 (for WRN), function of
linear and catapult phase. (c.f) A measured learning rate. Training diverges for learning rates in the shaded region.
In Section 3 we analyze the gradient descent dynamics of 2-layer linear networks, and find that they exhibit similar behavior that can be understood in the large width limit. When training with small learning rates, existing theory can describe the training dynamics of such networks. And, we present new theoretical results that explain the behavior large learning rates. To summarize our findings, in the large width limit we identify two separate learning rate regimes, or phases, with the following characteristics.
Lazy phase: 17 < 2/A0 For sufficiently small learning rate, the curvature A training step remains constant throughout training, and the model becomes equivalent to linear model (Jacot et al., 2018; Lee et al., 2019). The model converges to a nearby point in parameter space, and this behavior is sometimes called lazy training (Du et al., 2019; Zou et al., 2018; Allen-Zhu et al., 2019; Chizat et al., 2019).
Li & Liang, 2018;
Catapult phase: Merit < 1) < Maax At large learning rates the loss grows be of order the width over a number of training steps that is of order log(n). During the same period, the curvature decreases until reaching a flat minimum (relative to that reached in the lazy phase). The gradient descent dynamics in this phase are visualized in Figure and Figure SI.
below 2/1j. Beyond this point the loss decreases and training converges, ultimately
The maximum learning rate 7max (beyond which training no longer converges) depends the setup. In our theoretical model the maximum learning rate is 7max 4/ >g. For ReLU networks, we find empirically that Maax Rs 12/Ag-
1.3 LIMITATIONS
Our empirical results and our theoretical analysis focus on the case of MSE loss and on training with vanilla SGD, and do not extend case of cross-entropy loss, or to training with other optimizers such momentum. Additionally, our theoretical analysis involves a 2-layer linear network, and does not apply the case of networks with non-linearities such as ReLU.
1.4 RELATED WORKS
Our work builds on several existing results, which we now briefly review.
3

---

Under review as a conference paper at ICLR 2021
The existing theory of infinite width networks is insufficient to describe large learning rates. A recent body of work has investigated the gradient descent dynamics of deep networks in the limit of infinite width (Danicly, 2017; Jacot et al., 2018; Lee et al., 2019; Du et al., 2019; Zou et al., 2018; Allen-Zhu et al., 2019; Li & Liang, 2018; Chizat et al., 2019; Mei et al., 2018; Rotskoff & Vanden- Snjden-Finnden, 2018; Sirignano & Spiliopoulos, 2018; Woodworth al., 2019; Naveh al.; Xiao et al., 2019). Of particular relevance is the work by Jacot et al. (2018) showing that gradient flow in the space of functions is governed by a dynamical quantity called the Neural Tangent Kernel (NTK) which is fixed at its initial value in this limit. Lee et al. (2019) showed this result is equivalent to training the linearization of model around its initialization in parameter space. Finally, moving away from the strict limit of infinite width by working perturbatively, Dyer & Gur-Ari (2020); Huang Yau (2019) introduced an approach to computing the finite-widht corrections to network evolution.
Despite this progress, in many practical deep learning settings, the neural network has finite width and evolves nontrivially, with a large change its associated Neural Tangent Kernel. Depending on the architecture and hyperparameters, such networks may give superior performance. Prior work has compared the performance of finite-width, SGD-trained deep networks with the infinite-width kernels derived from the networks (Lee et al., 2018; Novak al., 2019; Arora et al., 2019). Performance gaps are observed in some cases, notably in convolutional networks, implying that existing anfinite-widie theory is insufficient explain the performance of deep networks in such settings where the network evolves nontrivially.
Large learning rate SGD improves generalization. SGD training with large initial learning rates often leads to improved performance over training with small initial learning rates (see (Li et al., 2019; Leclere & Madry. 2020; Xie et al., 2020; Frankle et al., 2020; Jastrzebski et al., 2020) for recent discussions). It has been suggested that one of the mechanisms underlying the benefit of large learning rates is that noise from SGD leads to flat minima, and that flat minima generalize better than sharp minima (Hochreiter & Schmidhuber, 1997; Keskar et al., 2016; Smith & Le, 2018; Jiang et al., 2020; Park et al., 2019) (though see Dinh et al. (2017) for discussion of some caveats). According to this suggestion, training with a large learning rate (or with small batch size) can improve performance because it leads to more stochasticity during training (Smith & Le, 2018; Mandt et al., 2017; Smith et al., 2017; Smith al., 2018).
We develop a connection between large learning rate and flatness of minima in models trained via SGD. Unlike the relationship explored in most previous work though, this connection is not driven by SGD noise, but arises solely as a result of training with a large initial learning rate, and holds even for full batch gradient descent.
EXPERIMENTAL RESULTS
In a variety of deep learning settings, we find clear evidence of the different phases introduced in Section 1. The experiments all use MSE loss, sufficiently wide networks, and vanilla SGD with learning rate v). Parameters such as network architecture, choice of non-linearity, weight parameterization, and regularization, do not significantly affect this conclusion.
In these experiments, we define the curvature À as the maximum eigenvalue of the Fisher Information Matrix or, equivalently, as the maximum eigenvalue of the Neural Tangent Kernel (NTK). Given network function f Rd -> R with model parameters 0 E R", and training set {(r., Vax) ] Ja=1> the NTK θ Rdd x 1d + R is defined by Ozz, x²) == m Σp (of(x)) oo (a'). We denote the curvature at time by Ao equal to the maximum eigenvalue of e. Another common measure of local curvature is the maximum Hessian eigenvalue; at large width we expect these measures to agree (Dyer & Gur-Ari, 2020), and we verify the agreement Appendix D.6.
Building on the observed correlation between flat minima and generalization performance (Keskar et al., 2016; Jiang et al., 2020), we conjecture that optimal performance occurs in the large learning rate (catapult) phase, where optimization converges to a low curvature minimum. For a fixed amount of computational budget, we find that this conjecture holds in all cases we tried. Even when comparing different learning rates trained for fixed amount of physical time tphys = t 77, find that performance of models trained in the catapult phase either matches or exceeds that of models trained with learning rates below /jrrit.
4

---

Under review as a conference paper at ICLR 2021
2.1 EARLY TIME CURVATURE DYNAMICS
Here we present empirical support for the lazy and catapult phases of training described in the previous section. Additional experimental results are presented in the appendix. We find that the lazy phase is characterized by small changes to the curvature and loss during training, while the catapult exhibits large deviations.
Figure 2 shows the curvature during the early part of training for two deep learning settings that involve sufficiently wide networks. The results are compared against the prediction of a phase transition lJerit 2/ Xn- For learning rates 7/ Meriit. (lazy phase), the curvature is independent of the learning rate and is approximately constant throughout training. For //crit < 7) < 7/max we find that the eurvature decreases during training to below 2/n--
Figure 2 also shows the loss initially increasing before converging for large learning rates, a signature of the catapult effect. This transient behavior is very short, taking less than 10 steps to complete. Because of this, the training curve for the test loss is very similar and also shows the catapult effect. In these and other experiments involving ReLU networks, we find that Twax ~ 12/2o is good predictor of the maximum learning rate (in Appendix E.4 we discuss other nonlinearities). We conjecture that this is the typical maximum learning rate of networks with ReLU non-linearities.
2.2 GENERALIZATION PERFORMANCE
We now consider the performance of trained models in the different phases discussed in this work. Keskar et al. (2016) observed a correlation between the flatness of a minimum found by SGD and the generalization performance (see Jiang al. (2020) for additional empirical confirmation of this correlation). In this work, we showed that the minima SGD finds are flatter in the catapult phase, as measured by the top kernel eigenvalue. Our measure of flatness differs from that of Keskar et al. (2016), but we expect that these measures to be correlated.
We therefore conjecture that optimal performance often obtained for learning rates above Berit and below the maximum learning rate. In this section we test this conjecture empirically. We find that performance in the large learning rate regime always matches or exceeds the performance when < Merrt. For fixed compute budget, we find that the best performance is always found in the catapult phase.
Figure 3 shows the performance of a convolutional network and a Wide ResNet (WRN) trained on CIFAR-IO. The experimental setup, which we now describe, was chosen to ensure a fair comparison of the performance across different learning rates. The network is trained with different initial learning rates, followed by a decay at a fixed physical time t 17 to the same final learning rate. This schedule is introduced in order to ensure that all experiments have the same level of SGD noise toward the end of training.
We present results using two different stopping conditions. In Figure 3a, 3c, all models were trained for a fixed number of training steps. We find a significant performance gap between small and large learning rates, with the optimal learning rate above Meriit and close to 7/max. Beyond this learning rate, performance drops sharply.
The fixed compute stopping condition, while of practical interest, biases the results in favor of large learning rates. Indeed, in the limit of small learning rate, training for fixed number of steps will keep the model close to initialization. To control for this, Figure 3b,3d models were trained for the same amount of physical time t- For the CNN of figure 3b, decaying the learning rate does not have a significant effect on performance and we observe that performance is flat up to 7mnax and there is no correlation between our measure of curvature and generalization performance. Figure 3d shows the analogous experiment for WRN. When decaying the learning rate toward the end of training to control for SGD noise, we find that optimal performance is achieved above Merit. In all these cases, 7aax is a good predictor of the maximal learning rate, despite significant differences in the architectures. Notice that by tuning the learning rate to the catapult phase. we are able to achieve performance using MSE loss, and without momentum, that is competitive with the best reported results for this model (Zagoruyko & Komodakis, 2016).
5

---

Under review as a conference paper at ICLR 2021
0.74
0.70
Test accuracy
Test accuracy
0.65
0.68
0.00
2/2b
0.64
12/
12/Ao
decry
No docay
Decay
2
77
(a)
0.90
0.06
0.90
0.90
Test accurag) ())
0.855
Test accuracy
0.84
0.89
2/Nr
12/Ap
12/Aa
0.91
No
0.035
0.
0.90
10
10-1
10"
(C)
(d)
Figure 3: Models perform best with a large learning rate. Test accuracy vs Icarning rate for (a.b) a CNN trained CIFAR-10 using SGD with batch size 256 and La regularization MMeiit 10- and (c.d) WRN28- 10 trained CIFAR-10 using SGD with batch size 1024, L2 regularization, and data augmentation (JJeiit 0.14); see Appendix A for details. (a,c) have fixed compute budget: (a) 437k steps and (b) 12k steps. (b,d) have been evalved for a fixed amount of physical time: (b) was evolved for 475/1 steps (purple) and evolved for 50k more steps at learning rate 2 10 (red) and (d) was evolved for 3360/n steps with learning rate ry (purple) and then evolved for 4800 more steps learning rate 0.035 (red). In all cases, optimal performance is achieved above Yeritt and close to the expected maximum learning rate, in agreement with our predictions.
In Appendix D.2, we present additional results for WRN CIFAR- with similar conclusions. The fact that optimal performance happens the catapult phase can also be observed for simple models like a fully-comnected ReLU network trained on a subset of MNIST (see the Appendix).
3 GRADIENT DESCENT DYNAMICS OF WIDE, 2-LAYER LINEAR NETWORKS
We now turn to theoretical analysis of the gradient descent dynamics of two-layer linear network at large but finite width. While such a setting omits complexities such as depth and nonlinearity, theoretical treatment already reveals the existence of three phases described in Section 1.2 with signatures that match our experiments.
Let the network function be given by =(=)= in the hidden layer), u. IR" are the model parameters (collectively denoted f), and IR is the training input. At initialization, the weights are drawn from (0, 1). We prove the following
= -1/2, ua. Here n is the width (number of neurons
Theorem 1. Consider 2-layer linear network of width n, trained with MSE loss and learning rate n) The training data has a single sample with (x, y) = (1, 0). Choose initial values fo + 0 and X > for the function fxx and curvature X. Let Tjerit := 2/ No and Mmmax A/>g, and choose 8 > O.
1. Lazy phase: If n < Merit then gradient descent achieves loss < in (nn) steps, and the final curvature >3 obeys [>g Aul O(n
2. Catapult phase: Merit <. 1) TJmax then gradient descent achieves loss L <, 5, the final curvature obeys Xg < 2/m, and during optimization the loss grows to be (2n ((n))o--nn)). (n)).
3. Divergent phase: If 2 Thaax then gradient descent does not converge to a global minimum.
The proof can be found in the appendix. We will complement the theorem with intuition about the dynamics of this model. In the appendix, we also generalize to the case of networks with arbitrary

---

Under review as a conference paper at ICLR 2021
input dimension. The gradient descent equations at training step are
u(+1 u, Tpn 1/2 feve, vc+1 2/_ - mp 1/2 feu..
(1)
The update equations in function space can be written in terms of the Neural Tangent Kernel. For this model, the kernel evaluated on the training set is a scalar which is equal to A, its top eigenvalue, and is given by ( 1) = À n=1 (Ilu11 + ||u| () At initialization, both f and À scale as n2 with the width n. The following update equations for and at step t can be derived from equation 1.
fi+1 = (1 -1- nd
+ 7 5²
Dft (n) - -4)
fe, A--1 = X
It is important to note that these are the exact update equations for this model, and that no higher-order terms were neglected. We now analyze these dynamical equations assuming the width is large.
Two learning rates that will be important in the analysis are Tjerit 2/A0 and Mmaa = 4/20-
Lazy phase.
Taking the strict infinite width limit, equations equation 2 become
fi+1 (1 - m7a) fe, A = d
(3)
When 1) < Merit. remains constant throughout training. This is a special case of NTK dynamics, where the kernel is constant and the network evolves linear model (Lee al., 2019). The function and the loss both shrink to zero because the multiplicative factor obeys |1 mptel < This convergence happens in (nn)) ((1) steps.
Catapult phase. When Merit <. Tmmax the loss diverges in the infinite width limit. Indeed, from equation we see that the kernel is constant the limit, while f receives multiplicative updates where 11 - n,,,1 > 1. This is the well known instability of gradient descent dynamics for linear models with MSE loss. However, the underlying model not linear its parameters, and finite width contributions turn out to be important. We therefore relax the infinite width limit and analyze equations (2) for large but finite width, n >> 1.
First, note that 77 4 < 0 by assumption, and therefore the (additive) kernel updates are negative for all f. During early training, [fll grows (as in the infinite width limit) while A remains constant up to small Onn updates. After ~ log(n) steps, lfel grows to order n²/22_ At this point, the kernel updates no longer negligible because f2/ /n is of order no. The kernel X receives negative, non- neggnegli updates while both f. and the loss continue to grow. This continues until the kernel effect the catapuli effect. Beyond this point, | - A < holds, ||x shrinks, and the loss converges to a global minimum. The dependence of the steps until optimization converges is log (n))
sufficiently small that the condition d
is met.² We call this survature-reductio
It is important for the analysis that we take a modified large width limit, in which the number of training steps grows like log(n) as becomes large. This is different than the large width limit commonly studied in the literature, in which the number of steps is kept fixed as the width is taken large. When using this modified limit, the analysis above holds even the limit. Note as well that the catapult effect takes place over (og(n) steps, and for practical networks will occur within the first 100 steps or so of training.
In the catapult phase, the kernel at the end of training is smaller by an order n' amount compared with its value at initialization. The kernel provides local measure of the loss curvature. Therefore, the minima that SGD finds in the catapult phase are flatter than those it finds in the lazy phase. Contrast this situation, which the kernel receives non- -egligiie updates, with the conclusions of Jacot et al. (2018) where the kernel is constant throughout training. The difference is due the large learning rate, which leads to a breakdown of the linearized approximation even at large width.
Completing the analysis of this model, when 1) Tmma the loss diverges because the kernel receives positive updates, accelerating the rate of growth of the function. Therefore, Mmax = Ao is the maximum learning rate of the model.
Thhe bound is exact because of the we neglected.
7

---

Under review as a conference paper at ICLR 2021
3.1 NON-PERTURBATIV PHASE TRANSITION
The large width analysis of the small learning rate phase has been the subject of much work. In this phase, at infinite width, the network map evolves linear random features model, sz+1
f(())
f - -er where f(0) is the function of the linearized model. At large but finite width (which we denote by 12), corrections to this linear evolution can be systematically incorporated via a perturbative expansion (Taylor expansion) around infinite width Dyer & Gur-Ari (2020); Huang & Yau (2019),
ft f + + +.... The evolution equations of the solvable model are an example of this. At large width and in the small learning rate phase, the O(n terms are suppressed for all times.
In contrast, the leading order dynamics of f! diverge when 7) > Merit. and so the true evolution cannot be described by the linear model. Indeed, the logits grow to O(n!/3) and thus all terms in equation 3 are of the same order. Similarly, the growth observed empirically in the catapult phase for more general models cannot be described by truncating the perturbative series at any order, because the terms all become comparable.
DISCUSSION
Previous work (Jacot al., 2018; Lee et al., 2019; Chizat al., 2019) has studied the lazy phase of deep neural networks, which is known to occur for sufficiently small learning rates. In this work argued that the lazy phase exists for learning rates smaller than Merit 2/20. where Ao is the curvature at initialization. This critical learning rate corresponds to where a linear model, constructed from deep network about its initialized parameters, would diverge under MSE loss. We pointed out the existence of the catapult phase in deep networks, corresponding to the learning rate regime 'eerit < 17 < Mmaxx Its unique empirical signatures include the early-time growth of the loss and convergence to a flat minimum. Empirically, the existence and properties of the catapult phase can be observed across variety network architectures and datasets. At yet larger learning rates r) >> 7/max. SGD dynamics are unstable.
A novel analysis illustrating the catapult phase. Through our analytical treatment of a two-layer linear network, we are able to clarify the dynamics behind the catapult effect. Among these, we (i) derived the quantitative changes in loss and curvature and the time scales over which they occur; (ii) derived an expression for Tmmax in terms of the curvature initialization; (iii) specified the manner in which the lazy and catapult regimes are distinct phases, which we elaborate on below, in a novel modified -iitt, infinite-time limit; and (iv) illustrated the dynamical mechanism stabilizing the catapult phase. Our approach reduces to analyzing two coupled difference equations relating the loss and curvature, which we hope may inspire a full treatment of deep networks with nonlinearities.
The change in behavior upon sweeping the learning rate from the lazy to catapult phase is reminiscent of phase transitions that commonly appear physical systems such as ferromagnets or water, as one changes parameters such as temperature. Indeed, in Appendix C this connection is made concrete, with the change in behavior sharpening as width is increased. In particular, these transitions are non-perturbative: Taylor series expansion of the lincarized model that takes into account finite width corrections is not sufficient to describe the behavior beyond the critical learning rate.
Catapult dynamics often improve generalization. Our results shed light on the regularizing effect of training at large learning rates. The effect presented here is independent of the regularizing effect of stochastic gradient noise, which has been studied extensively. Building on previous works, we noted the observed correlation between flatness and generalization performance. Based on these observations, we expect the optimal performance to often occur for learning rates larger than 'eerit, where the linearized model is unstable. Observing this effect required controlling for several confounding factors that affect the comparison of performance between different learning rates. Under a fair comparison, and also for a fixed compute budget, we find that this expectation holds in practice.
One outcome of our work is to address the performance gap between ordinary neural networks, and linear models inspired by the theory of infinite-wine widht networks. Optimal performance is often obtained at large learning rates which are inaccessible to linearized models. In such cases, expect the performance gap to persist even at arbitrarily large widths. We hope our work can further improve the understanding of deep learning dynamics and performance.
8

---

Under review as a conference paper at ICLR 2021
REFERENCES
Zeyuan Allen-Zhu, Yuanzhi Li, and Zhao Song. A convergence theory for deep learning via over- parameterization. In Kamalika Chaudhuri and Ruslan Salakhutdinov (eds.), Proceedings of the 36th International Conference Machine Learning, volume 97 of Proceedings of Machine Learning Research, pp. 242 252, Long Beach, California, USA, 09-15 Jun 2019. PMLR.
Shun-Ichi Amari, Hyeyoung Park, and Kenji Fukumizu. Adaptive method of realizing natural gradient learning for multilayer pereeptrons. Neural computation, 12(6):1399-144) 2000.
Sanjeev Arora, Simon $ Du, Wei Hu, Zhiyuan Li, Russ Salakhutdinov, and Ruosong Wang. On exact computation with an infinitely wide neural net. In Advances in Neural Information Processing Systems, pp. 8139-8148, 2019.
James Bradbury, Roy Frostig, Peter Hawkins, Matthew James Johnson, Chris Leary, Dougal Maclau- rin, and Skye Wanderman-Milne JAX: composable transformations of Python+NumPy programs, 2018. URL http: g/thbb.. com/ goog//ooogeeeeo e/ jax.
Lénaic Chizat, Edouard Oyallon, and Francis Bach. On lazy training in differentiable programming. In H. Wallach, H. Larochelle, A. Beygelzimer, F. d'lchééBucc E. Fox, and R. Garnett (eds.), Advances in Neural Information Processing Systems 32, pp. 2933-222222222 2943. Curran Associates, Inc., 2019. URL http::/papprs..atttt// nips wwtppppppppppettpptc.ccccct..ett 8559 -tap:///---------------aaaaa--a- lat---aaiaiinnnii.......... .ddf.
di
Amit Daniely. Sgd learns the conjugate kemnel class of the network. In Advances in Neural Information Processing Systems, pp. 2422-2430, 2017.
Laurent Dinh, Razvan Pascanu, Samy Bengio, and Yoshua Bengio. Sharp minima can generalize for deep nets. In Proceedings of the 34th International Conference on Machine Learning-Voles Volume 70, pp. 1019-1028. JMLR. org, 2017.
Simon S. Du, Jason D. Lee, Haochuan Li, Liwei Wang, and Xiyu Zhai. Gradient descent finds global minima deep neural networks. In Proceedings of the 36th International Conference on Machine Learning, ICML 2019, 9-15 June 2019, Long Beach, California, USA, pp. 1675-1685, 2019. URL http:// http:::/ppccceedinge. www. press/pees/eee v7 22997/du duul 9c. html.
Ethan Dyer and Guy Gur-Ari. Asymptotics of wide networks from feymman diagrams. In International Conference on Learning Representations, 2020. URL ptttts::/oppnzeeet ntt.:w/.eeeeeee.eeeeeee.eeeeeeee foormm? id-S1gFvANKDS,
Jonathan Frankle, David Schwab, and Ari S Morcos. The early phase of neural network training. arXiv preprint 200. 2020.
Sepp Hochreiter and Jürgen Schmidhuber. Flat minima. Neural Computation, 9(1):1-42, 1997.
Jiaoyang Huang and Horng- Tzer Yau. Dynamics of Deep Neural Networks and Neural Tangent Hierarchy. arXiv e-prints, art. arXiv:1909 156, Sep 2019.
Arthur Jacot, Franck Gabriel, and Clement Hongler. Neural tangent kernel: Convergence and generalization in neural networks. In S. Bengio, H. Wallach, H. Larochelle, K. Grauman, N. Cesa- Bianchi, and R. Garnett (eds.), Advances in Neural Information Processing Systems 31, pp. 8571-8580 Curran Associates, Inc., 2018.
Stanislaw Jastrzebski, Maciej Szymczak, Stanislav Fort, Devansh Arpit, Jacek Tabor, Kyunghyun Cho, and Krzysztof Geras. The break-even point on optimization trajectories of deep neural networks. arXiv preprint arXiv:2002.009522.00 2020.
Yiding Jiang, Behnam Neyshabur, Dilip Krishnan, Hossein Mobahi, and Samy Bengio. Fantas- tic generalization measures and where to find them. In International Conference on Learning Representations, 2020. URL https: http//pppeevv/eeeeeeeeee
n mntt/forum htt.://rummmmo// 2o-
Ryo Karakida, Shotaro Akaho, and Shun-ichi Amari. Universal Statistics of Fisher Information in Deep Neural Networks: Mean Field Approach. arXiv e-prints, art. arXiv: 10Xiv:186.0011 806.01316, June 2018.

---

Under review as a conference paper at ICLR 2021
Nitish Shirish Keskar, Dheevatsa Mudigere, Jorge Nocedal, Mikhail Smelyanskiy, and Ping Tak Peter Tang. On large-batch training for deep learning: Generalization gap and sharp minima. CoRR, abs/ Sess/660..48820 04036, 2016. URL http: httpaaxxvvvvv......a. org/ abs 1609. 04836.
Guillaume Leclere and Aleksander Madry. The two regimes of deep network training, 2020.
Jachoon Lee, Yasaman Bahri, Roman Novak, Sam Jeffefrrez,,,ffree) Pennington, and Jascha Sohl- dickstein. Deep neural networks as gaussian processes. In International Conference on Learning Representations, 2018. URL https: owwnreview net B00
Jaehoon Lee, Lechao Xiao, Samuel Schoenholz, Yasaman Bahri, Roman Novak, Jascha Sohl-Dickstein, and Jeffrey Pennington. Wide neural networks of any depth evolve as linear models under gradient descent. In H. Wallach, H. Larochelle, A. Beygelz- imer, F. d' Alché-Buc, E. Fox, and R. Garnett (eds.), Advances in Neural Informa- tion Processing Systems 32, pp. 8570-8581. Curran Associates, Inc., 2019. URL http:/ /ttt::/ppapee pwwprsappeeseeeeee nips..p..... cc..c///ppae pec/paper/99663 1e000-0wwd--euuraae -nteerrss-- any -eo-ddeph--volllee as- linear mode ls Aunderr----r gradient descent.... pdf.
Yuanzhi Li and Yingyu Liang. Learning overparameterizeed neural networks via stochastic gradient descent on structured data. Advances in Neural Information Processing Systems, pp. 8157-8166, 2018.
Yuanzhi Li, Colin Wei, and Tengyu Ma. Towards explaining the regularization effect of initial large learning rate training neural networks. In H. Wallach, H. Larochelle, A. Beygelzimer, F. d' Alché Buc, E. Fox, and R. Garnett (eds.), Advances in Neural Information Processing Systems 32, pp. 11669-11680. Curran Associates, Inc., 2019.
Stephan Mandt, Matthew D Hoffman, and David M Blei. Stochastic gradient descent as approximate bayesian inference. The Journal of Machine Learning Research, 18(1):4873-4907, 2017.
Robert May. Simple mathematical models with very complicated dynamics. Nature, 261 (5560): 459-467, 1976.
Song Mei, Andrea Montanari, and Phan-Minh Nguyen. A mean field view of the landscape of two-layer neural networks. 11563):E7666- 2018. doi: 10. 10.1073/pnas. 10003//paa..100579111... 1806579115.
Naveh, Ben-David, Sompolinsky, and Ringel. to published.
Roman Novak, Lechao Xiao, Yasaman Bahri, Jachoon Lee, Greg Yang, Daniel A. Abolafia, Jeffrey Pennington, and Jascha Sohl-dickstein. Bayesian deep convolutional networks with many channels gaussian processes. In International Conference on Learning Representations, 2019. URL https: /ttp://////eeeeeeee net forum?dd=B1130 j0qE7.
Roman Novak, Lechao Xiao, Jiri Hron, Jaehoon Lee, Alexander A. Alemi, Jascha Sohl-Dickstein, and Samuel S. Schoenholz. Neural tangents: Fast and easy infinite neural networks in python. In International Conference Learning Representations, 2020. URL https: /tt:://giit /github...///ggg//g com/ google honggl/e/eeaalllageennnn
Daniel S. Park, Jascha Sohl-Dickstein, Quoc V. Le, and Samuel L. Smith. The effect of network width on stochastic gradient descent and generalization: an empirical study. CoRR, abs/ 1907003776..
2019. URL http: waaaaaivvv.......... .ow.iv.orrgggeoe. abs 905. 03776.
Grant Rotskoff and Eric Handen-Eijnden, Parameters as interacting particles: long time convergence and asymptotic scaling of neural networks. In Advances neural information processing systems, pp. 7146-7155, 2018.
Justin Sirignano and Konstantinos Spiliopoulos. Mean field analysis of neural networks. arXiv preprint / 2018.
Samuel L. Smith and Quoc V. Le. A bayesian perspective generalization and stochastic gradient descent. In International Conference on Learning Representations, 2018. URL https:/ openreview http:/iwwwweeeeeeeeee feremt//oorreeeeeee :o
10

---

Under review as a conference paper at ICLR 2021
Samuel L. Smith, Pieter-Jan Kindermans, Chris Ying, and Quoc V. Le. Don't 't Decay the Learning Rate, Increase the Batch Size. arXiv e-prints, art. utXiv:1711.00488 Nov 2017.
Samuel L Smith, Daniel Duckworth, Semon Rezchikov, Quoc V Le, and Jascha Sohl-Dickstein. Stochastic natural gradient descent draws posterior samples in function space. arXiv preprint arXiv:1806.09597, 2018.
Blake Woodworth, Suriya Gunasekar, Jason Lee, Daniel Soudry, and Nathan Srebro. Kernel and deep regimes in overparametrized models. arXiv preprint /900/99976077 2019.
Lechao Xiao, Jeffrey Pennington, and Samuel S. Schoenholz. Disentangling trainability and general- ization in deep learning, 2019.
Zeke Xie, Issei Sato, and Masashi Sugiyama. A diffusion theory for deep learning dynamics: Stochastic gradient descent escapes from sharp minima exponentially fast. arXiv preprint arXiv:2002.03495, 2020.
Sergey Zagoruyko and Nikos Komodakis. Wide residual networks. CoRR, wwwp///0...oe 2016. URL httpt:////axxi.. org/ abs/1 0.146.
Difan Zou, Yuan Cao, Dongruo Zhou, and Quanquan Gu. Stochastic gradient descent optimizes over-parameterized deep relu networks. arXiv preprint arXiv:1811.08820 2018.

---

Under review as a conference paper at ICLR 2021
A EXPERIMENTAL DETAILS
We are using JAX (Bradbury al., 2018) and the Neural Tangents Library for our experiments (Novak et al., 2020).
All the models have been trained with Mean Squared Error normalized as C( ())) where l is the number of classes and yf are one-targets.
26 (f' y²)²,
+ E. a (x)(()) so that In a similar way, we have normalized the NTK as (z, = E|в the eigenvalues of the NTK are the same as the non-zero eigenvalues of the Fisher information:
B[ L
In our experiments we measure the top eigenvalue of the NTK using Lanezos' algorithm. We construct the NTK on a small batch of data, typically several hundred samples, compute the top cigenvalue, and then average over batches. this work, we do not focus on precision aspects such as fluctuations in the top eigenvalue across batches.
All experiments that compare different learning rates use the same seed for the weights at initialization and we consider only one such initialization (unless otherwise stated) although we have not seen much variance in the phenomena described. We let ow, Ty denote the constant wwitt--nddependen coefficient of the standard deviation of the weight and bias initializations, respectively.
Here we describe experimental settings specific to figure.
Figure 2a,2c,2e. Fully connected. three hidden layers 2048, ReLU non-linearity trained using SGD (no momentum) on MNIST. Batch size= 512, using NTK normalization, /2, o = 0.
Figures 2b,2d,2f. Wide ResNet 28-18 trained on CIFARIO with SGD (no momentum). Batch size of 128, LeCun initialization with ow v2, o. 0, L = 0.
Figures 54,517 Fully connected network with one hidden layer and ReLU non-linearity trained on 512 samples of MNIST with SGD (no momentum). Batch size of 512, NTK initialization with Tw /2, 0b. 0.
Figures 3a,3b. The convolutional network has the following architecture: Conv, (320) - ReLU - Convg(320) - ReLU + MaxPool((22).. 'VALID') + Conv (320) - ReLU + 2onv (2.2) 'VALID') - Flatten() - Dense(256) + ReLU + Dense(10), Dense denotes a fully-connected layer with output dimension n. Conv (n), (n)v denote convolutional layers with 'SAME' or VALID' padding and n filters, respectively; all convolutional layers use (3, (3 filters. (2.2), 'VALID') performs max pooling with 'VALID' padding and a (2,2) window size. LeCun initialization used, with the standard deviation of the weights and biases drawn as o., = v2. 0 = 0.05, respectively. Trained on CIFAR-10 with SGD, batch size of 256 and 1.2 regularization 0.001.
Figures 1, 3c,3d. Wide ResNet on CIFARI0 using SGD (no momentum). Training on v3-8 TPUs with a total batch size of 1024 (and per device batch size 128). They all use L2 regularization- 0.0005, LeCun initialization with a~ crop and mixup. With softmax classification, these models can get test accuracy 0.965 if one uses cosine decay, using JAX's implementation of Batch Norm which doesn't keep track of training batch statistics for test mode evaluation. We have not hyperparameter tuned for learning rates nor T.2 regularization parameter.
0. There is also data augmentation:
use flip,
don't 't observe big performance decay due to using MSE. Furthermore,
Figures $5,86. Wide ResNet on CIFARI00 using SGD (no momentum). Same setting figure 3c, 3d except for the different dataset, different L2 regularization 0.000025 and label smoothing (we have subtracted 0.01 from the target one-hot labels).
Figure $12. Two hidden layer, ReLUU network for one data point x = 1,y 1.
Figure $15. Fully connected network with two hidden layers and tanh non-linearity trained on MNIST with SGD (no momentum). Batch size of 512, LeCun initialization with ow, = 1, σ, 0.
Figure S13a. Two-hidden layer fully connected network trained MNIST with batch size 512, NTK normalization with o.
V2, m 0. Trained using both momenta 7 0.9 and vanilla SGD

---

Under review as a conference paper at ICLR 2021
(a)
(b)
(c)
(d)
Figure S1: Visualization of training dynamics in all three phases. In the lazy phase, the network is approximately linear in its parameters, and converges exponentially to a global minimum. In the catapult phase, the loss initially grows, while the weight norm and curvature decrease. Once curvature is low enough, optimization converges. In the divergent phase, both the loss and parameter magnitudes diverge. (a)-(d) Loss surface and training dynamics visualized The network has single hidden layer with width
2d linear subspace. 500, lincar activations, and is trained with MSE loss on a single ID sample x = with label y = 0. The parameter subspace defined by [dim1] r [dim2] s, U [dim1] [dim2] 8, where and
orthonormal vectors, u, TR are the weight vectors, and [diml], [dim2] are the coordinates in the subspace. initialized in this 2d subspace, W, and e, remain in the subspace throughout training, and so training dynamics can be fully visualized with two dimensional plot. (e) Visualization of the loss surface and training dynamics terms of nonlinear reparameterization, providing interpretable properties: x-axis correlation between weight vectors, y-axis curvature X. The trajectory shown is identical to that in (c), and in Figure 1.
for three different non-linearities: tanh, ReLU and identity (no non-linearity). The learning rate for each non-linearity was chosen to correspond to = x
Rest of appendix figures. Small modifications of experiments in previous figures, specified in captions.
2

---

Under review as a conference paper at ICLR 2021
B THEORETICAL DETAILS
B.1 THEBORE
For ID input & with label y. network [(x) -2 uxr learning rate 7), and loss L (f(x) /2. We s = and y = O. The curvature is À = n (1]u2 |
Theorem 1. Consider a 2-layer linear network of width trained with MSE loss and learning rate q. The training data has a single sample with input x = and label y = Choose the initial values fo * 0 and X > 0 for the value of the function fz) and curvature X. Let Meris := 2/ do and Imax == 4/ 4/20, and choose
1. Lazy phase: n < Merit then gradient descent achieves loss < O(n)) steps, and the final curvature X5 obeys [Ag- Ao O(n
2. Catapult phase: If Merrt. < 1) < then gradient descent achieves loss L < 6, the final curvature obeys Ay S 2/m, and during optimization the loss grows to be (2 log loollog'( (n)).
3. Divergent phase: If Tmax then gradient descent does not converge to a global minimum.
Proof. Under gradient descent, the function and curvature update equations given by
-nide n2St Afc
fr, AA: nf. (nde 4).
(SI)
n
Here we denote f f(11, Aft = fet1 - St, and AA Aet1 - Xt-
Convergence when 17 < 7mmax mf Mmax then Ad, S 0 and S No for all t. Therefore,
7 (ndo 4) Σ fe
A++1 do
Ad λo
(S2)
n
Since Ar+1 2 0 we then have
T
nd
f2 <
(S3)
114 17A0)
From the monotone convergence theorem, the sum on the left-hand side converges when taking T - oc, and therefore lim;- Jt 0, proving convergence when 17 < Jmmax-
Lazy phase. Next, assume that 1) < Merit- We will first show that, when taking the infinite n limit, optimization converges in a finite number of steps and the curvature is constant throughout optimization. We then show that the finite n solution is close to the infinite n one up to O(n -1) corrections, which is enough establish statement (I).
Taking n + oo, the update equations equation become
Aft -n
A 0.
(S4)
Here we denote by f A the function and curvature at step t the infinite width limit. For given (n independent) initial conditions fu, Au, in this limit the loss drops below 8 in O(n") steps, while the curvature remains constant.
Next, compare the optimization trajectory (fey Ax) of the equations at finite n, with the trajectory (J X²2) of the infinite n equations. We set the initial conditions to be the same, fo = f and Ag Ag Next, let us show by induction on it that the following holds at every step.
fel On)), > =n0),,) Ise 1: O(n 1 2;
O(n :0n=1((n1))
(S5)
3

---

Under review as a conference paper at ICLR 2021
At f 0 these relations hold by assumption. We now show the induction step. For the function trajectory,
1fi+1 f241
ft (7ra)) l
71
+ m// 52 - -v.sel n²² _11س
$ 1/2 12-
n
+ 1 18 - A + n1 - Ar1 115
S 1/4 18
O(n
(S6)
For the curvature trajectory.
d + T1?
Aet1
10t1l
(1)2e - -4) 270
n1
$ 1 1
+ "
(I 4)1
n
= (n
(S7)
The remaining relations in equation S5 follow from the fact that 1 and X are O(nº). Since the infinite n trajectory leads to L < 5 in O(n)) steps, this proves statement (1).
Catapult phase. We now show that any global minimum (f, A) with f = and À > 2/n is repulsive fixed point of the update equations. Assume that X 2 + for some d >> 0, and that 0 ffl < for some For sufficiently small 6,
fe+1 ft
2 f2 = --1 n
225
e72c2 > 1.
y +
1+ on
(S8)
We see that fi grows in absolute value, proving that optimization does not converge. We showed that gradient descent converges in this case, and therefore the final curvature X5 < 2/n
left to show that the loss grows to be S(n log (n)) during optimization. We choose e > 0.
If it enough show that if sufficiently small, then [ft] grows larger than log(n)
step
( - (n))
From now on let us make the following assumptions. For some > 0 (to be chosen),
log(n) O<L<< t log(n).
(S9)
Notice that these hold initialization for sufficiently large n.
Using the fact that X I An for all t, the curvature update is
A nff (np) - 4) n log(n) (ndr 4) log(n)
17
ncc
(n7o0 - 4).
(S10)
n
With the assumptions above, we have
Dr Aol
Σ Adp < <nrel/Ao - 4| #=1
(S11)
We find that with the above assumptions, we have
77, = 1 77i0 + 7/>, Apl = 17do + O((e).
(S12)
Here we used again that X < An In particular, since - noo = -1 46 for some 5 > 0, then
1 T/ -1 28 for sufficiently small
Next, consider the ft update.
Je+1 1 ni
rf -1 28 +
n/c
(S13)
ft
n
log(n)
4

---

Under review as a conference paper at ICLR 2021
Again setting sufficiently small, we get equation S9 hold,
f1 f
> + 6. We find that as long as the assumptions
M > (1 + Ioo
(S14)
Let us now choose >> log (1 + 8). Then for 1 =log(n)/2 and sufficiently small 6,
Ifll > n log ( [fo
vol
71
(S15)
log(n)
This shows that the first assumption in equation S9 is the first to be violated when increasing t from zero. Therefore, we find that L3 ne/ log(n) ) Qn log(n)) some point during training.
Divergent phase. Finally, suppose that 2 7mmax Then AA, 2 0 and X 2 Ao 2/17 for all I. We showed that gradient descent can only converge to minima with 2/17. and therefore gradient descent will not converge in this case. This concludes the proof of statement (3).
B.2 EMPIRICAL PLOT OF SIMPLE MODEL
Figure S2 illustrates the dynamics in the catapult phase. For learning rates eexit Mmax observe the catapult effect: the loss goes up before converging to zero. The curvature exhibits the expected sharp transitions function of the learning rate: it is constant in the lazy phase, decreases in the catapult phase, and diverges for v > 1/max..
2=3.04
2/
0 315
2-0.015
4/
<--1 31
9-111
-1
25
0
10
1 20
10
(b)
(c)
Figure S2: Empirical results for the gradient descent dynamics of the warmup model with n 10 for which Merit & 1. (a) Training loss for different learning rates. (b) Maximum NTK eigenvalue as a function of time. For 7) 1, X decreases rapidly to fixed value. (c) Maximum NTK eigenvalue at t= = 25/n. The shaded area indicates learning rates for which training diverges empirically. The results presented as function of t (rather than t) for convenience.
B.3 FULL MODEL ANALYSIS
Here provide additional details the theoretical analysis of the full model We introduce the notation fa = fxxa) for the function evaluated training sample, fa fa Ya for the error, and 0, O (Na,, 2r,)) for the kernel elements. We will treat f, evaluated the training set as vectors m whose elements are fa. fa. .Te gradient descent update equations are
++1 Uio
11 v.Jafaa of
Usa Fa
(S16)
Vnm
v num
and
Ogp nm
(10) 2 Dg + x u uxz))
(S17)
5

---

Under review as a conference paper at ICLR 2021
The update equations for the eπror and kernel evaluated on training set inputs are
f++1 = (dog - 1/aass) oF (IToTT))
(S18)
nmm
= OnB 1
m (1 + (f + 2 x
03
n
+ n²m
(1 (f (x))(())=)))
(S19)
Where ç= = fa e Rd. We now consider the dynamics of the kernel projected onto the direction, which is given by
(S20)
Let analyze the phase structure of equation S18 and equation S20. For now, we neglect the last term on the right-hand side of equation S18 (at initialization this term is of order n and negligible large width). Let λo be the maximal eigenvalue of the kernel initialization, and let max E Rm be the corresponding eigenvector. Notice that f projected onto the top eigenvector evolves as
maxx set1 = (1 max) O(
(S21)
Lazy phase. When daip < 2, we see that me f shrinks during training. The kernel updates are of order n while convergence happens in order n steps. Therefore the kernel does not change by much during training. This a special case of the NTK result (Jacot al., 2018). Effectively, the model evolves as a linear model in this phase.
Catapult phase. When 2 < n7 < 4, lIflla grows exponentially fast, and it grows fastest in the emax direction. Therefore, the vector becomes aligned with emaax after number of steps that of order no" Also, f itself grows quickly while the label is constant, and so we find that f maxT max after a similar number of steps. When these approximations hold, notice that fT N x- 1
Regarding the evolution of the top eigenvector itself, it is easy to from equation S19 that the kernel updates preserve the two separate subspaces of cigenvectors with X; < 2/1 and X; > 2/n, there is only one vector with nd; > 2, then the mmax direction stays constant during training. If there are multiple such vectors, we expect that ones with larger eigenvalues grow exponentially faster, and thus we expect that the eigenvectors are approximately preserved.
Now, from equation equation S20 we can then derive an approximate equation for the evolution of the top NTK eigenvalue.
A+11 RJ x CTTc 4)
(S22)
While if grows exponentially fast, will Ç. When Se becomes of order 1/2, 2, the updates to the top eigenvalue become of order n (and negative), causing X; to decrease by a non-negligible amount. This will continue until X₁ < 2/1, at which point f. will start converging to zero. Eventually, after a number of steps of order log(n), gradient descent will converge global minimum that has lower curvature than the curvature at initialization.
The justification for dropping the order -1 term in equation $21 explained in the warmup model: While this term may affect the details of the dynamics, eventually the maximum kernel eigenvalue must drop below 2/1j for the component m f of the error (and therefore for the loss) to converge zero.
Divergent phase. When Ti > 4, both 1l and À will grow, and optimization will diverge. Therefore, 4/20 is the maximum learning rate for this model.

---

Under review as a conference paper at ICLR 2021
C MODEL DYNAMICS CLOSE TO THE CRITICAL LEARNING RATE
Here we consider the gradient descent dynamics of the model analyzed in Section 3, for learning rates 7) that are close the critical point feerit = 2/Ao. The analysis reveals that the gradient descent dynamics of the model are qualitatively different above and below this point. For example, the loss decreases monotonically during training when 17 < 7eritt but not when rj > Meeiit- In this section we show that the transition from small to large learning rate becomes sharp once we take the modified large width limit, in the following sense: certain functions of the learning rate become non-analytic at 'erit in the limit. This sharp transition bears close resemblance to phase transitions of the kind found in physical systems, such as the transition between the liquid and gaseous phases of water. In particular, case involves dynamical system, where the dynamics descent equations. These dynamics undergo a phase transition as a function of the learning rate- an external parameter. We point to the logistic map May (1976) well-known example of dynamical system that undergoes phase transitions as a function of an external parameter.
governed by the gradient
C.I NON-PERTURBATIV DYNAMICS
phase transition is drastic change in a system's behavior incurred under a small change in external parameters. Mathematicallyy it is a non-analyticity in some property of the system as function of these parameters. For example, consider the property X. (1). the curvature of the model at the end of training as function of the learning rate. In the modified large width limit, X. (17) is constant for 17 < Merit. but not for 77 > Meiit. Therefore, this function is not analytic at Merit. Notice that this statement is true in the limit but not necessarily at finite width, where the final curvature may be an analytic function of the learning rate even at Merit- It is well known in physics that phase transitions only occur in limit where the number of dynamical variables (in this case the number of model parameters) is taken to infinity. One immediate consequence of the non-analyticity at Merit is that the large learning rate phase is inaccessible from the small learning rate phase via a perturbative expansion. In other words, we cannot describe all properties of the model for some 1) > MJorit by doing Taylor expansion around a point T/o < ferit and keeping a finite number of terms.
Dyer& & Gur-Ari (2020); Huang & Yau (2019) developed a formalism that allows one to compute finite-widthh corrections to various properties of deep networks, using a perturbative expansion around the infinite width limit. We have argued that the usual infinite width approximation to the training dynamies is not valid for learning rates above Teerit. and that full analysis must account for large finite-width effects. One may have hoped that including the perturbative finite-width corrections discussed in Dyer & Gur-Ari (2020); Huang & Yau (2019) would allow us to regain analytic control over the dynamics. The results presented here suggest that this is not the case: For 1) > Meerit. we expect that the perturbative expansion will not provide a good approximation to the gradient descent dynamics at any finite order in inverse width.
C.2 CRITICAL EXPONENTS
When the external parameters are close to phase transition, one often finds that the dynamical properties of the system obey power law behavior. The exponents of these power laws (called critical exponents) are of interest because they are often found to be universal, in the sense that the same set of exponents often found to describe the phase transitions of completely different physical systems.
Here we consider i. (z)). the number of steps until convergence, as function of the learning rate. We will now show that t, exhibits pawer-law behavior when 7/ is close to Teiit. For simplicity we consider the warmup model studied in Section 3. First, suppose that we are below the transition, setting iyyu = -- e for some small > O. From the update equation, f(+1 R (1 (1- r/):). & -(1-e)Je we see that ft will converge to some fixed small value f. after time f, e log(f. ~ e-! Here assumed that X₂ is constant in t, which is true as long as t, is independent of n (namely we fix and then take large). Therefore, the convergence time below the transition scales as f. (Merit - -1, and the critical exponent is -1.
Next, suppose that πiλu 2 c with > 0. Now the update equation reads f(+1 -(1 + f),i. This approximation holds early during training, when the curvature updates are small. Initially, [fl will grow until it is of order Vni at which point the updates X become of order n². This will happen in time L i~v - - log Vn. Following this, the optimizer will converge. At this point d is no longer
7

---

Under review as a conference paper at ICLR 2021
tuned be close to the transition, and the convergence time measured from this point will not be sensitive to e. Therefore, for small the convergence time will be dominated by the early part of training, namely t, s f ~~c c verification of this behavior.
The critical exponent is again -1. Figure S3 show empirical
experiment (Un) below (0.95) ft. above 1.0.3))
300
250
stape til convargence
200
150
100
0.925
0.350
0.975
1.000
1.025
1.050
1.075
leaming rate
(a)
Figure S3: The convergence time diverges when the learning rate is close to the critical value Meiit indicated by the solid green line. The measured exponents (shown in parentheses) are close to the predicted value of -1. Experiment involves the warmup model of Section 3 with width 16,000.
D EXPERIMENTAL RESULTS: LATE TIME PERFORMANCE
D.1 SIMPLE MNIST EXPERIMENT
Figure S4 shows the accuracy as function of the learning rate for a fully-connected ReLU network trained subset of MNIST. We find that the optimal performance is achieved above Teerit and close to Mmax 12, the expected maximum learning rate.
D.2 CIFAR-100 PERFORMANCE
We also repeat the performance experiments for CIFAR-100 and the Wide ResNet 28- 10 setup. In this case, using MSE and SGD we require to evolve the system for longer times, which requires smaller Le regularization. We didn't tune for it. but found that 2.5 10 works. With only one decay we can get within 3% of the Zagoruyko & Komodakis (2016) performance that used softmax classification and two learning rate decays. However, evolution for longer time is needed: we found that different learning rates converge at s 2000 physical epochs. Similar to the main text experiments, we observe that if we decay after evolving for the same amount of physical epochs, larger learning rates do better. See figure S5.
D.3 DIFFERENT LEARNING RATES CONVERGE AT THE SAME PHYSICAL TIME
We can also plot the test accuracy versus physical time for different learning rates to show that for vanilla SGD, the performance curves of different learning rates basically top of each other if we plot them in physical time, which is why we find that the fair comparison between learning rates should be at the same physical time.
We have picked a subset of learning rates the previous WRN28-18 18 CIFAR100 experiment of Appendix D.2. In figure S6, how even if the curves slightly different they converge
8

---

Under review as a conference paper at ICLR 2021
Accuracy
2/he
0.2
12/Ag
accurncy Test accuracy
0.0
10"
30'
10'
17
Figure S4: Final accuracy versus learning rate for fully-connected 1 hidden layer ReLU network, trained 512 samples of MNIST with full-batch gradient descent until training accuracy reaches or 700k physical steps (see Appendix A for details). We used subset of samples to accentuate the performance difference between phases. The optimal performance obtained when the learning rate is above /eriit. and close to Mmax
0.8
0.7
0.75
0.6
300uro0y
0.76
Text
0.4
0.722
13/Aa
12/
No decas
doday
0.01
9
(a)
(b)
Figure $5: Test accuracy vs learning rate for WRN28-10 10 and CIFAR100 with vanilla SGD, Ta regularization, data augmentation. label smoothing and batch size 1024. The critical learning rate is Merit & 0.4. (a) Evolved for 38400 steps. (b) Evolved for 96000/n steps with learning rate 77 (blue) and then evolved for 7200 more steps learning rate 0.01 (red).
roughly the same accuracy. The only curve which is slightly different is 2.5 which is rather high learning rate (close X²) ).
D.4 COMPARISON OF LEARNING RATES FOR DIFFERENT L2 REGULARIZATION FOR WRN28-10 ON CIFARIO
Even if in the main section we have considered model with fixed Le regularization, the effect without L2 or with a different value. In these two examples, we will be considering the same setup figures 3c.3d.
study
Without Ly regularization, we see that the larger learning rate does better even in the absence of learning rate decay, although training takes really long time. In our experience, comparing this setup with state of the art, Lg = 0 regularization makes the experiment take longer before convergence but does not influence performance much.

---

Under review as a conference paper at ICLR 2021
Max 0.78
1.2.
Max 200 0.75
0.8
Mox scr 0.75
Mbx
Test 3c0/00)
0.0
000
1000
1500
2000
Figure S6: Test accuracy physical time for different learning rates in the WRN CIFAR100 experiment of the previous section D.2
0.2
0.02
0.92
Text
0.88
0
1.00
200
2000
430
500
600
Physical spochs
Figure S7: WRN28-10 10 on CIFARIO without L2. Same setup as 3d but evolved for longer times.
In the presence of L2 regularization we picked the particular value L2 0.0005 in order to make sure that our conclusion not dependent on the choice of L2, the only hyperparameter (other than 7) have considered larger L2 0.001. We that the optimal performance in physical time is also peaked the catapult phase, although the difference here is smaller.
D.5 TRAINING ACCURACY PLOTS
The training accuracies of the previous experiments are shown in figure $9.
10

---

Under review as a conference paper at ICLR 2021
0.95
0.90
0.90
0.8G
0.70
2/
0.95
3/Xa
0.91
12/
12/Ao
decay
Desay 1 0.01
0.90
10P 0.90
10
10
10P
(a)
(b)
Figure S8: Test accuracies for larger L2 CIFAR10 experiment like that of the main section. (a)
WRN CIFAR-10 10 7200 steps as in figure 3c. (b) WRN CIFARIO R10 2400 physical steps and then 4800 more steps at learning rate 0.01 as in figure 3d.

---

Under review as a conference paper at ICLR 2021
0.95
0.90
Teen accuscy
Train accurdcy
0.83
0.75
0.80
0.70
k/
0.65
27 7m
12/Sa
No.
Deay
0. 0.085
0.50
0.70
10a
10
(a)
(b)
0..55
0.8
0.7
20curary
0.85
0.6
Trais
0.80
0.3
0.75
2/N0
0.4
27 Av
32/ As
12/ A6
0.3
0.70
10-1
10a
10-1
10'
(c)
(d)
Figure S9: Training accuracies for the performance experiments. Smaller learning rates have higher training accuracy when compared in physical time. However, they still perform worse for a fixed number of steps. (a) WRN CIFAR-10 12000 steps in figure 3c. (b) WRN CIFARIO 3360 physical steps as in figure 3d. (c) WRN CIFAR 100 38400 steps as in figure Sd) WRN CIFAR 100 96000 physical steps in figure S5b.
12

---

0.8 Under review as a conference paper at ICLR 2021
D.6 EQUIVALENCE Ào AND THE HESSIAN EIGENVALUE
In figure S10 we run an experiment for shallow Wide Resnet with 10 layers (1 block), without data augmentation and measure of the maximum eigenvalues of the NTK and Hessian for different widths. We see that they track each other fairly well for all times. There is intrinsic noise because different samples are used compute these cigenvalues.
1.00
WR -0.6
WRN 10-2 10-2.1A -...
WRNID-SS.rAs -3A.
0.95
WAN 10 10-5.n/.. -3.
NTK
acurccy
0.86
0.80
0.70
0.65
10000
20000
30000
40000
60000
10000
20000
30000
40000
50000
Staps
Figure $10: NTK vs Hessian eigenvalues for shallow Wide Resnet trained on CIFARIO during training. These values are normalized by An
D.7 LATE TIME BEHAVIOUR EIGENVALUES
We can compute the eigenvalues at late times for the experiment of figure 3d Figure S11 shows the final step of training (before decaying). We that the dependence of the curvature on the learning rate is preserved at late times.
2/X
Ague-decer
100
109
0.1
0.2
0.3
0.4
0.5
0.6
Figure S11: vs 17 late times for the experiment of figure 3d.
13

---

Under review as a conference paper at ICLR 2021
E EXPERIMENTAL RESULTS: EARLY TIME DYNAMICS
E. RELU ACTIVATIONS FOR THE SIMPLE MODEL
In the main text we have been using ReLU con-linearities. Compared with the simple model with no non- -nellinariiiee, ReLU networks have a broader trainability regime after _. It looks like these
Ao
12 This is generic feature of deep ReLU networks and can be already observed for non-linear generalization of the single sample model in Section 3 with a target = 1, two hidden layers and a ReLU non- -oneanit:: f = u. LU (w. RU figure S12). In this single sample context for r 2 X. the loss doesn't 't diverge but the neurons die and end up giving the trivial f = 0 function. For deep networks with more than one hidden layer and multiple samples, as discussed in the main text, we observe that the loss diverges after 12
networks generically will train until 1j Ao
as shown
X
22 364
2/50
4/5g
12/Au
7/9
0.50
0.25
(a)
(b)
Figure S12: Simple model ReLU non-linearity (Meii. 2.54). (b) is evaluated physical time 100.
E.2 MOMENTA
The effect of the optimizer also affects these dynamics. If we consider a similar setup with momenta, first we expect that a linear model converges linearities, we observe that for 17 < 2. the A is constant. However this is not true for ReLU, see figure S13a. In fact, for ReLu networks, we observe that there is small learning rate, roughly ceff Mezis below which the time dynamies of X is similar (but non-constant). However, for 7 > Neff.erit: there are strong time dynamics, we illustrate this in figure S13b with hidden layer ReLu network.
broader range 7 < Ao (1 + +)) For smooth non-
E.3 EFFECT OF L2 REGULARIZATION TO EARLY TIME DYNAMICS
We don't expect La regularization to affect the early time dynamics, but because of the strong rearrangement that goes on the first steps, it could potentially have a non-trivial effect; among other things, the Hessian spectrum necessarily is decaying. We how the dynamics that drives the rearrangement roughly the same, even in the maximum cigenvalue at early times is decreasing slowly.
E.4 TANII ACTIVATIONS
We observe that for Tanh activation, T/max is closer to the simple model expectation - see figure S15.
14

---

Under review as a conference paper at ICLR 2021
01.5
Mentity
Tash
350
300
(b)
(a)
Figure S13: (a) Evolution of the normalized curvature X A for 2 w 2048 FC connected networks evolved with momenta (same networks with SGD with dashed line for reference) evolved L. We observe that ReLU networks evolved with momenta doesn't have a constant kernel in the naive 'lazy" phase. (b) Verrit = 6.96, Ferrit,fff 0.69 Same setup as the FC network figure with momenta
for 7 = x
0.9: fully connected, three hidden layers = 2048, ReLU non-linearity. Territ is slightly different due to variations at initialization.
0-08
2/
12/Ao
m-0
10
2/9
2
(a)
(b)
Figure S14: Same WRN as figure 2d,f with T.2 regularization= 0.0005. Dynamics in physical steps of the d and d vs n. Ferit 0.18 λe b) X at physical time 25
E.5 WRN NTK NORMALIZATION
As illustrated in the text in figures 2c, 2e also this behaviour for NTK normalization. For completeness we include the WRN model with NTK normalization. From the linearized intuition, we expect the phases also be determined by the quantity nd, independently of the normalization. Figure $16 has the same setup as in figure
F RESTORATION OF LINEAR DYNAMICS
One striking prediction of the the catapult mechanism is that after period of excursion, the logit differences settle back to O(1) values, the NTK stops changing, and evolution is again well approxi- mated by linear model with constant kemel large width.
15

---

Under review as a conference paper at ICLR 2021
2/
4/30
000
2/9
0:005
30
2 20
90
10
0.05
0.10
0.15
(a)
(b)
Figure S15: Maximum NTK eigenvalue À at early times for 2 hidden layer fully connected network with tanh non- non-linearity trained on MNIST, with Merit 0.06. (a) Early time dynamics of the curvature for learning rates in the linear and catapult phase. (b) measured at 17t 3.
v--8.
2/
qu200.
4/30
0.06
40.0
12/Ao
2/9
qu-00.0
2
0.02
(a)
(b)
Figure S16: Same as Figures 2d,2f but with NTK normalization. a.b) Wide Resnet 28-10. Merit 31.47,A vs physical time 4000
We speculate that the return linearity and constancy of the kernel may hold asymptotically in width generally for a range of learning rates above '/erit- We test this by evolving the model for order (ng(n) steps until the catapult effect over, linearizing the model. and comparing the evolution of the two models beyond this point. Figure S17 shows an example of this. At fixed width, the accuracy of the linear and non-linear networks match for range of learning rates above the transition up 4/ Aoo We present additional evidence for this asymptotic linearization behavior in the Supplement.
F.1 ADDITIONAL EVIDENCE FOR LINEARIZATION IN THE CATAPULT PHASE.
Here we present some more detailed evidence for the re-emergence of linear dynamics in the catapult phase. Figure S18 show results for models trained on subsets of MNIST with learning rates Marit In figure Figure S18a we see that for mee-hddeen-layer fully connected model trained on 512 MNIST images, the performance of the full non-linear model and model linearized after 10 steps track closely.
16

---

Under review as a conference paper at ICLR 2021
1.0
Test accuracy
0.6
0.4
4/
0.2
12/Ag
1::10
Nos linear
0.0
10'
10'
10"
(a)
Figure S17: Evidence for linear dynamics after the catapult effect is over. Here show the same model as in Figure S4 with the addition of models linearized at step 0 and another linearized at step 10. We observe that the model linearized after 10 steps tracks the non-linear performance in the catapult phase up to 4/2o-
Models evolve linear models when the NTK is constant. In Figure S18b we give evidence that as networks become wider, the change in the kernel decreases.
0.6
0.6
(0(*) - (tuul)
Linarizal at 5tap=0 Unvarized at 10 Non inear
Tel
0.3
10
Ihi onn [tank), nu 09 (tank),
Dil (edu), Ru. 0.40
2v (nela), fit -047
(b)
(a)
Figure S18: Evidence a return of linear dynamics after 4in- (a.b) Show the same model as in figure S4 with the addition of linearized models step 0 and 10. We observe that the linearized model after 10 steps tracks the non-linear performance in the 'catapult' phase up to 17 ~ Ao (c) The change in the NTK between fim = 50 steps and = 1000 steps decreases as the width increases. Here we consider 2-class MNIST with 100 samples per class.
1
17


Reviews:
[{"review_id": "qKGtK8i7MHg", "reviewer": "AnonReviewer4", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "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": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "The paper studies the effect of the learning rate's magnitude when training neural networks. I believe this to be an extremely relevant problem since large learning rates are widely adopted in practice due to the their positive impact on the model's generalization, even though we don't understand the reason behind this.\n\nThe main contributions come from the analysis of a 2-layer linear network trained on a single data point, which shows the existence of three regimes that are governed by the learning rate's magnitude and the local curvature at initialization. Unlike prior work that also identifies different regimes governed by the learning rate, this submission does not point at noise as the main responsible for the regimes.\n\nPresented experiments are interesting and clearly show the different regimes and how they are connected to the learning rate, but do not seem to provide significant new insights compared to previous work. I believe the plots showing how the curvature behaves as the learning rate varies are novel to the best of my knowledge, but are not surprising given the previous discussions on how the learning rate affects the solution's sharpness.\n\nWhile the theoretical result does provide a precise characterization of the phenomena, it is restricted to a very constrained setting (2-layer linear net trained on one data point), making it hard to evaluate how the result translates to more complex settings (although the authors do a reasonable job at exploring this via experimental analysis).\n\nI also find the discussion on the presented theory to be somewhat lacking. For example, while the paper seems to downplay the importance of noise in characterizing the three regimes (compared to prior work), the curvature inevitably changes if noise is added to the sample gradients, e.g. adding isotropic noise will decrease \\eta_crit and shift the phase transitions, which is also expected to occur if as we decrease the batch size when doing mini-batch SGD. The fact that different batch sizes have been adopted across experiments (even for the same dataset/model) makes the experimental results harder to evaluate when compared to the developed theory. Except for a footnote on page 2, there is no discussion on how the initialization scale affects the phase transitions, which would be useful to better understand how the theory relates to prior work that focuses on the init scale instead of noise or width.\n\nOverall I believe the problem to be very relevant and the approach to be promising, but the theoretical results are overly limited and the contributions are below what I'd expected to suggest acceptance. Since extending the analysis to non-linear networks would seem to be a considerable technical leap, the authors could instead extend it to a less limited training setting where gradient variance is non-zero, in which case \\eta_crit would have a dependence on the batch size assuming mini-batch SGD updates.\n\n----\n\nI have read the response and am keeping my score. I agree that the simplicity of the results/model is valuable, but additional theoretical results (even extensions to Theorem 1, with more involved but stronger claims) would greatly improve the paper and make its contributions closer to what is expected of a ICLR submission. Extending the result to consider noise should be straightforward and yield a reasonably simple claim, which can be further verified empirically by adding synthetic noise to the gradients, adding label noise, and/or adopting extremely small batch sizes. As it stands now, the submission is still lacking in terms of contributions.", "extra_scores": {}, "extra_text": {}}, {"review_id": "L7gebsreHY", "reviewer": "AnonReviewer3", "summary": "", "questions": "", "limitations": "", "rating": 4, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "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": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "The paper is a detailed account of how large of a learning rate a given mode can take when it is trained by constant step size SGD. Many papers investigated the effect of learning rate (and batch size and width etc…) to the final accuracy. In general, the findings indicate increased accuracy up to a certain threshold (where the algorithm doesn’t converge any more). Such findings are abundant in the literature. However, the present paper investigates the path that SGD takes to get to find those ‘good’ performing points in the weight space. Two phases emerge: monotonic decay in loss vs catapult regime and the latter appears to perform the best. \nThe paper has interesting findings but also some shortcomings, in more detail:\n- Similar curve as in Figure 14 of https://arxiv.org/abs/1905.03776 as Fig 1 of the present paper, in that sense the findings of large lr better are certainly not novel, but the authors are well aware of this\n- The paper defines the curvature as the max eigenvalue of the Fisher information matrix. The large learning rate values are determined according to that value. Fig 2 c and f show the cross over of the regimes. What would make the cross over sharper? Increasing width? Larger batch size? \n- The large increase in loss value is huge! It is surprising that training doesn’t just diverge from there. What controls the growth and the subsequent decay? How consistent is it? Would the result hold in all random seeds or what fraction would still converge? \n- What’s the max tr loss for the wide res net? (Why not consistent plots for fully connected and resnets in fig2?) (Also, for clarity, Fig 2 label of B would be Max tr loss)\n- Fig 3 would be clearer if the test error axis of (a,b) would be aligned, same for (c,d). It appears as if decay actually makes the search for the correct large lr rate redundant. (Perhaps at the expense of searching for the right decay rate). In general, what can we do to maximize the range of hyperparameters that works for the model (to minimize search)? The paper proposes a heuristic based on the initial local curvature yet the findings for practice seem somewhat preliminary. \n- The comparison between lazy training and how large learning rate deviates from lazy training is a very interesting and active area of research. I think the paper would be more complete if the experiments would verify it in a more robust way. Measuring the deviation of the activations or the movement of the kernel would be helpful. Here is one of the references with clear observables to measure in the experimental framework https://arxiv.org/abs/1906.08034 \n\nThe characteristic difference in the early time dynamics is claimed to be related to the performance gains. There is a growing body of evidence that points to the role of the initial dynamics of SGD when it comes to identifying the final performance. The paper is further evidence in that direction. Therefore I think it’ll be useful for the ICLR community. Overall, the paper has very interesting bits and pieces but it fails to come together as a coherent whole to provide a consistent story.", "extra_scores": {}, "extra_text": {}}, {"review_id": "H5lsQgDLMX", "reviewer": "AnonReviewer2", "summary": "", "questions": "", "limitations": "", "rating": 3, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "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": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "This paper analyzes the effect of choosing a large step-size on the generalization of deep networks. They suggest that starting with a large learning rate, the loss initially increases before converging to a flatter minimum with improved generalization (catapult effect). When the learning rate is above a certain threshold, the authors observed this catapult effect along with a decrease in the curvature of the landscape. These observations were empirically demonstrated through various experiments and analytically studied for a two-layer linear network.\n\nPros:\nThe paper tackles an interesting problem in a vibrant field of research.\n\nSeveral experiments that demonstrate the authors claims were presented in the paper. These results can be further used to propose algorithms designed to converge to flatter solutions. \n\nThe related material is referenced and well-discussed in the paper.\n\nCons:\nThis paper lacks theoretical evidence for supporting the raised claims. The model studied in Theorem 1 is very simple. Moreover, the effect of choosing the step-size and batch-size on the sharpness of the computed solution has been observed in the literature. For instance, Keskar et al. (2017) justified the improved generalization achieved when using small batch sizes by the ability of such methods to converge to flat minima. The choice of the batch-size was formally related to the choice of the step-size in the work of Patel (2017) which provides a learning-rate lower bound threshold for the divergence of batch SGD. This threshold is also a function of the curvature. The higher the curvature around a critical point, the lower the threshold required for divergence. Hence choosing a large step-size tend to escape sharp minima and potentially converge to flatter minimizers. The last paper was not cited in the submitted manuscript.\n\nThe paper does not provide any intuition on how to compute the thresholds $\\mu_{\\mbox{crit}}$ and $\\mu_{\\max}$, and how is this related to the choice of non-linearity, structure of the network, …\n\nComments:\n1.\tFor the gradient descent update rule in equation (1): why do we have $f_t$? \n\nMinor Comments:\n1.\tWhat is meant by ``compute budget’’?\n2.\tIn the definition of $\\Theta$ in Page 4: I think it should by $\\sum_{\\alpha = 1}^m$ instead of $\\sum_{\\mu=1}^p$.\n3.\tIn Figure S1, the images are small and not very clear.\n4.\tExpression (S12) is missing a term.\n5.\tS(16), what is $u_{ia}$.\n6.\tFigure S10, the colors in the label do not match the plots.  \n\n", "extra_scores": {}, "extra_text": {}}]

Author rebuttal:

Meta-review: I agree with the reviewers that said that this paper has valuable insights. However, all reviewers ultimately recommended rejection. I think the main reason was that the reviewers did not feel these insights don't accumulate together to a message that would justify a paper. I hope the authors can address these concerns and resubmit. There were additional concerns, like the fact a very simplistic toy model is being used, but I agree with the authors that it makes sense to first explore such phenomena in the simplest model that produces them.
Decision: Reject