Quadratic models for understanding neural network dynamics
Authors: Libin Zhu
Chaoyue Liu
Adityanarayanan Radhakrishnan
Misha Belkin
Venue: Submitted to ICLR 2023 (2023)
Keywords: quadratic models
wide neural networks
catapult phase
optimization dynamics

Abstract: In this work, we show that recently proposed quadratic models capture optimization and generalization properties of wide neural networks that cannot be captured by linear models.  In particular, we prove that quadratic models for shallow ReLU networks exhibit the "catapult phase" from Lewkowycz et al. (2020) that arises when training such models with large learning rates. We then empirically show that the behaviour of quadratic models parallels that of neural networks in generalization, especially in the catapult phase regime. Our analysis further demonstrates that quadratic models are an effective tool for analysis of neural networks. 

Under review as a conference paper at ICLR 2023
QUADRATIC MODELS FOR UNDERSTANDING NEURAL NETWORK DYNAMICS
Anonymous authors Paper under double-blind review
ABSTRACT
In this work, we show that recently proposed quadratic models capture optimization and generalization properties of wide neural networks that cannot be captured by linear models. In particular, we prove that quadratic models for shallow ReLU networks exhibit the "catapult phase" from Lewkowycz al. (2020) that arises when training such models with large learning rates. We then empirically show that the behaviour of quadratic models parallels that of neural networks in general- ization, especially in the catapult phase regime. Our analysis further demonstrates that quadratic models are an effective tool for analysis of neural networks.
INTRODUCTION
recent remarkable finding on neural networks, originating from Jacot et al. (2018) and termed as the "transition to linearity" (Liu et al., 2020), is that, as network width goes to infinity, such models become linear functions in the parameter space. Thus, linear (in parameters) model can be built to accurately approximate wide neural networks under certain conditions. While this finding has helped improve our understanding of trained neural networks (Du et al., 2019; Nichani et al., 2021; Zou & Gu, 2019; Montanari & Zhong, 2020; Ji & Telgarsky, 2019; Chizat et al., 2019), not all properties of finite width neural networks can be understood in terms of linear models, is shown in several recent works (Yang & Hu, 2020; Ortiz-Jiménez et al., 2021; Long, 2021; Fort et al., 2020). this work, show that properties of finitely wide neural networks in optimization and generalization that cannot be captured by linear models are, in fact, manifested in quadratic models.
The training dynamics of linear models with respect to the choice of the learning rates' are well- understood (Polyak, 1987). Indeed, such models exhibit linear training dynamics, i.e., there exists a critical learning rate, Merit such that the loss converges monotonically if and only if the learning rate is smaller than Teii (see Figure la).
Linear Dynamics
Non-linear Dynamics
Convergence
Divergence
Convergenoe
Divergence
Catapult
Monotonic
Phase
Merit
Learning rate
0
7crit
mmax
Learning rate
(a)
(b)
Figure I: Optimization dynamics for linear and non-linear models based on choice of learning rate. (a) Linear models either converge monotonically if learning rate is less than Merit and diverge otherwise. (b) Unlike linear models, finitely wide neural networks and NQMs Eq. (2) (or general quadratic models Eq. (3)) can additionally observe catapult phase when Mert < < Maaxx-
Unless stated otherwise, always consider the setting where models are trained with squared loss using gradient descent.

---

Under review as a conference paper at ICLR 2023
Linear Dynamics
Catapuit Dynamics
Parameter space
0.55
Linear Dysamis
Catsuult Dy mynamics
Loos
0.80
Best Tedl Logs
Iteration
Iteratioa
0,70
We
[[]]]
[IKE
Constant Yangent Narnel Blw, R)
0.00
Berit
Globel minina
I
Itoration
Optiiizaiion pash
Sub-eriizui Lika Saper ottied ottioo LRs
Iieration
Learning
(a) Optimization dynamics for f (wide neural networks): linear dynamics (b) Generalization performance for and catapult dynamics.
f. fun and fquad-
Figure (a) Optimization dynamies of wide neural networks with sub-critical and super- critical learning rates. With sub-critical learning rates (0 < <, Meritt) the tangent kernel of wide neural networks is nearly constant during training, and the loss decreases monotonically. The whole optimization path is contained in the ball B (Wwo, Rt) == {w : ||w woll S R} with finite radius R. With super-critical learning rates (eoit < 1j < 7maax), the catapult phase happens: the loss first increases and then decreases, along with a decrease of the norm of the tangent kernel. The optimization path goes beyond the finite radius ball. (b) Test loss of fquad: f and füin plotted against different learning rates. With sub-critical learning rates, all three models have nearly identical test loss for any sub-critical learning rate. With super-critical learning rates, f and Sauad have smaller best test loss than the one with sub-critical learning rates. Experimental details are in Appendix J.4.
Recent work Lee al. (2019) showed that the training dynamics of a wide neural network f(w; x) can be accurately approximated by that of a linear model fun (w; x)::
fün (w: x) f (Wo; x) + (w wo) T (Wo;
where V fwoo æ) denotes the gradient² of with respect trainable parameters
ini-
tial point wo and input sample æ. This approximation holds for learning rates less than Merit R 2/ f(/o 2)13, when the width is sufficiently large.
However, the training dynamics of finite width neural networks, f. can sharply differ from those of linear models when using large learning rates. A striking non-linear property of wide neural networks discovered in Lewkowycz et al. (2020) is that when the learning rate is larger than Merit but smaller than certain maximum learning rate, 7mmax, gradient descent still converges but experiences "catapult phase." Specifically, the loss initially grows exponentially and then decreases after reaching large value, along with the decrease of the norm of tangent kernel (see Figure 2a), and therefore. such training dynamics are non-linear (see Figure lb).
As linear models cannot exhibit such a catapult phase, under what models and conditions does this phenomenon arise? The work of Lewkowycz et al. (2020) first observed the catapult phase phenomenon in finite width neural networks and analyzed this phenomenon for a two-layer linear neural network. However, theoretical understanding of this phenomenon for general non-linear neural networks remains open. In this work, we utilize a quadratic model as a tool shed light on the optimization and generalization discrepancies between finite and infinite width neural networks. We call this model Neural Quadratic Model (NQM) as it is given by the second order Taylor series expansion of f [w:: x) around the point Wo:
1
fquaa = -[(w) + (w - wo +-(w -wo Hy) - wo).
(2)
finnt))
Here in the notation we suppress the dependence on the input data and Hy (Wo) is the Hessian of with respect to w evaluated at Wo.
Indeed, we note that NQMs are contained in a more general class of quadratic models:
g w: x) T (xx) +,,"">""">>
> Σ(x)w,
For non-differentiable functions, e.g. neural networks with ReLU activation functions, we define the gradient based the update rule used in practice. Similarly, we H, denote the second derivative of f in Eq. (2).
2

---

Under review as a conference paper at ICLR 2023
where trainable parameters and æ is input data. We discuss the optimization dynamics of such general quadratic models in Section 3.3 and show empirically that they exhibit the catapult phase phenomenon in Appendix J.3. Note that the two-layer linear network analyzed in Lewkowycz et al. (2020) is special case of Eq. (3), when (xx:) = (See Appendix I).
Main Contributions. We prove that NQMs, fquad- which approximate shallow fully-connected ReLU activated neural networks, exhibit catapult phase dynamics. Specifically, we analyze the optimization dynamics of fquad by deriving the evolution of Squad and the tangent kernel during gradient descent with squared loss, for a single training example and multiple uni-dimensional training examples. We identify three learning rate regimes yielding different optimization dynamics for fquad- which are (1) converging monotonically (linear dynamics); (2) converging via catapult phase (catapult dynamics); and (3) diverging. We provide a number of experimental results corroborating our theoretical analysis (See Section 3).
We then empirically show that NQMs, for the architectures of shallow (see Figure 2b as an example) and deep networks, have better test performances when catapult dynamics happens. While this was observed for some synthetic examples of neural networks Lewkowycz et al. (2020), we systematically demonstrate the improved generalization of NQMs across a range of experimental settings. Namely, we consider fully-connected and convolutional neural networks with ReLU and other activation functions trained with GD/SGD multiple vision, speech and text datatsets (See Section 4).
To the best of our knowledge, work is the first to analyze the non-linear wide neural networks in the catapult regime through the perspective of the quadratic approximation. While NQMs (or quadratic models) were proposed and analyzed in Roberts al. (2022), our work focuses on the properties of NQMs in the large learning rate regime, which has not been discussed in Roberts et al. (2022). Similarly, the following related works did not study catapult dynamics. Huang Yau (2020) analyzed higher order approximations to neural networks under gradient flow (infinitesimal learning rates). Bai & Lee (2019) studied different quadratic models with randomized second order terms and Zhang et al. (2019) considered the loss in the quadratic form, where no catapult phase happens.
woll < R} with constant radius R > O, the transition to linearity of wide neural network (with linear output layer) is continuous in the network width m. That is, the deviation from the network function to its linear approximation within the ball can be continuously controlled by the Hessian of the network function, i.c. Hy, which scales with (Liu et al., 2020):
Discontinuity in dynamics transition. In the ball B(wo, R) == {w
||
sup II[(w) füi (w)ll wEB(
||H = 0(1
(4)
Using the inequality from Eq. (4), we obtain [[fquad - fiall (1//mm), hence fquad transitions to linearity continuously well in B(w,, R)². Given the continuous nature of the transition to linearity, one may expect that the transition from non-linear dynamics to linear dynamics for / and fquad is continuous in m as well. Namely, one would expect that the domain of catapult dynamics, [erit. 7maax]. shrinks and ultimately converges to a single point, i.c., eeeiit = Tmax. as m goes to infinity, with non-linear dynamies turning into linear dynamics. However, as shown both analytically and empirically, the transition is not continuous, for both network functions f and NQMs fquad- since the domain of the catapult dynamics can be independent of the width m (or 7). Additionally, the length of the optimization path of ff in catapult dynamics grows with m since otherwise, the optimization path could be contained
ball with a constant radius independent of m, in which I can be approximated by Stim Since fúi diverges in catapult dynamics, by the approximation, s' diverges as well, which contradicts the fact that f can converge in catapult dynamics (See Figure 2a).
2 NOTATION AND PRELIMINARY
We use bold lowercase letters to denote vectors and capital letters to denote matrices. We denote the set {1, 2,... .n by [n]. We use || || to denote the Fuclidean norm for vectors and the spectral norm for matrices. We use O to denote deiment-wise multiplication (Hadamard product) for vectors. We use Amax (A) and Amin(A) to denote the largest and smallest eigenvalue of matrix A. respectively.
For general quadratic models in Eq. (3), the transition lincarity is continuously controlled by 7-
3

---

Under review as a conference paper at ICLR 2023
Given model f(w; 2), where is input data and are model parameters, W to represent the partial first derivative ( x) /O/ When clear from context, we let Vf := w for case of notation. We use H and He to denote the Hessian (second derivative matrix) of the function f(w x) and the loss L(w) with respect to parameters w, respectively.
In the paper, consider the following supervised learning task: given training data {( y/ )== with data æ; E R4 and labels y/ E IR for E [n]. we minimize the empirical risk with the squared loss [(w) = 2 Σ== (f(:: æ. - y²)² Here f(w:-) a parametric family of models, e.g., a neural E RR. We use full-batch gradient descent to minimize the loss, and we denote trainable parameters w at iteration & by w(t). With constant step size (learning rate) 7), the update rule for the parameters is:
network kernel machine, with parameters
w(t + 1) = w(t) df(w)
(t), Vt 2. 0. dw
Definition (Tangent Kernel). The tangent kernel ( of f(w: is defined as
K K(w œ. =) (V[(w; m) (s(w; Vx, z E Rdd.
(5)
In the context of the optimization problem with n training examples, the tangent kernel matrix K R"*" satisfies Ki,((W) (w) = (w;;,,, æy), i,j [n]. The critical learning rate for optimization is given as follows.
Definition 2 (Critical learning rate). With an initialization of parameters wo, the critical learning rate of f(w; :) is defined
nerit = Amna (H)
(6)
learning rate yj is said to be sub-critical if < 1 < Merit or super-critical if feriit < v) < Mmax Here Vmaa is the maximum leaning rate such that the optimization of L(w) initialized at wo can converge.
Dynamics for Linear models. When f is linear in w, the gradient, V,, and tangent kernel are constant: K(w(t)) ((w)). Therefore, gradient descent dynamics are:
F(w(t: 1)) y (T ()oo wo where F(wo) = (w) (woo)) with (wo) = =[==[[n f(wo;
y). 14 0,
Noting that Hc(wa) (F(w.)) F(wo)) and that tangent kernel K(Wo) V VF((w)VF(ww))
share the same positive eigenvalues, we have X (Hc(Wo)) = Amax ( K (wo)). and hence,
Merit = 2/>max(X' ((w))).
(8)
Therefore, from Eq, equation 7, 0 < 1 < Meiit, the loss decreases monotonically and if > Meriit. the loss L diverges. Note that the critical and maximum learning rates are equal in this setting.
OPTIMIZATION DYNAMICS NEURAL QUADRATIC MODELS
In this section, we analyze the gradient descent dynamics of the NQM that approximates two-layer fully connected ReLU activated neural networks up to the second order. We show that the extra quadratic term in NQMs allows for catapult convergence: the loss increases early stage and then converges afterwards. Note that this type of convergence happens with super-critical learning rates and cannot happen for linear models. Interestingly, the top eigenvalues of the tangent kernel typically decrease after the catapult phase, while they are nearly constant when training with sub-critical learning rates, where the loss converges monotonically.
Neural Quadratic Model (NQM). Consider the NQM that approximates the following two-layer neural network:
flu, v;
Σ UgO
1 Vdd
where u, E Rd, v; E R for i E [m] are trainable parameters, x E Rd is the input, and a(-) is the ReL.U activation function. Each parameter is initialized i.i.d. following the standard normal
4

---

Under review as a conference paper at ICLR 2023
distribution, (0, 1). Letting (u, v: == fquad (u, v; this NQM has the following expression (See the full derivation in Appendix A):
g(u, v; f(uo, vo:x) +
vo uo æ1l a20) +
1 md
Σ (t) 20,)) (uZ, a)
md
0 vo - uo.i) (oos)
x1 (000e11(11 (ug
(10)
Vind
minimize the empirical risk with the squared loss [(w) E (ol w: mx) )² using GD with constant learning rate v). Throughout this section, we denote (lu(t),((); by (() and its tangent kernel ((u()), v()) by ((),, where f the iteration of assume the width of is much larger than the input dimension and the data size n, i... m7 > max [d, n). Hence, d and n can be regarded small constants. In the whole paper, we the big-O and small-o notation with respect to the width m. Below, we start with the single training example case, which already showcases the non-linear dynamics of NQMs.
Given training data {zg, 14 -
GD. We assume | O(1) and 13,1 = ((1) for i [ni], and
3.1 CATAPULT DYNAMICS WITH SINGLE TRAINING EXAMPLE
In this subsection, we consider training dynamics of NQMs with single training example (æ, v) where E I and y R. In this case, the tangent kernel matrix K reduces to scalar, and we denote K by À to distinguish it from a matrix.
By gradient descent with step size v, the updates for g(t) y and X(t), which we refer to as dynamics equations, can be derived as follows (see Appendix B.1):
Dynamics equations.
g(t + 1) - = 1 - n)(t)) + |
md 1f (g(()) (gg(t)
(o(t) - -0) : (o)) - y),
(11)
(t(
X(1 + X(1) + |
g(()
md (g(1) - y)² (A(11) g(t)
14 2 0.
(12)
Rx(1)
Note that as the loss is given by L(()= = =())=1/2(((0) - y) to understand convergence, it suffices to analyze the dynamics equations above. Compared to the linear dynamics Eq. (7), this non-linear dynamics has extra terms (t) and (t), which are induced by the non-linear term in the NQM. We will see that the convergence of gradient descent depends on the scale and sign of R2 and Rx(t), (t). For example, for constant learning rate that is slightly larger than Meii (which would result in divergence for linear models), Rt(t) stays negative during training, resulting in both monotonic decrease of tangent kernel X and convergence of the loss.
can show that with high probability over random initialization, AAo (1)) (see Appendix D). As A(t) = Xo + Σt= (T()), to track the scale of |μ(()], we will focus on the scale and sign of (t) and (t) in the following analysis. We start by establishing monotonic convergence for sub-critical learning rates.
For the scale of Åg. which is non- negative by Definition
sub-critical learning rates (n < 2/X(0)). The key observation that when gg()) is small, i.e., of the order of /m), and A(t) = 0(1), (t)()) and (t) are of the order o(1 (see Proposition 3 in Appendix F). Then, the dynamics equations approximately reduce to the ones of linear dynamics:
Monotonic convergence:
g(l + 1) - y= (1- + o(1) (9(1) - v),
X(t + 1)) + o(1)
Note that at initialization, with high probability over random initialization, the output satisfies le(0)) O(1) if ||=||| O(1) (Jacot et al., 2018), and have shown X(0) @(1). Applying
5

---

Under review as a conference paper at ICLR 2023
Proposition 3, with the choice of 17. we have for all 2 0. lu(t)] = |1 - -1t)) + +[1]] < 1; hence, [g(t) - y/ decreases monotonically. The cumulative change on the tangent kernel will be o(1), Le.. Ee (t) = o(1), since for all t, [Rx(t)] = o(1) and the loss decreases exponentially hence E (1) = o(1) log 0(1) = o(1).
Catapult convergence: super-critical learning rates (2/x(0) <1 < 4/x(0), The training dynamics are given by the following theorem.
Theorem (Catapult dynamics on a single training example). Consider training a NQM, Eq. (10), with squared loss on single training example by gradient descent. If the learning rate is super-
critical, i.e., 2/ < 17 < A/, then there exist T1, T2, Ts. T4 such that 0 < Ti < T2 < T3 < Ti and the training dynamics exhibits:
(i) Increasing phase: I. E Ti). In this phase, L(() = d(m). The loss grows exponentially and the tangent kernel is nearly constant, iee. [A(l) X(0)| o(1).
(ii) Peak phase: t E [T2, Ts). In this phase, [(t) ((m). The tangent kernel decreases significantly; A(t + 1) - X(t) < and [Xt +1 --01))--)]------------ =
(iii) Decreasing phase: L. E [Ta, x)). In this phase, the loss satisfies L(I) = o(m) again and decreuses. The tangent kernel is nearly constant until convergence: [x(1) - X(x) = o(1).
Furthermore, Ti = =o(log T2 = +llg m), T3- - TT = (1) and T4 = Ollg m))
Proof of Theorem 1. We will analyze the training dynamics in each phase sequentially. (i) Increusing phase. At the beginning. (g(l)] grows exponentially following linear dynamics. Specifically, since [0()) = (1). by Proposition 3, we have RRg (0)) o(L) and (0) o(1).
The loss grows exponentially at iteration O: by the choice of the learning rate, (() satisfies
|p(0)] = [1 m)) + Ry (0) (10) o(1) > 1.
-
Therefore (g(1) > (0) -y. and the tangent kernel almost does not change: X(1) = +o())
We can recursively apply this argument for the following steps as long as lg(()) ol vm) according to Proposition 3. Note that the increasing phase, the loss grows exponentially due to [u(t)] > 1.
Therefore, the tangent kernel is nearly constant since the cumulative change is E [[)(()] = En O (g(t)²/m) o(m). And can get Ti o(log m).
o(1), where we use the fact that g(t)² grows exponentially the order of
Furthermore, it is not hard to see that until the loss grows to the order of ((m), the tangent kernel does not change much hence the loss keeps increasing exponentially.
(ii) Peak phase. The key observation we use here is that when [g(()) is large, i.e., of the order O(Vm), (t) and (1) will be of the order (1)) Proposition 4 in Appendix F), which can lead to the decrease of the loss.
In the peak phase, we have 19(()) O @(vm), then by Proposition 4, the scale of [Rx(t)] is @(L), and (t) < 0 since A(t) < 4/n (in the increasing phase, A(t) almost does not change, hence we have X(t) Re X(0) < 4/n before the peak phase, and will decrease significantly). Consequently, by Eq. (12), A(t) will have significant decrease as [t)) = (1)) and A(t = A(t) + RA(t) < A(t), which is further smaller than 4/n-
Similarly, when [a()) = @v/m), [t) = O11) and (t) > 0 by Proposition 4. Then we can see the increase of loss slows down compared to that in the increasing phase:
lm(1) = 11 - A)t() + (1) < 11 17)(0) + (0) = 11 - -)(()) + o(1) =2 p(0).
In general, the loss grows exponentially prior to the peak phase hence Tr = (log m). And the peak phase only lasts (1) steps during training, i.e., [Tg- Tzl = (11). since the decrease of X(t) is O(1) at each step and the training dynamics will enter into the decreasing phase once [n(() < 1, which will happen if A()) is sufficiently small.
(iii) Decreasing phase. The peak phase ends when |a(t)] < 1, then lø(t) y/ starts to decrease. We note that our analysis implicitly assumes that the optimization path will not arrive at the saddle point, i.e., |p(t]| = 1, and it is generally true with discrete steps.
Recall that at the peak phase, A(t) decreases, and when [u(t)] > 1, [g(()) increases which causes the increase of R.y (1) since R. (1) scales with [g(()) As result [u(t) = [1 -nX(t) + (t())) decreases

---

Under review as a conference paper at ICLR 2023
and will be less than 1 ultimately. Note that though (a(t) yl starts to decrease, it is still of the order @Vmm) which makes X(t) decrease significantly as RA (t) = (1) by Proposition 4.
The decrease of A(t) will stop once (g()) decreases to the order of o(m), as Rx(t) = o(1) again by Proposition 3. At that moment, Ti- T3 = e (log m) as the loss decreases exponentially and the linear dynamics dominate again. Therefore, similar to the increasing phase, starting from such that [g(()] = m), the loss decreases exponentially hence the change of X(t) until convergence is o(1).
Divergence (1) Mmax = 4/) Initially, it follows the same dynamics with those in the increasing phase in catapult convergence: the loss increases exponentially and the tangent kernll is nearly constant. However, when [a(()] grows to the order of @(vm), corresponding to the peak phase in catapult convergence, X(1) does not decrease but increases significantly. Specifically, since 7 > 4/)(0),
approximately have 17 4/X(t) at the end of the increasing phase by the same analysis in catapult convergence. By Proposition 3, Rx(t) > O, then M(t) increases as A(t + 1) X(t) (t) > X(t). Larger Xt) leads to the faster increase on X(t), hence |n(()] becomes even larger. As result, (g(() - y grows faster, therefore the loss diverges.
3.2 CATAPULT DYNAMICS WITH MULTIPLE TRAINING EXAMPLES
In this subsection we show the catapult phase will happen for NQMs Eq. (9) with multiple training examples. We assume unidimensional input data, which is common in the literature and simplifies the analysis for neural networks (see for example Williams et al. (2019); Savarese al. (2019)).
Assumption 1. The input dimension d = 1 and not all x; is 0, i..., Σ | > 0.
Since x; is a scalar for all E [n]] with the homogeneity of ReLU activation function, we can compute the exact eigenvectors of (t) for all 2 0. To that end, we group the data into two sets S+ and S_ according to their sign:
S+ = fi : x4 2 0, E [n]] s- x4 < 0, E [n]]-
Now we have the proposition for the tangent kernel K(the proof is deferred to Appendix C):
Proposition 1 (Eigenvectors and low rank structure of K). For any u. E RTM. rank(K) <
Furthermore, P1. P2 are eigenvectors of K, where P1,; = x, (tES.) J. P2.: = (rEs. f.r E [n].
Note that when all x; are of the same sign, rank () = 1 and K only has one eigenvector (either Pi or P2 depending on the sign). It is in fact simpler setting since we only need to consider one direction, whose analysis is covered by the one for rank (K) = 2. Therefore, in the following we will assume rank = 2. We denote two eigenvalues of K K(t) by (t) and 22 corresponding to Pı and Pa respectively, i.e., K(t)p: = (t)pi, K(t)p2 = (2) Without loss of generality, we assume (0) 2 (0).
We similarly analyze the dynamics equations for multiple training examples (see Eq. (14) and (15) which are update equations of g(() - y and K())) with different learning rates. And we formulate the result for the catapult dynamics, which happens when training with super-critical learning rates,
into the following theorem:
Theorem (Catapult dynamics on multiple training examples). Consider training a NQM Eq. (10) with squared loss on multiple training examples by gradient descent. Under Assumption 1, if the learning rate is super-critical ie., (0) <n < min (2/)2 (//)
then there exist Tí, T2, T3, T4 such that 0 < T1 < To T3 < Ta and the training dynamics exhibits:
(i) Increasing phase: t E [0. Ti). In this phase, L(t) = om.). The loss grows exponentially; both eigenvalues are nearly constant i.e. (Ax(t) >v (0) o(1) for k 1,2.
(ii) Peak phase: t E [T2, Ta). In this phase, L(t) (m). For the tangent kernel:
(a) If2/ (0) < (0), both eigenvalues decrease significantly, i.e., (1 < X2 +1) >x(()) O and both difference are of the order (1). (b) If 2 > 4/x1 (0), only (t) decreases significantly.
(ii) Decreasing phase: t [Ta, oc). In this phase, the loss satisfies L(t) o(m) again and decreases. Both eigenvalues are nearly constant until convergence: (Ax(t) - >x = o(1). for k- 1,2.
7

---

Under review as a conference paper at ICLR 2023
Furthermore, T o(log m). T2 (llog m.) T3 T2 (11) and Ta @ m).
The proof idea is similar to the case for single training example; due to the low rank structure the tangent kernel, we only need to analyze directions pi and p2. It turns out that we can analyze the dynamics separately for each direction as we show the training dynamics in these two directions are almost independent to each other. We defer the proof of Theorem 2 and the analysis of other two parts of non-linear dynamics: monotonic convergence and divergence to Appendix H.
We verify the theoretical results above for multiple training examples via the experiments in Figure 3.
330
Spcgree langes egenvauu
byp
300
190
Traling 83c
LR= 0.20
LR 0.20
LR 0.80
LR 0:58
-00
0.5.
byt
LH 0:45
LH 0.45
150
LRt 0.45
LR 0.60
LR 0.60
LR - 0.80
100
80-
100
-3
50
100
10
10
10
175
Iteration
lleration
Iteration
(a) Training loss
(b) Largest cigenvalue of tangent kernel
(c) Sccond largest eigenvalue
tangent kernel
Figure 3: Training dynamics of NQMs for multiple examples case with different learning rates. By our analysis, two critical values are 2/ 2/)1(0) = 0.37 and 2/A2(0) = 0.39. When v) < 0.37, linear dynamics dominate hence the kernel is nearly constant; when 0.37 < 1) < 0.39, the catapult phase happens in P1 and only (4) decreases; when 0.39 < 7maax the catapult phase happens in P1 and p2 hence both (t) and (t(t) decreases. The experiment details can be found in Appendix J.1.
3.3 CONNECTION WIDE NEURAL NETWORKS AND GENERAL QUADRATIC MODELS
Wide neural networks. We have seen that NQMs, with fixed Hessian, exhibit the catapult phase phenomenon. Therefore, the change in the Hessian of wide neural networks during training is not required to produce the catapult phase. In our analysis, we show that the catapult phase arises because the eigenvectors of the tangent kernel "align" with the Hessian's 's spectrum, i.e., Hg, for S+ proportional with coefficients P1. and the same holds for Hg. for i E S_ with coefficients p2- = H if j. E S4. We believe this idea can be used to analyze the catapult dynamies in wide neural networks with changing Hessian. A similar behaviour of top eigenvalues of the tangent kernel with the one for NQMs is observed for wide neural networks when training with different learning rates (See Figure 5 in Appendix J).
General quadratic models. As mentioned in the introduction, NQMs are contained in a general class of quadratic models of the form given in Eq. (3). We show that the two-layer linear neural network analyzed in Lewkowycz et al. (2020) is special case of Eq. (3), and we provide a more general condition for such models have catapult dynamies Appendix I.
Furthermore, we empirically observe that broader class of quadratic models can have catapult dy- namies simply by letting ((x) and Σ be random and assigning small value to y (See Appendix J.3).
QUADRATIC MODELS PARALLEL NEURAL NETWORKS IN GENERALIZATION
In this section, we empirically compare the test performance of three different models considered in this paper upon varying learning rate. In particular, we consider (1) the NQM, fquadi (2) correspond- ing neural networks, f; and (3) the linear model, fin.
We implement our experiments on vision datasets: CIFAR-2 (a 2-class subset of CIFAR- 10 (Krizhevsky et al., 2009)), MNIST (LeCun et al., 1998), and SVHN (The Street View House Numbers) (Netzer et al., 2011), 1), speech dataset: Free Spoken Digit dataset (FSDD) (Jakobovski, 2020) and 1 text dataset: AG NEWS (Gulli, 2005).
In all experiments, we train the models by minimizing the squared loss using standard GD/SGD with constant learning rate 7). We report the best test loss achieved during the training process with
8

---

Under review as a conference paper at ICLR 2023
2-ayer FC on FSDD with
2-ayer FC on AG NEWS with GD
2-ayer FC on MNIST with GD
0LEE
Beel teaf Lase
Beel Twee Lowe
OrticsiL9
Of
Criicl
Mi
1.30
Nh
Guadiati
Gaaker
Quodit
Liver
000
Learning
Learning
CIFAR 2 SGD
GMN CIFAR-2
SVHN with
0860
Cricaal UF
0.8-
0.26
NY
Oucokidoo near
Beel Teet Loee
1001
0.79
Beet
Beet
Urt
Oucadralo
Liwer
066
03
100
2.3
Learning rate
Learnnng rate
Learning rate
Figure 4: Best test loss plotted against different learning rates for f(w), Jiin (w) and Squad (w) across variety of datasets and network architectures.
each learning rate. Experimental details can be found in Appendix J.4. We also report the best test accuracy in Appendix J.5. For networks with 3 layers, Appendix J.6. From the experimental results, we observe the following:
Sub-critical learning rates. In accordance with our theoretical analyses, we observe that all three models have nearly identical test loss for any sub-critical learning rate. Specifically, note that as the width m increases, f and fquad will transition to linearity in the ball (wo, Rt):
|f fiall O(//mm) Ilfquad - fuall Ó(1/mLL))))))
where R >> is constant which is large enough contain the optimization path with respect to sub-critical learning rates. Thus, the generalization performance of these three models will be similar when is large, as shown Figure 4.
Super-critical learning rates. The best test loss of both f(w) and fquad (w) consistently smaller than the one with sub-critical learning rates, and decreases for an increasing learning rate in a range of values beyond Nerit. which was observed for wide neural networks in Lewkowyez et (2020).
As discussed in the introduction, with super-critical learning rates, both fquad and ff can be observed have catapult phase, while the loss of fün diverges. Together with the similar behaviour of fquad and f in generalization with super-critical learning rates, we believe NQMs are better model to understand f training and testing dynamics, than the linear approximation fln.-
In Figure 4 we report the results for networks with ReLU activation function. We also implement the experiments using networks with Tanh and Swish (Ramachandran al., 2017) activation functions, and observe the same phenomena in generalization for J. fúi and fquad (See Appendix J.7).
CONCLUSIONS
In this paper, we use quadratic models tool to better understand optimization and generalization properties finite width neural networks trained using large learning rates, Notably, we prove that quadratic models exhibit properties of neural networks such the catapult dynamics that cannot be explained using linear models, which importantly includes linear approximations to neural networks given by the neural tangent kernel. Interestingly, show empirically that quadratic models mimic the generalization properties of neural networks when trained with large learning rate, and that such models perform better than linearized neural networks.

---

Under review as a conference paper at ICLR 2023
REFERENCES
Yu Bai and Jason D Lee. Beyond linearization: On quadratic and higher-order approximation of wide neural networks. International Conference on Learning Representations, 2019.
Lenaic Chizat, Edouard Oyallon, and Francis Bach. On lazy training in differentiable programming. Advances in Neural Information Processing Systems, 32, 2019.
Simon Du, Jason Lee, Haochuan Li, Liwei Wang, and Xiyu Zhai. Gradient descent finds global minima of deep neural networks. International Conference on Machine Learning, pp. 1675- 1685, 2019.
Stanislav Fort, Gintare Karolina Dziugaite, Mansheej Paul, Sepideh Kharaghani, Daniel M Roy, and Surya Ganguli. Deep learning versus kernel learning: an empirical study of loss landscape geometry and the time evolution of the neural tangent kernel. Advances in Neural Information Processing Systems, 33:5850-5861, 2020.
Antonio Gulli. AG's corpus of news articles, 2005. http://groups di www.piiii.gpuulll.... AG AG__orpus__oe _f news hwwsswwwwwleeeeeeeeeeeeeeeeeesee artiicassis..... html.
Jiaoyang Huang and Homg- Yau. Dynamics of deep neural networks and neural tangent hierarchy. In International Conference on Machine Learning, pp. 4542-4551. PMLR, 2020.
Arthur Jacot, Franck Gabriel, and Clément Hongler. Neural tangent kernel: Convergence and generalization in neural networks. In Advances in neural information processing systems, pp. 8571-8580, 2018.
Jakobovski. Fire-Spooen-DDiiit-aaatsss 2020. free-spoken-digit--ee
Ziwei Ji and Matus Telgarsky. Polylogarithmic width suffices for gradient descent to achieve arbitrarily small test error with shallow relu networks. International Conference on Learning Representations, 2019,
Alex Krizhevsky, Geoffrey Hinton, et al. Learning multiple layers of features from tiny images. 2009.
Yann LeCun, Léon Bottou, Yoshua Bengio, and Patrick Haffner. Gradient-based learning applied to document recognition. Proceedings of the IEEE, 86(1 86(11))2278-2344 1998.
Jachoon 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 Advances in neural information processing systems, pp. $570-8581 2019.
Aitor Lewkowycz, Yasaman Bahri, Ethan Dyer, Jascha Sohl-Dickstein, and Guy Gur-Ari. The large learning rate phase of deep learning: the catapult mechanism. arXiv preprint arXi::2003.022118 2020.
Chaoyue Liu, Libin Zhu, and Mikhail Belkin. On the linearity of large non-linear models: when and why the tangent kernel is constant. Advances in Neural Information Processing Systems, 33, 2020.
Philip M Long. Properties of the after kernel. arXiv preprint arXiv:2105.10585 2021.
Andrea Montanari and Yiqiao Zhong. The interpolation phase transition in neural networks: Memo- rization and generalization under lazy training. arXiv preprint arXiv:2007. 12826 2020.
Yuval Netzer, Tao Wang, Adam Coates, Alessandro Bissacco, Bo Wu, and Andrew Y Ng. Reading digits in natural images with unsupervised feature learning. 2011.
Eshaan Nichani, Adityanarayanal Radhakrishnan, and Caroline Uhler. Increasing depth leads to U-shaped test risk in over parapmmtmttri convolutional networks. In International Conference on Machine Learning Workshop Over-parameterizatton::: Pitfalls and Opportunities, 2021.
10

---

Under review as a conference paper at ICLR 2023
Guillermo Ortiz-Jiménez, Seyed-Mohsen Moosavi-Dezfooli, and Pascal Frossard. What can linearized neural networks actually say about generalization? Advances in Neural Information Processing Systems, 34, 2021.
Boris T Polyak. Introduction to optimization. Optimization Software, Inc, New York, 1987.
Prajit Ramachandran, Barret Zoph, and Quoe V Le. Searching for activation functions. arXiv preprint arXiv:1710.05941, 2017.
Daniel A Roberts, Sho Yaida, and Boris Hanin. The principles of deep learning theory. Cambridge University Press, 2022.
Pedro Savarese, Itay Evron, Daniel Soudry, and Nathan Srebro. How do infinite width bounded norm networks look in function space? In Conference Learning Theory, pp. 2667-2690. PMLR, 2019.
Francis Williams, Matthew Trager, Daniele Panozzo, Claudio Silva, Denis Zorin, and Joan Bruna. Gradient dynamics shallow univariate relu networks. Advances Neural Information Processing Systems, 32, 2019.
Greg Yang and Edward J Hu. Feature learning in infinite-width width neural networks. arXiv preprint 2.X5::0011.1422. 2020.
Guodong Zhang, Lala Li, Zachary Nado, James Martens, Sushant Sachdeva, George Dahl, Chris Shallue, and Roger Grosse. Which algorithmic choices matter which batch sizes? insights from noisy quadratic model. Advances in neural information processing systems, 32, 2019.
Difan Zou and Quanquan Gu. An improved analysis of training over- parameteeede deep neural networks, Advances in Neural Information Processing Systems, pp. 2053-2062, 2019.

---

Under review as a conference paper at ICLR 2023
A DERIVATION OF NQM
We will derive the NQM that approximate the two-layer fully connected ReLU activated neural networks based on Eq. (2).
The first derivative of f can be computed by:
0f Uu:
0f 1 Vm
VVa uf (vZuf(1)
1
U, 1 (uf (uf
x)
Vi E [m].
V Vmdd"
So,
And each entry of the Hessian of f. i.e. H,, can be computed by
Øʃ 0,
021 1
0,
Vmd (uf2203z²
Vi e [m].
Ou7
Now we get fquad taking the following form
NQM: squad (u, v; = f(uo, x) + Vmd
Σ i=1
{ng
to,i
vm Σ (va
(n
uo
m
iu,
uo x1 =20) (2z Uo.s)
(13)
md
B DERIVATION OF DYNAMICS EQUATIONS SINGLE TRAINING EXAMPLE
The NQM can be equivalently written as:
g(u, v; g Ug. Vo: + uo Vug(u,
vo. Vvg(u.v:2
u=w.v
0 pg(u v;
vo) (v- vo)
Ug
Oudv
u=
since 04g 0 and
8 gV
g
And the tangent kernel A(u, v: x) takes the form
8 gu, v; Juôv
X(u, v;
V Vug(u, v;
(v vv)
u=uo.v=Vo
4=
FF
Ou v;
+ Vug(u, v;
+ (u -uo))))² n0)
Oudy
v=vo
uu
Here
m
Vn,g(u, v;
Nos {u 220 [m),
lu md
Veg(u, v;
(u
md
12

---

Under review as a conference paper at ICLR 2023
In the following, we will consider the dynamics of g and À with GD, hence for simplicity of notations, we denote
Vug(0)= == Vug vi
V=Vo
Vvg(O)= Veg(u, v:
V=Vo
8g(O)
dudv = 8 (u, v;
duδv
V=Vo
By gradient descent with learning rate 17, at iteration t, we have the update equations for weights and v:
1) u(t) - n(g(t) y) (ug(()= + dg(())
audv (v(t) -v()))
v(t + 1) = - -(t())(((t)))) - 3)) Vvg(O)) + (u(t) u(0)) ( g()
Oudv
Then plug them in the expression of A(t + 1) and get
A(t + 1) [ug(O) + g2((0)
dudv (v(t + - + | + (u(t + 1) uu())² 02 9()) 2
Oudv
Vug(O) + 009(0)
dudv v(() n(9(t) v) (VrS(O) + (u(t) u()) r 00
Fudv - v()) |/3
+ Vog(O)+ (u(t) ((((t)) v) (va)(0) + 0g((0)
dudy (v(t) -(()) -u())
80
dudv
x(t)+ (7((t))) - y)² 8 (0) dudr
Vog(0) + (u(t) - u()))² 0 90)
dudv
(vag(0= + pg(())
8
g0(())
(0t(t( v)²
Oudr (v(t) - -v(0))
duôv
Vng(O) pg((0)
02((0)
Vvg(0) + (u(t) - =()))) 0.
2n( 8)
Oudv (v(t) v(0)),
Oudv
dudv
Vvg(0) + (u(t) u(0))² 029(0)
Vng(O) + 8g(0))
Judy (v(t) -(0)) ogg(0)
2n( 9)
dudv
Judv
8² [0) Judv we have
Due to the structure of
89(())
(Vrs(()) + (u(1) u())² O 9(0)
2 l|-
| 1
Vug(0) + (u(1) - u(0))² Pg(0)
Oudv
Oudv
md
Judv
| ||
[vg((t)1
md
and
(( (u)(()++ + P2(())
Judr (v(t)= v(0)) T pg)(0)
||x|
Vug(O) + pg(())
dudr (v(t) (v(t)-v(0)]
dudr
md
x 2 md
Vng(E)
13

---

Under review as a conference paper at ICLR 2023
Furthermore,
Vug(O) + 8g(()) Oudv (v(t) - v0)). 89((0)
(vrg(0) + (u(t) - u())) 7 (0)
Oudv
||||| md (v(1) v(0), Vvg(O)) + |
(Vug(0), u(t) - u())) + vn)( Pg(() Oudv Von(0))
md
( (8Pg(0)
Oudv (v(t) - -v(0), 80(())
dudv (u(t) u()))² 0 g()
dudv
||01| (v(1) v(0), Vvo(O)) + |
(Vug(0), u(() u()) + z(0) + | 12
v(1) - v(0). 0²((0)
Ōudv (u(1) -u())²)
md ==(t)]l==|///m
md
md
Similarly, we have
Vvg(0) + (u(t) u(0))² 8 Oudv'
(Tuoro)+ 8'9(0)
dudr (v(t) ((())-((1)) )²
89(())
==(t]l/=11//m
duôv
As a result,
x 2
"1 (g(() y)) 4/|2012
A(t + 1) M(t)
md n(g(t) - 4)g(t)
md
A(t) 7 ar 12 (g(() y)² (n((t) g(t) -
g(()
md
For g. we plug the update equations for u and v in the expression of g(t + 1) and we can get
g(t + 1) = g(0) + (u(t + 1) u(0), Vw)()) + (v(1 + - v(0), Vvg(O)
1) + (ult+ - -u)), 8g((0)
duδy (v(t +
1)) -v-
-(0) + (u(O) - -(((( -1) (vur()) + 8g
(v(() -v(0)) - -u(), Vuo(o))
dudv
v()) n(((t) - v) (vo(()) + (u(t) - =()))² O² ((0))
Oudv -()) FvO(0)
+ u(t)) ((((t) - 9) (vuo(0)+ + 80((0)
dudv (v(t) v(0))
u(0)
09)(0) (v(0) ((((() v) (vg(()) (u(() u()))" (0)(0)
v())
ōudv
duôv
gt) n(((t) 4) (v (ug(()) + d9((0)
dudv (v(t) v(0)), Vuo(O)
((9(() 5n (V + (u(()
8 g(0)
Vvg(0)
Oudv
+ ( 8)² (a)(()= + 89((0)
Pg(0)
(Vro(0) + (u(t) u()) Po(()
duδy (v(t) v(0),,
duδv
duōv
-(((() - v) u(t) u()), 89(())
Judv (vrg()) + (u(t) - u(0))² d 02
dudv
-ng(() 8) Vao(0)= + 0g((0)
oudr (v(t) - v()),, 0² 00)(0)
dudv (v(t) - v(0))
=((t) - n(((() - v))(()
+ ( 3)² V.g(0) + pg((0)
gg((0)
(vr((0) + (u(t) u(0))² Pa((0) Judr
Juey (v(t) - v())).
Oudv
= g(() (( v)a(t) |
md 11 (g(() 2)²g(t)
14

---

Under review as a conference paper at ICLR 2023
Therefore,
1) - y (1(1->>(0) +
md xf(o(t)-))s(t)) (o(t) - v)-
B.2 MULTIPLE TRAINING EXAMPLES
We follow the similar notation on the first and second order derivative of with Appendix B.1. Specifically, for E [n]. we denote
Vug:(0) == V Vu v3
V=
V. (0) vug(u, v;
u=
8 gh (0) 0 g(u. v; Oudv
V=Vo
By GD with learning rate y), we have the update equations for weights and at iteration t:
1
(V (Vagi(()= gg)
1) = n(() 7
Σ (t)
yn=)
(v(t) v(0))
Oudv
h=1
v(t 1) v 17 Σ Con (t) yn) (v)gx(0)+ + (u(() u())²² 0 (0)
n
dudv
We consider the evolution of ((t) first.
(ig(1+1) V Fngg(()) + (0)
Oudv (v(1 vv0), -()) (0); + Pg2(1)
(v(1+ 1)-v())))
Oudv
82 (0(
02 (0)
(ogi(0) (u(t 11)++(0)))
V.g, (0) (n(t: 1) u(D))²
Ouôr
ōudv
8 (0) dudr
Σ (ge(t) yba) (Vrge((0) + (u(t) u())) 82 (1.(0))
(ty(t)
ouer)
8 g (v(t) -v(0))
V09,(0)
Oudv
8 (0) 1 Juôv
n (t)(t) y1b)
(Vrg((0) + (u(() u(0) 82 g0)
Vugi(0) 8a((0)
(v(t) v(O))
δuôv
Qudt
(0),(0) 11 dudv
(1)((1) - ya)) (v.nx(0) + (u(1) - u(0)) 8 (9)((0)
δudv
82 (0)
Vvgh(0) + (u(t) -o)) 8.gg.gggl gx((0)
11 Ouδv
(g)(t) - -16)
duδv
(0)
(un + 8 (0)
Ouδv (v(t) -v()) Vvg.(0) + (u(() un0)) T 8 (0)
" Ouδv
(t)(t) yk))
Ōuôv
(0)
V (0g)(0) + (0)
Vvg( + (u(t) u()) r. 8 (0) tuev
1 tudv
(t) - yk)
Judv (v(t) v(0))
(0)
V (0) + g9
+
(t) - Uk)
Judy (v(t) v(0)
dudv
(7,,(0))
V (ugi(0) (0)
(g)(t) yx)
Judv (v(t) v(0))
Judv
15

---

Under review as a conference paper at ICLR 2023
We separate the data into two sets according to their sign:
S+ = (i : ai 2 0, E [m]], S_ == (i: 0, [n]]
We consider two scenarios: (1) x and xy have different signs; (2) 2, and 2j have the same sign.
With simple calculation, we get if and zj have different signs, i.., i E St.j E S_ or i E S_,j E S+. (0) 8 g, = 0, iuûv dudv
(1)
( (ug;(0)= =0,
(0) Vo, 0.
dudv
dudv
Without lose of generality, we assume i St. E S.. Then we have Kt + = -K,s(t).
(2) If a; and xy have the same sign, i.e., i,j E S+ or i,j E
(0) 0² (0) 1 ()) 099.(0)
Oudv V Vug,(0) (0) = 1
Pg.(0) Oudv
v0g 1 Vm
Vuni()
Vvgi(0)zg.
Oudv Oudv vm dudv zj
mm
For Fo St, we have
2n
Σ (t)(t) yk)zz (v)g(()) + (u(t) u(())² 8g (0)
(t+ (t)(t)
- m
Oudv
kES+
vng;(0) + 89 g0)
Judr (v(1) - v(0))
2m
(u)g((0)+ 0 g (0) dudv (v(t) v(0)), (vg + (u(t)
u(0))r (0 Oudv
(g)(t) - uk)a,
mm
KESt
112
Σ (gx(1) - yk) (0) + (u(() -(()) 02 (0)
m
audty
KES+
212
Σ (gx(t) - yk) (V + (g) (v(t) - -(()))
+ m
ouδv
NES.
4n m
Σ (gx(1)) (1)) + n La (g((1) - y) 2 m,)" (()) ((g()) - - m+)
K1,(1)
KES+
4n
(t)
2 ((g(t) yy) 3 m+) (g(t) ① m+)
m
2 n,æ (gg(t) y) m+) ) K(t) ((g(t) y) + m.).
m
Similarly, for i, E S_,, we have
41 nx, ((g(t) y) 3 m_ (g())mm))))))
Ki + = Ki,(())
+ n²
may ((g(t) y) の m ((()(((())) y)0m).
m
Combining the results together, we have
1 K(t ( K(t) ((g(t) y) o m,,)) ((() ((g(t) - y) m.) PiPl
m
2 ((g(t) - y) o m )r ((t) ((g(t) - y) o
m
4m
((g(1) y) m T (g(1) mm pPPp
4n
((g()) y) ② m 1d {g(() m p2p.
16

---

Under review as a conference paper at ICLR 2023
Now derive the evolution of g(t) y. Suppose E S+. Then we have
g:(l + 9: (u(t + 1) u()), (0)) + - Vvg:(0)) +
u(t + 1) - -u()), 0 (0)
(v(t + 1)
Oudr
Σ (gx(t) ykk) (u) gg)(()
(v(t) v(01)) vuo:(0)
= g:((
Judr
n
(v)gx()) + (u(() - u(0))² 0 (n((0)
(1) - yn))
vvgi(0)
audv
(u))(0) + 9 (0)
82 g()
(gx(t) - (/))
dudv (v(t) v(0))
(v(t) - v(0))
Oudv
(1)(1) - -ya) (o.)(()) + (u(1) - u())²² 8 g0)
(u(() - u())² 8g (0) δndv
Ondv
Σ (9) (t) -3h) (ugn(()= + gg (0)
+ x
(v(t) v(0))
Oudv
P.:(())
0 gu (0)
(1) - yk.) Vegn (0) + (u(t) - m())²²
dudv
Oudv
(9x(1)= yk) (t. 1/2
(t) 11
Σ
(ge(1) (0) y) )o; (1)
KES+
KES. JES+
Similarly, for i S_,, we have
yk (f. + o2
gi(t 1) = 1)= (t) -
Σ
(g)(t)=
Σ
Σ (ax(1) - (k)(()()) ys v))0s(()es (1)
RES
DES_fJES.
Combining the results together, we have
n² 10-1-111 (1-n((t) + ((g(t) y) o mmmm))) ) o m+)ppPT
g(t + - y
m
+ 1 (gg(t) y) 2 m_ ) (g(t) O m pppp (g(t) y).
C PROOF OF PROPOSITION 1
Restate Proposition 1: For any u, v E RT, rank(K) 2. Furthermore, P1, P2 eigenvectors of K, where P1, = x; (ieS,) ) P2,5 = x,l((ies.)) f,oo i [n].
Proof. By Definition 1,
Kijj
2 + u , 1l (1uz (wazy20p i,j E [n]. cu
k=1
17

---

Under review as a conference paper at ICLR 2023
By definition of eigenvector, see
Ki Ki
(uf + 2 (uiz) I (1azz (58,
k=
Σ
n
# (JE8))
1
Σ
(uL# (u
m
k=1
a Σ 2f
1
(mES))
Σ (vé + u (112,200)
m
j=1
where use the fact that if 2æ < 0, Kuj =0.= =0
As Pia Iil(4) and 25=1 2 (jES+ H m 2m + L (u,z,220 does not de- pend on i, we can see Pt is an eigenvector of K with corresponding eigenvalue À Ej=1 zÿ (jes+) 1 Dm ( + +2 (
The same analysis can be applied show P2 is another eigenvector of K with corresponding d = 2j=1 (JeS E kk=1 (o) - (u,,,20 20].
For the rank of K, not hard to verify that K AIPIPT + >22ppp hence the rank of K is at most 2.
D SCALE OF THE TANGENT KERNEL FOR SINGLE TRAINING EXAMPLE
Proposition 2 (Scale of tangent kernel). For any 5 E (0.1). if m 2 d (og((/6) where is an absolute constant, with probability least 1- 5, | S X(0) 3||20| /(2d).
Proof. Note that when / t= O,
x(0) = md
1
T x1 (uT,
)²+ mdd Σ (m || (1(
According to NTK initialization, for each i [m]. Vo,s NN( 1) and ug,s (O. We consider the random variable
:== -Z, x {u] 1 w20) : vo [u] #20)
it is not hard to see that Sr and & are sub-guassing since - x and vo, are sub-gaussian. Specifically, for any & 2 0.
Pílsi 2 t < P( P( 2 s 2 2cxp(-f/(21=13))
P( 2 1) S P( 2 21 S 2exp(-12/2).
where the second inequality comes from the definition of sub-gaussian variables.
Since & is sub- gaussian, by definition, 2² is sub-exponential, and its sub-exponential norm is bounded:
| ll,, I 11 SC
where C is absolute constant. Similarly we have ||11 Cl|
By Bernstein's 's inequality, for every 2 O, we have
t2
P
exp
-emin
1
max; |1E?
where e is an absolute constant.
Letting ( = m/4, we have with probability at least 1 - 2exp (-m/c),
m
3m
<
4
4
==
18

---

Under review as a conference paper at ICLR 2023
where c' = c/(AC).
Similarity, we have with probability at least 1 - 2 exx (-m/c').
m
-
-
S
3m
x 12
4
4
As result, using union bound, we have probability at least 1 - 4 exp (-m/c/)
||0||
3
3=
- x(0)
-
2d
2d
E SCALE OF THE TANGENT KERNEL FOR MULTIPLE TRAINING EXAMPLES
Proof. As shown in Proposition 1, P1 and p2 are eigenvectors of K, hence we have two eigenvalues:
(0) pf K (0)p1 Take (0) as an example:
p (()
A2(0)
p1 1
P2
(0))ppi:1 Σ
m 1 (2,20) 1 (2,20) Σ
(n + (10
iJ=1
Σ (n + vg,e) (L(m0.20)) k=1
)²
Σ
x (z,20) 4 (2,20}
1,=1
Similar to the proof of Proposition 2, we consider & := 1o (20, which is a sub-gaussian
random variable. Hence & is sub-exponential so that | I C where C > 0 is an absolute constant. By Bernstein's inequality, for every t 2 0, we have
m
-emin
P
2t
exp
Σm
1I?
max; |I
where с is an absolute constant.
Letting m/4, have with probability at least 1 - 2 eexp (-m/d),
3m
4
where d' = c/(AC).
The same analysis applies Ca 4o 1 2u well and have with probability at least 1 2 cxp((m/c))
3m
As a result, we have probability at least
4exp (-m/c"),
(1)0)plp1112
( + +g,x) (1 (11m)(0)201))
Σ
202
i=k
1,j=1
Σ N (2.20) (2,20);
3
Σ
2 (2,20) 1 (2,20)
1==1
1,f=1
19

---

Under review as a conference paper at ICLR 2023
Applying the same analysis to A2()), we have with probability - 4exp (-m/c),
m Σ
f + vH,x) (1 (1
?
Σ
2
(0)
m
i=
1==1
E
2 1 12 suys
3
Σ
12
i,j=1
6,=1
The largest eigenvalue is max (A1 (0), 22 (0)) Combining the results together, have with proba- bility at least 4 exp (-m/c) c')),
2タ M S ||(( - M.
Ei.-1 (ft
En (0,00
where M max
Σi Ei=1 (m:20)
F SCALE ANALYSIS FOR RA AND Rg
Proposition 3. Let = lg(()// /mm Assume m >> 1 and [X(I] < C' for some constani C > 0. If p 1. Le.. 19(t) mmm then
(i) (R((t) & | 1/ p² +
dym p p o(1), o(1),
(ii) 2t) (1-(0) 4n) ( p² +
A(1)+ m=1 ( p dym
1111=11302x(1) o(1).
dm
Proof of Proposition 3. According to dynamics equations, i.e., Eq. (11) and (12),
(1) |
(g(() -))((1,
md a
(o(1) y)² 11((1)= - g(1)
(x
md
g(t) y
Let o(1). Then with simple application of the triangle inequality. we have
[Ry(1) 11) 7 md
11p d p² d/m
lyl
o(1)
md
and
||ا0 (g(1) - y)² - 4||=|13n ((()((()) - v)
[[x(()]= =
mdd
md
[=>(1)||==11 - 4n) 41)12=1 1| g(t)2 + 4| ny - 2/|20| 2 yyt)
g(t) + 2 137?2?
md
md
md
(n))(( - |0=12 122 +
2n 12/ ((t) + |
n²² | 2 y)t(() o(1)
d
d/m
p
dm
Proposition 4. Let p = |g(())/ v and assume m > 1. If p E [C1, C2) for some constants C2 0, i.e. lo(t) ((mm)) then
(i) (t) E ||01|n000 | C2
L C
|x|| [i +
(ii) i 4/1, (t) E otherwise, Rt E
[ [
| 2=0
20

---

Under review as a conference paper at ICLR 2023
where c 0L1/mm).
Proof of Proposition 4. According to dynamics equations, i.e., Eq. (11) and (12),
Rt(t) || | 17 (g(() - y)g(t),
md
2 (u(1) y)²² (TA(C) - g(t) -49((t)
Rx(()= md
((t) y
Let Ci S p C2. Then with simple application of the triangle inequity, we have
| (1)((1)))))
|<=:)31/99(0) md
||=2||/1)²
| | 21,"" v md P-
(4)
0²
md
d
Then
x 112 C7
| 102" /mdd
|,02|3112
> 132 C1. Vmd
C2 S Rt({) S
C2
And
Rt(t)= = 1,7 (g(t) - -0))
4||x 12
(g() - yy
md
md
172 {(t |0212 - 41/11211
4 <| 2 21 11
1 A(t) md
p²
P
md
If (()) S 4/n, we have
Rx(t) 2 )²A(t) 4n||0
4||=||| ?ny + 2||30|| ytt)
C2
Cz.
d
/md
17 (
||= 4n|1
4||==|| my + 2||20| || y)(()
|| 1|=113r7222((t))
Rx(t) <
C7 +
C1 +
d
/md
md
If X(I) 2 4/n, have
[1x 4n|12:11
4||a=|/2y my + 2||x0| (11)))((1))
Rx(1)= N
C²
C2.
d
/md
12 4nl|
C² + 4||=||2ny + 2|| 1/722 y(t(
RA(t) <
Ci +
d
/idd
md
Picking c = max (2ll=l
|| ||=
||m]|3 Cz).
C2 +
we have the result.
/mdd
md
G SCALE ANALYSIS FOR R& AND Rg
Proposition 5. Let p := (((II//m Assume m >> and (1 S C. for some constant C > 0. 1f P 1, i.e. [[g(()] =((mm), then
(i) [Rg(t)]] < 2p2 -Iyllp/ -,a o(1)
(i) [Rx(1]] s 20/ + |lylle/ E, Sn(p2 E
21

---

Under review as a conference paper at ICLR 2023
Proof of Proposition S. According to dynamics equations for multiple training examples, ie.. Eq. (14) and (15), we have
112 (g(t)
Rg(()
(g(() m -pppp (g(t) mm) o )ppp
m
Rx
'
((g() - y) & m+) ) ((t) ((g(1) y) o m+)
m
T ((g(t) - y) o m-) > ((t) ((g(t) y) の m m
((g(1) y) の m )" (g()) 0 m piPl
4n
((g(t) y) ) (g() m pap
m
Let o(1).
With simple application of the triangle inequality, we have
(t)l - (p² + ltyllp//m) Σ =
since
| 2p S Σ 2 and ||2| s Σ 2²?
And
= + /yylp/ Σ + +8y(p² + ltllp/ oς
Proposition 6. Let := (lg(t)))t///))) Assume m > 1 and rj < (t), 1f E C2),foo some
constants C1, ie.. [g( @vvm), then
(i) [[g(())] E m P1 PC
+ ma +cmmc//p//PPP////////))1++1++++1
(ii) [Rx( E [- min (1p. P Ip (()))) - 4n/l| - max p ( - 4n/7 +
where e =
Proof of Proposition According to dynamics equations for multiple training examples, i.e., Eq. (14) and (15), we have
(t) 2 (g(t) - y) C m (g(t) m (g(t) y) (g(() @ m- m_)pppp
m
v?
Rx(()
(((()) - y) R m+ 3 ((t)( - o 0
n ((g(t) y) o m ((() (((() y) m P2p
An
((g(t) y) o mm (g(t) 3 m PIPf
4nj
((g(t) - y) o > (g(t) O m PaP.
Note that g(t) m, + g() ④ m g(t). We further denote p+ == llg(t) m+1l/ym and
P- = [lg(() m /1/ m. Then is not hard to see that 12 1² = =p² And we have
(2(()) = /2/m g() 3 m_ 2 P + "² /m/g(()) c m_ 1'2 -1 /2/mmy mm) c m )pPP
- - m x" (g(() o m P2pp
Therefore
11221 +0
< [[(((()) s max
I| || |C2)" +0
mm
22

---

Under review as a conference paper at ICLR 2023
For Rx (t), since the top eigenvalue of (t) 4nI is negative by our assumption, we have
[R](1)] S min || [| [l2ll - 4n/ -0//m)
| (t)l 2 | - -0(1//m)
H PROOF OF THEOREM AND ANALYSIS ON OPTIMIZATION DYNAMICS FOR MULTIPLE TRAINING EXAMPLES
By Eq. (5), the tangent kernel K at step t is defined as:
(t) (Vvgi(t), (t) + (t), Vug)(t))
m
((nx + (cx()))² 2,2, 20) (m (m)(0)a200) Vi.j [n].
m
Similar to single example case, the largest eigenvalue of of tangent kernel is bounded from 0:
Proposition 7. For any 8 (0.1), if m 2 d' c'log(4/8) where d' is an absolute con-
stant, with probability En (15 E. (2,
least 8, M/2
I Amax Amax(K(0)) S 3M/2 where M
n7 Ei 21-1 1 (x;20)
x7 14, (2y
max
The proof can be found in Appendix E. For the simplicity notation, given p, m R" we define the matrices and Qp.m
Kpm (1) = ((g(1) - y) m) K(I) ((g(1) y) @ m)) pp" Qp. (t) = ((g(t) y) c m) (g(t) @ m) pp"
It is not hard to see that for all t. Kp.m and Qp.m are rank-1 1 matrices. Specially, is the only cigenvector of Kp. and Qp.m-
With the above notations, we can write the update equations for g(t) and ((f) during gradient descent with learning rate rp.
Dynamics equations.
1/2 (0p.m, (1) + Qp (8)) (g(() y),
g(++1) - y I-
(14)
m
(y(())
K(t+1) K(t)+ ?? (Kprm (t) + Kprm (f)) 4n
(Qpr, (t) + 2p.mm (t)), (15)
m
m
Rx(()
where m., m_ I are mask vectors:
m 1 (1ES+))
m- (EES))-
Now ready discuss different three optimization dynamics for multiple training examples case, similar to single training example case in the following.
Monotonic convergence: sub-critical learning rates (n < 2/>x We use the key observation that when [[g(()] is small, i.e., o(Vm), and [[]]] is bounded, then || R2 (t)l and R| (t) are of
23

---

Under review as a conference paper at ICLR 2023
the order o(1) (see formal statement Proposition 5 in Appendix G). Then the dynamics equations approximately reduce to the ones of linear dynamics for multiple training examples
g(t + 1) - y= = (1 (I ((t 1)+ k(t)+ o1))
(t( + ((1) (g(t) y),
At initialization, ||g(0)] O(1) with high probability over random initialization. By Proposition 5, when [|]((]] = of m.) the optimization follows linear dynamics. By the choice of the learning rate, will have for all i. 2 0, || (!)ll 2, hence [lg[() - -yl decreases exponentially. The cumulative change on the norm of tangent kernel is o(1) since [[x((t] = o(1) and the loss
decreases exponentially hence Σ | [()]l o1) log O(1) o().
Catapult convergence: super-critical learning rates (2)A1 < 1 < minn(2/2(0),4//((0))
Restate Theorem 2: Consider training a NQM Eq. (10) with squared loss multiple training ex- amples by gradient descent. Under Assumption 1, if the learning rate is super-criticad i.e., (0) < " min((222((0,)))) 4/21) (0)), then there exist Ti, T2, Ta, T4 such that Ti < T2 T3 T4 and the training dynamics exhibits:
(i) Increasing phase: t E [i]. In this phase, [(t) = o(m). The loss grows exponentially;
both eigenvalues nearly constant iee,, (bx(t) As (0)] o(1) for k = 1,2.
(ii) Peak phase: I. E [T2, Ta]. In this phase, L(I) = ((m). For the tangent kernel:
(a) / A2 (0) S 4/21 (0), both eigenvalues decrease significantly, iee, A (1+11)) (1++1)-21()) <
0, A2 + 1) -xe(t) < O and both difference of the order (())
(b) (0) > 4/\1 only (t) decreases significantly.
(iii) Decreasing phase: t E [T', oc). In this phase, the loss satisfies C(t) = o(m) again and decreases. Both eigenvalues are nearly constant until convergence: (tx du (c0) = o(1),
for k- 1,2.
Furthermore, Ti= o(log m). T2 +(log m). T3- T2 = 0(1 and T4 m).
Proof of Theurem 2. We assume 2/)2(0) < (0), since in this scenario, the catapult phase happens in both directions Pl P2. i.e., g() has significant projection in both directions at the peak of the loss. In fact, as p1 is orthogonal to p2. the training dynamics in two directions are almost independent to each other due to the special structure of the Hessian. If instead 2/ 20) > (0), the catapult phase mainly happens in the direction P1 since in the direction of p2 the linear dynamics dominate. In this simpler setting, the analysis will be implied by our following analysis.
(i) Increasing phase. Initially llg(t) yll grows exponentially following linear dynamics by Proposi- tion 5. By the choice of the learning rate, we will have for in the increasing phase,
IpT - ([[In((()) (1)))) /Ilpi 2|2
= = - nyi = 2 - (0) + o(1) 11 1.
(22//)/2/lllppl)) > well. Therefore, ilg(t) y| increases in the direction pl and pa (If instead 2/)2(0) > 4/A, (0). pp (I-n) (4))(l)
And following the same analysis, have
I therefore llg(() (is(()yy]) -yl only grows in the direction P1). Note that when both ppp (I-nk (1) and Ip (I - nk (t(())pp (t)) |222 are smaller than 1, the loss stops increasing hence enters the peak phase.
Up to the iteration when [gg(t)]| = ol mm) the cumulative change of (t) and (2 from initializa- tion will be o1) since [g(())] increases exponentially. Specifically, for A1(t)
T2 θ (1lg())l2/m) = o(1)
(t) (z = Σ p Rx (ttpl///ll
Σ 1=0
The same analysis can be applied A²(t) well which gives (Aa(t) - >2 o(1)
Furthermore, it is not hard to see that until the loss grows to the order of ((m), the tangent kernel does not change much hence the loss keeps increasing exponentially.
24

---

Under review as a conference paper at ICLR 2023
(ii) Peak phase. We will use the key observation that when llg(t)] is large, i.e., of the order m))), | (t)] and R| (t)]l will be of the order (1) (See the formal statement Proposition in Appendix G), which can lead to the decrease of the loss.
When [t]( increases to the order of @ m since [|g(()) only grows in direction p1 and P2. we see that g(t) is mostly aligned with P1 and p2, i.e., g(t) p1/ P + g(t) P2/ P2ll = [lg(()]l-
Here by our assumption lyyl (1) which is small compared to the scale of [lg()) hence we can
omit it. Since yy is initially chosen to be in the interval (2/\1 (0), 4/>, and Ai(() X (0) + o(1) when [lg(() = o( mm), we in general have (1) < 4/1 at the begining of the peak phase. Therefore, by Proposition 6, the tangent kernel decreases significantly in direction p1 i... (t) decreases. Specifically, we have
K[(++1)pp =P] - -x(Rx(0p < <pI
where IoT RI p1 = (1) and p. (x))x < <.. The same analysis works for the direction Pa then we have P2 ((L 1)p2 < p2 K()p2, i., A2(1 +
X2()).
Similarly, as [[g(()] increases, (()l increases as well by Proposition 6. Therefore, the factor (t) + (t) on g()) decreases which slows down the increase the loss.
In general, the loss grows exponentially prior to the peak phase hence T2 = (llgg And the peak phase only lasts (1)) steps during training, i.e., [Tg T2l (1), since the decrease of X (t) and (2 is (11) at each step.
(iā) Decreasing phase. With the tangent kernel ((t) decreasing and llg(t) - yl increasing in the direction P1 and P2, the factor / - K (t) + (t) will ultimately smaller than 1 in both directions, which makes the peak phase ends since (ig(t) yl starts to decrease. Again, due to the large scale of (t)]. i.c O(ym), the tangent kernel still decreases significantly.
Similar to the single training example case, the decrease of (1 and A2(t) stops when [[g(
decreases to the order o( || (t)l o(1). In general we have ||z - nK (t)l smaller than 1 in both directions. Hence by Proposition 5, the training dynamics become linear that the loss decreases monotonically and the tangent kernel is nearly constant until convergence.
Divergence: (1) 7max (0))) Similar to the increasing phase in the catapult convergence, initially [[g(t) - y| increases in direction pi and P2 since linear dynamics dominate and the learning rate is chosen to be larger than Vesit Also, we approximately have r > 4/>z(t) at the end of the increasing phase. by similar amalysis for the catapult convergence. We consider the evolution of K (t) in the direction P1- Note that when [g(())] increases the order of ((ym). g(t) @ m, will be aligned with Pl, hence with simple calculation, we approximately have
[gg())]]|lpr (A(( 4n) > 0.
d Rx (1p1
m
Therefore, (t) increases since pf K(t 1)P; pf K())pi + (1)p > pf K(())) As a result, || - nk + Rg(()] becomes even larger which makes [|g(t) - yll grows faster, and ultimately leads to divergence of the optimization.
I SPECIAL CASE OF QUADRATIC MODELS WHEN ((x) = 0
In this section we will show under some special settings, the catapult phase phenomenon also happens and how two layer linear neural networks fit in quadratic model.
We consider one training example (xx y) with label y = and assume the initial tangent kernel X(0) ((1). Letting the feature vector of = 0, the quadratic model Eq.(3) becomes:
g(w) = 2 E(x)w.
For this quadratic model, have the following proposition:
25

---

Under review as a conference paper at ICLR 2023
Proposition 8.
With learning rate x n X(o [(x) =|=11? g(w) exhibits catapult
phase.
Proof.
With simple computation, we get
g(l + 1) (1 n)((1) + ( s(), X(L - =1==)x00 X(1) +/)))=/((((()))))))))))))))))))
1)=
4).
We note that the evolution of and À is almost the same with Eq. (11) and Eq. (12) if we regard ? 1/m. Hence we can apply the same analysis to show the catapult phase phenomenon.
It is worth pointing out that the two-layer linear neural network with input e I&d analyzed in Lewkowyez et al. (2020) that
f(U, v; = mm
UZu,
where E RT, U E mmd is special case of our model with [Vec(U)", v²] . 1/Vm and
Im 6 x
Σ = I & x!
E mmdd+m
J EXPERIMENTAL SETTINGS AND ADDITIONAL RESULTS
J.1 VERIFICATION NON-LINEAR TRAINING DYNAMICS NQMs, LE., FIGURE 3
We train the NQM which approximates the two-layer fully-connected neural network with ReLU activation function on 128 data points where each input drawn i..dd from (N-2, 1) if the label is -1 or (2, 1) if the label is 1. The network width is 5, 000.
J.2 EXPERIMENTS FOR TRAINING DYNAMICS OF WIDE NEURAL NETWORKS WITH MULTIPLE EXAMPLES.
We train two-layer fully-connected neural network with ReLU activation function on 128 data points where each input is drawn i.i.d. from ( 1) if the label is -1 or (22,, 1) if the label is 1. The network width 5, 000. See the results in Figure 5.
390
3000
10"
LR 0.20
200
LR 0.80
LR - 0.80
LR38
LA 0.38
007
LR 0.45
LR 0.60
331
(4
LR - 0.65
LR 0.80
Decone
LR 0.80
LR - 0.45
20
60
50
100
40
10
100
40
80
379
Iteration
leeraiin
Iteration
(a) Training loss
(b) Largest eigenvalue of tangent kernel
(c) Sccond largest eigenvalue tangent kernel
Figure Training dynamics of wide neural networks for multiple examples case with different learning rates. Compared to the training dynamics of NQMs, i.e., Figure 3, the behaviour of of top eigenvalues is almost same with different learning rates: when y) < 0.37, the kernel nearly constant: when 0.37 7 0.39, only X (t) decreases: when 0.39 17 < Tmax both X₃ (t) and At(t) decreases. See the experiment setting in Appendix J.2.
26

---

Under review as a conference paper at ICLR 2023
J.3 TRAINING DYNAMICS OF GENERAL QUADRATIC MODELS AND NEURAL NETWORKS.
As discussed at the end of Section 3, a more general quadratic model can exhibit the catapult phase phenomenon. Specifically, consider general quadratic model:
g w: = w ((x) + 2 >>"I(x)w E(x)w.
2
We will train the general quadratic model with different learning rates, and different y respectively, to see how the catapult phase phenomenon depends on these two factors. For comparison, also implement the experiments for neural networks. See the experiment setting in the following:
General quadratic models. We set the dimension of the input 100. We the feature vector 6(x) z x/||0|| where (O 1) ii.d. for each i E [d]. We let 2 be a diagonal matrix with Ei E [- randomly and independently. The weight parameters w are initialized by W Ia). Unless stated otherwise, 7 10 and the learning rate is set be 2.8.
Neural networks. We train two-layer fully-connected neural networks with ReLU activation function on 20 data points of CIFAR-2. Unless stated otherwise, the network width is 10 and the learning rate is set to be 2.8.
See the results in Figure
U1-23
L8-33
(A)
Nom to Homm
Milh
W With- 1a3
LRx22
LR-28
L5-
L5-02
(B)
Nenators
Sarsoon
fevadoos
(a) Loss (log scaled) vs. (b) Tangent kernel norm (c) Loss vs. learning rate (d) Tangent kernel norm 7/widthh
7/width
learning rate
Figure 6: General quadratic models have similar training dynamics with neural networks when trained with super-critical learning rates. Panel (A): experiments on general quadratic models. Smaller ry or larger learning rates lead to larger training loss at the peak. Larger learning rates make tangent kernel decrease more. Panel (B): experiments on two-layer neural networks. Larger width (corresponding to smaller ?) and larger learning rates have similar effect on the training loss at the peak and decrease of tangent kernel norm with quadratic models. Note that width or ? seems have no effect on the tangent kernel norm at convergence.
J.4 TEST PERFORMANCE OF f, fün AND fquad, I.E., FIGURE 2(B) AND FIGURE 4
For the architectures of two-layer fully connected neural network and two-layer convolutional neural network, we set the width to be 5, 000 and 1, 000 respectively. Specific to Figure 2(b), we use the architecture of a two-layer fully connected neural network.
Due to the large number of parameters in NQMs, choose small subset of all the datasets. We use the first class (airplanes) and third class (birds) of CIFAR-IO, which we call CIFAR-2, and select 256 data points out of it the training set. We use the number O and 2 of SVHN, and select 256 data points as the training set. We select 128, 256, 128 data points out of MNIST, FSDD and AG NEWS dataset respectively the training sets. The size of testing set is 2, 000 for all. When implementing SGD, we choose batch size to be 32.
For each setting, report the average result of independent
27

---

Under review as a conference paper at ICLR 2023
J.5 TEST PERFORMANCE OF f, fiin AND fquad IN TERMS OF ACCURACY
In this section, report the best test accuracy for J. Jin and Squad corresponding to the best test loss Figure 4. We use the same setting as Appendix J.4.
2-layer FC on MNIST with GD
2- 2-layer FC on AG NEWS with GD
11.66
NN
0.84
Acouracy
Quadratic
Best Test Accuracy
Linear
0.62
0.89
0.38
Critical LR
NN
6.78
Quadratic
0.24
Lingar
20
Learning rate
Learning rate
Figure 7: Best test accuracy plotted against different learning rates for Squads I and fim- Left panel: 2-layer FC on MNIST trained with GD. Right panel: 2-layer FC on AG NEWS trained with GD.
2-layer FC on CIFAR-2 with SGD
2-layer CNN on SVHN with GD
0.785
Critical
0.80
NN
Accuracy
0,750
Quadratic
Best Test Accuracy
Linear
0.79
0,771
0,760
0.78
Orifical LR
0.756
Guadratic
Linear
0.730
0.5
3.0
3.5
2.9
3,0
3.5
Learning rate
Learning rate
Figure 8: Best test accuracy plotted against different learning rates for fquads f. and fin- Left panel: 2-layer FC on CIFAR-2 trained with SGD. Right panel: 2-layer CNN on SVHN trained with GD.
2-layer FC on FSDD with GD
2-layer CNN on CIFAR-2 with GD
0.500
Cntical
0.78
NN
Accuracy
0.885
Quadratic Linear
Best Test Accuracy
Quadratic
Linear
0,580
0.75
0.575
0.75
0.570
0.905
0.74
103
1.25
150
Learning
Learning rale
Figure 9: Best test accuracy plotted against different learning rates for fquads f. and fin- Left panel: 2-layer FC on FSDD trained with GD. Right panel: 2-layer CNN on CIFAR-2 trained with GD.
28

---

Under review as a conference paper at ICLR 2023
J.6 TEST PERFORMANCE OF f, fiin AND fquad WITH ARCHITECTURE OF 3-LAYER FC
In this section, we extend our results for shallow neural networks discussed in Section 4 to 3-layer fully connected neural networks. In the way, we compare the test performance of three models, f. /in and Lquad upon varying learning rate. We observe the same phenomenon for 3-layer ReLU activated FC with shallow neural networks. See Figure 10 and 11.
We use the first class (airplanes) and third class (birds) of CIFAR-IO, which we call CIFAR-2, and select 100 data points of it the training set. We the number 0 and 2 of SVHN, and select 100 data points as the training set. We select 100 data points out of AG NEWS dataset as the training For the speech data set FSDD, select 100 data points in class and the training set. The size of testing set is 500 for all.
For each setting, we report the average result independent runs.
layer FC on
layer
Oritica
Oritical LR
0.84
NN
0.76
NN
Best Test Loss
Lineer
0.71
Lineer
0.78
Bost
0108
0.78
822
2.5
0.06
Learning
Learning
Figure 10: Best test accuracy plotted against different learning rates for fquad, f, and fin- Left panel: 3-layer FC on CIFAR-2 trained with GD. Right panel: 3-layer FC on SVHN-2 trained with GD.
3-lyyer on
3-layer
Gritical LH
Best Test Loss
0.17
Lineer
00.09
Bost
Crinial
NN
0.30
Liner
Leerning rate
Learning
Figure 11: Best test accuracy plotted against different learning rates for fquad and fli- Left panel: 3-layer FC on FSDD-2 trained with GD. Right panel: 3-layer FC on AG NEWS trained with GD.
J.7 TEST PERFORMANCE WITH TANH AND SWISH ACTIVATION FUNCTIONS
We replace ReLU by Tanh and Swish activation functions to train the models with the setting Figure 4. We observe the same phenomenon as we describe in Section 4.
29

---

Under review as a conference paper at ICLR 2023
layer
layer
Critical un
NN
0.0003
111000)
0,0373
Best Test Loss
Lineer
000353
0,0123
010002
MN
Linar
0/857
0.067
Learning
Learning
(a) Swish activation function
(b) Tanh activation function
Figure 12: Best test loss plotted against different learning rates for fquad> f, and fiin We choose 2-layer FC as the architecture and train the models on AG NEWS with GD.
2-layer FC on FSDD with GD
2-aayer FC FSDD with GD
Critiaal
Qusoratic
Loss
0/773
Beat Test
0.072
Criscal LR
0.002
Quadratic
20
100
1.30
40
601
10
101
Learning rate
Learning rate
(a) Swish activation function
(b) Tanh activation function
Figure 13: Best test loss plotted against different learning rates for fquads fa and fim- We choose 2-layer FC the architecture and train the models on FSDD with GD.
2 2eleyer CNN on CIRAR with GD
2. 2-layer CNN on CIFAR with GD
0.56
0.83
Best Test Loss
0.90
11.33
0.75
Criticl LA
Criical
Quadratic
Cuadratic
0,70
Linaar
Linear
Leerning rate
Learning rate
(a) Swish activation function
(b) Tanh activation function
Figure Best test loss plotted against different learning rates for Squads f, and fim- We choose 2- 2-layer CNN the architecture and train the models CIFAR-2 with GD.
30

---

Under review as a conference paper at ICLR 2023
2-ayer CNN SVHN with GD
2-layer CNN on SVHN with GD
1/340
Best Teat Loss
Teat Less
0,233
Critial
Criscll Ln
Quadratic
Quadratic
1/013
15
201
25
3/9
3.5
+31
6.5
Learning rate
Learning rate
(a) Swish activation function
(b) Tanh activation function
Figure 15: Best test loss plotted against different learning rates for fquads fr and fim- We choose 2-layer CNN as the architecture and train the models on SVHN with GD.
31


Reviews:
[{"review_id": "savugbNQT2", "reviewer": "Reviewer_Ctfq", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 5, "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": "This article analyzes the dynamics of GD for quadratic approximations to sufficiently wide one layer ReLU network in which the network function is replaced by its second order Taylor expansion with respect to the parameters. The main result of this article is that for rather simple data (either one datapoint, or multiple datapoints in 1D), unlike linear approximations, quadratic approximations exhibit the so-called catapult phase. Specifically, for moderately large learning rates the training loss is unimodal, first rapidly growing and only then decaying, and the top eigenvalue of the NTK becomes decreases significantly by the end of training.\n\nUPDATE AFTER REBUTTAL AND READING OTHER REVIEWS\nI share the concerns of the other reviewers regarding the rigor of the proof of the main results. Originally I was kind of OK with it because I saw that I could fill in the details myself. But ultimately I when I read the other reviews and looked again, I became of the opinion that I quite like the results but that a paper of this kind should really have a higher standard for providing precise mathematical details in the proofs. So I am revising my assessment of this paper down from an 8 to a 6. ", "summary_of_the_review": "This article gives a novel analysis of GD with moderate learning rates in a second order approximation to wide one layer ReLU networks. The main result of this article is that for simple data (either one datapoint, or multiple datapoints in 1D), unlike linear approximations, quadratic approximations exhibit the so-called catapult phase. As mentioned in the strengths section, it is rare to give a more or less complete characterization of GD dynamics in non-linear models at relatively large learning rate. I think the basic observations and techniques in this article will be useful for further developing theory and also that the results themselves are interesting.", "strength_and_weaknesses": "Strengths:\n\n- It is rare to be able to more or less fully analyze any kind of non-linear learning dynamics. That the authors manage to do this, even if it is for somewhat stylized datasets, in a regime of large learning rate is really nice.\n- The fact that the top eigenvectors for the NTK and the Hessian provably align in the catapult phase is an interesting and important contribution. Empirically, this is often known to happen but as far as I am aware this is the first time this has been proved to occur in a non-linear model. \n\nWeaknesses:\n\n- The restriction to very simple datasets, either with one example, or inputs in 1D is quite severe. While the authors write “The assumption of d = 1 is common in the literature and simplifies the analysis for neural networks,” this assumption drastically simplifies (at least as far as I know) the structure of the NTK, as evidenced by Proposition 1. Perhaps the authors could comment on what is expected to occur in higher dimensions.\n- To a lesser extent, the restriction to one layer networks is undesirable. Perhaps the authors could comment on what is expected to occur in deeper networks.\n- There is a slightly strange apparent inconsistency between the theory developed in this article and the original work on the catapult phase. Specifically, figures 3 and 4 in that article seem to show that for shallow ReLU networks the catapult phase is delineated approximately by taking  the learning rate to be between $2 / \\lambda_{max}$ and $12 / \\lambda_{max}$. However, the authors here seem to prove that the largest convergent learning rate is $4 / \\lambda_{max}$. Of course this is not a direct contradiction: quadratic models are not necessarily high enough fidelity to capture the constants controlling the catapult phase. Nonetheless, it would be nice if the authors could comment on this point. ", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "4: The contributions are significant, and do not exist in prior works.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "Hdn9qGb5Pbj", "reviewer": "Reviewer_JdaH", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": 2, "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": "The authors present the theoretical analysis of a neural quadratic model as an approximation of a two-layer neural network.\nThe neural quadratic model exhibits an interesting non-linear aspect of training dynamics that can be seen with two-layer neural networks in practice: The catapult phase.\nThe catapult phase occurs when the learning rate is high enough to be beyond locally critical but low enough so as to not globally diverge and leads to a jump into a potentially better optimum with lower\nneural tangent Kernel norm.\nThe authors' theoretical analysis makes use of insightful estimation possible due to considering the wide network limit but the theoretically predicted phenomena are relevant in practice and the behavior in training loss and tangent kernel largest eigenvalue are also demonstrated using experiments. ", "summary_of_the_review": "I recommend to accept the paper as I perceive it to analyze an interesting practical phenomenon of neural networks training using and extending the toolbox of modern theoretical analysis tools for neural\nnetwork training dynamics.", "strength_and_weaknesses": "Strengths:\n- The paper analyzes an interesting practical phenomenon using powerful theoretical tools that will enable better understanding of training dynamics for neural networks\n- Having the main theoretical analysis demonstrated for the single training example case first helps make it more accessible\n- The empirical validation in Figure 3 backs up the theory\n\nWeaknesses:\n- NQMs appears in the figure caption for Figure 1 (b) which is referenced before the abbreviation (for Neural Quadratic Model) is first introduced in the text\n- Does the theoretical analysis in principle extend for squared to cross-entropy loss? (Makes the dynamics less elegant but would still be interesting to know how this could work)", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "4: The contributions are significant, and do not exist in prior works.", "empirical_novelty_and_significance": "4: The contributions are significant, and do not exist in prior works."}, "extra_text": {}}, {"review_id": "AKgiHFwDjod", "reviewer": "Reviewer_uqyZ", "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": "This paper proposes an approximation of Shallow ReLU networks in the large width regime to study their optimization dynamics. More particularly, they show that for uni-dimensional inputs, this approximation exhibits the 'catapult phase phenomenon', which corresponds to a regime of learning rate larger than the largest step size advocated by the theory on linear dynamics. ", "summary_of_the_review": "This paper provides some exciting contributions. However, I have some concerns regarding the formalism of the proof. The results feel correct to me, but, at this stage, the paper only contains sketches of proof. Moreover, I am slightly uncomfortable with how a critical assumption is presented. This critical assumption could be missed by quickly reading the paper. It should be mentioned earlier in the text, and it should be more visible. \n\nFor all these reasons, I am currently evaluating this paper under the acceptance threshold. However, if the authors address my concern, I am eager to significantly increase my score. \n\n", "strength_and_weaknesses": "## Strengths:\n- This paper tries to tackle a difficult problem which is to understand the non-linear dynamics of learning neural networks. \n- The model they are proposing is simple enough to be analyzed but complex enough to model non-linear learning dynamics that seem to correspond (to a certain degree) to what happens with two-layers ReLU networks. In particular, it exhibits a 'catapult phase' for the step-size.\n- The intuition of the main elements of the proofs is well communicated. \n- The experiments back up the theory pretty well.\n\n## Weaknesses:\n- My first main complaint would be regarding the critical assumption regarding the fact that the input has to be of dimension 1. This is very restrictive and relatively buried in the paper (see my comments in the section about clarity)\n- My second main complaint would be with the proof of the main theorem. (In particular, Theorem 1). What I  thought was a proof sketch was actually supposed to be the complete proof (please let me know if I missed the detailed proof). The statement in the \"proof of Theorem 1\" (page 6 and 7) are not formal enough to be verified. I give a very good idea of the structure of the actual proof, but it is not enough. We cannot formally have $g(t) = o(\\sqrt{m})$ for $t<T_{incr}$ and then $g(T_{incr }) = \\Theta(m)$ (even though I understand the intuition).\n- The argument that $\\sum_{t=0}^{T_{incr}} \\Theta(g(t)^2/m) = o(1)$ is not 100% clear to me since $g(t)$ will be close to $\\Theta(m)$ when $t$ is close to $T_{incr}$, moreover even if we had $g(t)^2= o(\\sqrt{m})$ we would have $\\sum_{t=0}^{T_{incr}} \\Theta(g(t)^2/m) = o(T_{incr}) = o(\\log(m))$.\n", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "2: Several of the paper’s claims are incorrect or not well-supported.", "technical_novelty_and_significance": "4: The contributions are significant, and do not exist in prior works.", "empirical_novelty_and_significance": "4: The contributions are significant, and do not exist in prior works."}, "extra_text": {}}, {"review_id": "1HRvDh5VJ8L", "reviewer": "Reviewer_LYg1", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 2, "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": "This paper studies the learning dynamics of deep neural networks using the quadratic models.  Specifically, the \"catapult phase\" is identified under this model when the learning rate is high, while it is not for linear models. Numerical simulations are provided to support the theoretical analysis. \n\n\n\n", "summary_of_the_review": "A solid theoretical work on the analysis of the learning dynamic of NN by using the quadratic models, as opposed to previously used linear models. The most interesting result is that, it  successfully captures the interesting \"catapult phase\" from Lewkowycz et al. (2020), while linear networks fail. This demonstrates the usefulness of the quadratic models in understanding NNs. The only concern is that the current analysis is restricted to the output dimension d=1 and it is not clear whether the results or methods still apply to high-dimensional outputs, or even deeper hidden layers. \n\n", "strength_and_weaknesses": "Strenghths:\n\n1. It is a solid theoretical work on the learning dynamics of deep neural networks. The adopted quadratic model successfully captures the interesting \"catapult phase\" from Lewkowycz et al. (2020), which provides insights and thus help understanding of the sophisticated dynamics of NN. \n\n2. Numerical results support their theoretical analysis. \n\nWeaknesses:\n\n1. The output dimension d is limited to dimension 1 and it is unclear what happens in the case d>1, and whether similar results still hold, which I think is the biggest weakness, since in high-dimensional case, some counter-intuitive phenomena might happen. \n\n", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}]

Author rebuttal:

Decision: Reject