Tight conditions for when the NTK approximation is valid
Authors: Enric Boix-Adserà
Etai Littwin
Venue: Accepted by TMLR ()

Abstract: We study when the neural tangent kernel (NTK) approximation is valid for training a model with the square loss. In the lazy training setting of Chizat et al. 2019, we show that rescaling the model by a factor of $\alpha = O(T)$ suffices for the NTK approximation to be valid until training time $T$. Our bound is tight and improves on the previous bound of Chizat et al. 2019, which required a larger rescaling factor of $\alpha = O(T^2)$.

Published in Transactions on Machine Learning Research (11/2023)
Tight conditions for when the NTK approximation is valid
Enric Boix-Adsera
eboin edu
MIT Electrical Engineering and Computer Science Apple
Etai Littwin
elitwwn Capple. com
Apple
Reviewed OpenReview: https: openreviee. net/ formm? id= qNTJJBBRRO
Abstract
We study when the neural tangent kernel (NTK) approximation is valid for training a model with the square loss, In the lazy training setting of Chizat et al. (2019), we show that rescaling the model by factor of a OT) suffices for the NTK approximation to be valid until training time T. Our bound is light and improves the previous bound of Chizat et al. (2019), which required a larger rescaling factor of 0
1 Introduction
In the modern machine learning paradigm, practitioners train the weights of large neural network model fw Id + Rooe via gradient- pradient-based optimizer. Theoretical understanding lags behind, since the training dynamics non-linear and hence difficult to analyze. To address this, Jacot et al. (2018) proposed an approximation to the dynamics called the NTK opprozimation, and proved it was valid for infinitely-widd networks trained by gradient 1escent The NTK approximation has been extremely influential, leading to theoretical explanations for a range of questions, including why deep learning can memorize training data (Du al., 2018; 2019; Allen-Zhu et al., 2019a;b; Arora al., 2019a; Cao & Gu, 2019; Lee al., 2019), why neural networks exhibit spectral bias (Cao et al., 2019; Basri et al., 2020; Canatar et al., 2021), and why different architectures generalize differently (Bietti & Mairal, 2019; Mei et al., 2021; Wang al., 2021). Nevertheless, in practice the training dynamics of neural networks often diverge from the predictions of the NTK approximation (see, e.g., Arora et al. (2019b)) Therefore, it is of interest to understand exactly under which conditions the NTK approximation holds. In this paper, we ask the following question:
Can we give tight conditions for when the NTK approzimation is volid?
1.1 The "lazy training" setting of Chizat et al. (2019)
The work of Chizat et al. (2019) showed that the NTK approximation actually holds for training any differentiable model, as long as the model's outputs are rescaled so that the model's outputs change by a large amount even when the weights change by small amount. The correctness of the NTK approximation for infinite-width models is a consequence of this observation, because by the default the model is rescaled the width tends to infinity; see the related work in Section 1.3 for more details.
Rescaling the model Let h : + JF be porottllypprrmmeerizee model, where F is separable Hilbert space. Let a > 0 be a parameter which controls the rescaling of the model and which should be thought of as large. We train the rescaled model ah with gradient flow to mininize smooth loss function R : F - R..²
1 Under a specific scaling of the initialization and learning rate as width tends to infinity. We use the Hilbert space notation as in Chizat et al. (2019). We can recover the setting of training a neural network fur Rd + R on a finite training dataset ((11.11)..... yı)...., (Øn, (x)) C Rd R with empirical loss function L(w) = 1. " -E= (((o(a)))
follows. Let R2 be the Hilbert space, let h(w)) [S (w)))) (oi)...... (oo and let R(w)= E A(v..y))

---

dw 1 Published in Transactions on Machine Learning Research (11/2023)
Namely, the weights wtt) RP are initialized at w(() wo and evolve according to the gradient flow
dt a² (u
(1)
NTK approximation Define the linear approximation of the model around the initial weights wo by
h(w)= h(wo) + Dh(wo)(w wo).
(2)
where Dh is the first derivative of h in w. Let un(t) be weights initialized at (i( wwo that evolve according to the gradient flow from training the rescaled linearized model ah:
dw dt a
Vw R(ah(w(()))
(3)
The NTK approximation states that
ah(w(tt) ~ ah(w(()))
In other words, it states that the linearization of the model h is valid throughout training, This allows for much simpler analysis of the training dynamics since the model linear its parameters, and so the evolution of h(w) can be understood via a kernel gradient flow in function space.
When is the NTK approximation valid? Chizat al. (2019) proves that. if the rescaling parameter a is large, then the NTK approximation is valid. The intuition is that the weights do not need to move far from their initialization in order to change the output of the model significantly, so the lincarization (2) is valid for longer. Since the weights stay close to initialization, Chizat et al. (2019) refer to this regime of training as "lazy training." The following bound is proved. 3 Here
Ro R(ah(wo))
is the loss at initialization, and
T Lip(Dh) Ro
is a quantity that will also appear in our main results.
Proposition 1.1 (Theorem 2.3 of Chizat et al. (2019)). Let R(y) lly y²11 be the square loss, where y" E F are the target labels. Assume that h is Lip(sht-ipsshitzz and that Dh is Lipsshitt in ball of radius p around Wo. Then, for any time 0 < T S app (Lip(h) /Ro),
[(a((w(T)) aT (T))l TLip(h) (/1h))////// /Ro
Ro
(4)
Notice that as we take the rescaling parameter a to infinity, then x goes to 0, so the right-hand-sidee of (4) is small and the NTK approximation is valid.
1.2 Our results
Our contribution is to refine the bound of Chizat al. (2019) for large time scales. We prove:
Theorem 1.2 (NTK approximation error bound). Let R(y) = 1 lly 1|2 be the square loss. Assume that Dh is
in ball of radius p around Wn. Then, at any time 0 < T < a² p² /Ro:
[[ah(w(T)) - ah( S min(Gn Ro, /Ro)
(5)
Furthermore, the converse is true. Our bound is tight to a constant factor. Theorem 1.3 (Converse to Theorem 1.2). For any a, T, Lip(D), and Ro, there is a model h + R, an inilialization wo R. and a target y² E such that, for the risk R(y) 1 - y )², the initial risk is R(a((wo)) = Ro, the derivative map Dh.
and
[[ah(w(T) ah(w(T)) 2 min( F Roo Ro) See Section 1.3 for discussion on the other results of Chisat et (2019).
2

---

Published in Transactions on Machine Learning Research (11/2023)
Comparison to Chizat et al. (2019) In contrast to our theorem, the bound (4) depends the Lipschitz constant of h, and incurs an extra factor of (.)) So if (i),))), Lip(h), and Rg are bounded by constants, our result shows that the NTK approximation (up to O(e) error) is valid for times T = O(ae), while the previously known bound is valid for T = O(vac). Since the regime of interest is training for large times T > I, our result shows that the NTK approximation holds for much longer time horizons than previously known.
1.3 Additional related literature
Other results of Chizat et al. (2019) In addition to the bound of Proposition 1.1 above, Chizat et al. (2019) controls the error in the NTK approximation in two other settings: (a) for general losses, but a must. be taken exponential in T, and (b) for strongly convex losses and infinite training time T, but the problem must be "well-conditooned." We work in the sctting of Proposition 1.1 instoad, since it is more aligned with the situation in practice, where we have long training times and the problem is ill-conditioned. Indeed, the experiments of Chizat et al. (2019) report that for convolutional neural networks on CIFAR10 trained in the lazy regime, the problem is ill-conditioned, and training takes a long time to converge.
Other works on the validity of the NTK approximation The NTK approximation valid for infinitely- wide neural networks under a certain choice of hyperparameter scaling called the "NTK parametrization Jacot et al. (2018). However, there is another choice of hyperparamcter scaling, called the "mean-field parametrization", under which the NTK approximation is not valid at infinite width (Chizat & Bach, 2018; Rotskoff & Vanden- Vanden-Eijnden, 2018; Sirignano & Spiliopoulos, 2022; Mei et al., 2018; 2019; Yang & Hu, 2021). It was observed by Chizat et al. (2019) that one can interpolate between the "NTK parametrization" and the "mean-field parametrization" by varying the lazy training parameter a. This inspired the works Woodworth et al. (2020); Geiger et al. (2020; 2021), which study the effect of interpolating between lazy and non-lazy training by varying a. Most work points towards provable benefits of non-lazy training (Allen-Zhr & Li, 2019; Bai & Lee, 2019; Bai al., 2020; Chen et al., 2020; Nichani et al., 2022; Ghorbani et al., 2020; Malach al., 2022; Damian et al., 2022; Bietti et al., 2022;
et al., 2021; Abbe al., 2022; 2023; Mousavi-Hossein Ba et al., 2022) although interestingly there are settings where lazy training provably outperforms non-lazy training (Petrini et al., 2022).
Finally, our results do not apply to ReLU activations because we require twice-differentiabllii of the model as in Chizat et al. (2019). It is an interesting future direction to prove such an extension. One promising approach could be to adapt a technique of Cao Gu (2020), which analyzes ReLU network training in the NTK regime by showing Lemma 5.2 that around initialization the model is "almost" linear and "almost" smooth, even though these assumptions are not strictly met because of the ReLU activations.
2 Application to neural networks
The bound in Thoorem 1.2 applies to lazy training of any differentiable model. As a concrete example, we describe its application to neural networks (a similar application was presented in Chizat et al. (2019)). We parametrize the networks in the mean-field regime, so that the NTK approximation is not valid even as the width tends to infinity. Therefore, the NTK approximation is valid only when we train with lazy training.
Let fm : Rdd - IR be a 2-layer network of width m in the mean-field parametrization Chizat & Bach (2018); Rotskoff & Vanden Vanden-Eijnden (2018); Sirignano Spiliopoulos (2022); Mei et al. (2018; 2019), with activation function 0 IR- R,
fw( vm Σ a. (Vm u.))
The weights are w (a, U) for a [a,...., am] and U = [u,,... unl. These are initialized at wo with i.i.d. Unifl-1/Vm, 1/ entries. Given training data (æ1, y/1 ), ... (In, yn), we train the weights of the network
3

---

Published in Transactions on Machine Learning Research (11/2023)
with the mean-squared loss
((w) = Σ (fw() vi), ((a,b) 2 (a - b)²
(6)
In the Hilbert, space notation, let H R", so that the gradient flow training dynamics with loss (6) correspond to the gradient flow dynamics (1) with the following model and loss function
h(w)= = / (x1), (z) E R"; R(v) - 11 vn
y
2
Under some regularity assumptions the activation function (which are satisfied, for example, by the sigmoid function) and some bound the weights, it holds that Lip(Dh) is bounded.
Lemma 2.1 (Bound on Lip(Dh) for mean-field 2-layer network). Suppose that there is constant K such that (i) the activation function a is bounded and has bounded derivatives ||r [|o'llooo ||o" I|o" S K, (ti) the weights have bounded norm ||a|| + ||U|| S K, and (ūä) the dala points have bounded norm max, || S K. Then there is constant K' depending only K such that
Lip(Dh) S K'.
Proof- See Appendix C.
Note that our bounds hold any finite width of the neural network, because we have taken initialization uniformly bounded in the interval [-1/ 1/mm Since the assumptions of Theorem 1.2 met, obtain the following corollary for the lazy training dynamics of the 2-layer mean-field network.
Corollary 2.2 (Lazy training of 2-layer mean-field network). Suppose that the conditions of Lemma 2.1, and also that the lobels are bounded in norm ll I VnK. Then there constonts c, C depending only K such that for any lime 0 < T < ca?,
[ah(w([)) ah(w(T)) s C'min(T/a,1).
Notice that training in the NTK parametrization corresponds to training the model Vmfws where fwr the network the mean-field parametrization. This amounts to taking the lazy training parameter a = / /m in the mcan-field setting. Therefore, under the NTK paramctrization with width m, the bound Corollary 2.2 shows that the NTK approximation valid until training time O(m) and the bound is O(T/
3
Proof ideas
3.1
Proof ideas for Theorem 1.2
Proof of Chizat et al. (2019) In order to give intuition for our proof, we first explain the idea behind the proof in Chizat et al. (2019). Define residuals r(1), F(() E F under training the original rescaled model and the linearized rescaled model as r(t) = = - at) and r(t) = yy - ah(w(t)). is well known that these evolve according
dr Ker and dF di
-KoF,
for the time-dependent kernel K₁ F F which is the linear operator given by K; == (h(w(t)) To compare these trajectories, Chizat et al. (2019) observes that, since Ko is p.s.d..
1 d 71 -(r - F, Ker Kor) < -(y (K: Ko)r),
2 d/
4

---

Published in Transactions on Machine Learning Research (11/2023)
which, dividing both sides by ||r F|| and using that |||| s VRg implies
dt || Fl1 - K; Kollllrl| & 2Lip(h)Lip( - wolll//wwwooolllllloooooooo Ro.
(7)
Using the Lipschitzness of the model, Chizat et al. (2019) furthermore proves that the weight change is bounded by ||w(t) woll s tv R Plugging this into (7) yields the bound in Proposition 1.1,
[ah(w(T))) - = (|((T) - F(T) < (Dap()) Roo '." tdt
=
First attempt: strengthening of the bound for long time horizons We show how to strengthen this bound to hold for longer time horizons by using an improved bound on the movement of the weights. Consider the following bound on the weight change.
Proposition 3.1 (Bound on weight change, implicit in proof of Theorem 2.2 in Chizat al. (2019)).
woll S wwol [ww(T)
TRa/a and (T) woll S VTRo/a.
(8)
Proof of Proposition 3.1. By (a) Sauchy-Schwaar and (b) the nonnegativity of the loss R,
/"
dw (a)
1/
dw
1 V a² 6-" di
[wo(T) (0()) <
|| |
||tt
S
V
T
R(h((w()))d
di
di
(b)
TRo// R.
The bound for in is analogous.
This bound (8) has the benefit of V dependence (instead of linear dependence), and also does not depend on Lip(h). So if we plug it into (7), we obtain
[ah(w(T)) (T)(w((())))))) (T)) (iip(h) L'" Vidt
4 (iP
3
This improves over Proposition 1.1 for long time horizons since the time dependence scales 73/2 instead of as T However, it still depends on the Lipschitz constant Lip(h) of h, and it also falls short of the linear in T dependence of Theorem 1.2.
Second attempt: new approach to prove Theorem 1.2 In order to avoid dependence on Lip(h) and obtain a linear dependence in T, we develop a new approach. We cannot use (7), which was the core of the proof in Chizat al. (2019), since depends on Lip(h). Furthermore, in order to achieve linear T dependence using (7), we would need that |/w woll = 0(1) for constant that does not depend the time horizon, which is not true unless the problem is well-conditioned.
In the full proof in Appendix A, we bound (Ir(T) = [[ah(w(T)) - añ (w())) which requires working with a product integral formulation of the dynamics of r to handle the time-varying kernels K, (Dollard & Friedman, 1984). The main technical innovation in the proof is Thcorem A.8, which is a new, general bound on the difference between product integrals.
To avoid the technical complications of the appendix, we provide some intuitions here by providing proof of simplified theorem which does not imply the main result. We show:
Theorem 3.2 (Simplified variant of Theorem 1.2). Consider r'(t) F which is initialized as r' r(0) and evolves de' KTr'. Then, dt
[I((()) - -(())l & min(3k Ro, V (8)))
(9)
5

---

Published in Transactions on Machine Learning Research (11/2023)
It is at the final time = T that the kernel K2 can differ the most from Ko- So, intuitively, if we prove in Theorem 3.2 that rT()) and F(T) are close, then the same should be true for r(T) and F(7) in Theorem 1.2. For convenience, define the operators
Dh(w))) and B ()(ww())) Dh(wo))
Since the kernels do not vary in time, the foosed-form solution is
r' (1) = (A+B)
and F(I) = e Al
We prove that the time evolution operators for r' and F are close in operator norm. Lemma 3.3. For any 2 0, we have ||e (A= (A+B)) e e
e llll s B] Vt.
(A (A + (B) (A+ By the fundamental theorem of calculus Proof of Lemma 3.3. Define Z(C) = -
||e (A=)))) - -1 11ll ||e - - - S/" d AC) S sup [Ef0.1] dc
d 2
dc
Using the integral representation of the exponential map (see, e.g., Theorem 1.5.3 of (Dollard & Friedman, 1984)),
de 7(C)
=11 (1 e1-T))((( d
((S))e drl
dç
d
1|t (1--))Z()) (AT AABL B + B A+ 2BB B)e B)e 2(()
2() drl|+ /' ((--))(()) BT (A+ (B)e 7(() dr|l
$ || (1--))((() (A + (B) Be 7(0)
7
(Term 1)
(Term 2)
By symmetry under transposing and reversing time, (Torm 1) = (Torm 2), so suffices to bound the first
term. Since 2()) < 1,
(Term 1) s1 [e(1-))2(0(
(A (A+CB) (IIBB)IIBB
|||| 2(()) |dr -
s B||
[|e0)1)2((1 (A + (B)
BI| /."" (1 [0
sup V 2(1-
X20
v / (1/(2e(1--))
B1I
where in the third-to-last line we use the Courant-Fischer- Fischer- Weyl theorem and the fact that Z({) is negative semidefinite. Combining these bounds ||e (A+BB) (A- (A-B)(
A
-- A At
2 2 BI.

---

Published in Transactions on Machine Learning Research (11/2023)
Finally, let combine Lemma 3.3 with the weight-change bound in Proposition 3.1 to prove Theorem 3.2. Notice that the weight-change bound in Proposition 3.1 implies
B|| & Lip woll S (ip((D))TRR///
So Lemma 3.3 implics
(T) F(T)I RLpp(DN)T/Raa (1 = 20]]]
2|)r -2v2Ro implies (9). Thus, have shown Theorem 3.2, which is the result of Theorem 1.2 if we replace by r'. The actual proof of the theorem handles the timo-varying kernel K,, and is in Appendix A.
Combining this with (T) S [T)++ [F(T)
3.2 Proof ideas for Theorem 1.3
The converse in Theorem 1.3 is achieved in the simple case where h(w) 2 bbw for /7 and Lip(Dh), and 0 and R(y) 2 /2R as we show in Appendix B by direct calculation.
4 Discussion
A limitation of our result is that it applies only to the gradient flow, which corresponds to SGD with infiitesimally small step size. However, larger step sizes are beneficial for generalization in practice (see, e.g., Li et al. (2021); Andrinshchenko et al. (2022)), so would be interesting to understand the validity of the NTK approximation in that setting. Another limitation is that our result applies only to the square loss, and not to other popular losses such as the cross-entropy loss. Indeed, the known bounds the setting of general losses require either a "well-conditoonig"" assumption, or taking a exponential in the training time I' (Chizat et al., 2019). Can one prove bounds of analogous to Theorem 1.2 for more general losses, with depending polynomially on T, and without conditioning assumptions?
A natural question raised by our bounds in Theorems 1.2 and 1.3 is: how do the dynamics behave just outside the regime where the NTK approximation valid? For models where Lip(h) and Lip(Dh) are bounded by a constant, can we understand the dynamics the regime where T 2 Ca for some large constant and a >> C, at the edge of the lazy training regime?
Acknowledgements
We thank Emmanuel Abbo, Samy Bengio, and Joshua Susskind for stimulating and helpful discussions. EB also thanks Apple for the company's generous support through the AI/ML fellowship.
References
Emmanuel Abbe, Enric Boix-Adsera, and Theodor Misiakiewicz. The merged-staircase properly: a necessary and nearly sufficient condition for sgd learning of sparse functions on two-layer neural networks. In Conference Learning Theory, pp. 4782-4887. PMLR, 2022.
Emmanuel Abbe, Enric Adsx-a, and Theodor Misiakiewicz. Sgd learning on neural networks: leap complexity and saddle-to-saddle dynamics. arXiv preprint 1000 2023.
Zeyuan Allen-Zhu and Yuanzhi Li. What can resnet learn efficiently, going beyond kernels? Advances in Neural Information Processing Systems, 32, 2019.
Zeyuan Allen-Zhu, Yuanzhi Li, and Yingyu Liang. Learning and generalization in overparameterized neural networks, going beyond two layers. Advances in neural information processing syslems, 32, 2019a.
Zeyuan Allen-Zhu, Yuanzhi Li, and Zhao Song. A convergence theory for deep learning via over- parameterization. In International Conference Machine Learning, pp. 242 252. PMT.R, 2019b.
7

---

Published in Transactions on Machine Learning Research (11/2023)
Maksym Andriushchenko, Aditya Varre, Loucas Pillaud-Vivien, and Nicolas Flammarion. Sgd with large step sizes learns sparse features. arviv preprinl or arXiv:2210.05.0 2022.
Sanjeev Arora, Simon Du Wei Hu, Zhiyuan Li, and Ruosong Wang. Fino-grained analysis of optimization and generalization for overparameterized two-layer neural networks. In Proceedings of the 36th International Conference on Machine Learnöng, pp. 322-332, 2019a.
Sanjeev Arora, Simon S Du, Wei Hu, Zhiyuan Li, Russ R Salakhutdinov, and Ruosong Wang. On exact computation with an infinitely wide neural net. In Advances in Newral Information Processing Systems, 2019b.
Jimmy Ba, Murat A Erdogdu, Taiji Suzuki, Zhichao Wang, Denny Wu, and Greg Yang. High-dimensional asymptotics of feature learning: How one gradient step improves the representation. arXiv preprint ar arXiv:2205.01443 2022.
Yu Bai and Jason D Lee. Beyond linearization: On quadratic and higher-order approximation of wide neural networks. Xrxiv preprint arXic:1910.01619, 2019.
Yu Bai, Ben Krause, Huan Wang, Caiming Xiong, and Richard Socher. Taylorized training: Towards better approximation of neural network training at finite width. arXiv preprint (56::0002.00111 2020.
Ronen Basri, Meirav Galun, Amnon Geifman, David Jacobs, Yoni Kasten, and Shira Kritchman. Frequency bias in neural networks for input of non-uniform density. In International Conference on Machine Learning, pp. 685-694. PMLR, 2020.
Alberto Bietti and Julien Mairal. On the inductive bias of neural tangent kernels. Advances in Newral Information Processing Systems, 32, 2019.
Alberto Bietti, Joan Bruna, Clayton Sanford, and Min Jae Song. Learning single-index models with shallow neural networks. Xice preprint orXiv:2210.155: 15651, 2022.
Abdulkadir Canatar, Blake Bordelon, and Cengiz Pehlevan. Spectral bias and task-model alignment explain generalization in kernel regression and infinitely wide neural networks. Nature communicctions, 12(1):2 2021.
Yuan Cao and Quanquan Gu. Generalization bounds of stochastic gradient descent for wide and deep neural networks. In Advances in Neural Information Processing Systems, 2019.
Yuan Cao and Quanquan Gu. Generalization error bounds of gradient descent for learning over-parameterized deep relu networks. In Proceedings of the AAAI Conference Artificial Intelligence, volume 34, pp. 3349933366 2020.
Yuan Cao, Zhiying Fang, Yue Wu, Ding- Xuan Zhou, and Quanquan Gu. Towards understanding the spectral bias of deep learning. arXiv preprint arXiv:1912.011900 2019.
Minshuo Chen, Yu Bai, Jason D Lee, Tuo Zhao, Huan Wang, Caiming Xiong, and Richard Socher. Towards understanding hicrarchical learning: Benefits of neural representations. Advonces in Neural Information Processing Systems, 33:22134 22145, 2020.
Lenaic Chizat and Francis Bach. On the global convergence of gradient descent for over-parameterizel models using optimal transport. Advances in neural information processing systems, 31, 2018.
Lenaic Chizat, Edouard Oyallon, and Francis Bach. On lazy training in differentiable programming. Advances in Neural Information Processing Systems, 32, 2019.
Alexandru Damian, Jason Lee, and Mahdi Soltanolkotabi. Neural networks can learn representations with gradient descent. In Conference on Learning Theory, pp. 5413 5452. PMLR, 2022.
John Day Dollard and Charles N. Friedman. Product Integration with Application Differential Equations. Encyclopedia of Mathematics and its Applications. Cambridge University Press, 1984. doi: 10.1017/ C100978107340701.
8

---

Published in Transactions on Machine Learning Research (11/2023)
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-16886 PMLR, 2019.
Simon S Du, Xiyu Zhai, Barnabas Poczos, and Aarti Singh. Gradient descent provably optimizes over- parameterized neural networks. arXiv preprint arXiv:1810 111.:0810.00054, 2018.
Mario Geiger, Stefano Spigler, Arthur Jacot, and Matthieu Wyart. Discntangling feature and lazy training in deep neural networks. Journal of Statistical Mechanics: Theory and Experiment, 2020(11):113001, 2020.
Mario Geiger, Leonardo Petrini, and Matthieu Wyart. Landscape and training regimes in deep learning. Physies Reports, 924:1-18, 2021.
Behrooz Ghorbani, Song Mei, Theodor Misiakiewicz, and Andrea Montanari. When do neural networks outperform kernel methods? Advances in Newral Information Processing Systems, 18:14820-14800, 2020.
Arthur Jacot, Franck Gabriel, and Clément Hongler. Neural tangent kernel: Convergence and generalization in neural networks. Advances in neural information processing systems, 31, 2018.
Jachoon Lee, Lechao Xiao, Samuel Schoenholz, Yasaman Bahri, Roman Novak, Jascha Sohl-Dickstein, and Jelfrey Penninglon. Wide neural networks of any depth evolve as linear models under gradient descent. Advances in neural information processing systems, 32, 2019.
Zhiyuan Li, Sadhika Malladi, and Sanjeev Arora. On the validity of modeling sgd with stochastic differential equations (sdes). Advances in Neural Information Processing Systems, 34:12712-12725, 2021.
Eran Malach, Pritish Kamath, Emmanuel Abbe, and Nathan Srebro. Quantifying the benefit of using differentiable learning over tangent kernels. In International Conference on Machine Learning, pp. 7379- 7389. PMLR, 2021.
Song Mei, Andrea Montanari, and Phan-Minh Nguyen. A mean field view of the landscape of two-layer neural networks. Proceedings of the National Academy of Sciences, 115(33):E7665 -7666-67771,66 2018.
Song Mei, Theodor Misiakiewicz, and Andrea Montanari. Mean-fiet theory of two-layers neural networks: dimension-free bounds and kernel limit. In Conference on Learning Theory, pp. 2388-288820 2464. PMLR, 2019.
Song Mci, Thcodor Misiakiewicz, and Andrea Montanari. Learning with invariances in random features and kernel models. In Conference on Learning Theory, pp. 3351 3418. PMLR, 2021.
Alireza Mousavi-Hosseini, Sejun Park, Manuela Girotti, Toannis Mitliagkas, and Murat A Erdogdu. Neural networks efficiently learn low-dimensional representations with sgd. arXiv preprint XrXiv:2209.1 2022.
Eshaan Nichani, Yu Bai, and Jason Lee. Identifying good directions to escape the ntk regime and efficiently learn low-degree plus sparse polynomials. arXiv preprint arXiv:2206.n 2.50:2206600688, 2022.
Leonardo Petrini, Francesco Cagnetta, Eric Vanden-Eijnden, and Matthieu Wyart. Learning sparse features can lead to overfitting in neural networks. arXiv preprint arXiv:2200. 12314, 2022.
Grant M Rotskoff and Eric Vanden-Fijnden. Neural networks as interacting particle systems: Asymptotic convexity of the loss landscape and universal scaling of the approximation error. stat, 1050:22, 2018.
Justin Sirignano and Konstantinos Spiliopoulos. Mean field analysis of deep neural networks. Malhemalics of Operations Research, 47(1):120-11) 152, 2022.
Terence Tao. Topics in random matrix theory. Graduate Studies in Mathematics, 132, 2011.
Sifan Wang, Hanwen Wang, and Paris Perdikaris. On the cigenvector bias of fourier feature networks: From regression to solving multi-scale pdes with physics-informed neural networks. Computer Methods in Applied Mechanics and Engineering, 384:1 384:113938, 2021.
9

---

Published in Transactions on Machine Learning Research (11/2023)
Blake Woodworth, Suriya Gunasckar, Jason D Lee, Edward Moroshko, Pedro Savarese, Itay Golan, Daniel Soudry, and Nathan Srebro. Kernel and rich regimes in overparametrized models. In Conference on Learning Theory, pp. 3635 3673. PMLR, 2020.
Greg Yang and Edward Hu. Tensor programs IV: Feature learning in infinite-width neural networks. Tn International Conference on Machine Learning, pp. 11727 11737. PMLR, 2021.
A Proof of Theorem 1.2
A.1 Notations
We let Ro == R(ah(wa)) denote the loss at initialization. We define the residuals ((t), E F under training the original model and the linearized model as
r(t) y" - ah((w(()), and F(() - at)(
Since the evolution of w and w is given by the gradient flow in (1) and (3), the residuals evolve as follows. We write Dh dh denote the adjoint of Dh dw
dr dr dw dt dw dt aDh(w)( (-VFa aD(w)( Dh(w)
T VR(ah(w)) Dh(w))
since VR(ah(u=)) -w" ah(w)) r. An analogous result can be derived for the residual under the linearized dynamies:
di = DD((w)heeee Dh(w) Fr.
dt
For any time 2 0, define the kernel K₁ F - F as
K, = Dh (h(ww(t)
Since Ko = (M(((()))D(((()))))
= Dh (h((((t) (h((()) we can write the dynamics in compact form:
dr Ker and di -Kor
(10)
A.2 Proof overview
Our proof of Theorem 1.2 is outlined in the flowchart in Figure 1. First, we use Lemma A.1 to argue that the change in the weights is bounded above during training. This implics several basic facts about the continuity and boundedness of the kernel K, over time, which are written as Lemmas A.2, A.3, A.4 A.5. These lemmas allow us to write the dynamics of r and F using product integrals Dollard & Friedman (1984) in Lemma A.6. Product integrals are a standard tool in differential equations and quantum mechanics, but known results on product integrals do not suffice to prove our theorem. Thus, in Theorem A.8, we prove new operator-norm bound between differences of product integrals. Theorem A.8 implies main Theorem 1.2 when specialize to the case of gradient-flow training of model, and combine it with the weight-change bound from Lemma A.1. The bulk of the proof is dedicated to showing Theorem A.8. The proof is an interpolation argument, where we interpolate between the two product integrals being compared, and show that the operator norm of the derivative the interpolation path is bounded. In order to show this, we require several new technical bounds most notably Claim 6, which shows that for any matrices Y ga and any t 2 we have
||e (X- +X +Y) -xx < <3|||]
10

---

Published in Transactions on Machine Learning Research (11/2023)
Theorem 1.2 (NTK approximation error bound)
Lemma A.6 (Product integral
Theorem A.8 (Bound difference between product integrals)
formulation of dynamics)
Lemmas A.1, A.2, A.3, A.4, A.5 (Weight change bounded) (Weights in ball around init) (Dh Dh bounded)
Finite- cinttedimeesiionl case of
Lemma A.9 (Reduction to fiillediimensional case)
Theorem A.8 (Bound difference between product integrals)
(K, bounded)
(t + K. contimuous)
Lemma A.10 (Bound derivative of interpolation between product integrals)
Claim 2
Claim 1 (Product integral is contraction)
(Bound on product integral times interpolant between kernels)
Claim 4
Claim 3
(Approximate product integral with exponential)
(Bound - xxX xxxxxlI) xII)
Claim 6
(Telescoping argument, usymmettic matrices)
Claim 5
('elescoping argument, symmetric matrices)
Figure Proof structure.
11

---

Published in Transactions on Machine Learning Research (11/2023)
A.3 Basic facts about boundedness and continuity of the kernel
We begin with series of simple propositions. Recall Proposition 3.1, which we restate below with slight strengthening in (11) which will be needed later on. Lemma A.1 (Restatement of Proposition 3.1). For any time t,
(t) -woll - vtRo/a and || -wol S tRoo/a
and furthermore
{ dw dt ||t VtRo/a and [ dw di dr
dr S VtRu/a.
(11)
Proof. The statement (11) is also implied by the proof of Proposition 3.1 in the main text.
The first implication is that the weights un(t stay within the ball of radius p around wo, during the time-span that we consider.
Lemma A.2. For any time 0 S TT, we have (1w(t) - -woll S.
Proof. Immediate from Lemma A.1 and the fact that T S p²aa Ro.
This allows to use the bounds Lip(h) and Lip(Dh) on the Lipschitz constants of h and Dh. Specifically, the kernels Ke and Ko stay close during training, in the sense that the difference between Dh and Dh during training is bounded.
Lemma A.3. For any time Stt T, we have
[Dh(ww(t) Dh(w(()) < Lip(Dh) vioo/a
Proof. Since (a) h is the linearization of at wo and (b) [[w - woll s min(p, /ERo/a) by Lemin A.2 and Lemma A.1,
[)((w(t)) (t)) @ [Dh(w(t)) (wwwl) Lip(Dh) Ltt
And therefore the kernel K, is bounded at all times 0 < 1 1.5 S in operator norm. Lemma A.4. For any time & ttsT STT, T, we have |K || < (h)) + Ra)
Proof. By triangle incquality and Lomma A.3,
|K Koll [I)((((())D(((()))) (M(w((0))Dh(w((())
< ()D)
<
< 3
And finally we note that the kernel evolves continously in time.
Lemma A.5. The map + K, continuous (in the operator norm topology) in the interval [0, T].
Proof. First, t w(t) continuous in time, since solves the gradient flow. Second, we know that wr(t) is in the ball of radius p around wa by Lemma A.2, and in this ball the map w + Dh(w) continuous because Lip( (h) < co. Finally, Dh + DhDhT T is continuous.
12

---

Published in Transactions on Machine Learning Research (11/2023)
A.4 Product integral formulation of dynamics
Now we can present an equivalent formulation of the training dynamics (10) in terms of product integration. For any S I < T, let P(y,x) IPP + F solve the operator integral equation
(y =1 S"d xt)(t,
(12)
A solution 2)y, is guaranteed to exist and to be unique:
Lemma A.6. The unique solution to the integral equation (12) is given as follows. For any define Sm.j (y z)mm/m) and d (y z)/m, and let Ply, x)) be the product integral
5 ST
P(,,)) x) = II K,ds
lim m=+00
II f=1
5K,,,
lim e 8K - e 8k.ze 8K.1 m=000
Proof. Existence, uniqueness, and the expression an infinite product guaranteed by Theorems 3.4.1, 3.4.2, and 3.5.1 of Dollard & Friedman (1984), since + K2 lies in (0, T), which is the space of "strongly integrable" functions on [0, T] defined in Definition 3.3.1 of Dollard & Friedman (1984). This fact is guaranteed by the separability of F and the continuity and boundedness of t -> Kg (Lemmas A.4 and A.5).
The operators P(y, r) are the time-evolution operators corresponding to the differential equation (10) for the residual error Tz. Namely, for any time 0 S t S T,
Tt P(t. 0)ro
On the other hand, the solution to the linearized dynamies (10) is given by Kotra.
since eot is the time-evolution operator when the kernel does not evolve with time.
A.5 Error bound between product integrals
To prove Theorem 1.2, it suffices bound Pt,,, - K-Kee],, the difference of the time-evolution operators under the full dynamics versus the linearized dynamics. We will do this via general theorem. To state it, must define the total variation norm of time-indexed sequence of operators:
Definition A.7. Let (Cr)el wi be a sequence of lime-bounded operators C. IRP + so that C. is continuous in the interval [z, y). Then the total variation norm of {Ci) is
Ap-1
V((Cthees.))
sup PEP
10 C..ill
where the supremum is taken over partitions P (P {a h 12 S S hmp-1 Shp f}} of the interval [z, v].
We may now state the general result.
Theorem A.8. Let F be separable Hilbert space, and let {Aiheeo, (Biheeoo,y be time-indered sequences of bounded operalors A, : IRP + F and B, Rp F such that + Ae and 1 + B, are contivuous in [0, T). Then,
T II D
T II
A.A.
B.B. do
||
sup ||, + 37 V([A, Be)ee(o,r)) [6(0.]
If we establish Theorem A.8, then may prove Theorem 1.2 follows.
13

---

Published in Transactions on Machine Learning Research (11/2023)
Proof of Theorem 1.2. for each 2 0 choose the linear operators A,, B, - A, B. ()))(()) so that A,Ag = K, and B2 B, Ko- We know that A,, Br are bounded by Lemma A.4, and that + A, is continuous by Lemma A.5. (Also t + Be is trivially continuous). So we may apply Theorem A.8 to bound the difference in the residuals.
Dh(w(()) and
We first bound the total variation norm {Ae - BeJeEO,TT, By (a) the fact from Lemma A.2 that w(t)) is ball of radius most p around wo where w - Dh(w) differentiable, since it solves gradient flow; and (c) Lemma A.1,
(b) the fact that t + w(() is
UP
V([A. BeheE(o))
sup
DD(w((++1)) D)(w(()) (O) Dh(wo(()) + Dh(w(()))
PEP
sup
|D Dh(ww(4))
PEP
(0)
sup I Dhp
PEP
(
Di)
du de lile
Tip(Dh)
When we plug the above expression into Theorem A.8, along with the bound ||A, Bell s LtRoDa from Lemma A.3, obtain:
1
Kede
II
Kods rrol
I (Tip( Dh) TRo - 3I Ro/o) 2Ro
= (2x + 3x²²) 3xx²)) /
where k LLip(Dh))RRg
Also note that
|rr Frll S ||r + |1rll (2R(ah(u(T))) + (2((ah(((T)))
(2
s 222,
since the gradient flow does not increase the risk. So
||r - Frll S min(2n (2x + 3xx 2) 2R & 6x Rg
It remains only to prove Theorem A.8.
A.6 Reduction to the diitensiinniional case
We first show that in order to prove Theorem A.8, it suffices to consider the case where F is a finite-dimensional Hilbert space. The argument is standard, and uses the fact that F has countable orthonormal basis.
Lemma A.9. Theorem A.8 is true for general separable Hilbert spaces F if it is true whenever F is Snile-dimensiioal.
Proof. Suppose that F is a separable Hilbert space and A,, B, R2 + F satisfy the hypotheses of Theorem A.8. Let (Ji)ee be countable orthonormal basis for F, which is guaranteed by separability of F. For any n, let.
14

---

Published in Transactions on Machine Learning Research (11/2023)
Pa : -F be the linear projection operator defined by
Pn((f) fi. 1sisn otherwise
0,
By (a) Duhamel's formula in Theorem 3.5.8 of Dollard & Friedman (1984), (b) the fact that ll le=, s
and ||e A. P.
because A.Ag and PRA,A,P, are positive semidefinite.
T II
A.A. ds
T
do
(A,A, - -PAA,A
II
PuA.A.TT P.
drll
e
T
T
I|
A, de II - PRA,A, n P.A.A.P delde
1 |4 - P
(13)
We have chosen Pn. so that for any bounded linear operator M : F + F, we have lim P,MP,, = M. Since 7 + A, is continuous in т, and A, is bounded for each T, the expression in (13) converges to 0 as n + cc. By triangle inequality, we conclude that
T II
T II
B.B. dell < sip IIII E TIE.AAAA A.A, de III BI..B.B
P. P.A.A do
P. Cc..BB.B
0
0
Notice that An and P, Be are bounded maps from Rº to span fu), and + PaA and E P,B. are continuous and bounded. So using the theorem in the case where F is finite-dimensional, the right-hand side can be bounded by
T II 0
lim limsupl| n-0000
Pa P.A.AIPT
II
P.B. B.B. dell
lim sup sup I
|n - P,, B., 3T V((PRAL- - -P (r
n+ tE(O,T]
= ( sup ||A, - (21))(2VT +3 + -.heeoor))dt).. (0
A.7 Bound for the mnitt-dimensional case
We conclude by proving Theorem A.8 in the finite-dimensional case, where we write A,, B, E Raxd as time-indexed matrices. In order prove a bound, we will interpolate between the dynamics we wish compare. Let
C, == B, - A e Radd
For any E [0, 1] and t [0, we define the kernel
Ke = + +ci + (C.)) Rnxn.
This interpolates between the kernels of interest, since on the one hand, Ke = A.A] and on the other K₁, = B,B, For any x, E [0, let
P(y, = II
Keedt E Raxn
z
15

---

Published in Transactions on Machine Learning Research (11/2023)
The derivative 0 Pyy.z;)) () exists by Theorem 1.5.3 of Dollard & Friedman (1984), since (i) Kig is continuous äc in t for each fixed 5, (ii) Kig is differentiable in Ç in the L' sense", and (iii) has partial derivative 0Kt.c that is integrable in t. The formula for 8ç Ply. x; is given by the following formula, which generalizes the integral representation of the exponential map:
aç
0 x(y, = J P(y.t:))
"
0KLs P(t, x; ()dt
(14)
oç
Our main lemma is: Lemma A.10. For all Ç [0, 1], we have the bound
aP(T.0: S sup [[C,I)(2VT++ 3T V((Cohetorr))) te(o,]]
oc
This lemma suffices to prove Theorem A.8.
Proof of Theorem A.8. Using the fundamental theorem of calculus,
T B,B. B..BB.BB
1) P(TT,0:011 f(² ōç
@P(TT, 0:5) |
III
c CTAAA
A.AJ de II e
which combined with Lemma A.10 proves Theorem A.8.
A.8 Proof of Lemma A.10
By a direct calculation,
oK,5 = (A, + (CCCC + C((A+ dç
so, by (14);
0
I' P(T.EC)((AA+ + CC(A++ (CA T
P(T,0:{) oç
P(T,E + P(t,O: P(,0E)) f" P(I,, + (C.) P(t, d)d.
Mi
M2
The arguments are similar for bounding M1 and M2, so we only bound Mi. We will need two technical bounds, whose proofs are deferred Section A.9.
Claim 1. For any 0 < t < T, we have Ptt 0; ()Il < 1. Claim 2. For any 0 S I T, we have [I (()) (C,)ll W 3V((Ca)se))m)
Using (a) Claim 1, and (b) Claim 2,
(a)
||C,I) /' [P(T, (Co)ldd
[M ||
sup
(15)
tE(D,1]
( 2VT-1" + + - +3)((Cshoettt)
sup sE[o,]]
|(C.1)
(16)
[[CI)(VV+ +3 S' V((CChaele.m)dt)
( sup [E[O.T]
(17)
5 1 Keg-Kas oKo lst = 0, since the matrices A,, B, are uniformly bounded.
'For any E [0, 11, lime
2-6
25
5The proof ls consequence Duhamel's formula. This tool used variety of contexts, including perturbative analysis of path integrals quantum mechanics Dollard & Friedman (1984).
16

---

Published in Transactions on Machine Learning Research (11/2023)
Lemma A.10 is proved by noting that, symmetrically,
[[+]))(VT +3 [ v((Cabeeuu))d))
|| S ( sup
1e[0,]]
and, for any E [0,
V((CA) + V((C,) = v((Cabeefor1).
A.9 Deferred proofs of Claims and 2
Claim 1. For any 0 s t s T, we have [[ 0: C)l - 1.
Proof. This follows from the definition of the product integral as infinite product, and the fact that each term e -8 in the product has norm at most because Kt,ç is positive semidefinite.
]
In order to prove Claim 2, we need two more claims:
Claim 3. For any X E Raxd and 2 0,
le xxxx S
1
2
Proof. Since X is positive semidefinite,
||e xxe-xx xx X = V xxxlexxxxxxxe-e
S suu - = = V 1
2/7
A20
X20
Claim 4. For any time-indezed sequence of matrices (Xi)tefo,r) in gaxd such that t + X, is continuous in [0, T), and any s < TT
b xxxxxxxa Xall S (V((s))sela.))
|le (b-a))-a))
This latter claim shows that we can approximate the product integral with an exponential. The proof is involved, and is provided in Section A.10. Assuming the previous two claims, we may prove Claim 2. Claim 2. For any 0 < 1 S T, we have [P, (17,((9)(AA)) (ci)ll S श + 3V((Ca))ee).))
Proof. By Claim 3, since KT.c (AT + +CCT) + (CT)T
1 |le ( +CTr]l 2/T-T -
By the triangle inequality, it remains to prove
|le Kr.(T -D (AT + (CT) P(T, + (C.) 3V((CC)
and this implied by Claim 4, defining X, A, (C₁ and b= T.
17

---

Published in Transactions on Machine Learning Research (11/2023)
A.10 Deferred proof of Claim 4
The proof will be by interpolation, using the integral representation of the exponential map provided in (14), similarly to the main body of the proof, but of course interpolating with respect to a different parameter. We begin the following claim, which we will subsequently strengthen.
Claim 5. For any symmetric X, E R and 2 0,
-x*x (ee (X +Y) - X S 3||||
Proof. For any T E [0, 1], define X()) = X + +Y, which interpolates between and Y. Then by (a) the derivative of the exponential map in (14), and (b) the fact that X'(T) Y and ||elll-- (7 s 1,
|e (x + Y) - x XI|
=11 J/- d
( (-x(()(((T)) drll
dr
= || ! dr
{ c
(-x2 x drll
d x( -
< sup
dr
rE/0.1]
(a)
||e ---(()1//2X x2(7)4/2
sup rE(0.1]
-11/2) (2/2) X - ( x2
(t/2) e XXX
(x2/2) dsll
(
|| + (1(122) /' +=0=)))) X X (3(711/2 dsll
< sup
rE(0,1]
T.
[l(//2)
and by (a) ozrdns dy 121/2 = Vet, and (b) [le-M1 < 1 if M is p.s.d.,
T < (t/2)
(6)
$ (t/2) /'
(=-1=))((22(( x()))) (T ds
Vet
$ V1
||
x + ||e -1--))x (xx2())/22 |ds
2/e1
(4) Vt
lle-
(1-s))
v(1 8)t ||llle- (+)))/ Vst
2e
(9)
||
1
ds
2e
Jo
V8
V(1 -8)
||]]
(2V5 2
8)6
2e
< ||
Similarly, T2 |||||-
18

---

Published in Transactions on Machine Learning Research (11/2023)
Claim 6. For any X, E Rxxd (not necessarily symmetric) and 2
|e |e- ((++)) -e -e-xxexxxx XX| S 3||
Proof. We will use Claim 5, combined with general method for lifting facts from symmetric matrices asymmetric matrices (see, e-g., Tao (2011) for other similar arguments). Assume without loss of generality that d, since otherwise we can pad with zeros. Define ñ 2n and
X 0
O Y
Y
X
R and Y
E
R
X
0
These are symmetric matrices by construction. Furthermore,
X O
XX and (X 2)2= = [X [X+ (X+))((+X)
o (X (x+
X
Because of the block-diagonal structure of these matrices,
e xx 0
(X+))
82
and e -8+88) x-xxx
So
x-xXXXx
-
e-xxxxx
0
Similarly,
(x (X+P)= (
0
+Y
For any matrix M Raxn, we have ||M||| supess
|| So for any matrices M1, M2 RRxx we have
Mi
0
Mi
v
> sup
sup -ES
| = M2
M2
M2
0
w
This means that, using Claim 5,
le (x (X (X - -xx
-
+8) -x 53|81]. S ||e (x+x)x))) -XXTx
Finally, using the symmetry of Y and the block-diagond structure of y
Y o
1/2
|| = 1/2
max J|Y 1// 11//)- |А]]-
0 YYT
We conclude by using these results to prove Claim with telescoping argument. Claim 4. For any inneiindzed sequence of matrices (Xi)tefo,ry in gaxd such that t + X, is continuous in [0, T), and any S STS
b
|le (b-a) Xo II
e X.xJass ds Xall S 33(((x))seja.)
19

---

Published in Transactions on Machine Learning Research (11/2023)
Proof. We will do this by approximating the product integral by finite product of m matrices, and taking the limit - oc. For any 2 0 and j E {O..... m), and /. E [0. 7), let (m.j = (b a)(i/m) + and define the finite-product approximation to the product integral for any l E m!.
Pm,k =
II j=++1
Qms
Smik
where
Qm. = -Ximi Xim. (b-a)/m
Notice that for k = 0 we have
Pm, II
-Xim X km (b-a)/m
and so by continuity of 8 + X, for 8 E [0, and boundedness Ns, we know from Theorem 1.1 of Dollard & Friedman (1984) that
b
lim Pm. = II m-
X,XX
a
This means that
[
6 c x,xTdo do
[[
|le X Xo (b-a) Xb
Xal = ||P X Xtmm
exexxXx7a do
Xall
lim m-000
Pm Xtm Xemmm PmmuXemoll
Finally, telescoping and (a) using 1lQm.j || S 1 for all j, and (b) Claim 6,
|P Ptm.X.m. Xtm..
||
Pmm Pm. Xtmm all
km
II 1++
Qms
Qmi (a (Om Xe Xem.k-ill
1(m,,)) 1-1 X, N.m.n (Qm, 1k-- 1 Xemiiril
31m Xem.e-ll
$ <3(((X)
The lemma follows by laking m - xc, since the bound is independent of m.
B Proof of Theorem 1.3
The converse bound in Theorem 1.2 achieved in the simple case where h. R R is given by h(w) aw + =bw² for / and b Lip(Dh). We also let wo 0 and v2Ro so that all conditions are satisfied. The evolution of the residuals r(t) = y² ah(w(t)) and F(() = y² ah(m((t) is given by dr
= Ker and dF dt
di = Kor. -of,
where K, = (h(w(t)) DM(((()) = (a + b and Ko a² Since 1y" a(aw + hw?),, we express
the evolution of the residuals and as:
dr di.
and dr - (a² + 2b(y" r)/r)r
27
(18)
db
20

---

Published in Transactions on Machine Learning Research (11/2023)
Since b(y" -r)/a at all times, must have a² 2b(y" - r)/a a². This means that at all times
r(1) F(I) - -a't
So, at any time 12 7/2,
r(t) (()))) e - y 11211""
By plugging this into (18), for times 2 T/2,
dr dt
- + - - (110D/)0/000 S -aa² +1.++.11 1.1x/T)r.
So, time T, assuming that S without loss of generality,
r(7) S (T(T/2)e
(TT/2)e-112--000 1/2-0.55m < -
-0.55m Sy¹e e
- 0..kk).
[ah(w(TT) ah(w(T)) = (r(T) F(T)) 1e- (1 -0.4k)e vy 2 0.4kc 0.AKe-1/2R '2RRo 2 v Ru/5.
C Deferred details from Section
The bound on Lip(Dh) for 2-layer networks is below.
Lemma C.1 (Bound on Lip(Dh) for mean-field 2-layer network). Suppose that there is a constant such that (i) the activation function o bounded and has bounded derivatives ||a]]oes ||o ooo ||a" ||a" S K (88) the weights have bounded norm ||a]] + ||U|| K, and (ti) the data points have bounded max; ||= K. Then there is constant K' depending only K such that
Lip(Dh) I K".
Proof. Let m + md be the number of parameters of the network. Then Dh e R2xp is
Df
Dh 1
:
Vn
(x
Lip(Dh) max Lip( :5(n]
21

---

Published in Transactions on Machine Learning Research (11/2023)
So for the 2-layer network,
- [ip((o(/m(u),, o(Vm(um::il)) xil]))
Di) S max
xe m
1 Lip( a
am (ym(uum ai) o0
m
o' mi))aa o'(ym((u,, x1))
0
o.
1
0
= max =E[n] vm
o' ((mmm/umm (0i
|| |c
Lip( Jajo'( ( (yml mm(uu, xxi))
I'm (ym(
mm
1 1 1l x
[ |l
+
Idiag( [a'( (o'(vvi (vm((n, ai))
d'(mm(um,z;))
vm
m
(1o"(/m((u,, T
+ ||0:||
Om Om (um x.))
|10' 1o |aa;
+ ||a"|l 2| ||aa | |||
m
22


Reviews:
[{"review_id": "7WS5OBiEF9", "reviewer": "Reviewer_t2p5", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strengths:\n\n* In terms of the results, this paper shows that the NTK approximation can be achieved by setting $\\alpha=O(T)$, which improves the previous result that requires $\\alpha=O(T^2)$.\n\n* This paper develops a more refined proof technique, which can be used in other theoretical analyses for dynamic kernels.\n\nWeaknesses:\n\n* Although the main goal of this paper is to improve the conditions made in [Chizat et al. (2019)], it still needs some more explorations, at least how this condition affects the conditions for the neural network width and the convergence rate.\n\n* Following the previous comment, the authors should also derive the convergence rate under the improved conditions.\n\n* The authors only consider the general NTK model, it would be better to consider the applications to certain neural network models such as multi-layer ReLU networks.\n\n* The comparison with [Cao and Gu, 2019] should be added, in which the authors also prove a bound on the linearization approximation error (see their Lemma 5.2).\n\nCao and Gu 2019, Generalization Error Bounds of Gradient Descent for Learning Over-parameterized Deep ReLU Networks, https://arxiv.org/pdf/1902.01384.pdf.\n", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "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": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "This paper establishes a tighter bound on the condition required for NTK approximation. ", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "See the weakness section.", "broader_impact_concerns": "No concern.", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "CG3OqEgqCD", "reviewer": "Reviewer_7ztJ", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "The strong points of the paper are:\n- Clear improvement over current SOTA (Chizat et al.)\n- General and novel approach that can probably be further refined/extended\n- Thorough derivation, missing only a good recap (ideally a proof flowchart) for the extensive appendix\n- Corresponding lower bound\n\nThe weaknesses:\n- The result only holds for continuous flow, and is tight only under very general assumptions while real-world data might satisfy stronger assumptions\n- At the same time the authors do not discuss how realistic the assumptions they inherit from Chizat et al. are. This is particularly important when these assumptions are playing an important role in the proof (e.g. assuming that $Dh$ is $\\text{Lip}(Dh)$-Lipschitz in a ball of radius $\\rho$ makes it much easier to prove boundedness of the operator)\n- The proof sketch in Sec. 2 gives an intuition of the main result by saying that the $K_T$ process bounds any of the $K_t$ intermediate processes. However in the appendix the proof mainly relies on a $\\sup_{t \\in [0,T]}$ analysis, and while I find it plausible that that $K_T$ is a bound on this sup I cannot find a clear point where it is shown and the derivation is not immediate.\n\nMinor:\n- This is a theory paper so empirical validations are not strictly necessary. However the result would be strengthened by showing that it capture the real-world behaviour of NN optimization (the original goal of NTK). For example the authors could show empirically how close (or not, for \"easy\" data) real-world flows are to their bounds, how much finite steps make a real-world system deviate from their analysis. Finally, since the authors provide a constructive lower bound, it would have been nice to provide some real-world simulation to see if the designed model keeps close to the theoretical behaviour under finite steps.\n- Section 2 is slightly confusing w.r.t. how novel the proof blueprint is. The authors start with Chizat et al.'s algebra-based proof, then introduce an intermediate algebra-based lemma, then switch to their integral-representation based proof but eventually return to the lemma. The flow can be greatly improved by clarifying which parts of the approach are novel and which are off-the-shelf.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "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": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "The paper introduces a tighter analysis of the NTK evolution in the\nmost general setting introduced by Chizat et al. (2019).\nIn particular, their new bound holds for a much longer time horizon\n($\\alpha^2\\rho^2$ vs $\\alpha\\rho/\\text{Lip}(Dh)$) and is tighter ($\\kappa$ vs $\\kappa T\\text{Lip}(Dh)$). Framing this in terms of initial rescaling and a fixed error budget epsilon, the time horizon is squared from Chizat et al.'s $O(\\sqrt{\\alpha\\varepsilon})$ to $O(\\alpha\\varepsilon)$.\n\nTo achieve this the authors propose a new proof blueprint based on a product integral formulation of the dynamics of the residual, leveraging in a few points results from Chizat et al.\n\nThe main technical challange seem to be showing continuity and operator-boundedness guarantee for the main quantities involved in the process, which in turns makes it possible to leverage classical results from Dollard & Friedman (1984) for the evolution of bounded continuous operators. Once the main reformulation is achieved, the rest of the proof is a more straightforward (but still highly technical) sequence of bounding to recover a quantity comparable to Chizat et al.\n\nOnce they establish their tighter bound, the authors also provide a worst case example that achieves it, which prove that the bound is indeed tight under this specific set of assumptions.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "Add a short discussion of the necessity and impact of the assumptions, and a clarification of the $\\sup_t$ to $K_T$ bound.\n\nBeyond that, I think the paper as is achieves overall the threshold for acceptance in TMLR. However a broader discussion of the assumptions, more structured presentation of the proof, and an illustrative empirical validation would strengthen the contribution.\n\nMinor:\n- $\\kappa$ is defined as part of Prop. 1.1 while it is a shared quantity used across theorems and should be independently defined as $\\alpha, R_0$ and other common quantities are\n- adding a proof flowchart/graph would help legibility of the appendix a lot.", "broader_impact_concerns": "N/A as this is a theory paper", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "XxjRb9LgsZ", "reviewer": "Reviewer_jka3", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strengths\n\nThe result is a clear improvement over existing work and is shown to have a matching lower bound, thereby providing a very satisfying conclusion to this problem. \n\nIn fact, there is very little to cherry pick, as the contributions are quite clear and easy to understand. ", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "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": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "The authors tightened an NTK approximation result of Chizat and Bach (2019) based on time rescaling, and showed this result is tight with a matching lower bound. \n\nI believe this submission exactly fits the purpose of TMLR: to publish correct and useful technical results, without judging for the result's \"impact\" or other subjective criteria. I would recommend accepting this submission as it is. ", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "N/A ", "broader_impact_concerns": "N/A ", "additional_comments": "", "extra_scores": {}, "extra_text": {}}]

Author rebuttal:

Meta-review: The paper is reviewed by three experts. Although two of the reviewers are positive about the paper, one reviewer raised concerns which authors only partially addressed. 

In particular,  the reviewer would like to see a discussion on global convergence, which I think authors can briefly mention in the final version (does not have to be a technical statement if it is beyond the scope). 

I think the paper will be a good addition to TMLR and its contributions  will be appreciated by its audience.
Decision: Accept (poster)