A note on regularised NTK dynamics with an application to PAC-Bayesian training
Authors: Eugenio Clerico
Benjamin Guedj
Venue: Accepted by TMLR ()

Abstract: We establish explicit dynamics for neural networks whose training objective has a regularising term that constrains the parameters to remain close to their initial value. This keeps the network in a lazy training regime, where the dynamics can be linearised around the initialisation. The standard neural tangent kernel (NTK) governs the evolution during the training in the infinite-width limit, although the regularisation yields an additional term appears in the differential equation describing the dynamics. This setting provides an appropriate framework to study the evolution of wide networks trained to optimise generalisation objectives such as PAC-Bayes bounds, and hence contribute to a deeper theoretical understanding of such networks.

Published in Transactions on Machine Learning Research (04/2024)
A note on regularised NTK dynamics with an application to PAC- MAC-Bayesian training
Eugenio Clerico
eugenio. cla.mio.ceercco@mma.. com com
Universitat Pompeu Fabra, Barcelona.
Benjamin Guedj
b.guedj&uucl.com uk
Centre for AI and Department of Computer Science, University College London & Iaria London
Reviewed on OpenReview: https: // openreview. net/ fot/fo id= 2la55Bekuy
Abstract
We establish explicit dynamics for neural networks whose training objective has regular- ising term that constrains the parameters to remain close to their initial value. This keeps the network in a lazy training regime, where the dynamics can be linearised around the initialisation. The standard neural tangent kernel (NTK) governs the evolution during the training in the infinitewwidth limit, although the regularisation yields an additional term that appears in the differential equation describing the dynamics. This setting provides an appropriate framework to study the evolution of wide networks trained to optimise general- isation objectives such as PAC-Bayes bounds, and hence contribute to a deeper theoretical understanding of such networks.
1 Introduction
The analysis of infinitely wide neural networks can be traced back to Neal (1995), who considered this limit for a shallow (I-hidden-layer) network and showed that, before the training, it behaves as Gaussian process when its parameters are initialised as independent (suitably scaled) normal distributions. A similar behaviour was later established for deep architectures, also allowing for the presence of skip-connections, convolutional layers, etc. (Lee et al., 2018; 2019; Arora et al., 2019a; Novak et al., 2019; Garriga-Alonse et al., 2019; Yang, 2019; Hayou et al., 2021). Tee et al. (2019; 9; 2020), among others, brought empirical evidence that wide (but fimite-size) architectures are still well approximated by the Gaussian limit, while finite size corrections were dcrived in Antognini (2019) and Basteri & Trevisan (2022).
Although the previous results hold only at the initialisation (as the Gaussian process approximation is only valid before the training), Jacot al. (2018) established that the evolution of an infinitely wide network can still be tracked analytically during the training, under the so-called neural langent kernel (NTK) regime. In nutshell, they showed that the usual gradient flow on the parameters space induces the network's output to follow a kernel gradient flow in functional space, governed by the NTK. A main finding of Jacot et al. (2018) is that although for general finite-sized networks the NTK is random at the initialisation and evolves during the training, in the tifdh limit it becomes a deterministic object that can be exactly computed, and it stays fixed throughout the training. Later, Lee et al. (2019) provided a new proof of the convergence to the NTK regime, while Yang (2019) established similar results for more general architectures, such as convolutional networks. Chizat et al. (2019) exlended the idea of linearised dynamics to more general models, introducing the concept of lazy truining and finding sufficient conditions for network to reach such regime, They also pointed out that this lincarised behaviour may be detrimental to learning, as also highlighted by Yang & Hu (2022) who showed that the NTK dynamics prevent network hidden layers from effectively learning features. However, the NTK has been a fruitful tool to analyse convergence (Allen-Zhu et al.,
Work mostly done while affiliated at the Department Statistics, University Oxford, UK.

---

Published in Transactions on Machine Learning Research (04/2024)
2019b; Du et al., 2019) and generalisation (Allen-Zhr Zhu et al., 2019a; Arora al., 2019b; Cao & Gu, 2019) for over-parameterried settings under (stochastic) gradient descent.
The standard derivation of the NTK dynamics (Jacot et al., 2018; Lee et al., 2019) requires the network to be trained by gradient descent to optimise an objective that depends on the parameters only through the network's output. This setting does not allow for the presence of regularising terms that directly involve the parameters. Yet, in practice, a network reaches the NTK regime when its training dynamics can be linearised around the initialisation. This happens if the network parameters stay close enough to their initial value throughout the training, the defining property of the laay training regime. It is then natural to expect that a regularising term that enforces the parameters to stay close to their initialisation will still favour linearised dynamics, and so bring a training evolution that still can be expressed in terms of a fixed and deterministic NTK. This is the focus of the present paper, where we discuss the evolution of a network in the NTK regime trained with an e-regularisstion that constrains the paramcters to stay close to their initial values. We remark that similar ideas are also present in Hu et al. (2020), where the authors study the evolution of linearised approximation of a neural network under the (-rreuularisation, without however proving the convergence of the original network's dynamics to those of the linearised model.
We note that regularisers centred at the initialisation typically appear in PAC-Bayes-inseee tives, where the mean vector of normally distributed stochastic parameters is trained via (stochastic) gradient descent on a generalisation bound (an approach initiated by the seminal work of Langford & Caruana, 2001 and further explored by Alquier et al., 2016; Dziugaite & Roy, 2017; Neyshabur et al., 2018; Letarte et al., 2019; Nagarajan & Kolter, 2019; Zhou et al., 2019; Nozawa al., 2020; Biggs & Guedj, 2021; 2022; Dziugaite et al., 2021; Pérez-Ortiz et al., 2021a;b; Pérez-Ortiz et al., 2021; Chérief-Abdelle
training objec-
et al., 2022; Lotfi et al., 2022; Tinsi & Dalalyan, 2022; Clerico et al., 2022; 2023a; Viallard et al., 2023). As these training objectives yield generalisation guarantees, we conjecture that the exact dynamics that we derive could be a starting point to obtain generalisation bounds for more general kernel gradient descent algorithms.
As final remarks, we note that Chen et al. (2020) considers a NTK regime that allows for regularisation, but they only consider the mean-field setting of two-layer neural networks (introduced by Chizat Bach, 2018; Mei et al., 2018, and further explored by Mei et al., 2019; Wei al., 2019; Fang et al., 2019), where the network is not initialised with scaling yielding Gaussian process. Finally, mention that also Huang et al. (2022) atlempled the analysis of PAC-Bayesian dynamics via. NTK. However, their approach dilfers from ours, it does not highlight the effect of the regularisation term, and the whole analysis deals with a simple shallow stochastic architecture where only one layer is trained.
Outline. We present our framework and notation in Section and then treat the unregularised NTK dynamics in Section 3 a starter. We then move on the more interesting case of regularisod dynamics in Section 4, first for the simple case of (rreeularisaaion and then for a more general regularising term. We instantiate analysis to the example of least square regression in Section 5 and illustrate the merits of our work with application to PAC-Bayes training of neural networks in Section 6. The paper ends with concluding remarks in Section 7 and we defer technical proofs to Appendix A.
2 Setting and notation
We consider collecconeeeted feed-forward neural network of depth L, which we denote as F X R where X' C IR is compact set. We denote as n, the width of the I-th hidden layer of the network, as p the input dimension and ny q the output dimension. We consider the network
nh w ((T(())))
1
F(x)=U (x)
U](x)
W3ay;
VTu
Vno
where Φ is the network's activation, acting component-wise. The network's prediction in the label space y is (xx) [(F())), for some IR9 + y. We denote as W the parameter space where the weights lie, and
2

---

Published in Transactions on Machine Learning Research (04/2024)
W the parameters of the network. We consider the infinite-width limit, where all the hidden widths (1 11=1,.... T. 1) are taken to infinity".
Data consists of pairs instance-label 2 = (x, v) E Z = X x V/ with I E X' and y E V. We consider a non-negative loss function & W x 2 + [0, 2o). For a dataset 8 E z", we define the empirical loss C, as the average of on 8, namely
1 Σ ((W
(W)
m =E.
As will often encounter empirical averages, we define the following handy notation
(o(Z)) = (x(x,)), 1 Σ g(r.v).
(x.D)
that (W) = (e(W, Z)) (),. We assume that l depends on W only through the network's output F, i.c.,
there exists a function i such that we can rewrite
= ((F()))
The network training follows the gradient of a learning objective C.,, namely
(t) = (W,))==CC(W(()) (W(()),,
where V denotes the gradient with respect to the parameters. We assume that C. can be split into two terms, the empirical loss, and rogularisation term that depends directly on the parameters (without passing through the net work's output F) and does not depend on 8:
C, == L. XR,
(1)
for some À 2 O. We let P: R a strictly increasing differentiable function such that p(0) = 0, and we consider the case of a regulariser in the form
R(W)= = (( 11I -(0))
(2)
with (0) denoting the value of the parameters at the initialisation, and with || denoting the square of the Frobenius norm on W, namely |IW 1/ Ei,
For conciseness, we write (t()) for g(t) g(0), where is any time-dependent term. Moreover, we introduce the following notation
aUg(1;;) oW!
di'' Rig
(x;t) (x;t)
ij
(x, = Σ
bild' (x: v/ (z':t);
(3)
P=1
uk (zi Aw" (t),
where denotes the component-we wise inner product between matrices (or vectors) of the same size. O is the so-called neural tangent kernel (NTK) of the network. We remark that if = represent the linear approximation of the variation of U with respect to changes in the parameters.
Following the approach of Jacot (2018), we consider the case where this limit taken recursively, layer by layer (that also have n (nj + 0).
3

---

Published in Transactions on Machine Learning Research (04/2024)
3 The unregularised dynamics
Without regularisation (namely when we set 0) the standard NTK dynamics hold (Jacot et al., 2018; Lee et al., 2019). We briefly cover this case, where the width of the network is taken to infinity.
At the initialisation, the network's output behaves as centred Gaussian process, whose covariance kernel can be computed recursively. More preciscly, we have that all the components of the output are i.i.d. Gaussian processes defined by
Fk GP(
where
x x E (x,
= = -CG.00--r1OSSa2220000(0001111111
(4)
no
Moreover, the neural langent kernel 04, defined in (3), tends to a diagonal deterministic limit (with respect to the initialisation randomness), and stays constant during the training. In particular, we have
(z, x';t) (z, z')
where O6 can be computed recursively as follows:
= (x,')= 1 St
(5)
with é denoting the derivative of
During the training, the network's output obeys the dynamics
OL(F (((F()))).))
((((xx))((F(())
(l(F(x))
0x m (2.0)Ex
Σ (x(x,
(6)
OF
0F
We remark that the network being in the NTK regime simply means that the model is linear around the initialisation. Indeed, neglecting terms of order [[]AW(())) one has
F(x;t) ~ F(x; 0) + Jw [[( (W()),
with Jw denoting the Jacobian with respect to the parameters. Assuming that this linear approximation holds we also have
Jw (
Now, up to terms of order (()I)()))I).
(J (Jw[F(X:))
V = (J ),
a- -VC =
Jacot et al. (2018) showed that in the infinite-width limit this lincar approximation becomes exact, and
[F'F' 0)] Jw (F(; lends in probability to the O(x,z') Id, leading to the dynamics (6), as
0,F(z;t) ~ [ (Jw[F(X:e]] V=((F((xt). ( (0
4 The regularised dynamics
In this section, we discuss the impact of the regularisation term on the NTK dynamics. We first particularise this to the specific case of l feregularisation, then move to the treatment of more general regularisers.
Intuitively, we note that the regulariser that we are studying tends to keep the walues of the parameters close to their initialisation. Hence, we still expect that the dynamics can be linearised. If this assumption indeed holds, we get (in the simpler case p piid) id)
a,W, - -J]][[[(XO)] ((L( - -AA(()).
4

---

Published in Transactions on Machine Learning Research (04/2024)
Keeping in mind that for the linearised model have AF(x;t) = Jw (F now have
0.F(x: - (OX,))=-(((G...NV)000Y..
AAF(x;
which a regularised version of the NTK evolution. We will establish this more rigorously in the next sections, showing that under regularised dynamics the linearised model a valid approximation of the neural network in the infinitewwidth limit.
4.1 Simple e &rregularisation
We first consider the case p= id in (2), so that R(W) = 2ll and
C.(W) Lo(W) + 2 AWl.
(7)
By just applying the chain rule to (L) (VC)(W()), we find that
nL
Oc(F (((F(X;t), (t)
(t -WC, -
-
wotd (x::)))))
(8)
OFk
This translates into the following functional evolution of the network's output
(=;1))))- Σ (e, al(F ol(F(X;),
(x;
OF
Our goal to prove that the infinite width limit the next two properties hold:
1. The NTK is constant at its initial value, which coincides with the standard deterministic NTK in (5), and more generally for all layers
+ x':t) SE O'
2. The term _² is exactly AF, and more generally
=((:t)= Ax) Ux((x) (0(:::)))((xx0))))
We note that these two properties are actually rather intuitive: the first one tells us that the Jacobian of the output stays fixed during the training, as if the dynamics where linear; the second one that we can indeed linearise around the initialisation. We recall that the standard NTK regime (with no regularisation) holds when the parameters do not move too much from their initial values, so that the dynamics can be lincarised. The addition of regularising term does actually enforce the parameters to stay close their initial value. Hence, adding regularisation does not hinder the linearisation of the learning dynamics.
From the two properties above, we conclude that the NTK evolution of F is given by the following.
Theorem 1. In the infinite-width limit, taken recursively layer by layer, the network's output evolues as
0,Fx(r;1)= = Σ 0² σ) di(F(z;t). A(Fa(e;t) Fz(x;)),
OF
where L is defined in (5).
Note the term (x; F& (z; 0)), which constrains the network's output stay close its initialisation, is present in the standard NTK dynamics (6).
5

---

Published in Transactions on Machine Learning Research (04/2024)
Evolution of the training objective. As side remark, see how the training objective C2 evolves during the training. First,
O.C,(W(O) (V)RF((:0)))) -.F(:)))
=(0(XXX"V - - ->(V=((F((:0))) -FF(XX
where the notation (g(Z, (3)) sas denotes term have that
Exes Eve, On the other hand, for the regularising
O, = - -V -F(X))))))) - 2X)(W(()).
(9)
Thus, overall we gel that
(W((()()))
-- (=((F(X:))))'))
22 (F -FF(X:))) -2X²R(W(t)).
As side remark, note that if & convex in F, then we always have that
Vp((F(x;(,.)) (F(;; 2 ((P()=)))))) ((F(x;)),v)
In particular we get that
a.R(W(t) s < (AI(F(;;)))))) 2XR(W(H)) = - - XR)
Since t ++ (W((())) is non-decreasing, we obtain that
Rt)(t)) s ( (C((v()))CC(((())) - (1
from which follows that (t) can be controlled by the variation in L. (t) as
1 - At.
XR(W(()) < 2 (c)()))))
4.2 General regulariser
We consider the case of more general regularising term, which still leads to tractable training dynamics. Let R + R+ be a differentiable strictly increasing function. Define
D(t) 1 (t)
and let
R(W(t)) = p(D(t))
Theorem 2. Consider a function p [0, 0c) + [0, oo) as above and assume that ē is locally Lipschilz. Assume that the dynamics are given by
O.W(t) = VC, VL.(W(()) - XR(W(()).
Then, in the infthite-width limit (taken recursively layer by layer) we have
0.Fx(1;1)= = (0(0; X) a((F(x;1))))
(0((()))
OF
a,D(t)= - Σ AF((X;t) (((F(Xx)).
- (()
OF
k=1

---

Published in Transactions on Machine Learning Research (04/2024)
Proof. The proof is based on the following result, which establish in Appendix A.1.
Proposition 1. Fir a time horizon T > O and depth L. Assume that for t E [0, T) and for all I E [1 L]
n O, = Σ
(X:1)Vi(Z:0) - Xr(D())AW,
for some mappings : 2 x + Ru and r: (0, + R. Then we have that Uh obeys the dynamics
RL
(x;)) Σ XI: -M(
Moreover, if D(t) ((1 for all t [O,T],
[. (IV(Z),))),, dt 0(1) and [' (((D)))))dd 0(1),
if & is "Lipshhshttzz and Bg smooth, then (in the infinite- width limát taken starting from the layer 1 and then going with growing index), for all E [1 : L}, Ol is constant during the training, and E' = AU
To prove Theorem 2, need to check that the assumptions of the above proposition hold when setting r(D;t) = (((())) and V(z;t) = V Vp((F(x;t),)) y). First, notice that C,(t) S C,(0) (,) for all t 2 0, as we are following the gradient flow. Now, with arbitrarily high probability (on the initialisation), we can find a finite upperbound for (,),), which holds when taking the infinito-width limit (as the network output becomes Gaussian). particular, we have that since independent of the width and we can write J (1) (where the O notation referred to the infinite width limit). Now, we have that
R(t) S J/A O()).().
Since p is invertible, in particular for all t we have that
D(t) p (J) 0(1).
We prove in Lemma 1 (Appendix A.1) that AF(() O(), and so we know that with arbitrarily high probability we can find a radius J' such that [F t)ll S J' for all t > 0. In particular, the regularity of implies that ()z; = Vr((F((; v) is uniformly bounded for all > 0. So, f- (Z;()) di = 0(1).
Finally, we have that the t(:;t) in the previous statement is (D(())). Since D(t) is bounded throughout the training and p locally Lipschitz can bound the integral over and apply Proposition 1.
Finally, the dynamics for D()) follow from the chain rule.
Clearly Theorem is just particular instance of Theorem 2, as only need to set p the identity.
5 The example of least square regression
As simple application of what established, study the evolution of network under least square regression, namely when we have ((ū.v) v) = [0 2) The dynamics of the training are given by
0,F(z;t) = (0(a,
Y)). F(;;;))))((;))
This is a linear ODE and can be solved exactly. For convenience we introduce the following notation. We let θ denote the NTK Gram matrix whose entries 0(z,x'//m, with and z' ranging among the instances of the training sample s. Similarly F(t) and y are the vectors made of the network's output and labels, for the datapoints in We thus have
a.F(1)= = -Ō((() - -(F(()) F(0)),
7

---

Published in Transactions on Machine Learning Research (04/2024)
which brings
= = F(0) + -Td- -V - -F(O))
where VA AId + न, which is always invertible for À > 0.
Note that asymptotically for large we have that
F(I) -()+F(())) = (ld P(0) + VY
which is exactly what would obtain optimising CN in (1). Clearly, for small values of the labels are almost perfectly approximated, as we have
F(co) = Y + Λ - 1Y 16- - Y) Y) - O(A).
On the other hand, if X is very large then F(x)) = F(0), as
① A Id
θ X (Y F(O)) F(0) + 0(1A))
F(oo)=
Once the evolution of on the training datapoints has been computed, one can directly evaluate the value of F(x;t) t) for any input that not in the training sample. Defining
T(;;t) () (((XX))))))))..)
we have
=(x;)))((r)))+++(())))) L'' ((0;8)e de'.
6 An application to PAC-Bayesian training
A motivation to study dynamics in the form of (8) comes from the PAC-Baycsian literature. We consider a dataset 8 made of i..dd draws from a distribution j on Z. We are seeking for parameters with a small population loss
(z [e,])(()))
As is unknown, we rely on the empirical loss Lg as a proxy for Lz.
The PAC-Bayesian bounds (introduced in the seminal works of Shawe-Taylor Williamson, 1997; McAllester, 1999; Seeger, 2002; Maurer, 2004; Catoni, 2004; 2007 we refer to the recent surveys from Guedj, 2019; Alquier, 2021; Hellström et al., 2023 for a thorough introduction to PAC-Bayes) are generali- sation guarantees that upperbound in high probability the population loss of stochastic architectures, in case networks whose parameters W are random variables (this means that every time that the network sees an input I, it draws W from some distribution Q, and then evaluate F(x) for this particular realisation of the parameters). The PAC-Bayesian bounds hold in expectation under the parameters law Q (commonly referred to as posterior). Here is a concrete example of this kind of guarantecs. Fix a probability measure P and 17 > 0, for a bounded loss f [0,1] 1] we have (see, e.g., Alquier, 2021)
(++ KL(QIP) + log)))/))((/0)))
S Ew-o (w)) 1-5
(10)
8m
1
where the inequality holds uniformly for every probability measure Q, in high probability (at least 1 - 5) with respect to the draw of s. The probability measure P (typically called prior) in (10) is arbitrary, as long as it is chosen independently of the particular dataset used for the training.
Several studies (sce Section this line of work started with Langford & Caruana, 2001, was reignited by Dziugaite & Roy, 2017 and then followed by a significant body of work by many authors) have proposed
We remark that T(2;;) 1) can computed: when averaging on 8, F(X;1) only takes values the components of F(1).
8

---

Published in Transactions on Machine Learning Research (04/2024)
to train stochastic neural networks by optimising PAC-Bayesian bound. A possible approach consists in considering the case when all the parameters of the net work are independent Gaussian variables with unit variance (namely Wiy ~ ((mi;: 1)). The training then usually amounts to tune the means m. Typically, the imitial values of the means are randomly initialised (we denote them as m(0),, and P can be chosen as the distribution of the networks parameters at initaaiiiooonn (Dziugaite & Roy, 2017); namely a multivariate normal with the identity as covariance matrix and m(0) as mean vector. In this way, one gets that
(L(()P)- = (llm(t) m())l
Defining (m) Bww.Wm.ia)L... (W)),, we see that using the bound (10) as training objective is equivalent to optimise
C. (m) (m) 1 m m
(11)
2n/ 8m
which is exactly in the form of (7) with À 1/6nv//mm More generally, many PAC-Bayesian bounds are not linear in the KL therm. However, they can still fit in our framework (with a general reguliriser p in Section 4.2) as long as they are in the form
SwwolSa(W))
& Ew-olf. ((KL(Q)P)) 1-8
for some strictly increasing and differentiable This is for instance the case for the training objective used for the PAC-Baycsian training by Dziugaite & Roy (2018).
While significant experimental work has focused on PAC-Bayesian training methods and achieved promising results, to our knowledge the literature lacks of rigorous theoretical studies of these training dynamics. Since the NTK formulation has already been successfully used for the study of gradient descent in the unregularisod case, we anticipate that the closed-form expression for the network's evolution that we derived could help study various properties (such as rates of convergence, convergence to global/local minima, etc.).
Training wide and shallow stochastic networks
Clerico et al. (2023b) has recently shown that, when considering the infinite-width limit for a single-hidded hidden- layer stochastic network, a close form for (m) can be computed explicitly, and one can actually see the stochastic network deterministic
(with different activation function), where the means trainable parameters. We summarise and rephrase the results of Clerico et al. (2023b), and show explicitly how exact continuous dynamics can be established via our results, in the infinite-widtle limit.
the
We focus on binary classification problem (i.e., y = (+1)), where we assume that all the inputs x are normalisod so that ||ж|| = /no (i.e., X' = 5mo- (Vro), the sphere of radius vno in RTo). We consider a stochastic network with single hidden layer and one-dimensional output, and Lipschitz and smooth activation 6,
F(x) = Vn Σ (x)(
U7(x)
W TTo j=l
1
The prediction of the network is the sign of the output. The stochastic parameters can be rewritten as W'
w'=5++m m'
where ç is matrix of the same dimension of W whose components are all independent standard normals (resampled every time that new input is fed to the network) and m' is matrix of trainable parameters (the means of W'). We assume that the components of m' are all randomly initialised as independent draws from a standard normal distribution. In this setting, Clerico et al. (2023b) showed that in the infinite-width limit n + 00
F(x) ((MT(()) 3Tee limit distribution with respect the intrinsic stochasticity of the ('s, and in probability with respect the random initialisation.
(12)

---

Published in Transactions on Machine Learning Research (04/2024)
where M and Q² are the output mean and variance in the limit.
When the training objective is in the form (11), the network is constraint to remain close to its initialisation. Interestingly, in this lazy training regime Q² stays constant to its initial value, which is deterministic (with respect to the initialisation randomness) and independent such that (²xx;) t) = o², for all x E X' and t. We refer to the Appendix A.2 for details.
when X is a sphere. We can thus define a > 0
Following the derivation of Clerico et al. (2023b) we note that M² can actually be seen the output of neural network with parameters m, whose activation function is Φ, defined Gaussian convolution of ø,
(()) = I (E--
Concretely, this means that
1 m² m (M M'(x) Vn
1
Mxxx
m'r.
vno
Now, for a loss function ((F,2), we can define the expected loss
E(M, 2) = Ec- + M,
In this way, we have that the expected empirical loss (m) appearing in (11) is given by
(.(m) Σ
(x)
Hence, optimising the PAC-Bayes bound (10) induces the dynamics
(x; = - (0(11; X) (x;))(
(M;
(13)
0M2
n/8m 8m
where ((x,x') = (x; V M2x; is given by x' Ō x1xxx B((())(())
E(x,x') =(((())(()
no
na with (5,5")~ ~ N (1, x-zz" mxn) )) and v denoting the derivative of sb.
Misclassification loss. A common choice is to set ((F,2) 1 if sign F(x) # y. and 0 otherwise, that is the so-called misclassification loss. In such case we can easily derive that.
7 2) PPcwN(..) (<> (M(())
1 1 erf (yM(z) 2
o
ov2
Tt follows that (13) now reads
(x; Yō(x; (M((X:))/(23
1 AM (x;11;).
0
7/88m
This does not have simple close-form solution, and can only be solved using numerical integrators.
Quadratic loss. In order to obtain simpler dynamics we consider the loss ((F,z) = (1 - yz))) Note that this quadratic loss unbounded, and so a generalisation bound such as (10) is not guaranteed to hold. However, the quadratic loss is always greater than the misclassification loss, so the RHS of (10) is still valid generalisation upperbound for the misclassification population loss. This choice of é yields the dynamics
(xM2(z;t) = 22
Y)).
AM² (x:
n/8m
10

---

Published in Transactions on Machine Learning Research (04/2024)
Proceeding as in Section 5, and again using tilde denote vectors and matrices indexed the training sample 8, we get that
(t) (0) + (ld e 22V2222 V (0))
where Va/2 AId/2+ θ and X = 111 8m). Asymptotically, for large t, M will approach
(x) +V M2)))).
For large the empirical loss approaches
Lx 11=-1117 11 - YIl
m
where for a positive definite matrix A we define ||0113 =TTAu Av. We note that the eigenvalues of (A/2) are the form
>²/4
04
(X/2 + 0,)
where the 's are the eigenvalues of Ô. In practice (f.e., max, 02 ~ 1) when the sample size m grows to infinity (Murray et al., 2023). Since À ~ 1/ 1/ym, we get that the network will be able reach a small empirical loss (of order X2 1/m) if 2 Y lies in
expect the largest eigenvalues of θ to be of order
cigenspaces of θ where the eigenvalues 0 ~ 1 > m/.
On the other hand, for large the regularising term R will approach Roo; which from (9) must satisly
2 (00) (M(0)) - v) = --2XR
From this we can derive that
Rx 2m
M| -V11 ,
Here, the eigenvalues of V2 are of the form
0: +
8:
If we are in the regime where the datapoints are such that (0) - Ý again lies where 0 1, then expect Re 1. This means that if we are able learn well the labels while optimising the PAC-Bayeste Bayesian bound, we are ensured that the KL term is of order 1, and so the objective (11) will be of order 1/ 1/mm. resulting in a non-vacuous bound. We argue that this is what happens when data comes from a reasonable underlying distribution, matching the implicit regularisation induced by the NTK.
We finally note that the 8 are small also when < 1/ /mm However this is due to the fact that these directions are not promoted by the NTK dynamics and so if M2 - Y completely lies in eigenspaces with very small eigenvalues of θ, then the network will essentially stay fixed to its initial configuration, keeping small penalty, but also not improving its performance.
7 Conclusion
We established explicit dynamics for infinitely wide fully connected networks trained to optimise regularised objective, where the regularisation pushes the parameters to stay close to their initialisation. Under this regime we show that the model undergoes linearised dynamics during the training, which turns out to be a regularised version the standard NTK evolution.
Our analysis follows similar ideas to the NTK convergence proof of Jacot al. (2018) and presents the first regularised NTK analysis that also be applied to PAC-Bayesian training. We conjecture that stronger
11

---

Published in Transactions on Machine Learning Research (04/2024)
follow-up results could be derived, for instance, following the approach of Lee et al. (2019) to show that the convergence holds when the infinite width limit is taken for all the hidden layers simultaneously, and to study discretised dynamics.
We also anticipate that further analytical and empirical studies of the induced PAC-Bayesian dynamics might be of interest to shed some light on generalisation-driven training of neural networks.
References
Z. Allen-Zhu, Y. Li, and Y. Liang. Learning and generalization in overparameterized neural networks, going beyond two layers. NewrIPS, 2019a.
Z. Allen-Zhu, Y. Li, and Z. Song. A convergence theory for deep learning via over-parameterization. 1CML, 2019b.
P. Alquier. User-friendly introduction to PAC-Bayes bounds. arXiv:2110.11211 2021.
P. Alquier, J. Ridgway, and N. Chopin. On the properties of variational approximations of Gibbs posteriors. Journal of Machine Learning Research, 17, 2016.
J.M. Antognini. Finite size corrections for neural network Gaussian processes. ICML Workshop, 2019.
S. Arora, S. Du, W. Hu, 7. Li, R. Salakhutdinov, and R. Wang. On exact. computation with an infinitely wide neural net. NewrIPS, 2019a.
S. Arora, S. Du, W. Hu, Z. Li, and R. Wang. Fine-grained analysis of optimization and generalization for overparameterized two-layer neural networks. ICML, 2019b.
Basteri and D. Trevisan. Quantitative Gaussian approximation of randomly initialized deep neural networks. Xxxx:2203.07877 2022.
F. Biggs and B. Guedj. Differentiable PAC-Bayes objectives with partially aggregated neural networks. Entropy, 23(10), 2021.
F. Biggs and B. Guedj. Non-vacuous generalisation bounds for shallow neural networks. ICML, 2022.
S. Boucheron, G. Lugosi, and P. Massart. Concentration Inequalities A Nonasymptotic Theory of Indepen- dence. Oxford University Press, 2013.
Y. Cao and Q Gu. Goneralization error bounds of gradient descent for learning over-paramotorized deep relu networks. AAAI Conference on Artificial Intelligence, 2019.
O. Catoni. Statistical Learning Theory and Stochastic Optimization. Ecole d'Eté de Probabilités de Saint- Flour. Springer, 2004.
O. Catoni. PAC-Bayesian supervised classification: The thermodynamics of statistical learning. IMS Lecture Notes Monograph Series, 2007.
Z. Chen, Y. Cao, Q. Gu, and T. Zhang, A generalized neural tangent kernel analysis for two-layer neural networks. NewrIPS, 2020.
B.E. Chérlef-Abcdellttif,,, Y. Shi, A. Doucet, and B. Guedj. On PAC-Bayesian reconstruction guarantees for VAEs. AISTATS, 2022.
L. Chizat and F. Bach. On the global convergence of gradient descent for over-parameterized models using optimal transport. NeurIPS, 2018.
L. Chizat, E. Oyallon, and F. Bach. On lazy training in differentiable programming. NeurIPS, 2019.
F. Clerico, G. Deligiannidis, and A. Doucet. Conditionally Gaussian PAC-Bayes. ATSTATS, 2022.
12

---

Published in Transactions on Machine Learning Research (04/2024)
E. Clerico, G. Deligiannidis, and A. Doucet. Wide stochastic networks: Gaussian limit and PAC-Bayesian training, ALT, 2023a.
E. Clerico, T. Farghly, G. Deligiannidis, B. Guedj, and A. Doucet. Generalisation under gradient descent via deterministic PAC-Bayes. arXin:2209. 02525, 2023b.
S.S. Du, J.D. Lee, H. Li, L. Wang, and X. Zhai. Gradient descent finds global minima of deep neural networks. ICML, 2019.
G.K. Dziugaite and D.M. Roy. Computing nonvacuous generalization bounds for deep (stochastic) neural networks with many more parameters than training data. UAI, 2017.
G.K. Dziugaite and D.M. Roy. Data-dependent PAC-Bayes priors via differential privacy. NeurIPS, 31, 2018
G.K. Dziugaite, K. Hsu, W. Gharbieh, G. Arpino, and D.M. Roy. On the role of data in PAC-Bayes bounds. AISTATS, 2021.
C. Fang, Y. Gu, W. Zhang, and T. Zhang. Convex formulation of overparameterized deep neural networks. wrXiv:191..07626, 2019.
A. Garriga-Alonso, C.E. Rasmussen, and L. Aitchison. Deep convolutional networks as shallow gaussian processes. ICLR, 2019.
B. Guedj. A primer PAC-Baycsian learning. Second congress of the French Mothematical Socicty, 2019. S. Hayou, E. Clerico, B. He, G. Deligiannidis, A. Doucet, and J. Rousseau. Stable ResNet. AISTATS, 2021.
F. Hellström, G. Durisi, B. Guedj, and M. Raginsky. Generalization bounds: Perspectives from information theory and PAC-Bayes. orXiv:2309.044 04381, 2023.
W. Hu, Z. Li, and D. Yu. Simple and effective regularization methods for training noisily labeled data with generalization guarantee. ICLR, 2020.
W. Huang, C. Liu, Y. Chen, Liu, and Richard Y. Da X. Demystify optimization and generalization of over-parameterized PAC-Bayese AAC-Bayesian learning. arXiv:2202.01955 019558 2022.
Jacot, F. Gabricl, and C. Hongler. Neural tangent kernel: convergence and generalization in neural networks. NewrIPS, 2018.
J. Langford and R. Caruana. (Not) bounding the true error. NeurIPS, 2001.
J. Lee, Y. Bahri, R. Novak, S. Schoenholz, J. Pennington, and J. Sohl-Dickstein. Deep neural networks Gaussian processes. ICLR, 2018.
J. Lee, Xiao, Schoenholz, Y. Bahri, R. Novak, J. Sohl-Dickstein, and J. Pennington. Wide neural networks of any depth evolve as linear models under gradient descent. NewrIPS, 2019.
J. Lee, S. Schoenholz, J. Pennington, B. Adlam, L. Xiao, R. Novak, and J. Sohl-Dickstein. Finite versus infinite neural networks: an empirical study. NewrIPS, 2020.
G. Letarte, P. Germain, B. Guedj, and F. Laviolette. Dichotomize and generalize: PAC-Bayesian binary activated deep neural networks. NewrIPS, 2019.
S. Lotfi, M. Finzi, S. Kapoor, A. Potapczynski, M. Goldblum, and A. Gordon Wilson. PAC-Bayes compres- sion bounds tight that they can explain generalization. NewrIPS, 2022.
Maurer. A note on the PAC Bayesian theorem. arXiv:0411099, 2004.
D.A. McAllester. PAC-Bayesian model averaging. COLT, 1999.
S. Mei, A. Montanari, and P.M. Nguyen. mean field view of the landscape of two-layer neural networks. Proceedings of the National Academy of Sciences, 115, 2018.
13

---

Published in Transactions on Machine Learning Research (04/2024)
S. Mei, T. Misiakiewicz, and A. Montanari. Mean-field theory of two-layers neural networks: dimension-fin -free bounds and kernel limit. COLT, 2019.
M. Murray, H. Jin, B. Bowman, and G. Montufar. Characterizing the spectrum of the NTK via power series expansion. ICRL, 2023.
V. Nagarajan and J. Zico Kolter. Deterministic PAC-Bayesian generalization bounds for deep networks via generalizing resiliece,,,,eee ICLR, 2019.
R.M. Neal. Bayesian learning for neural networks. Springer Science & Business Media, 118, 1995.
B. Neyshabur, Bhojanapalli, and N. Srebro. PAC-Bayesian approach spectrally-normaaized margin bounds for neural networks. ICLR, 2018.
R. Novak, L. Xiao, J. Lee, Y. Bahri, G. Yang, J. Hron, D. A Abolafia, J. Pennington, and J. Sohl-Dickstein. Bayesian deep convolutional networks with many channels are gaussian processes. ICLR, 2019.
K. Nozawa, P. Germain, and B. Guedj. PAC-Bayesian contrastive unsupervised representation learning. UAI, 2020.
M. Pérez-Ortiz, O. Rivasplata, E. Parrado-Hernandez, B. Guedj, and J. Shawe-Taylor. Progress in self- certified neural networks. NewrIPS workshop Bayesian Deep Learning, 2021.
M. Pétiz-Ortiz, O. Risvaplata, J. Slawe-Taylor, and C. Szepesvári. Tighter risk certificates for neural networks. Journal of Machine Learning Research, 22, 2021a.
M. Pérez- Ortiz, O. Rivasplata, B. Guedj, M. Gleeson, J. Zhang, J. Shawe- Shawe-Taylor, M. Bober, and J. Kittler. Learning PAC-Bayes priors for probabilistic neural networks. arXiv:2109. 10304, 2021b.
M. Seeger. PAC-Bayesian Generalization Error Bounds for Gaussian Process Classification. Journal of Machine Learning Research, 3, 2002.
J. Shawe- Taylor and R.C. Williamson. PAC analysis of a Bayesian estimator. COLT, 1997.
Shawe-Taylor
L. Tinsi and A.S. Dalalyan. Risk bounds for aggregated shallow neural networks using gaussian priors. COLT, 2022.
R. Vershynin. Introduction to the non-asymptotic analysis of random matrices, chapter 5 in Compressed Sensing: Theory and Applications, pp. 210-268. Cambridge University Press, 2012.
P. Viallard, M. Haddouche, U. Şimşekli, and B. Guedj. Learning via Wasserstein-based high probability generalisation bounds. NeurIPS, 2023.
C. Wei, J.D. Lee, Q. Liu, and T. Ma. Regularization matters: Generalization and optimization of neural nets v.s. their induced kernel. NewrIPS, 2019.
G. Yang. Scaling limits of wide neural networks with weight sharing: Gaussian process behavior, gradient independence, and neural langent kernel derivation. arXiu:1922.02 0960 2019.
G. Yang and E.J. Hu. Feature learning in infinile-width neural networks. urXim:2011.1455 2022.
W. Zhou, V. Veitch, M. Austern, R.P. Adams, and P. Orbanz. Non-vacuous generalization bounds at the Tea scale: PAC-Bayesian compression approach. ICLR, 2019.
14

---

Published in Transactions on Machine Learning Research (04/2024)
A Appendix
A.1 Proof of Proposition
Lemma 1. Assume that, when the infinite-width limit is taken recursively layer by layer, D(t) = O(1). Assume that o is -o-Lipschizz then we have that for each layer
(x; 0(1)
Proof. We have that
| (I;
|I O()) V
Then, recall that at initialisation the components of (z 0) are independent normal distributions, with mean 0 and variance (1,,,)), 2),, as defined in (4). So, we have that all the (x; ()))s))ss are independent and equally distributed, with finite variance (thanks to the Lipschitzness of Φ). By the standard CLT we get that
|o(U" - 1 Σ (U
~~~ 1 m
TH Σ -)
T2
and (x 0))1 Oly Now, using the Lipschitzness of o we get
1
||AU+1(x (x; 1)ll
(1)* 1)1 (10000/000|000000000000000000000
(0I))
ViI
0
1+
1+
n1+1
(z;+)))
T
where we used that || (0)1l
Ol , classical result in random matrix theory (Vershynin, 2012). Assuming that the limit is taken layer after layer, we have that n441 /nj -> and
(1:0)==000
by induction.
For the rest of this section, define (t) as
D') 2 1=1
| (t) 100-w(((0F (0)
Now, we restate and prove Proposition 1.
Proposition 1. Fix time horizon T > 0 and depth L. Assume that for E [0, T) and for all l [1 L]
ny
(,W(,(t= Σ
(x; (X; (xxt)) ():
Lnl
for any moppings V x R - L and :0,0 10, R. Then we have that Thh obeys the dynamics
n
(.;)( Σ (x)) S,0K,//
Moreover, if (1 0(1) for all [0,
[' (II(Z) di = 0(1) and [ ( 0(1),
if 6 is No- Lipschitz and Ba-sthhth, then (in the infinite-wied width limit taken starting from the layer 1 and then going with growing inder), for all E [1 I.), @l is constand during the training, and E AU
15

---

Published in Transactions on Machine Learning Research (04/2024)
Proof. The first statement follows directly from the chain rule. For the second statement, proceed by induction, taking inspiration in the original NTK proof of Jacot et al. (2018). For I= 1 the model is linear and the statement holds. Now, assume that the statement holds for networks of depth L, we want to show that it is true also for architectures of depth L + 1. We hence consider a network of depth L ++1 following the dynamiis
(1) = - Σ
Wh Whi (X; (X:1)..(2) - (0:++W0
dt and satisfying 1e (IV(2;0))),d==00) 1f0 (D)=))) We note that for 1 E [1: L]
(1) = 0(1), for all t E [0,
L+t Pii (x;t)= (
oU%+ (x; aw
nL
aUk (x;
UR (x:
1
nn
ouk.
VLL
===1
We define
(2;t)
Whe (0)0U((::)))))))
Vnn
k=1
and
F(D;t) -r(1---ww---- LII D;
so that we can rewrite the dynamics for E [1 L.] as
Σ (M
-M (1)
We have that
VuL ['" (IV(2;01)) de 7o / [IW++(0I (01I (1)
which is of order O(1) if [|)++(()]]/ is. This is indeed the case, as know that ||W+++()) (t)l < ||A O1), and ||WL+100] = Ol + vTL) (this is a classical result random matrix theory; see for instance Vershynin, 2012). Moreover, for each t we have that that in particular
=
so
((DD 1-" In(D T+ (t):t)ldt (+())++ O()).
We can hence apply the inductive hypothesis to the sub-network made of the first L layers, and obtain that, for E [1 L). O' stays constant during the training and E = AU', We also recall that at initialisation the kernels @ are diagonal, and so we have that for all t E [0 T]
O' (x,x'; x'; = Sس O (I,x').
Now, in order to conclude we need to check that the claim holds also for the last layer. We have
L+1 m
O (z, x';)
=
L+1N
(x:
aUL+ (a;
out
2/4
(x:t)
w/
++1:++1 (x/:t)
auf
ouy
1.d
aUt+ (x; 1) au 6 (x,
1/
b+1:L+1
(x;
(x:
w/
1+1:L+1 (x':t)
oUf
j=1
16

---

Published in Transactions on Machine Learning Research (04/2024)
81 (mst)
Let us denote as (2; the vector with components (x;((;;t)
We easily see that
|| t)ll
Ta Trs. &
(
1 (0)1 + Te Vnl.
To
(0]=+ O(1),
where we used that ( y and that the norm of the row of a matrix is always bounded by the Frobenius norm the matrix. On the other hand, we have that
A (2; t):
To |AW +t 1++ (+)++(++))) (t)l
+ So Vns
l | (O])<]|AU**)))
VnL
where we used that @ Y0-Tipschitz and fg-smooth. We know that |AWL+ O(1) as DL+1 O(1). Moreover [[ (0) loc ||W = Ovolog vi.)) (Boucheron et al., 2013). On the other hand, ||AU)(()) O(1) by Lemma We hence easily conclude that
behaves as the maximum of nt independent standard normals, that is
nL OUL+1 (a;t) - 0z (xx,z)) 0( ((vloo(m))/mm) a
auj
Now to show that the variation of the NTK vanishes during the training we only need to control (x; (x; 2 + (d):t)). First, notice that
A A(vf+I:L+1
1 1+ (x; 0)1 S
0(1).
Moreover, have that 1
+++1: 1: (x; Vns.
bis
(x; and
| ++1+++1 (x;
Ta 1001=)
Vn
We thus deduce that and
A A A (x;
1+11++11111
2/ ++++1++1 -
A x';() O 000==0((ooooo (og(mz)//z))
kk
which shows that in the infinite-width limit the NTK stays constant during the training.
Now are left with showing that EL+1 AUL+1 Recalling the notation uky (x; a0rf++(a;)) (ait) write
80
L+1
+x;t
E
1 Σ
++H AW d (2;1+(((;;)) AWL-I(t) + (x
1=1
Using the induction hypothesis EL = AUL, we get that
E (x; w +++:1+1 (x:t) AWL+ ux; AU*(x;t).
Using that a.U
++1:1+1 O,WL-1 un 0,UL, we can write
=x+1(a;t) AU
f/' Wh Cebthitt (x; - - ++++i++ (x;1)) a.W.++ + f." (un(o; -u ux(xr;))) 0,U"
In particular,
++ (x;
k
| ++1:++1 (x; 1')ll / 11,, (2. +2 sup (E
|| (z; (')l J [10,U4 (x;
22 2 sup (E(0,(]
17

---

Published in Transactions on Machine Learning Research (04/2024)
From what we have shown already, we know that
| 1o:::+++ (I; (')ll = (0100(1/)))) and
[ (') (((lgg())m))))))
sup r'(oad]
sup #e(oo.t
are left with checking that the last two integrals of order O(1) in order to conclude. First, we have
W - _
VnL (x - (O):)Aw) (t).
Since (x; O(1), we have |0((0" (x; =( O(v and we can define
K' sup [Z(0 0(1)
VTr r'Es
We have thus obtained
|0, 1.+ (t)ll S K' (IV(Z)))))), + A1/(D *
(t)ll
So,
L (()W++(()))ll s K S. (II((Z dt' + sup (AW-+((t')l J (()))a=)) = 0(1),
rejo.d
since both integrals in the RHS can be controlled by hypothesis.
Now we just need to control S 11, (')''"'. Defining K(x) = suparss (x,,,,),, we easily get that
|0,0 ((x) (1((2:01)) + XFDD"
We have established that Jo since by assumption Sn (f)) - 0(1), we have that indeed
(1 (IV dt' = 0(1) and supreeja. | = 0(1. particular,
[ I,, = 0(1)
With this last step we have shown that in the infinite width limit
Eat t) = Ax;+1(ztt))
which concludes the proof.
A.2 Variance of the wide stochastic network
From Clorico al. (2023b) (eq. therein, applied to a onc-dimensional output) we know that the output's variance Q2 of a shallow wide stochastic network is given by
1 n M [1 + (m)( (xxxt)=
Σ (M)())((MM(():))
where we define {(u) = B,-Nioo,)l((( + u) and we recall that (uu) = E + +v].
We now consider an initialisation where each component m' and m² is sampled independently from (0, 1). Then, since for any input x we have ||^||| = nog each component of M) (0) distributed (with respect the initialisation's randomness) as a standard normal distribution. Since m² is independent M we can easily derive that, the limit n > 00;
Q (x; 0) 2Rc (osovi.
18

---

Published in Transactions on Machine Learning Research (04/2024)
which is a deterministic value.
We now show here that Q² keeps constant during the training. We have that
(a; f)l - v
= (t)l=
Now, let ut(() = 1 + (t) We have that
Σ (1 (())))))))))) - u(() (z;)) + Au(t) ((;) + Au(t) (x;t) j=1
Let us start by the first term. We have that
(x; Ec- (26 (z: +) + +
and (recalling that we are assuming that ø is Co Lipschitz)
lu(0) (E(M((;;())
< 2
B Bc-N 40000000000000000000004400000000 +
5=1
1/2
n Σ (0)
s
+ (0:00
For large we have that
Σ (0)(0)(((M + O)I + 2E + = -()
(in probability with respect to the random initialisation) and ||u(0)] + Vnn. Since AaM = 0(1), we have that
(0))
=
Proceeding similarly get that
1/2 m m300(((A(<;;00))
AA(()) E(M (2; 0))1 2
Σ
1/2
< 2
Σ (x)0))((A((x000)
-(V)).
Finally, we have that
Au(t) AE(M (z;t)) = Σ
m3 (1000)))(()))) + (t)) (((M))
1=1
1=1
We have that
(A
1/2
22CgEc-Nio.)
+ m 0))
[tma(()]a (a;
+ (3]l AM =
19

---

Published in Transactions on Machine Learning Research (04/2024)
and
Am3(0)
2CC []l(( + "1 [M (c
With analogous reasoning, we obtain that
A
(m (A)))))))))) - -0(1/m))
==1
and so conclude that
namely the output's variance is constant the deterministic value o² throughout the training, n oc.
20


Reviews:
[{"review_id": "1UTfU1jHWz", "reviewer": "Reviewer_qo7i", "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* Overall the paper is well-written and easy to read (modulo minor issues listed in requested changes).\n\n* Connection to PAC-Bayesian training appears interesting, but I'm not well-familiar with this field to properly understand its importance.\n\n## Weaknesses:\n\n* The case of $\\rho \\equiv 1$ (i.e. $\\ell^2$ weight decay to initialization) seems to have been considered already in https://arxiv.org/abs/1905.11368 (Theorem 4.1).\n\t\n    * Impact and/or motivation for generalizing to other monotonic $\\rho$ isn't obvious -- the paper doesn't consider examples of other $\\rho$, and from my understanding the dynamics in this case do not admit closed-form solutions.\n    \n* Implications from the connection to PAC-Bayesian training are also not very clear, and would benefit if motivation / envisioned consequences were more specific and spelled-out.\n\n* The paper has no experiments, and to my understanding no straightforward/immediate applications.\n\nFor the above reasons I am on the fence regarding whether the paper is of high-enough interest to the TMLR audience.", "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 analyzes the infinite width (lazy; Neural Tangent Kernel (NTK) regime) training dynamics with a regularizer equal to [a monotonic function $\\rho$ of] the $\\ell^2$ distance of weights to their initialization (not to $0$). The paper:\n\n* proves that in the infinite width limit the NTK stays constant throughout such regularized training;\n\n* derives closed-form dynamics for $\\rho \\equiv 1$ (i.e. $\\ell^2$ weight decay to initialization) and shows that it's equivalent to regularizing the network outputs (to initialization);\n\n* shows that a certain setting of PAC-Bayesian training in the infinite width limit yields similar training dynamics (NTK with a regularization to init term).", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "## Major:\n\n* Please outline novelty / impact of your results relative to https://arxiv.org/abs/1905.11368\n\n* Could you elaborate on the importance / impact of section 6, describe specific use-cases / follow-ups to it?\n\n  * (Overall, the paper would benefit from explaining its impact / importance more)\n\n## Minor:\n\n* Under equation 3: would be good to add a brief comment on what $\\Xi$ represents.\n\n* Section 4.1, \"Evolution of the training objective\": when reading the paper linearly the purpose of this subsection isn't clear, would be good to motivate it / put it more into context.\n\n* Section 4.2 mentions a \"... general regularizing term which still leads to tractable training dynamics\". Could you elaborate in what way is it tractable? Aren't equations in Theorem 2 not solvable in closed form? \n\n* I assume $F(x; t)$ in section 5 (immediately before section 6) can be solved in closed form? If so, would be good to mention it / write down the solution.\n\n* Some light proof-reading could be useful, e.g.\n\n\t* Abstract: \"yields an additional term [which?] appears\"\n    * Intro: \"NTK dynamics prevent~~s~~\"\n    * Intro: \"approach ~~approach~~\"\n    * Section 2: \"Data consists ~~in~~[of]\"\n    * Before equation 5: \"~~Moreover,~~[Where]\"", "broader_impact_concerns": "No ethical issues -- a theory paper.", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "6oF3gQRe6O", "reviewer": "Reviewer_phPi", "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": "A main strength is that there is no way to study regularization in the framework of NTK yet. While this work does not make too much a progress, it is a decent progress forward I think. (Also, see the requested change section for why I think the progress is limited)\n\nIn terms of weakness, I do not have much to complain about. I think the theory is straightforward, but this is not a problem.\n\nGiven the novelty of the paper and there is no major concerns in its results I do lean towards acceptance", "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 works extends the framework of neural tangent kernel to the case when regularization is used (though a special type), and applies the theory to understand PAC bayes learning.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "I think the authors should add the following discussion to the paper. Essentially, the result of the paper can be understood through a simple Taylor expansion.\n\nIn the NTK regime, the model is linear around the initialization:\n$$f(\\theta) = f(\\theta_0) + \\nabla^T f(\\theta_0) \\Delta \\theta$$\n\nThis means that the loss function is also essentially quadratic in $\\theta$.\n\nBut the regularization term under consideration is also quadratic in $\\theta$. Therefore, the result studied by the authors can be simply understood by observing that the effective loss function is nothing but:\n$$L = L0 +   V^T \\Delta \\theta + (\\Delta \\theta)^T (H + \\lambda)\\Delta \\theta,$$\nwhere $\\lambda$ is the regularization strength, and $V$, $H$ is a constant vector/matrix.\n\nNo wonder the resulting dynamics is still linear. Adding this discussion to the paper can make the results of the paper much easier to understand for ordinary readers and improves the paper in my opinion.", "broader_impact_concerns": "NA", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "WIXGkF9ttq", "reviewer": "Reviewer_Gxk2", "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- The problem is well-motivated and focuses on a task that is relevant to the TMLR audience\n- The paper is well organized and follows a clear structure. \n\n## Weaknesses\n- The paper is rather dense with respect to its notation, although that probably can't be avoided.\n- Question: I am probably overlooking something, but where does the regularizing $\\lambda$ appear from in (9)? \n\n### Minor\n- The manuscript seems at certain places unnecessarily hesitant/humble for the results the authors discuss. E.g., in the abstract \"potentially contribute\", or in the conclusion \"appears to be a regularized version\".\n- Minor notation. It is somewhat overlapping that F is both the output as well as the Frobenius norm\n\n### Typos\n- Equation above (9) should be $\\hat \\ell$ instead of $\\ell$ unless I am mistaken\n- Sec 5, 2nd line: $\\hat \\ell(\\hat y, y) = \\frac{1}{2}(\\hat y - y)^2$", "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 consider the question of how neural networks evolve in an infinite-width setting and extend current results from the neural tangent kernel (NTK) literature to regularized objectives, as long as these regularize the parameters to remain close to their initialization. \n\nIn the second part of their manuscript, they apply these results to PAC-based learning objectives, extending results by Clerico et al. (2023).", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "None", "broader_impact_concerns": "None", "additional_comments": "", "extra_scores": {}, "extra_text": {}}]

Author rebuttal:

Meta-review: The reviewers have made a number of concrete suggestions for improvement but the paper has not been revised yet to incorporate them. Therefore, a revision is required as a condition for the paper's full acceptance. Specifically, I would like to see the following incorporated in a revision:

- All comments made by Reviewer qo7i under "Requested changes" (both "Major" and "Minor").
- All comments made by Reviewer phPi under "Requested changes".
- All comments made by Reviewer Gxk2 under "Minor" and "Typos".
Decision: Accept (poster)