Quadratic models for understanding catapult dynamics of neural networks
Authors: Libin Zhu
Chaoyue Liu
Adityanarayanan Radhakrishnan
Mikhail Belkin
Venue: ICLR 2024 poster (2024)
Keywords: quadratic models
wide neural networks
catapult phase
optimization dynamics
Primary area: optimization
TL;DR: Quadratic models capture properties of wide neural networks in both optimization and generalization.

Abstract: While neural networks can be approximated by linear models as their width increases, certain properties of wide neural networks cannot be captured by linear models. In this work we show that recently proposed Neural Quadratic Models can exhibit the "catapult phase" 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.

Published as a conference paper at ICLR 2024
QUADRATIC MODELS FOR UNDERSTANDING CATAPULT DYNAMICS OF NEURAL NETWORKS
Libin Zhu 12 Chaoyue Liu?, Adityanarayanan Radhakrishnan' Mikhail Belkin
1 Department of Computer Science, UC San Diego
2 Halicioolu Data Science Institute, UC San Diego Harvard & Broad Institute of MIT and Harvard [ (libinzhu, ch1212, mbelkin (ucssd edu 3 aaradha edu
ABSTRACT
While neural networks can be approximated by linear models as their width in- creases, certain properties of wide neural networks cannot be captured by linear models. In this work we show that recently proposed Neural Quadratic Models can exhibit the "catapult phase" (Lewkowycz et al., 2020) that arises when train- ing such models with large learning rates. We then empirically show that the behaviour of neural quadratic models parallels that neural networks in general- ization, especially in the catapult phase regime. Our analysis further demonstrates that quadratic models can be 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 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, as is shown in several recent works (Yang & Hu, 2020; Ortiz-Jiménez et al., 2021; Long, 2021; Fort al., 2020). In this work, we 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.c. there exists a critical learning rate, Merit. such that the loss converges monotonically if and only if the learning rate is smaller than Terii (see Figure la).
Recent work Lee et 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 fün (w; æ)::
fiin (w; x) = f(Wo; x) + (w wo) T V (w)
where V [ m) denotes the gradient2 of with respect trainable parameters
ini-
tial point wa and input sample æ. This approximation holds for learning rates less than Meritt Ro 2/1 f x) )) when the width sufficiently large.
However, the training dynamics of finite width neural networks, f. can sharply differ from those of línear models when using large learning rates. A striking non-linear property of wide neural
Unless stated otherwise, always consider the setting where models trained with squared loss using gradient descent.
FFor non-differeee differentiable functions, neural networks with ReLU activation functions, we define the gra- dient based on the update rule used in practice. Similarly, we use H, to denote the second derivative of in Eq. (2).

---

Published as a conference paper at ICLR 2024
Linear Dynamics
Non-linear Dynamics
Convergence
Divergence
Convergence
Divergence
Catapult
Monotonic
Phase
erit
Learning rate
7crit
max
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, funitely wide neural networks and NQMs Eq. (2) (or general quadratic models Eq. (3)) can additionally observe catapult phase when Morit < <
Mmax
Linear Dynamies
Catapult Dynamies
Parameler space
0.55
Lincee Dysumis
Casspuli Dynamics
Loss
Bost Test Lawe
Heradoo
Ittration
0,700
We
[K]]
[KK
Constant 'anngent Karnel B(wy R)
0.00
Bexit
Globel minina
Iteration
Optimisation path ouper
Sub-eriiedl LRa
Iieration
Learning
20
(a) Optimization dynamics for f (wide neural networks): linear dynamics and catapult dynamics.
(b) Generalization performance for
f. fun and fquad-
Figure 2: (a) Optimization dynamics of wide neural networks with sub-critical and super- critical learning rates. With sub-critical learning rates (0 < 7) < Aeritt) 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(wo, R) == f ||w woll < R} with a finite radius R. With super-critical learning rates (Meitt < 1) < 7max)) the catapult phase happens: the loss first increases and then decreases, along with decrease of the norm of the tangent kernel The optimization path goes beyond the finite radius ball. (b) Test loss fquads f and Jü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 fquad have smaller best test loss than the with sub-critical learning rates. Experimental details are in Appendix N.5.
networks discovered in Lewkowyez et al. (2020) that when the learning rate larger than Merit but smaller than certain maximum learning rate, Mmax gradient descent still converges but experi- ences a "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 1b).
As linear models cannot exhibit such 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, a theoretical understanding of this phenomenon for general non-linear neural networks remains open. In this work, we utilize a quadratic model as tool shed light on the optimization and generalization discrepancies between finite and infinite width neural networks. We define Neural Quadratic Model (NQM) by the second order Taylor series expansion of f (w; a)) around the point Wg:
NQM:
Squaa (w) = f(wo) + (w wo) (V)o))((W) + [w
H,, - wo).
(2)
Here in the notation we suppress the dependence the input data æ, and Hy(wg) is the Hessian of with respect to evaluated Wo- Note that fquad (w) fim (w) + (ww- wo) (wo)()) wo).
2

---

Published as a conference paper at ICLR 2024
Indeed, we note that NQMs are contained in a more general class of quadratic models:
General Quadratic Model: x(w;)) = w 6(x) + 2 2 (rw"((/)ww
(3)
where are 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 N.4. Note that the two-layer linear network analyzed in Lewkowyez et al. (2020) is special case of Eq. (3), when (xx) = O (See Appendix M).
Main Contributions. We prove that NQMs, fquad, which approximate shallow fully-connected ReLU activated neural networks, exhibit catapult phase dynamics. Specifically, we analyze the opti- mization dynamics of fquad by deriving the evolution of /quad and the tangent kernel during gradient descent with squared loss, for a single training example and multiple uni-dimensional training ex- amples. We identify three learning rate regimes yielding different optimization dynamics for fquad. which are (1) converging monotonically (linear dynamics); (2) converging via a catapult phase (cat- apult 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, our 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), work focuses 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 dynamies. 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. A recent work showed the existence of the catapult phase in two-layer, homogenous networks Meltzer & Liu (2023).
Discontinuity in dynamics transition. In the ball B(wo, R) := {w ||w woll S R} with constant radius R > O, the transition to linearity of a wide neural network (with linear output layer) 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, ie. Hy, which scales with m (Liu et al., 2020):
ll((w) füaa(w (w)ll < sup WEB(Wo.R)
[H R = Ō(1vmm)
(4)
Using the inequality from Eq. (4), we obtain []. - fiall = 0(1 (m), hence fquad transitions to linearity continuously as well in B(wo, R) Given the continuous nature of the transition to linearity, one may expect that the transition from non-linear dynamics to linear dynamics for f and fquad is continuous in m as well. Namely, one would expect that the domain of catapult dynamics, Meritt "maax] shrinks and ultimately converges to a single point, i.e., /erit 7maxx as m goes to in- finity, with non-linear dynamics 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 ?). Additionally, the length of the optimization path of f in catapult dynamics grows with m since otherwise, the optimization path could be contained in ball with a constant radius independent of m, in which can be approximated by fin- Since the optimization of fíin diverges in catapult dynamics, by the ap- proximation, the optimization of f diverges as well, which contradicts the fact that the optimization of f can converge in catapult dynamics (See Figure 2a).
'For general quadratic models in Eq. (3), the transition lincarity is continuously controlled by
3

---

Published as a conference paper at ICLR 2024
2 NOTATION AND PRELIMINARY
We use bold lowercase letters to denote vectors and capital letters to denote matrices. We denote the set 2,, 2,. , by [rx]. We use -I Il to denote the Euclidean norm for vectors and the spectral norm for matrices. We use @ to denote element-wise multiplication (Hadamard product) for vectors. We use Amax (A) and Amin (A) denote the largest and smallest eigenvalue of matrix A, respectively.
Given a model ((:: ar), where æ is input data and w are model parameters, we use / to represent the partial first derivative 0f(w: /0w When clear from context, let V == / for of notation. We use H and He to denote the Hessian (second derivative matrix) of the function fww x) and the loss C(w) with respect to parameters w, respectively.
In the paper, we consider the following supervised learning task: given training data {(z;, yi )== with data 2; E Rd and labels y/ E R for i E [n]] we minimize the empirical risk with the squared (w;
loss L(w) マ Σi= (J ..) ya)². Here ((w;:) is parametric family of models, e.g., neural network or a kernel machine, with parameters w E IR.. We use full-batch gradient descent to minimize the loss, and we denote trainable parameters w iteration by w(t). With constant step
size (learning rate) the update rule for the parameters is:
dC(w) (1), VI 0.
w(t 1) w(1) 11 dw
Definition (Tangent Kernel). The tangent kernel K (w):s)) of f(w; -) is defined as
K K(w; 8. z) (V](w; V[(w; z)), Va, z E Rd.
(5)
In the context of the optimization problem with n training examples, the tangent kernel matrix K E R=x= satisfies Kij (w) = K (w; æz), E [n]. The critical learning rate for optimization given follows.
Definition 2 (Critical learning rate). With an initialization of parameters Wo. the critical learning rate of f( -w:)) is defined as
Mexis == 2/ 2maxxx (Hc(W)))
(6)
A learning rate yj is said to be sub-critical if < < qerit or super-critical if Marit < 1) < Tmmax- Here Tmax is the maximum leaning rate such that the optimization of C(w) initialized at Wo can converge.
Dynamics for Linear models. When f is linear w, the gradient, Vʃ, and tangent kernel are constant: K( K(w(t)) = K(wo). Therefore, gradient descent dynamics are:
F(w(l 1)) y (T n (wo) where F(wo) = [f (wo), .... (wo)) with f:(wo) f(wo: ..))
y): 14 2 0,
Noting that (wo) (F((w)) FWW))) and that tangent kernel K K(Wo) F (wo) WF))))())) share the same positive eigenvalues, we have Amaas (Wo(w)) = Amax [K (Wo)), and hence, Meriit = 2/Amax (ommax(((vr)) (W()))
(8)
Therefore, from Eq. (7), if < <7 < Merit. the loss L decreases monotonically and if Merit.. the loss L. keeps increasing. Note that the critical and maximum learning rates are equal in this setting.
OPTIMIZATION DYNAMICS IN NEURAL QUADRATIC MODELS
In this section, we analyze the gradient descent dynamics of the NQM corresponding to a two-layer fully-connected neural network. We show that, unlike a linear model, the NQM exhibits catapult dynamics: the loss increases at the early stage training then decreases afterwards. We further show that the top eigenvalues of the tangent kernel typically become smaller consequence of the catapult.
Neural Quadratic Model (NQM). Consider the NQM that approximates the following two-layer neural network:
1
1
f(u, v; =
Uio
T
(9)
m
4

---

Published as a conference paper at ICLR 2024
where u, E R 2/2 R for i E [m] are trainable parameters, z R is the input, and of-) is the ReLU activation function. We initialize u, ~ (o0 Ia) and v E Unif[[-1,1] for each i independently. Letting g(u, v: x) == fquaa (u, v; x)), this NQM has the following expression (See
the full derivation in Appendix A):
g(u, v; = fuu, vo + md
(u uo,i) a:ll {uā [ug +
/md Σ (2s v0 (u)
v vo "o.s)"==(("i
x1 ( m20
(10)
Vmd
Given training data fær, ye 1 we minimize the empirical risk with the squared loss L(w) -i= (g( (:; (2:))))))0 yy using GD with constant learning rate 1/. Throughout this section, we denote g(u(t), v();:)) by g(t and its tangent kernel K (u(t), v()) by K(t), where is the iteration of GD. We assume ||> (1) and ly.1 = (1) for i e [n]] and we assume the width of f is much larger than the input dimension d and the data size i.c., m > max(d, nt.. Hence, d and n can be regarded as small constants. In the whole paper, we use 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.
3.1 CATAPULT DYNAMICS WITH A SINGLE TRAINING EXAMPLE
In this subsection, we consider training dynamics of NQM Eq. (10) with single training example (a, y) where æ E Rd and y E IR. In this case, the tangent kernel matrix K reduces to a scalar, and we denote K by À to distinguish it from matrix.
By gradient descent with step size 7,, the updates for g(1) -y y and X(1), which we refer to as dynamics equations, can be derived as follows (see the derivation in Appendix B.
Dynamics equations.
| md "i (g(()) w)9(1)
glt + 1 1
m)(() 1-n2(1-
(9(1) - v) = x()(((l)) -)),
(11)
Ry(()
X(1 *(+1) x()) 1 ||00|
g(()
md (o(1) - w)² 4 g(1) y
mx(1)
14 2 O.
(12)
(1)
Note that as the loss is given by L(t) (g(() 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 (l) and RA (1). For example, for constant learning rate that slightly larger than Merie (which would result in divergence for linear models), (t) stays positive during training. resulting in both monotonic decrease of tangent kernel À and the loss.
As A(t) = Au A'= (T),). to track the scale of [p()),, we will focus on the scale and sign of Rgy() and R, (t) in the following analysis. For the scale of Xg which is non-negative by Definition 1, we can show that with high probability over random initialization, Aaal = @(1) (see Appendix 1). And lg(0)] O(1) with high probability well Du et al. (2018). Therefore the following discussion is with high probability over random initialization. We start by establishing monotonic convergence for sub-critical learning rates.
Monotonic convergence: sub-critical learning rates (n 2/2o Ment). The key observation is that when g(t) = O(1), and X(t) = @(1), [t) (t)) and [t)l are of the order o(1). Then, the dynamics equations approximately reduce to the ones of linear dynamics:
g + 1) - y = (1 - -)t + o(1)) (g(t) - v). A(t + 1) x(t) o(1)
5

---

Published as a conference paper at ICLR 2024
Note that at initialization, the output satisfies |g(0)] ( and we have shown Au = ((1). With the choice of 77, we have for all t 2 0. [x(t))
n)((t) + o(1) < 1; hence, (g(t) y/ decreases monotonically. The cumulative change the tangent kernel will be o(1), i.e., E, (1) = o11), since for all I, (t)
((/m) and the loss decreases exponentially hence Σ (t) O(1 (0/m) log O(1) o(1). See Appendix C for detailed discussion.
Catapult convergence: super-critical learning rates (7erit = 2/A0 < 7 < 4/A0 = Tmaxx). The training dynamics are given by the following theorem.
Theorem (Catapult dynamics
single training example). Consider training the NQM Eq. (10) with squared loss on a single training example by GD. With a super-critical learning rate E the catapult happens: with high probability over random initial-
24 1-c
where c= 0 log ( ).
Ao
Ao
ization, the loss increases the order of
minrAg -2)" then decreases to O(1).
log m
Proof of Theorem I. We use the following transformation of the variables to simplify notations.
|
md n² (g(t) v)², w(t) ||=||
a(t)
md " (g(() y)y, v(t) m)((t).
Then the Eq. (11) and Eq. (12) are reduced to
u(t + 1) = (1 v(t) + u(t) + (t)) = x(t)u(t),
(13)
u(t 1) (tt) - (a))( vtt) 4w(().
(14)
At initialization, since |(0) = ( we have u(0) = O ( and w(0) 0 (+)). Note that by
definition, for all 2 0. u(t) 2 and we have v(t) 2 since A(t) is the tangent kernel for a single training example.
In the following, we will analyze the above dynamical equations. To make the analysis more under- standable, we separate the dynamics during training into increasing phase and decreasing phase. We denote == (1 Merit) = 11 - 2.
Increasing phase. In this phase, [u(4) increases exponentially from (++) to θ
g²
and
log
8 (u(t) v00] O logg ).) This can be shown by the following lemma.
u(() increases exponentially with infrejo.r) k((t) 2 1 - 0 log (m)mm) > 1 and supre(o.ry (o(t) - e(0)] = = log
Lemma 1. For T > 0 such that supre[o,T] ul() = log
Proof. See the proof in Section D.
After the increasing phase, based on the order of u(t) we can infer the order of loss is θ
mn52
log
Decreasing phase. When u(t) is sufficiently large, v(t) will have non-negligible decrease which leads to the decreasing of s(t), hence turn making u(f) decrease as well. Consequently, we have:
Lemma 2. There exists T 0 such that u(T')) O (1)
Proof. See the proof in Section E.
Then accordingly, the loss is of the order O(1).
Once the loss decreases the order of ((1), the catapult finishes and we in general have 1) < 2/A(() where |Rg (t)) O(C(t)/m) O(1/m). Therefore the training
as [t(t) [1) RR (t)l < dynamics fall into linear dynamics, and we can use the same analysis for sub-critical learning rates for the remaining training dynamics. The stableness of the steady-state equilibria of dynamical equations can be guaranteed by the following:
Theorem 2. For dynamical equations Eq. (11) and (12), the stable steady-state equilibria satisfy g(t) = y (ose,,oss is 0), and X(t) [c.2/n c) with Olog m/ /m).

---

Published as a conference paper at ICLR 2024
Divergence (1) 1/uax 4/).). Initially, it follows the same dynamics with that in the increasing phase in catapult convergence: (g(t) - y) increases exponentially as |p(t)] > and the X(t) almost does not change as (t) is small. However, note that R2 (t) > 0 since 1) g(t) -0) y) = 1 when g(4) becomes large and 2) y > 4/A(1). Therefore, X(t) keeps increasing during training, which consequently leads to the divergence of the optimization. See Appendix G for detailed discussion.
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 et al. (2019)).
Assumption 1. The input dimension d = and not all is 0, i.e., Σ | > 0.
We similarly analyze the dynamics equations with different learning rates for multiple training ex- amples (see the derivation of Eq. (22) and (23) in Appendix) which are update equations of g(() yy and (4). And similarly, we show there are three training dynamics: monotonic convergence, cata- pult convergence and divergence.
In the analysis, we consider the training dynamics projected to two orthogonal eigenvectors of the tangent kernel, i.e., P1 and P2. and we show with different learning rates, the catapult phase can occur only in the direction of Pt or occur both directions. We consider the case where 2/A2(0) < 4/)1(0) hence the catapult directly obtained from our results. We denote the loss projected to P; by I1,L := 1/ (g y. p,) for i = 1,2. We have (() = O()) with high probability over random initialization of weights.
can be occur in both directions. The analysis for the other
We formulate the result for the catapult dynamics, which happens when training with super-critical learning rates, into the following theorem, and defer the proof of it and the full discussion of training dynamics to Appendix K.
Theorem 3 (Catapult dynamics on multiple training examples). Supposing Assumption holds, consider training the NQM Eq. (10) with squared loss on multiple training examples by GD. Then, with high probability over random initialization we have
1. with 17
E
2+c 2- (x(
the catapult only occurs in eigendirection P1 11,C increases to
Xa(O)
the order of
(0)n)1(0)-21)
then decreases to O(1);
log m
2. with
2tc A2
the catapult occurs in both eigendirections P1 and P2: II, for
X1
1,2 increases to the order of m(nn:(()) 2)²
then decreases to O())
log 77
where e == θ Iog m
vmn
We verify the results for multiple training examples via the experiments in Figure 3.
10"
350
laggen egeneabie
10P
Tralking loso
LR 0.20
LR -..0
UR 0.20
LR 0.38
LR 0.58
LR 1R5-0.38 0.38
bot
0.60
UR - 0.60
LR - 0.60
50
lteration
(a) Training loss
(b) Largest eigenvalue of tangent
(c) Sccond largest eigenvalue of tangent kernel
kernel
Figure 3: Training dynamics of NQMs for multiple examples case with different learning rates.
By our analysis, two critical values are 2/ ( (0) = 0.37 and 2/ (0) = 0.39. When 1j < 0.37, linear dynamics dominate hence the kernel is nearly constant; when 0.37 17 0.39, the catapult phase happens in and only X₁ (t) decreases; when 0.39 < < Tmaxx the catapult phase happens in P1 and p. hence both d (t) and Az(t) decreases. The experiment details can be found Appendix N.I.

---

Published as a conference paper at ICLR 2024
3.3 CONNECTION TO GENERAL QUADRATIC MODELS AND WIDE NEURAL NETWORKS
General quadratic models. As mentioned in the introduction, NQMs are contained in general class of quadratic models of the form given in Eq. (3). Additionally, we show that the two-layer lin- neural network analyzed in Lewkowycz al. (2020) is a special case of Eq. (3), and we provide more general condition for such models to have catapult dynamics in Appendix M. Furthermore, we empirically observe that broader class of quadratic models g can have catapult dynamics simply by letting <xx) and V be random and assigning a small value to 7 (See Appendix N.4).
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. We will discuss the high-level idea of analyzing the catapult phase for a general NQM with large learning rates, and empirically show that this idea applies to neural networks. We train an NQM Eq. (2) fquad on n data points {(x, v:) IR R with GD. The dynamics equations take the following form:
1 I I-nK(t) + 2
fqudd (t + 1) =
(Equad (t)-)) y))
(15)
Reyua (8)
1
K(L+ 1) = 1) K(()
4" (AG(1) equad (1) (((())))))
(16)
(x(
where (t, = (fquad(t) y'" (Quuuaa)))))
E 1= for E [n]]
In our analysis for fquad which approximates two-layer networks in Section 3.2, show that cata- pult dynamics occur in the top eigenspace of the tangent kernel. Specifically, we analyze the dynam- ics equations confined the top eigendirection of the tangent kernel p, (i.e, II, L and x (t)). We show that p Ri and p RxPi scale with the loss and remain positive when the loss becomes large, therefore Pf Kp, (i.c., Amax (K)) as well as the loss will be driven down, and consequently we yield catapult convergence.
We empirically verify catapults indeed happen in the top eigenspace of the tangent kernel for addi- tional NQMs and wide neural networks in Appendix N.3. Furthermore, 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 in Appendix N).
4 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, fquad: (2) corre- sponding neural networks, f: and (3) the linear model, flin-
We implement our experiments on 3 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) and text dataset: AG NEWS (Gulli).
In all experiments, we train the models by minimizing the squared loss using standard GD/SGI with constant learning rate 7). We report the best test loss achieved during the training process with each learning rate. Experimental details can be found in Appendix N.5. We also report the best test accuracy in Appendix N.6. For networks with 3 layers, see Appendix N.7. 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 B (Wo R):
|f fuall 0(L//m), Ilfquad - fuall Ó(1//m)
where R > 0 is constant which is large enough to contain the optimization path with respect to sub-critical learning rates. Thus, the generalization performance of these three models will be similar when m is large, as shown in Figure 4.
8

---

Published as a conference paper at ICLR 2024
FC FSCD
Blaayer NEWS GD
NNIST
(000
0110
Bast beet Lose
Bost Loos
8072
Critiaa
006
OfficaLn
CrticiL
Learning
Learning
Leariig (ele
Blayer FC on CIFAR-2 SGD
Bllayer Chil on CIFAR-2 with
2layer Chi on SVHN with
0007
Bal Teat Loee
Griical Ln
Begt Test Loce
Best
NB
MN
Quediatk
Quashalic
Uneer
Linour
Learnig
Learnig
Learring rele
Figure 4: Best test loss plotted against different learning rates for f(w), fun(w) and fquad (w) across variety of datasets and network architectures.
Super-critical learning rates. The best test loss of both f(w) and Squaa(w) is consistently smaller than the one with sub-critical learning rates, and decreases for an increasing learning rate in range of values beyond Merit. which was observed for wide neural networks in Lewkowyez et al. (2020).
As discussed in the introduction, with super-critical learning rates, both fquad and can be observed to have catapult phase, while the loss of fum diverges. Together with the similar behaviour of fquad and f generalization with super-critical learning rates, we believe NQMs are better model to understand f in training and testing dynamics, than the linear approximation ftin-
In Figure we report the results for networks with ReLU activation function. We also implement the experiments using networks with Tanh and Swish (Ramachandran et al., 2017) activation functions, and observe the same phenomena generalization for f. fiin and fquad (See Appendix N.8).
5 SUMMARY AND DISCUSSION
Summary. In this paper, generalization properties of finite width neural networks trained using large learning rates. Notably, we prove that quadratic models exhibit properties of neural networks such as the catapult phase that cannot be explained using linear models, which importantly includes linear approximations to neural networks given by the neural tangent kernel. Interestingly, we 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.
quadratic models tool better understand optimization and
Future directions. As quadratic models are more analytically tractable than finite width neural networks, these models open further avenues for understanding the good performance of finite width networks in practice. In particular, one interesting direction of future work is to understand the change the kernel corresponding model with large learning rate causes a decrease in the eigenvalues of the neural tangent kernel, and it would be interesting to understand the properties of this changed kernel that correspond with improved generalization. Indeed, prior work Long (2021) I) has analyzed the properties of the "after kernel" corresponding finite width neural networks, and it would be interesting observe whether similar properties hold for the kernel corresponding to trained quadratic models.
trained quadratic model. As we showed, training a quadratic
Another interesting avenue of research is understand whether quadratic models can be used for representation learning, Indeed, prior work Yang & Hu (2020) argues that networks in the neural tangent kernel regime do not learn useful representations of data through training. As quadratic models trained with large learning rate can already exhibit nonlinear dynamics and better capture generalization properties of finite width networks, it would interesting to understand whether such models learn useful representations of data well.

---

Published as a conference paper at ICLR 2024
ACKNOWLEDGEMENTS
We thank Boris Hanin, Daniel A. Roberts and Sho Yaida for the discussion about quadratic models and catapults. A.R. is funded by the George F. Carrier fellowship at Harvard School of Engineering and Applied Sciences. We are grateful for the support from the National Science Foundation (NSF) and the Simons Foundation for the Collaboration
the Theoretical Foundations of Deep Learn- ai/) through awards DMS 1 and #814639 and the TILOS institute (NSF CCF-2112665) This work used NVIDIA V100 GPUs NVLINK and HDR IB (Expanse GPU) at SDSC Dell Cluster through allocation TG-CIS220009 and also, Delta sys- tem at the National Center for Supercomputing Applications through allocation bbjr-delta-gpu from the Advanced Cyberinfrastructure Coordination Ecosystem: Services & Support (ACCESS) pro- gram, which supported by National Science Foundation grants #2138259, #2138286, #2138307, #2137603, and #2138296.
ing (https:/ / deeepfoundatinn
REFERENCES
Yu Bai and Jason D Lee. Beyond lincarization: On quadratic and higher-order approximation of wide neural networks. In International Conference on Learning Representations, 2019.
Nicoletta Bof, Ruggero Carli, and Luca Schenato. Lyapunov theory for discrete time systems. arXiv preprint arXiv:1809.05200 2018.
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. In International Conference on Machine Learning, pp. 1675- 1685, 2019.
Simon S Du, Xiyu Zhai, Barnabas Poczos, and Aarti Singh. Gradient descent provably optimizes over-parameterized neural networks. arXiv preprint 1rX1v:1810.02044, 2018
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. http:://tg///// .aspps.ddd dii www.uuiiiii..uuuu... corpuss_f__nes__rrtti
html.
Jiaoyang Huang and Horng-Tzer Yau. Dynamics of deep neural networks and neural tangent hier- archy. In International Conference on Machine Learning, pp. 4542-4551. PMLR, 2020.
Arthur Jacot, Franck Gabriel, and Clément Hongler. Neural tangent kerne:: Convergence and gen- eralization in neural networks. In Advances in neural information processing systems, pp. 8571- - 8580, 2018
Jakobovski.
Pree-Spoken-Diii--ataee gree-spoken-diie dataset.
httpss///gtthubb.cm//aatoooot
Ziwei Ji and Matus Telgarsky. Polylogarithmic width suffices for gradient descent to achieve ar- bitrarily small test error with shallow relu networks. In 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(11):2278-2324 1998.
Jaehoon Lee, Lechao Xiao, Samuel Schoenholz, Yasaman Bahri, Roman Novak, Jascha Sohl- Dickstein, and Jeffrey Pennington. Wide neural networks of any depth evolve linear models under gradient descent. In Advances in neural information processing systems, pp. 8570-8581, 2019.
10

---

Published as a conference paper at ICLR 2024
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 arXiv:2003.02218, 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 Long. Properties of the after kernel. arXiv preprint arXiv:2105. 2021.
David Meltzer and Junyu Liu. Catapult dynamics and phase transitions in quadratic nets. arXiv preprint arXiv:2301.077 0.7737, 2023.
Andrea Montanari and Yiqiao Zhong. The interpolation phase transition in neural networks: Mem- orization and generalization under lazy training, arXiv preprint arXiv:2007.12800 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, Adityanarayanan Radhakrishnan, and Caroline Uhler. Increasing depth leads to U-shaped test risk in over-parame perepaaraettriiee convolutional networks. In International Conference Machine Learning Workshop on over-paramettrizaatin:: Pitfalls and Opportunities, 2021.
Guillermo Ortiz-Jiménez, Seyed-Mohsen Moosavi-Dezfooli, and Pascal Frossard. What can lin- earized 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: 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 on Learning Theory, pp. 2667-2690.
PMLR, 2019.
Francis Williams, Matthew Trager, Daniele Panozzo, Claudio Silva, Denis Zorin, and Joan Bruna. Gradient dynamics of shallow univariate relu networks. Advances in Neural Information Process- ing Systems, 32, 2019.
Greg Yang and Edward J Hu. Feature learning in infinite-width neural networks. arXiv preprint 20Xi:22111.14522. 2020.
Guodong Zhang, Lala Li, Zachary Nado, James Martens, Sushant Sachdeva, George Dahl, Chris Shallue, and Roger B Grosse. Which algorithmic choices matter at which batch sizes? insights from noisy quadratic model. Advances in neural information processing systems, 32, 2019.
Difan Zou and Quanquan Gu. An improved analysis training werrparameterized deep neural networks. In Advances in Neural Information Processing Systems, pp. 2053-2062, 2019.
11

---

Published as a conference paper at ICLR 2024
APPENDIX
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:
Øf 1
0f 1 de,
Vi E [m]] (v2
ou; V Vmd" U; (u. m20 x²,
vm
And each entry of the Hessian of f, i.c., Hy, can be computed by
0 0,
01 1
0,
(uf (1/020)=" Vi e [m]
du
Oc?
ou
Vmd
Now we get Squad taking the following form NQM: fquad (u, v;
m
1 m
f(uo, vo:
u uo Txl #20
vo
vo
(va 2b
md
mm
[u,
uo xl 220 Vo,)).
(17)
md
B DERIVATION OF DYNAMICS EQUATIONS
For simplicity of notation, we denote fquad by g. Note that at initialization, the first and second derivatives of with respect to parameters are the same those of
B. SINGLE TRAINING EXAMPLE
The NQM can be cquivalently written as:
g(u, v; x) = g(u,, Voi + u uo Vug(u, v;
+ v vo. Vug(u.via
u=
u=
8 gu, v;
+
Un:
(v vo)
Ouôv
u=
since 84g 0 and
8 2
And the tangent kernel A(u, v) takes the form
8 Lglu, v; Judr
X(u, v; = Vng(u, v:
(v Vn)
v=
4=
F
0 8gg( (u, v:
Vug(u, v;
(u- uo)
ōudv
u=uo
V=
Here
m
Vn,g(u, v;
Nos {ug x200 [m),
lu md
Veg(u, v;
(u
md
12

---

Published as a conference paper at ICLR 2024
In the following, we will consider the dynamics of g and À with GD, hence for simplicity of nota- tions, we denote
V Vug(O) == Vug(u, v:
V=Vo
vvg(O)= == Vug(u, v;
V=Vo
8g((0)
0g ( Ouδv
Oudv
V==Vo
By gradient descent with learning rate n7 at iteration t, we have the update equations for weights and v:
u(L + 1) = u(() - v(((() - y) Vung()) + 8'(())
Oudv (v(1) - v(0))
v(t + 1) = v(t) - m(g(t) - y) Vvg(()) + (u(t) - u(0)) S g(0)
Oudv
Then we plug them in the expression of {(t 1) and we get
V Vug(O) + 8'((0)
dudv (v(t= + 1) - + V + (u(t 1) u()) (0)
A(t + 1) =
dudv
Vuo(O)= + -g((0)
ouey (v(1) m((()) -)) (V.S(O) + (u(() -u()) 8 gg)(0)
- v(0) 1
Oudv
+ Vog(O)+ + (u()) - ( - v) (vno(O) + 09((0)
dudy (v(1) -(0) -u())
8 89((0)
Ouδv
x(t) +n(((())) 2)² 18 (0)
Vvg(0) + (u(t) u()) T 0 (0)
dudv
dudv
Vu)(()= + 89(0))
8 90)
+g)(t(t) v)²
Oudv (v(t) v(0))
δuδv
Vng(O) pg((0)
09((0) Oudv
Vog(0) + (u(t) - u(())² 8 (0)
2n(o(t) v)
ouδv (v(t) v(0)).
δuδv
Vvg(O) + (u(t) (())) 89((0)
(vug(0) + 8g(()
dudy (v(t) -v(0))
00
2n(o(t) v)
duôr
Due to the structure of 2 o(0) Buōv we have
(²)(0)
(vug(0)+ + (u(t) - ( p 9(0)
2 l|- md
| 1|
Vvg(0) + (u(() - u(0) pg
Oudv
Judv
dudv
||2: ||
[
md
and
( (Vnol()+ + P2(())
Judr (v(t)= v(0)) T pg)(0)
||x|| md
Vug(O) + pg(())
dudr (v(t) (v(t)-v(0)]
dudr
x 2 md
Vng(E)
13

---

Published as a conference paper at ICLR 2024
Furthermore,
Vug(O) + 8g(()) Oudv (v(t) - v0)). 89(())
(vrg(0) + (u(t) - u(0)) 7 (0)
Oudv
|,0| md (v(1) - v(0), Vvg(O)) + md
|
(Vug(0), u(t) - u())) + (ong()) rg((0) Oudv voo(0))
( (8Pg(0)
Oudv (v(t) - -v()), g0(
dudv (u(t) u()))² 0 g()))
dudv
md (v(1) v(0), Vvo(o)) + | ==(t)]]==l//mm)
||0|1
| 12
v(() - v(0). 8²((0)
(Vur(0), u(() u()) + z(0) +
Ōudv (u(1) -u())²)
md
md
Similarly, we have
Vvg(0) + (u(t) u(0))² 8 Oudv'
(Tuor0)+ 8'9(0)
dudr (v(t) ((())-((1)) )²
89(())
==(t]l==1l//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)
(v(t +
1)) -v-
duδy
-(0) + (u(O) - -(((( -1) (vur()) + 8g
(v(() -v(0)) - -u(), Vuo(o))
dudv
v()) n(((t) - v) (vo(()) + (u(t) - =()))² O² ((0))
Oudv -0)) 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(1) 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((t) 2)²g(t)
14

---

Published as a conference paper at ICLR 2024
Therefore,
1t -
y= (0=111200) +
x
(o(t) - v)- - -(o((t)-))(tt)
md
B.2 MULTIPLE TRAINING EXAMPLES
We follow the similar notation the first and second order derivative of g with Appendix B.1. Specifically, for E [n]. we denote
V Vng == Vng(u, v;
V. (0) V Vvg(u, v;
u=
8 gh (0) Oudv
0 g v;
V=Vo
By GD with learning rate y,, we have the update equations for weights and at iteration t:
1
gg)
+ u(t) 7 Σ (t)
m1=) (vag((()++
Oudv (v(t) - -((0))
k=1
v(t 1) v 17 Σ (t) yk) (vrgx(0)+ + (u(t) - u())²² 0² 8gg gx(x0))
n
dudv
We consider the evolution of ((t) first.
(ig(1+1) V Vngi + (0)
(v(1 1) v()) v,g,(() + Pg2(1)
(v(1+ 1)-v())))
Oudv
Oudv
(ogi(0) (u(t 11- 1))))())) 82 (0(
02 (0)
Ouôv V.g, (0) (n(t: 1) u(D))²
Oudv
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 gx)
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)
Vvgi(0) + (u(t) -o)) 8.gg.gggl gx((0)
11 Ouδv
(g)(t) - -16)
duδv
(0)
(u + 8 (0)
Ouδv (v(t) -v()) Vvg.(0) + (u(() un0)) T 82 (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) - Uk6
Judy (v(t) v(0)
dudv
(7,,(0))
V (ugi(0) (0)
(g)(t) yx)
Judv (v(t) v(0))
Judv
15

---

Published as a conference paper at ICLR 2024
We separate the data into two sets according their sign:
S+ = (i : ai 2 0, E [n]], S_ == (i: 0, E [n]]-
We consider two scenarios: (1) xy and 2 have different signs; (2) 2, and zy have the same sign.
With simple calculation, we get if aj and Hy have different signs, i.., E St,j E S_ or i E S_,j E S++
(1)
(0) 8 (0) iuûv dudv 0,
g (0) =
Pg
(o) Vo,(0) 0.
duàv
dudv
Without lose of generality, we assume i e St, j 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) 1 ()) 0²),(0)
Oudv V Vug,(0) (0) = 1
Pg.(0) Oudv
v0g 1 Vm
(nni()
Vvg.())zg.
Oudv Oudv
Vm dudv zj
vm
For Fo St, we have
2n
Σ (t)(t) yk)zz (v)g(()) + (u(t) u()))² 8g (0)
(t (t+++)) (t)(t) - m
Oudv
kESt
vng;(0) + 99
Judv (v(1) - v(0))
2n
(ung(()++ 8 (0) dudv (v(t) v(0)), (vg + (u(t)
u(O)) p gj (0) Oudv
(g)(t) - uk)a,
m
KES+
112
Σ (gx(t) - yk) (Vrgee()) + (u(() (vrge(0)+(u(())
-(()) 02 (0)
m
audty
KES+
Σ (gx(t) - yk) (v + (g)
ouδv (v(t) - -(()))
+ m
NES.
4n m
K()(()
Σ (9)(1) (1) + Lag, (g((1) - y) 2 m,)" K(1) ((g()) - yy の m+)
KES+
4n
(t)
2 ((g(t) -y) 3 m+)) (g(t) ① m+)
m
1 x, (gg(t) y) m+ )" K(t) ((g(t) y) + m.).
m
Similarly, for i,j E S_,, we have
41 n., ((g(t) y) 3 m (g())mm)
Ki + = Ki,(!)-
+ n²
may ((g(t) y) 1 m ((()(((())) y)0m).
m
Combining the results together, we have
1 K(t (1++1) K(t) ((g(t) y) o m,,m) ((() ((g(t) - y) m,) pIP
m
2 ((g(t) - y) o m )r ((t) ((g(t) - y) o
m
4m
((g(1) y) m, T (g(1) mm P
4n
((g()) y) ② m 1d (g(() ② m. )
16

---

Published as a conference paper at ICLR 2024
Now derive the evolution of g(t) y. Suppose E S+. Then we have
g: + +) (0) + (u(t + 1) - u(0), (u); + (v(++ 1) vvg:
+ u(t + 1) - 8 (0)
(v(t + +1)++(())))
Judv
Vugx + 8 g0) Oudv
(v(t) - -(0)) Vusi(0)
= 91 (t)(t) ym)
(gx(t) yn) (vgg(0) + (u(t) u()))² 8 Fgm(10)
δudr
(1) yv)) (0) + 8 9x (0)
(v(I) v(01)
0 g:
(v(1) v(0))
δudv
duδv
(ge(t) vkk Vogg((0) + (u(t) u(0))? 02 g))
(u(t) u(0)² 02 (0) Oudv
Oudv
Σ (1) - m/h) (Vnon(0)+ 8 (0)
(v(1) -(())
1)
Oudr
(0)
8 g (0)
Σ (1) - y1h)
-2h)
(vegk(0) + (u(t) - u(O))²
Oudv
δndv
kmml
(t) yk (b)))( + 17
(t) -
Σ
(gx(t) (0, (t)xx.r
KES
RES JESI
Similarly, for i E S_,, we have
ge(t + = (t) 17
Σ
(gu(t) - ya) + n²
Σ Σ (gx(t) (t) (t)
RES
kES
Combining the results together, we have
1/2 ((g(t) y) 0 m= (g(() o m
g(t + 1) -y
[ T- [T-nK(t) ((()
m
2 ((g(t) - 3 (g(1) m m- (g(())
p.Dk (g(1) y).
+ m
C OPTIMIZATION WITH SUB-CRITICAL LEARNING RATES
Theorem 4. Consider training the NOM Eq. (10), with squared loss on a single training example by GD. With sub-critical learning rate n [e, 2=1] with
θ log /mm
the loss decreases
exponentially with
0 (1 )) ((1) = (1 5 L(1),
C(1 1)++1) (1 1
where ö = min(1)0, 7/70).
Furthermore, sup, (A(t) x(0) 0 (1))
We use the following transformation of the variables to simplify notations.
|| u(1) md. (g(()) v)². w(I) md."
|| 11 (g(()) y)v, v(() = 1)\().
17

---

Published as a conference paper at ICLR 2024
Then the Eq. (11) and Eq. (12) are reduced to
u(t + 1) = (1 - (1-v(()- + u(() + =(()))((() = x())(())
v(t + 1) = =(t) (4 4w(t).
At initialization, since 1g(0) = O1). we have u(0) I Cu/m for some constant C. > 0. As |un(t) Cyu(t) where C == -| llv O, we have |av(0)) S C.C/m Note that by definition, vm
va
for all 2 0, u(t) 2 0 we have vt) 2 0 since A(t) is the tangent kernel for single training example. From the definition of 6, we can infer that 1.
In the following, we will show that if v(0) [e, c] with e = θ
log vm
then there exist constant
CR., C. > O such that for all t 2 O,
Cx
C2
x(t)= v 1 6
1, ((t) v(0)
mo)
mo
if C, 2 9C, + (Ca + CJC and m satisfies
12C₂
V6C,
C2
m > max
82
83/2
Given the condition m, we have (1 - 5 0 md
Cx )2 <1. 1.
We will prove the result by induction. When t = 0, we have
<(0) (1- v0) (()) + w()))²
1 - 8 +
Cu
CuC
<11-
1-
Ce md
3/2
where we use the assumption C, 2 (Cn, + CuC)ō.
Therefore the result holds =.
Suppose when & = T the results hold. Then at t = T 1, by the inductive hypothesis that u(t) decreases exponentially with n(t) < (1 + Cx )² we can bound the change v(T + 1) from v()):
|u(T + v(0)) = Σ (u(t(( - v(t)) + 4w(())
T
|u(1]|4- v(t)l + Σ 4|w((t]
t=
t=
< max lu(t) - -| u(() x(T) max,c stt)
w(O) wT)) (ax /k(t)
+
(E(O,T]
1 max |o, [o,r] k(t)
max Vi((t)
u(()
w()
<4.
max 16(1)
1 max,e 0.7] V x())
Cu/m
C.C/m332
4
8/2
8/2
9C,,
m8²
For the summation of the "geometric sequence" i.e., {u(0), u(1),... ,u(T)} where u(O)) and u(T) have the determined order but the ratio has an upper bound, we use the maximum ratio, i.c., max k(t) in the denominator to upper bound the summation.
18

---

Published as a conference paper at ICLR 2024
For (0,7] k(t), we use the bound that
(1 118 8+
Cc
2
C
2C,
2Cn
28
8²
mā
m 82
m
8
8
8
6
6
2
where we use the assumption on m. Furthermore,
v +
1))2
(1 =(0) + (0) vT + - (T 1) + w( + 1))²
1 + 9Cμ
C₂ + 3/2
C.C²
md
C
2
+
md
Here we use the assumption C₁ 2 9Cn + (Cu, + C.C)6.
Therefore, finish the inductive step hence finishing the proof.
D PROOF OF LEMMA
We present the formal statement of Lemma 1:
Lemma 1. Consider constants Cu, C₁ Cv, CR2 > O which satisfies Cx. 2 28C + C +
C C where C = 7/|==/lul//V and C, 2 28C,, If m satisfies
C CC
20,, (20,+ +4CC
576C2
, exp(Cc)))
2 max
C.C.
then with high probability over random initialization of the weights, the following holds: for 5²
T > such that suPre(o.T] u(t)
log m u(t) increases exponentially with ratio inf k=()) 2 and suPre [v(() v(0) S
11 log m
Czb
C.6 log m
Proof. Due the random initialization of the weights, we have with probability, there exists
stant C₁ > 0 such that [u(0)| < Cu/m. As lw(t)] = Cvu(e) where := |||=|||_|
vd > 0, we have
CVC [w(0)] < m²/2
We prove the results by induction.
Recall that d = viAo 2 E |e, where E log mm, C² m/ /m] for some constant
C. < C.
When 0, as (00) n7u 8 + 2, we have
x()) = (1 v00) + (0) + 2
(1 (8 (8+ + u00) + w(0)²²
= (1 6 u(0) - =(0))²
19

---

Published as a conference paper at ICLR 2024
Based on the condition that 2 C.(C+1)) we have,
C. C,O 2
(1+6- w 2
1 1+ 8
m m3/2
Cu
CaC
1 + 8
m
m
C
2
IV
+ 8
log m
And by the condition m 2
2CC+ C.C.
we get
40 m3222 S C25/ logm.
[e(1) - [0) - [u(0)112 - 61 + 40w(00) S 2Cm/m +
Therefore the results hold at = 0.
Suppose when t = T the results hold. Then at t = T + 1, by the inductive hypothesis that u(t) increases exponentially with rate at least (1 (1+5 we can bound the change of ut):
C.6
from w(0) S Cu/m to u(T")
C²
log
log m
T
|o(T" + - v(0)]= =
(((((()) - 4) + 4wo(1)
[o(t) 4 Σ (tt) + Σ lw(t)l
< max rE[0,T"
t=1
u (T)) min k(t))
min x(()= 4 |mo(T)) mi k(t))
max rE(0,Ty)
[o(1)
mince(o,Ty) v /((()
&² log m
(1 + 8)2
max (u(() - v0)) (u(0) - 41) (E
(1 + 8
log m ))²
Ceb 2
CVC,
4
(1 +
/m log
8
Cia
log
C.5
90'6
24C CK
2
+
+
log
log m
Vmllg
(18)
log m
Here are the techniques we used for the above inequalities: for the summation of the "geometric sequence" i.c. {u(0), u(1)............ ,u(T)} where w0)) and u(T") have the determined order but the ratio has lower bound, we use the smallest ratio, i.c., inf x(t) to upper bound the summation. Specifically, we apply the following inequality to bound the summation:
T
u(T')
T"
min x(t))
Σ ((t))
u(T))
S "(T")
min k(t))
m0 a(t))
mine 0,29 k(() -1
Additionally, sine m 2 exp(4C,5)) we used the inequality
O56
C.5
1 1+ log m
1
1
82 +2
1
log m
log m
> 2
1
8 3 8,
log m
20

---

Published as a conference paper at ICLR 2024
and 1 [+
C. log
))
to bound the denominator of the summation of the geometric
sequence.
24C VCL
Ció log
And we further used the inequality < 8 < 2 and the final upper bound.
by the condition on m to get
vm log
Consequently, by the assumption C. 2 28C 8, we have lu(T' +1) -00) -
C.d. log m.
Now we bound the ratio K(T" 1). By our assumption, u(T" +1) S u(T) S similarly bound |wo(T" + 1)1 S
C log m
and we can
vm log m as (T("T + 1)) CVu(T++1)
C CVC.5
vm
And the rate R(T' + 1) satisfies
n(T" + 1) = (1 - -T + + + + w(T' + 1))²
(1 v0) (00) =(T' + 1) - u UT 1) + w(T + +1)
= (1 + 8 + (T'" + +) ((0 - (T' + - ( + +))²
Note that |u(T" + +1 v(0) S
2 4
log m by Eq. (18). By the assumption that m 2 exp(4Cx) and
C₂ 28C + C,, CVC... have 8 [T + 1) =(0) u(T' +) |T + 1)). Consequently, we can get
Kr <(T' + +) = (1 + 8 + u(T' + 1) v(0) - u(T' + 1) - w(' + 1))²
2 (1 + 5 |u(T' +1) v(0) uTT' + + - |w(T + 1)1)
2
28C
C2 CVC,
1+ 8-
log m
log loglmm m
C.d 2 log m
+ 8
Since m 2 exp(4Cx), we have (1+ -
o.
2 (1 +1+) 00 > 1.
log
Then we finish the inductive step hence finishing the proof.
E PROOF OF LEMMA 2
formal statement of Lemma 2 is as follows:
Lemma 2: Under the condition of Lemma 1, if we further assume that m satisfies
256C²
C C,, -CC 8 8 CH
m. max exp(2C.00, CC.-0) exp 5(C, + 4C exp
200C
where == 18C,, 2C,, and C, 2 4C CC, C.> > 20, then with high probability over random initialization of the weights, the following holds: there exists T such that u(T*)= = 0 (.))
Proof. The main idea of the proof is the following: as u(t) increases, ett) decreases since u(t)
v(() >> w(()= = @(vu(()) (()/ /m) in Eq. (14) and u())(4 - v()) < 0. Furthermore, the increase of utt) speeds up the decrease of vt)). However, v(t) cannot decrease infinitely u(t) > 0 by definition. Therefore, u(t) has to stop increasing at some point and decrease to a small value.
We first show that by the choice of the learning rate that 4 v(0) 2 where θ
log we Vm
will have 4 v(t) > 0 for all i in the increasing phase. Recall that 8 antn 2.
21

---

Published as a conference paper at ICLR 2024
2/3
48C /C Co
Proposition Under the condition in Lemma 1, if we further assume m. > exp
C.
then for such that suPreJo,T) u(t) S
log logm' m. we have u(T) <4- C. logm 2/m
See the proof in Appendix H.1
Given the constant C² in Lemma I, we define the end of the increasing phase by Ti, i...
Ti == sup { (
C²
u(())))) -
(19)
log m
We further show that there exists 72 2 such that e(T2) S
Note that indeed can show that there exists T2 such that v(T2) < C where C E (2, 4) is constant independent of m. Here for the simplicity of the presentation, we take C as 3. Furthermore, we note that T1, T2 depends m.
Before that, we present a useful result that controls the decrease of u(t):
4C Proposition 2. For t such that b(t) < 4, if m(4-v(1))z. then u(t + 1) < v(t).
See the proof in Appendix H.2.
Now ready to show the existence of T2 such that v v(T2) 3.
Proposition 3. Under the condition of Lemma 1, if we further assume that m satisfies
C ,ex + +Ce, -c), exx
[48C/Cl
2/3
16C
m max exp
C.C²2
C.
C
and C₁ 2 4C², there exists T2 2 Ti such that v(T2) S 3.
See the proof in Appendix H.3
Since vTT)) < 3 hence 4 v(Ts) 2 Simply using Proposition 2, we get
Proposition o(f) keeps decreasing after T2 until u(t) O (2.).
By definition vtt) nı(t) where A(t) 2 0, ett) will not keep decreasing for t co hence there exists 7 such that u(T**) O ( And it indicates that the loss will decrease the order of O(1).
F PROOF OF THEOREM 2
We compute the steady-state equilibria of Eq. (13) and (14). By letting u(t+1) = u(t) and u(t+1) = vt), we have the steady-statt equilibria (u", v*) satisfy one of the following:
(1) u² 0, o* E I; (2) 8² + u² + = u² v²) + 4w" = O.
As wt)2 CPu(t) where C' ==
lvl dd
we write w as a function of u for simplicity, hence
w² w(u").
As the dynamics equations are non- non-linear, we analyze the local stability of the steady-state equilibria. We consider the Jacobian matrix of the dynamical systems:
2(1 + w w1((1 (1+ + du)u + (1
u auwww w)2 -2(1 + a + w)ul
J [u, v)
de
4 4 due
1 1
22

---

Published as a conference paper at ICLR 2024
We analyze the stability of two equilibria separately.
For Scenario (1), we evaluate (u, a) at the steady-state equilibrium ("))))))))"))")))) then we get
(1 (1=0")2
0
,(",,"") v²) 10- - 4
day du
1
We get the two eigenvalues of J(u", v*) are 1 and (1 v*)a. We will show the Lyapunov stability of the equilibrium (u", v*). Specifically, we apply Theorem 1.2 in Bof et (2018). We find the domain
D ((u, (0 u S C1 1V 0-1 S min((C; v=1. C2 =-
where C1 (11/m)) and C2 ((1/vm) and the Lyapunov function V(u, v) (u 2 )² not hard to verify that is locally Lipschitz in D as V is continuous in a compact domain. Furthermore, can see that (u", ,²") with u 0, o² E [e, e] where Olog m/ vm) satisfies the condtions Eq. (3,4) in Theorem 1.2 in Bof al. (2018). Therefore, (u", x²) with u and E [e, el is a stable equilibrium point.
For Scenario (2), we again evaluate (u, v) at the steady-state equilibrium
then we get
-2 + CVT
2u
J(u (u',x")==
m
1 u²
where we replace v" by 4 4w" /u" based on the second equality in Scenario (2). Note that u²(4 0*) > since < during the whole training process, therefore we have w" achieve the equilibrium.
We can compute the eigenvalue of ,(") then we get
c mm
C2VU"
Cu" vm
A = +
2/m u
16
+
)3/21
2 2
m
Note that when Scenario (2) holds, there are only two possible cases
(2.1) u² @11mm, (w" = (1/mm) and x* = O(1);
(2.2) u² (1/mm, lw" ((1/m) and 8" @(1mm).
For (2.1), , by the first equality v² = u² + ww" E (1, Then plugging o" into the second equality
yields E 20
For (2.2), by the second equality that (4 - 0²) 4w" 0, we have u² C² + o(1mm)
By computing the modulo of AJ, we have
5C
1
1+
Vu" n² +
Vm
m
Therefore, for both (2.1) and (2.2) we have AAyl > 1 which indicates (u*, v²) is unstable.
G OPTIMIZATION WITH 1) 7max
Theorem 5. Consider training the NQM Eq. (10), with squared loss a single training example by GD. If the learning rate satisfies E
4+c
Nu o0) with
[ogm then GD diverges.
23

---

Published as a conference paper at ICLR 2024
Proof. We similarly the transformation transformation of the variables to simplify notations.
| o2 (g(() - y)², w(t)
||=:1 mdd"
u(t)
1 (g(() y)v, v(t) - m)(())
md
Then the Eq. (11) and Eq. (12) are reduced to
u(t + 1) (1 - ett) + u(t) + (u(t))((()))))) x(t)u(t)
v(t + 1) = e(t)
- - 4w(().
We similarly consider the interval [o. T] such that (0,7] u(t) = 0 log m [0, 27, u(t) increases exponentially with a rate ( kc()) > 9. We assume (w(t) > |u(t)(4
By Lemma 1, in
(()) for all [0, 7), which is the worst case as vtt) will increase the least. By Lemma 1, we have Ex [w())] = O Vmggmm which is less than Therefore, we have v(T) > 4.
Then at the end of the increasing phase, we have [u(Ti)(4 e(Ti))) = 201 vm) is of a greater order than (wii)) O(1/ /m log m.) hence v(t) will increase at Ti. Note that x(Ti) (1 + o(1))² = 9 + o(1), hence u((t) also increases at 7'
It is not hard ically, if (4(t) utt) = [4)(
that (1) will keep increasing unless u(() decreases
smaller order. Specif-
(. it requires u(t) to be of the order at least O /log/ologn m) (by letting cu(t) (ww(t) @ @(v((t (m)t))))), which will not happen as k(t) = x(t) +o 2 and it contradicts the decrease u(t).
Therefore, both ul) and v() keep increasing which leads to the divergence of GD.
H PROOF OF PROPOSITIONS
H.1 PROOF OF PROPOSITION 1
Proof. Note that 4 v(0 = 2 - log m by definition, where C. >> 0 is a constant. To show vm
4 - (T)
C. log sufficient condition is (TT) - v0)) < 2/m
C. log
2/m
CL
Specifically, we will prove for T such that sup u(t)= S
the following holds:
log m
(T) v00) (w(t)s 24C C%
Vm log
where C, C² are the same constants defined in Lemma 1. Then by the condition that >
2/3
488 C.
we have (T) v(0) < C. log m 2/m
exp
We will prove the result by induction.
When T = 0, the result holds trivially.
Suppose T T the result holds. When T = T" + 1, since v(T') -0(0) O. we have v(T") < 4. Therefore, by the update equation of u(t Eq. (14), we have
"(T' +1) = v(T") u(I') - v(T") - 4w((T)
I (TT) - 4wo(T')
$ v(T") + 4/wo(T)))
Then e(T' + - -00) = v(T'
+ 1) ((T) + u(T') v(0) S Et' (w((t))
24C /CL
By Lemma I, we have Σ 21=0 |w(l)] <
Indeed, this inequality holds for any T" + 1 such
Vm log
that (1) u(()
log m
24

---

Published as a conference paper at ICLR 2024
Therefore, finish the inductive step hence finish the proof.
H.2 PROOF OF PROPOSITION 2
Proof. A sufficient condition for v(I) to decrease is
(4 - v(t) > 4/w(t)] 40 vu(t))
vm
If u(t)
m(4-v(t))) then the above condition is satisfied.
1C
H.3 PROOF OF PROPOSITION 3
Proof. Note that for E [0, T1), the change of utt) satisfies sup, (m(() - -000 _
C.u log m
by Lemma 1.
For 8 C.á log m i..., v(0) 3 log m v(Ti) < 3. Therefore, the existence of T2 can be guaranteed by simply letting T2 = Tı.
C..
have ((i) v(0) ( v(0) 2 6 (0(0)
For Tog m i.c., v0) 2 Tog m
C.d.
C.d
we will show there exists 72 2 Ti which depends on m
such that (T2)) 3.
We prove the existence of T2 by contradiction. Suppose that for all f 2 Tj + have u(t) 2
For the simple case that if all u(t) m(1-v(0))² then by Proposition 2, vtt) keeps decreasing which will ultimately lead Ett) < 3.
4C
Suppose there is an iteration 2 Ti + 1 such that v(I) (0)= The following Proposition guarantees that e(f) will decrease to smaller value after t once such occurs. Therefore, we can find T2.
Proposition 5. Under the condition of Lemma 1, suppose m further satisfies
7682 C.
16C
C exp (2C), + c..).
m max
exp
where C²2 2 402.
Then if there T 2 0 such that u(T)) m[]]e(T]] and v(T)) 3, we have v('' + +2) < v(T) and u(T' + 1) (())) where T is the end of the increasing phase starting from T.
10
10
See the proof in Appendix H.4.
H.4 PROOF OF PROPOSITION
Proof. Since u(T) <
C² log m
by the assumption that m
16C²
C the training dynamics falls into the
increasing phase from T. We denote the end of the increasing phase starting from T by TV, i.e.
T' sup {
U: u(t) ul:(()=
S C2
T
log m
We will prove the result by induction.
Suppose 7" is the end of the first increasing phase, i.e., T" Tí. By proposition 1, v(Ti)
C. And the magnitude of u(Ti + 1) can be lower bounded by
2/m
C²
u(Ti 1 2
(20)
4 log m
25

---

Published as a conference paper at ICLR 2024
where we use 2 by the assumption on m that and plug ō C
logmi
Note that when m > exp (28C"), 8 z 1 is necessary for e(Ti) > 3 as (o(Ti) v(0) Inequality (18).
S
28C, à Tog
by
Furthermore, with the above bound on u(Ti + 1), we have
(Ti 1)++1==0
- -((Ti) - 42o(Ti)
< C.log m
C₁
Co m +T]w((Ti)
2/m
4 log m 2/m
C. mo
C²
2/CL
< 4
2/m
8ym
vm lollog m
C. log m
S 4
(21)
4/m
Consequently, when O2 2C and m exp (2C), + C.), we have
u(Ti + 1) (TT))u()) = (1 - x(T) + ((T)) + w((T)))
C. Colog m
C²
+ CVCL
C₁
1 - 4
4ym
loogm 2/mlogm
1
90%
4 log
Therefore, at Ti + 1, we have
4C'
VuVuuT +
"Ti + u(Ti 2 u(Ti + 1)(44)1()) + 1))
Vm
C₁ Cologm
SC
z 4 logm m 4ym
C.C.
60
16/m
Vmlog m
C²,
32.mm
where we use the assumption that m exp
2560 90.C7
Note that the increase is caused by the term w=(t) since for all t E [T, Ti+1] we have u(t)(4-((1)) 0. Then we have the maximum increase during [T,Ti + 1] be bounded by
24C /CL
u u + - v(T)
4|w(t)]
Vm m
where we use Eq. (18) in the proof of Lemma 1.
By the assumption on m that m > exp
C2 C.
we have v(Ti < v(T)..
If there is T such that u(T) <
m(4-v(T)) while uT) > there is another increasing
40
phase. Since u T(T) < (T), we can apply the same analysis under the same condition to show u(T, +2) (T), where Tj is the end of the increasing phase starting from T. Therefore, finish the inductive step hence finish the proof.
26

---

Published as a conference paper at ICLR 2024
H.5 PROOF OF PROPOSITION 7
Restate Proposition 7: For any u, v E R", rank(K) S 2. Furthermore, Pt, P2 are eigenvectors of K, where Pı x; (1ES+) P2,1 n; ( (iES_,, [n].
Proof. By Definition 1,
m
Ki,
m E (uî + 2g) (wea)20) (mm2220); (mmr,2))) i.j e [n].
By definition of eigenvector, we can see
K Kip
(0g + ( (12z,200 (58t)
k=1
Σ j=l m
2 11 (jES
1
Σ (ut + 2z (mmm;20) (mxa220)
k=1
Σ (of + u (m)z,200
=x, (m=E5 }
x (jEst)
where we use the fact that if n, < 0, Kig =0
As P1, Fil(mess)) and Σj=1 2 (ESS)) 1 E
u (m),, does not de-
pend on we can see Pj is eigenvector of K with corresponding eigenvalue A 2j=1 x (JESt) 2 (eg + +
The same analysis can be applied to show P2 is another eigenvector of K with corresponding A2 Ej= of wef (jes_) >
Ex=1
(vg uz) (ul2, 203-
For the rank of K, it is not hard to verify that K = >IPIPT + >appp? hence the rank of K is at most 2.
I SCALE OF THE TANGENT KERNEL FOR SINGLE TRAINING EXAMPLE
Proposition 6 (Scale of tangent kernel). For any 5 E (0,1), 1), if m absolute constant, with probability at least 1 - 5, /=2))/(2)) S A(0) S 3/|=0| 1((d).
e' (og(4/8) where d' is
Proof. Note that when O.
X(0)
1 md
Σ
(w x (ug, (1f, ) + md
Σ
(o
)² ||11 (1 (1(
2m1
According to NTK initialization, for each i E [m], 00 N'0, 1) and uo ~ AO(, We consider the random variable
: x0 x1 =20) & : Uo (u), #20)
it is not hard to see that &i and & are sub-guassian since _H and vo, are sub- sub-gaussian. Specifically, for any & 2 0,
[l 1 < P( P( 20 < 2exp(_(/(21)=13)
PIL. 2 t) < P( 0o 2 < 2 2xxp(-t/2).
where the second inequality comes from the definition of sub-gaussian variables.
Since & is sub-gaussian, by definition, _² is sub-exponential, and its sub-exponential norm is bounded:
| 1l,, < ||E:l12 < C,
27

---

Published as a conference paper at ICLR 2024
where C > 0 is a absolute constant. Similarly we have lIGell By Bernstein's inequality, for every 2 O, we have
C||2
m
12 E
max; |1E?11 ))
P
2 t
2cxp
cxp
-emin
1 p
where c is absolute constant.
Letting ( = m/44 we have with probability at least 1 - 2exp ( (-m/c'),
m
m
3m
4
4
where d' = c/CCC).
Similarity, we have with probability at least 1 - 2 exp -m//c).
m
3m
x
4
4
As a result, using union bound, we have probability at least - 4 exp (-m/c),
||
3
x
S x(0)
2d
2d
J SCALE OF THE TANGENT KERNEL FOR MULTIPLE TRAINING EXAMPLES
Proof. As shown Proposition 7, P1 and P2 are eigenvectors K, hence we have two eigenvalues:
A (0) pf (()p1
A2(0) = k0)(00)2
P 12
P2
Take A (0) as an example:
m
Σ
2, (2,20) 1 (z,20) Σ (ub, + x, (0000000000002200000000000
ij=1
m
Σ a + vo,e) (1100
Σ 1,=1
xn (2,20) L 12,207
Similar to the proof of Proposition 6, we consider &u == Vo 1 (uo, (wo,_(u
which is a sub-gaussian
random variable. Hence & is sub-exponential so that ||<2 llv C where C > 0 is an absolute constant. By Bernstein's 's inequality, for every i 2 0, we have
P
t
exp
-e
2
m
max; |K
where > 0 is an absolute constant.
Letting = m/4, we have with probability least 1 2exp (-m/c!),
m
3m
4
4
where c' c/(AC).
28

---

Published as a conference paper at ICLR 2024
The same analysis applies to Çs == 2oo 1 (v0 as well and we have with probability at least 1 - 2 (-m/d))
m
3m
4
4
i=1
As result, we have probability at least -eexp( (-m/c).
(0)1|p111
1
m
u + vb2x) (M(m)())20))
Σ
12
27 (m:20) (my20)
m
s,==1
E
2qa (a.20) 1 (m,20)
3
Σ
2 2
by (m.20) (z,20)
2
f,==1
1,1=1
Applying the same analysis to (0), we have with probability 1 4 exp (-m/d'),
m
x2(0))lp,11 = m
Σ
x + vo.x) (1 (I(m)(0)=)))*
Σ
27 (0: (0)50)
i=
1,==1
3
E
2,4 (0:50) 1 (m
Σ
2
(m,50) (2,50))
2
t,j=1
i,=1
The largest eigenvalue is max (0), A2 (0)). Combining the results together, we have with proba- bility at least 1 4 exp (-m/c) c'),
M < ||(]]]] <
2 M,
2
En (1
Ee (2 (
where M max
ER 20
K ANALYSIS ON OPTIMIZATION DYNAMICS FOR MULTIPLE TRAINING EXAMPLES
In this section, we discuss the optimization dynamics for multiple training examples. We will see that by confining the dynamics into each eigendirection of the tangent kernel, the training dynamics similar to that for single training example.
Since x; is a scalar for all E [m], with the homogeneity of ReLU activation function, we can compute the exact eigenvectors of k(t) for all 2 0. To that end, we group the data into two sets S4 and S_ according to their sign:
S = (i : x, 2 0, E [n]],
S_ (i: 1::;; < 0, [n])-
Now we have the proposition for the tangent kernel K(the proof is deferred to Appendix H.5):
Proposition 7 (Figenvectors and low rank structure of K). For any u, v E IRm rank K) < 2. Furthermore, Pt, Pa are eigenvectors of K, where pı, 2; (1ES,) } P2, 2; (irs_), E [n].
Note that when all 2; are of the same sign, rank(K) = and K only has one eigenvector (either Pj or Pa depending 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(K) 2. We denote two eigenvalues of (() by At(t) and (² corresponding to pı and pa respectively, i.c., K(t)p, = (1(t),, K(t)p2 \2(t)p,- Without loss of generality, assume )₂ 2 (0)
29

---

Published as a conference paper at ICLR 2024
By Eq. (5), the tangent kernel K step t is defined as:
(t)(t)= (Vvgi (t), (vg,(t)) + (Vagi(t), Vug;(t))
m ((ux(t))² + (t))²) 2 (m=(0)=.20) (uz(0)a;20) Vi,j e [n].
Similar to single example case, the largest eigenvalue of tangent kernel is bounded from 0:
Proposition 8. For any ò (0. if m 2 c'log(4/0) where c' is an absolute con-
stant, with probability at least 1 6, M/2
S ( I 3M/2 where M
Eii=1 [[i=]=]a]] (e, (2,200 E 2525 (1,500 (0) Σ 27 (0:20)
max
E 10,
The proof can be found in Appendix J. For the simplicity of notation, given p, E R we define the matrices and Qp.
K Kpm (t) = ((g(t) y) m) K(t) ((g(t) - y) o m) pp", Qp.m (t) = ((g(t) y) m) T (g()) m) pp"
It is not hard to see that for all 4, Kp.m and pp. are rank-1 1 matrices. Specially, is the only eigenvector of Kp.m and 2p.m.
With the above notations, we can write the update equations for g(t) and ((t) during gradient descent with learning rate yy.
Dynamics equations.
x (0p (t) + Qp (8)) (g(t) y),
g( 1 T-nk(t)
(22)
(t)))
1)++1) K(t)+ 7? (Kp,m (t) + pm (f) An
(Qp (t) + Pp. (t)),
(23)
m
m
(x(()
where m,, m_ Rn are mask vectors:
m 1 ()eS))
m- = (ies }.
Now are ready discuss different three optimization dynamics for multiple training examples case, similar to the single training example case in the following.
sub-critical learning rates (1 < 2/Ar We use the key observation that when [g(()]| is small, i.c.. O(1), and ||((t)) is bounded, then [t)( and [| are of the order o(1). Then the dynamics equations approximately reduce to the ones of linear dynamics for multiple training examples:
Monotonic convergence:
+ - - (I-n(()) o(1)) (g(() y),
-()+1)=(()) o
At initialization, [lg( = O(1) with high probability over random initialization. By the choice
of the learning rate, we will have for all > 0. (r)))( < 2, hence [[g(t) y| decreases exponentially. The cumulative change on the norm of tangent kernel o(1) since (t)] =
O(1/m) and the loss decreases exponentially hence Σ Rx (t)l = 0(1mm) log 0(1) = o(1).
Catapult convergence: super-critical learning rates (2/ (0) < (/A/ We summarize the catapult dynamics in the following:
30

---

Published as a conference paper at ICLR 2024
Restate Theorem 3 (Catapult dynamics multiple training examples). Supposing Assumption holds, consider training the NQM Eq. (10) with squared loss on multiple training examples Then,
eximples by GD.
1. with E
2+ 2-6 A1(0]'
the catapult only occurs in eigendirection P1 II increases to
the order of S2 200002/A000)" then decreases to O(1);
2. with 1 E
2++ X2(0) (o
the catapult occurs in both eigendirections P1 and P2: I1,C for
1, increases to the order of f
m(n-2/A.(0))
then decreases to O(1),
log m
where c = θ log m
The proof can be found in Appendix L.
For the remaining eigendirections P..... Pa. i.c. the basis of the subspace orthogonal to p1 and Po- show that the loss projected to this subspace does change during training in the following proposition. It follows from the fact that K, Rg (t) and Rx (t) are orthogonal to p.P. for
i 3, n..
Proposition 9. Vt 2 [1, 11, for = 3,
Once the catapult finishes as the loss decreases to the order of O(1), we generally have 1j > 2/A1 and 1) 2/X2. Therefore the training dynamics fall into linear dynamics, and use the same analysis for sub-critical learning rates for the remaining training dynamics.
Divergence: (1) >> 1/max = (0)) Similar the increasing phase in the catapult convergence, initially (3(t)-yl increases in direction P1 and P2 since linear dynamics dominate and the learning rate chosen increasing phase, by similar analysis for the catapult convergence. We consider the evolution of K ((t) in the direction P1. Note that when [g((t)] increases to the order of ((ym), g(t) @ mm will be aligned with Pl. hence with simple calculation, we approximately have
larger than Meii.. Also, we approximately have 1j > 4/A1(t) at the end of the
lg (g)())|3llp 11
o RK (t)p, ~
(t) 4n) > 0.
m
Therefore, (t) increases since k 1)p k(((t)p; + RT pt))t)) As result, || - (t(() + [t]l becomes even larger which makes [|g(t) - yll grows faster, and ultimately leads to divergence of the optimization.
L PROOF OF THEOREM 3
As the tangent kernel K has rank 2 by Proposition 7, the update of weight parameters is subspace with dimension 2. Specifically,
Ög OL w(t + 1) w(t) Ow 0g (t),
where 0g/Ow has rank 2. Therefore, to understand the whole training dynamics, it is sufficient to analyze the dynamics of the loss in eigendirection P1 and p2-
We will analyze the dynamics of the loss L and the tangent kernel K confined to P1 and P2. It turns out that the dynamics each eigen direction almost independent on the other hence can be reduced to the same training dynamics for single training example.
We start with cigendirection P1. For dynamics equations Eq. (22) and (23), we consider the training dynamics confined to direction Pı and we have
((t) (1 7 + Rg C(, = Ki ((t) A 1) A1 p.. Rx (t)pt,
31

---

Published as a conference paper at ICLR 2024
where we use the notation II ((t) (g(t) y.P1).
We further expand pl Rg (t)p, and p Rx (t)p: and we have
21/² IT, [()) + r m ((g(t) O m+ ④ m+),
p Rg((
2n (()((AA)))))) - 4)
4n (g(()) - y) o m o m+).
p Rx (1)P1
m
Analogous the transformation for Eq. (11) and (12) as we have done in the proof of Theorem 1, we let
2T un (1) IT, C(I,, wi (1)
T ((g(l) - y) @ m a mm+), vi(1) - (1)
Then the dynamic equations can be written as:
(t + 1) = (1 (t) + (1 - w1 (t)) (t), (t
(24)
=)) =)) un
(1)44)22(1)
- 4w1 (t).
(25)
Note that at initialization, [lg(()] (vV)) with high probability, hence we have un u2 = 0
and w/ (0) 0 ( (we omit the factor is constant). Furthermore, |wn (t) O
V (t)
vm
Therefore, both the dynamic equations and the initial condition are exactly the same with the ones for single training example (Eq. (13) and (14)). Then we can follow the same idea of the proof of Theorem 1 to show the catapult in eigendirection Pt
Similarly, when we consider the training dynamics confined to P2, we have
( + 1) (1 v2() + (t) + w2()))² (t), 02 ( = (1)) - (1) - v2(()) - 4wz(
(26)
(27)
where
2n² 1lt)()), (t)
o ((g(t) y) C y 22(() = 11)2(t).
u2
Then the same analysis with Theorem I can be used show the catapult in direction P2-
Note that when 20) > (0), the learning rate is only allowed to be less than (0) otherwise GD will diverge, therefore, there will be no catapult in direction P2-
M SPECIAL CASE OF QUADRATIC MODELS WHEN ((x) 0
In this section will show under some special settings, the catapult phase phenomenon also hap- pens and how two layer linear neural networks fit in our quadratic model.
We consider one training example (z, y) with label y 0 and assume the initial tangent kernel x(0) = 2(1). Letting the feature vector (xx) = 0, the quadratic model (3.(3) becomes:
g(w) = w T "(x)w.
For this quadratic model, have the following proposition: Proposition 10. With learning rate X(O) phase.
2
<
x if E(x)² = ||^=|| I, g(w) exhibits catapult
Proof. With simple computation, we get
g(t 1)= = (1 ()) (== ot),
X(t 1)= X(t) (gt)) mT)(t))
We note that the evolution of and À is almost the same with Eq. (11) and Eq. (12) if we regard 7 1/m. Hence can apply the same analysis to show the catapult phase phenomenon.
32

---

Published as a conference paper at ICLR 2024
It is worth pointing that the two-layer linear neural network with input R analyzed in Lewkowycz et al. (2020) that
S(U
viz)
Uæ,
mm
where v E R", U E Rnd is special case of our model with w [oe(UU) v 1/mm and
E = Im xx
Im ② x
E pmddm
N EXPERIMENTAL SETTINGS AND ADDITIONAL RESULTS
N.1 VERIFICATION OF NON-LINEAR TRAINING DYNAMICS OF NQMS, L.E., 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.i.d. from (V 1) if the label is -1 or (2, 1) if the label The network width is 5, 000.
N.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.dd from (-22, if the label is - or (22,, 1) if the label is 1. The network width is 5, 000. See the results in Figure
Saccoeei lorgent aggnnaa
Tranba tots
-R 0.20
LR 0.20
LR - 0.80
LR 0.45
200
LR 0.45
LR - 0.45
-H 0.00
UH 0.00
100
wn
100
251
101
350
300
Iteration
lteration
Iteration
(a) Training loss
(b) Largest eigenvalue of tangent kernel
(c) Second largest eigenvalue of tangent kernel
Figure 5: 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 cigenvalues almost the same with different learning rates: when < 0.37, the kernel is nearly constant; when 0.37 n 0.39, only A (t) decreases; when 0.39 1) 7maxx both Az and Az(t) decreases. See the experiment setting in Appendix N.2.
N.3 TRAINING DYNAMICS CONFINED TO TOP EIGENSPACE THE TANGENT KERNEL
We consider the corresponding dynamics equations (15) and (16) for neural networks: f(t + 1) y = [T - (t( + Re (t) (f(t) y), K(t 1) K(t) (t).).
(28)
(29)
Note that for NQMs, (1) and Rx (1) have closed-form expressions but generally for neural net- works they do not have.
We consider the training dynamics confined to the top eigenvector of the tangent kernel p, (t):
(p. y (1- (1= (p (t), -y),
P (t) K(t (t), (1() (t)
33

---

Published as a conference paper at ICLR 2024
We conduct experiments to show that P1 (t) Re(t)p (t) and p1 (x)) (t) scale with the loss and remain positive when the loss is large. Furthermore, the loss confined to P1 can almost capture the spike in the training loss.
In the experiments, we train two-layer FC and CNN with width 2048 and 1024 respectively on 128 points from CIFAR-2 (2 class subset of CIFAR-10) and SVHN-2 (2 class subset from SVHN-10). The results for NQM can be seen in Figure and for neural networks can be seen in Figure 7.
da
baiing
Tadg
1000
(a) FC CIFAR-2
(b) CNN CIFAR-2
(c) FC SVHN-22
(d) CNN SVHN-2
Figure 6: Training dynamics confined to the top eigenspace the tangent kernel for NQMs.
Sarig his
Baring
daving bis
(a) FC CIFAR-2
(b) CNN CIFAR-2
(c) FC on SVHN-2
(d) CNN SVHN-2
Figure 7: Training dynamics confined to the top eigenspace of the tangent kernel for wide neural networks.
N.4 TRAINING DYNAMICS OF GENERAL QUADRATIC MODELS AND NEURAL NETWORKS.
As discussed at the end of Section 3, more general quadratic model can exhibit the catapult phase phenomenon. Specifically, we consider general quadratic model:
1 w: x = w ((x(( + 2 y
T E(x)w.
We will train the general quadratic model with different learning rates, and different respectively, to see how the catapult phase phenomenon depends 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 d 100. We let the feature vector ((x)) = x </||==|| where x; ~ (o0, 1) i.i.d. for each E [d]. We let be a diagonal matrix with En, 1-1, randomly and independently. The weight parameters w are initialized by (0, Ia). Unless stated otherwise, 7 10-3, and the learning rate is set to be 2.8.
Neural networks. We train a two-layer fully-connected neural networks with ReLU activation function 20 data points of CIFAR-2. Unless stated otherwise, the network width is 104, and the learning rate is set to be 2.8.
See the results in Figure
34

---

Published as a conference paper at ICLR 2024
N0m
(A)
Iheatons
Nome of Sarget Herm
Trann
Nerrer SamrerBaegeeee
Mith Mach
Uh
Lh-23
Mach
(B)
(a) Loss (log scaled) vs. (b) Tangent kernel norm (c) Loss learning (d) Tangent kernel norm
7/width
vs. r//width
learning rate
Figure& 8: General quadratic models have similar training dynamics with neural networks when trained with super-critical learning rates. Panel (A): experiments on general quadratie models. Smaller y 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 the training loss the peak and decrease of tangent kernel norm with quadratic models. Note that width or ? seems to have no effect on the tangent kernel norm at convergence.
N.5 TEST PERFORMANCE f. Júin AND Squad, I.E., FIGURE 2(B) AND FIGURE
For the architectures of two-layer fully connected neural network and two-layer convolutional neural network, set the width be 5, 000 and 1, 000 respectively. Specific to Figure 2(b), 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-10, which we call CIFAR-2, and select 256 data points out of it the training set. We use the number 0 and 2 of SVHN, and select 256 data points as the training set. We select 128, 256, 188 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 runs.
N.6 TEST PERFORMANCE OF f, funn AND fquad IN TERMS OF ACCURACY
In this section, report the best test accuracy for f. fin and fquad corresponding to the best test loss Figure 4. We use the same setting as in Appendix N.5.
N.7 TEST PERFORMANCE OF fiin AND Squad WITH ARCHITECTURE OF 3-LAYER FC
In this section, we extend our results for shallow neural networks discussed Section 4 to 3-layer fully connected neural networks. In the same way, we compare the test performance of three models, f. fim and fquad upon varying learning rate. We observe the same phenomenon for 3-layer ReLU activated FC with shallow neural networks. See Figure 12 and 13.
We use the first class (airplanes) and third class (birds) of CIFAR-10, which we call CIFAR-2, and select 100 data points of it as 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 set. 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, report the average result of 5 independent
35

---

Published as a conference paper at ICLR 2024
2-layer FC on MNIST with GD
2-layer FC on AG NEWS with GD
0.75
Linear
NN
0.84
Best Test Accuracy
Quadratic
Best Test Accuracy
Ofttical LR
0.82
0.80
0,000
NN
Qundratic
0.73
Critical LR
201
10
Learning
Learning rate
Figure 9: Best accuracy plotted against different learning rates for fquads f, and flin Left panel: 2-layer FC on MNIST trained with GD. Right panel: 2-layer FC on AG NEWS trained with GD.
2-layer CIFAR-2 with SGD
2-layer CNN SVHN GD
0.786
NN
0.80
Quadratic
Test Accuracy
0,730
Linear Critical
Best Test Accuracy
0,770
0.79
0,770
0.78
0,700
Quadratic
0.756
0.77
Linear
Critical
0.700
0.5
1.0
1.5
2.0
1.0
2.
3.
Learning rate
Learning rate
Figure 10: Best test accuracy plotted against different learning rates for fquads f, and fli- 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 2-layer CNN on CIFAR-2 with GD
0.78
Quadratic
Quadratic
Accuracy
0.685
Linear
Best Test Accuracy
Linear
Critical LR
0.77
Critical LR
0.680
0.75
0.676
3.670
0.66
0.74
25
200
125
100
Learning
Learning rate
Figure 11: 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 CIFAR-2 trained with GD.
N.8 TEST PERFORMANCE WITH TANH AND SWISH ACTIVATION FUNCTIONS
We replace ReLU by Tanh and Swish activation functions to train the models with the same setting as Figure 4. We observe the same phenomenon as we describe Section
36

---

Published as a conference paper at ICLR 2024
3-layer FC on CIFAR-2 with GD
3-layer FC SVHN with GD
NN
Quaddwiic
Quadratic
Bost Teal Loss
0.52
Linear
0.70
Critical LR
Linear
0.50
66772
0.75
0.70
0:76
0.74
15
2.1
2.5
3:3
2.55
4.6
16
2.0
235
25.5 Learning rate
4.01
4.
530
Learning rale
Figure 12: Best test accuracy plotted against different learning rates for fquads 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-layer FC FSDD-2 with GD
3-layer FC AG NEWS with GD
0.18
0.45
Quadratic
Linear
Bost Test Lozs
0.37
Best
0.33
0:15
Quadmaic
Linear
8:11
Critical
15
200
2.5
2.0
4.5
Learning rate
Learning rate
Figure 13: Best test accuracy plotted against different learning rates for fquads f. and fin- Left panel: 3-layer FC FSDD-2 trained with GD. Right panel: 3-layer FC AG NEWS trained with GD.
2-ayer FC AG NEWS with GD
2-layer NEWS with GD
0.0000
0.1025
Quadratic
0.000
Linear
acoms
Bost lest Los
80060
Quadratc
Linear
(1)
150
130
.0
40
Z01
Learning rate
Learning rate
(a) Swish activation function
(b) Tanh activation function
Figure 14: Best test loss plotted against different learning rates for fauads f, and fii. We choose 2-layer FC as the architecture and train the models on AG NEWS with GD.
37

---

Published as a conference paper at ICLR 2024
2-aaer on
2-kyer
0..0%
0007
Quadrato
0170
Best Test Loss
Critical LR
Logs
0.77
0,074
0.602
9.77
0.003
01772
Quadratic
Critical
300
Learning rate
Learning rate
(a) Swish activation function
(b) Tanh activation function
Figure 15: Best test loss plotted against different learning rates for fquads f, and flim. We choose
2- 2-layer FC as the architecture and train the models on FSDD with GD.
2-layer CNN on CIFAR-2 with GD
2-layer CNN on CIFAR-2 with GD
0.55
Beat Test Loos
0.50
0,92
Best
0,81
0.75
NN
Quadmiic
Quadratic
Lincar
Linear
0.70
Critica
Critical
Learning rate
Learning rate
(a) Swish activation function
(b) Tanh activation function
Figure 16: Best test loss plotted against different learning rates for fquads f, and fin We choose 2-layer CNN as the architecture and train the models on CIFAR-2 with GD.
2 2-aayer CNN on SVHN with GD
layar CNN on SVHN with
1.50
0.55
1007
Best Test Lood
0.56
Best Loss
1102)
8.64
0.52
L.0g)
a.so
NN
NN
Quadrato
1.000
Quadratic
0.58
Critical UR
Criical
0.58
Learning
Leerning
(a) Swish activation function
(b) Tanh activation function
Figure 17: Best test loss plotted against different learning rates for Squad, S, and fim- We choose 2-layer CNN the architecture and train the models on SVHN with GD.
38


Reviews:
[{"review_id": "96Jf9BNAIo", "reviewer": "Reviewer_Unww", "summary": "This paper studies a quadratic approximation of two-layer neural networks to understand optimization behaviours that cannot be captured by linear models. More precisely they theoretically evidence the so-called catapult dynamics for respectively one and multiple training points, ie they show the existence of two critical values for the learning rate which delimits respectively exponential convergence to the minimum, catapult dynamics (ie first increase of the loss then convergence to low loss) and finally divergence. Finally they provides experiments evidencing catapult dynamics as well as the fact that the quadratic approximation of a neural networks shows similar behaviours that the neural network above the critical learning rate.", "questions": "1) I am curious if you could clarify the links, if they exist, between catapult dynamics and recent works on edge of stability analysis of neural networks (cf the paper \"Second-order regression models exhibit progressive sharpening to the edge of stability\" that is not cited but seems to me very related to your setting). It seems interesting because both studies explore the influence of a non-linear quadratic term on the coupled dynamics of the sharpness and the learning rate.\n\n2) It is written in the text that the results hold with high probability over initialization but I'm not sure if this is written in the statement of the theorems. Perhaps it could be added in the theorems themselves.\n\n3) I find the results interesting and they provide good contribution. I would increase my score if the authors address my concerns about the clarity of $O,o,\\Omega$ notations in the proof, especially the proof by induction.", "limitations": "", "rating": 6, "confidence": 1, "soundness": 2, "presentation": 4, "contribution": 3, "strengths": "This paper is well-written with clear figures and a good explanation of the setting and the results. This paper studies a very interesting phenomenon about the optimization of neural networks and has to deal with non-linear phenomena which are usually not well understood. They are able to evidence theoretically the catapult dynamics and the existence of two critical learning rates when dealing with neural quadratic models.\n\nEmpirical results shed light on the similarity of behaviours in terms of generalization between the NQM and the neural network function, evidencing the coherence of the quadratic model.", "weaknesses": "I find the proofs a bit hard to understand because of the use in the proofs of O,o, $\\Omega$ notations. Hiding constants behind such notations makes the proofs a bit cloudy to me. \n\n1)Especially I have some concerns about the proof of lemma 1 in appendix D: you use a proof by induction, and still use O,o notations. But summing $O$ terms remains $O$ only when the number of iterations is controlled. To still have $O$ in the end of the summation with respect to m, it should be checked that the summation index T is itself $O(1)$. However I am not sure such a result is proved (correct me if I'm wrong).\nEspecially it seems to me that it might not be the case by considering the fact that T must satisfy a relation of the form $(1+\\delta)^T\\sim \\log(m)$ with $\\delta \\sim \\frac{\\log(m)}{\\sqrt{m}}$. Indeed $u(t)$ goes from $O(1/m)$ to $O(\\log(m)/m)$. In that case, $T\\sim \\frac{\\log(\\log(m))\\sqrt{m}}{\\log(m)}$ which is not $O(1)$.\n\n2) Another ambiguity about the $O$ notation is for example when it is stated: $\\kappa(0)>(1+\\delta-O(1/m))^2>(1+\\delta-O(\\delta^2/\\log(m)))^2$ in appendix D. In full generality it seems wrong for general sequences as it depends on the constants in the $O$ itself and their sign. For example $(1+\\delta))^2<(1+\\delta-(-3\\delta^2/\\log(m)))^2$. I think this statement in its current form would perhaps need an additional study of the constants, their sign or if they are zero.\n\n3) The non-linearity that is used in this paper is a ReLU: it allows to compute easily the second derivative of the neural network function (it deletes the diagonal terms of the hessian of the neural network function). My only concern is regarding the generalization to non-linearities that are not piecewise linear and hence which make another term appear in the Hessian of the neural network function, which corresponds to the second derivative of the non-linearity (correct me if I'm wrong). Could the author comment about how to handle this and if they expect the analysis to be the same and the results to still hold?", "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": "", "extra_scores": {}}, {"review_id": "5LDCetm5Rb", "reviewer": "Reviewer_e4J5", "summary": "The paper studies the \"catapult phase\" during training models with large learning rates.\nImportantly, it proposes Neural Quadratic Models (NQMs) as a tool to study it and proves that when large learning rates are used they exhibit a similar catapult phase as modern neural networks, which is not the case of other tools such as linear models, a popular theoretical tool to analysis the learning of neural networks.\nThe paper presents these findings with proofs and suggests that NQMs can be useful tools to analyze neural networks in the future.", "questions": "Some recent findings such as the deep double descent are high-dependent on the size of the neural networks and other hyper-parameters.\nIs the \"catapult phase\" here sensitive to other hyper-parameters apart from learning rates?\nI don't see many conditions for the theorem presented in the paper.", "limitations": "", "rating": 8, "confidence": 1, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "**originality** The finding of the paper seems to be novel and the proposal of using NQMs can be new as well.\n\n**quality** The finding and theory from the paper seems to be sound.\n\n**clarity** The paper is overall well-written but can be hard to follow for people who are not very familiar with the field.\n\n**significance** The contribution of the paper seems to be significant and can enable many future work on analysis with NQMs, which could lead to more useful findings than the catapult phase.", "weaknesses": "The proposed term Neural Quadratic Models seems to be unnecessary as for linear models we don't call them Neural Linear Models.", "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": "", "extra_scores": {}}, {"review_id": "LWP3q1SbQE", "reviewer": "Reviewer_F9oB", "summary": "This paper proposes using neural quadratic models -- second order taylor expansion of any function f around initial parameters $w_0$ -- as a way to study neural network dynamics. The authors start by explaining that linear dynamics perspective of neural networks falls short in explaining some of the behaviors in neural network dynamics; specifically, catapult phase of learning rate. They approximate a two layer neural network using its second order expansion. They first show that for a single training example, where tangent kernel reduces to a scalar, monotonic convergence, catapult phase, and divergence are separated based on learning rate and inverse of kernel value at random initialization. They further extend their analysis to uni-dimensional multiple example setting where the analysis is driven by the eigenvectors of the kernel matrix. Finally, the authors empirically show that for wide neural networks, catapults happen in the top eigenspace of the kernel -- similar to the multiple example setting. Experimental results suggest that catapult phase results in lower test error compared to error with sub-critical learning rate for the quadratic model, mimicing the dynamics of neural networks more closely.", "questions": "Please see above for related questions as they are more meaningful within their respective contexts.", "limitations": "", "rating": 5, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "The paper is written well and in general easy to follow. Empirical results on top eigenspace of tangent kernel for analyzing general wide neural networks is convincing.", "weaknesses": "Several assumptions about underlying the theory are not clear and different from the practice.\n\n1. The neural network in Eq (9) is initialized as $u_i \\sim \\mathcal{N}(0, I_d)$ and $v_i \\sim Unif(-1, 1)$. Can you explain why $v_i$ is initialized different from $u_i$, also different from the typical inverse fanin/fanout initialization in practice?\n\n2. The assumption that width (m) is larger than data size (n) which is assumed to be a small constant is unrealistic. While in the limit where width goes to $\\inf$, this would be the case but for any finite width network, this does not hold in practice. Is this assumption crucial, can you still assume that $n/m$ does not necessarily go to zero?\n\n3. $p_1(t)$ is the top eigenvector of $K(t)$. Given that $p_1(t)$ is not necessarily equal to $p_1(t+1)$, it is not clear how you derived $\\lambda_1(t+1)=\\lambda_1(t)-p_1(t)^TR_K(t)p_1(t)$. Could you please explain in more detail?\n\n4. In Eq(12), $R_\\lambda(t)$ is defined without the minus sign. In the following paragraph, you mention that \"$R_\\lambda(t)$ stays positive and results in monotonic decrease of kernel\" which makes sense since $\\lambda(t)=\\lambda(t-1)-R_\\lambda(t)$. But in the next paragraph, you write down $\\lambda(t)=\\lambda(0)+\\sum_{\\tau=0}^{t} R_\\lambda(\\tau)$ which suggests that $R_\\lambda(t)$ should include the minus sign. Please clarify.\n\n5. Related to above, I think it should be $\\lambda(t+1)=\\lambda(0)+\\sum_{\\tau=0}^{t} R_\\lambda(\\tau)$ or $\\lambda(t)=\\lambda(0)+\\sum_{\\tau=0}^{t-1} R_\\lambda(\\tau)$\n\n6.  You mention in the decreasing phase section in page 6 that decrease in $v(t)$ would cause a decrease in $\\kappa(t)$. But in Eq (13), reducing $v(t)$ would increase inside of square which should lead to an increase in $\\kappa(t)$; unless, $u(t)+w(t)<0$ which is not clear if it holds. Please clarify.\n\n7. In Eq (10), $1/\\sqrt(d)$ is missing from $\\sigma(u_{0,i}^Tx)$. It is present in Appendix A.\n\n8. Page 27 in the Appendix, it should be $\\Pi_1 \\mathcal{L}(t)=K_1(t)\\Pi_1 \\mathcal{L}(t-1)$", "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": "", "extra_scores": {}}, {"review_id": "FEM3AOCjJx", "reviewer": "Reviewer_JHXp", "summary": "I was asked to review this paper at the last minute, so although I did not get the chance to work carefully through the proofs, I can comment on the broad content and contributions\n\nThe paper analyzes optimization and generalization properties of neural networks using quadratic models. It shows theoretically and empirically that quadratic models exhibit the \"catapult phase\" with large learning rates, explaining a property of neural networks not captured by linear models. The quadratic model and its (changing) tangent kernel is studied analytically for the case of a single training example and multiple uni-dimensional training samples. Three regimes are identified. When catapault effects occur, better generalization is observed in quadratic models. Experiments demonstrate quadratic models parallel neural networks better than linear models in generalization with large learning rates.", "questions": "Have you analyzed how the Hessian evolves during catapult phase? This could shed some light onto the phase transition.\n\nIt may be interesting to consider explicitly varying the feature learning parameter (e.g. $\\alpha$ in Chizat et al) and see how the quadratic model differs from the linear one. This would be related to going into $\\mu$ parameterization as in Yang. \n\nCan you relate observed properties of trained NQMs to improvements in generalization, perhaps through the lens of something like kernel-target alignment for the after-kernel of the quadratic model?", "limitations": "", "rating": 8, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "The paper rigorously analyzes the catapult phase in NQMs theoretically. Prior works have attempted similar analyses, but the presentation here is particularly readable. The experiments validating theory and demonstrating applicability are thorough. They successfully demonstrate that NQMs better capture neural network behavior than linear models. The experiments in the appendix in particular provides good additional support. The mathematical analysis appears solid.", "weaknesses": "The architectures and datasets used are still relatively limited. It would also be good to highlight what the incremental value of this paper is over prior work on quadratic models. A CNN and transformer experiment would be particularly welcome. \n\nFrom a cursory literature review, I have found this paper which has relatively substantial overlap in topic, and I believe should at least be cited (Meltzer and Liu https://arxiv.org/abs/2301.07737). A more thorough set of references on recent work around the catapult effect and edge of stability would also benefit the paper. If these changes are incorporated and a more thorough related works section is added, I will likely raise my score.\n\nFrom empirical work, I've also seen that quadratic model's validation curve only tracks the NN's curve in the early stages of training, but then diverges from it. It would be nice to give an example of the limitations of the quadratic model for understanding generalization. Being clear about potential limitations of quadratic models to fully model neural networks would be welcome.", "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": "", "extra_scores": {}}, {"review_id": "2rm13ZSK4P", "reviewer": "Reviewer_5b1y", "summary": "This work shows that quadratic models exhibit the catapult phase of neural networks.", "questions": "See weakness\n\n\n------\nI feel more positive about the paper with the new title. However, I still feel that the paper needs to clarify why or why not chaos is present in the model and clarify its connection to the studies in this line.", "limitations": "", "rating": 5, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 1, "strengths": "The strength of the work lies in the soundness of the theory and that the topics covered by the paper are all pertinent problems to the contemporary deep learning theory.", "weaknesses": "There are quite a few fatal weaknesses in my opinion.\n\n1. The paper covered too many topics that it feels that it does not achieve any point satisfactorily. For example, the paper feels mistitled. The main focus of the paper is the catapult mechanism -- which does appear in deep learning but cannot represent all types of \"neural network dynamics.\" In my opinion, the catapult mechanism is a rather special / specific type of dynamics and the title is an overclaim. If the authors change the title to a more proper one, the paper could be much easier to evaluate.\n\n2. Lack of discussion of a highly relevant problem. Essentially, within the framework of quadratic models, it appears to me that the catapult mechanism is nothing but what the academia traditionally calls \"chaos.\" For example, consider a width-1 quadratic model, and compare the dynamics of GD to that of a logistic map -- the dynamics is essentially identitical -- the loss of local stability leads to chaos in the logistic map, and the same here in the quadratic model. The authors need to discuss this point in my opinion.", "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": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (poster)