Convergence of Stochastic Gradient Langevin Dynamics in the Lazy Training Regime
Authors: Noah Oberweis
Semih Cayci
Venue: Accepted by TMLR ()

Abstract: Continuous-time models provide important insights into the training dynamics of optimization algorithms in deep learning. In this work, we establish a non-asymptotic convergence analysis of stochastic gradient Langevin dynamics (SGLD), which is an Itô stochastic differential equation (SDE) approximation of stochastic gradient descent in continuous time, in the lazy training regime. We show that, under regularity conditions on the Hessian of the loss function, SGLD with multiplicative and state-dependent noise (i) yields a non-degenerate kernel throughout the training process with high probability, and (ii) achieves exponential convergence to the empirical risk minimizer in expectation, and we establish finite-time and finite-width bounds on the optimality gap. We corroborate our theoretical findings with numerical examples in the regression setting.

Published in Transactions on Machine Learning Research (02/2026)
Convergence of Stochastic Gradient Langevin Dynamics in the Lazy Training Regime
Noah Oberweis
Deportment of Mathematics
RWTH Aochen University
Semih Cayci
Department Mathematics
RWTH Aochen University
Reviewed on OpenReview: https: // openreview. net/ net/forum? id= 35037DKmLo
Abstract
Continuous-tine models provide important insights into the training dynamics of optimiza- tion algorithms in deep learning. In this work, we establish a non-asymptotic convergence analysis of stochastic gradient Langevin dynamics (SGLD), which is Itô stochastic dif- ferential equation (SDE) approximation of stochastic gradient descent in continuous time, in the lazy training regime. We show that, under regularity conditions on the Hessian of the loss function, SGLD with multiplicative and state-dependent noise (i) yields a non- degenerate kernel throughout the training process with high probability, and (ii) achieves exponential convergence to the empirical risk minimiize expectation, and we establish finite unite-time and finite-width bounds on the optimality gap. We corroborate our theoretical findings with numerical examples in the regression setting.
1 Introduction
Stochastic gradient descent (SGD) has been the main workhorse of deep learning due to its low computational complexity and effectiveness. The cxisting thcoretical analyses of first-order optimization methods in deep learning predominantly study the deterministic batll- gradient descent, which neglects the impact of inherent stochasticity of SGD in training neural networks. Although there exist recent works that investigate stochastic dynamics (see, for example, Lugosi & Nualart (2024h), explicit finite-time convergence rates for the stochastic first-order methods under realistic and interpretable conditions still remain elusive, which motivates our work.
In this work, establish theoretical guarantees for stochastic gradient Langevin dynamics (SGLD), continuous-lime approximation of SGD, in the lazy training regime. One important application of theoretical analysis is finite-time and finite-width bounds the training loss for deep feedforward neural networks trained by SGLD in a supervised learning setting. The main contributions of our work include the following:
Exponential convergence rates in expectation: By leveraging tools from stochastic calculus, we derive stochastic Grönwall-type inequality for SDEs in the lazy training regime, which establishes an exponential decay rate for the expected optimality gap as long the parameters are within a certain neighborhood of the initialization. To that end, we explicitly characterize the impact of the curvature of the loss function in the parameter space and the output scaling factor on the effective noise during the training process. The resulting scaling law describes sufficient conditions on the curvature of a global minimum to prevent SGLD from escaping it, which is one of the key outcomes of our analysis.

---

Published in Transactions on Machine Learning Research (02/2026)
High probability bounds on the conservation of the lazy training regime: A key quantity in the lazy training analysis is the first-exit time at which the parameters leave certain neighborhood of their random initialization, which lower bounds the duration of descent throughout the optimization steps. Leveraging our stochastic error analysis, we derive high-probabiltt bounds on this first-cxit time; and show that, for sufficiently large output scaling, the parameters remain in this neighborhood indefinitely with high probability, thereby completing the convergence argument in the stochastic setting.
Applications in deep learning: We apply our theoretical analysis to derive explicit non- asymptotic (finite-time and finito-width) bounds on training loss for deep neural networks trained by SGLD in the lazy training regime. To the best of our knowledge, this is the first analysis of stochastic gradient Langevin dynamics in this regime.
Our work builds on and extends the deterministic lazy training analysis in Chizat al. 2020)) which analyzes the convergence of gradient flow for overparameterized neural networks in the lazy training (or the so-called kernel) regime. In Chizat et al. (2020), it was demonstrated that the lazy training phenomenon arises from an appropriately chosen output scaling factor, which ensures that the neural network paramcters remain in a small enough neighborhood of their random initialization. We note that Chizat et al. (2020) considers deterministic (full-batch) gradient flow in the lazy training regime, and does not address the stochasticity that stems from stochastic subsampling in SGD, which is crucial for learning with large datascts. In this work, we extend this framework by analyzing a stochastic continuous-tim approximation of the SGD,
1.1 Related Work
Convergence of overparameterized neural networks in the kernel regime. Overparameterized neural networks have been shown to achieve interpolation, which implies zero training error when they are trained with first-order methods despite the highly non-convex optimization landscape. At the same time, they achieve impressive generalization performance even with noisy data (Belkin et al. 2019 Bartlett et al. 2020 Zhang et al.] 2021). Additionally, works such as Oymak & Soltanolkotabi (2019) have demonstrated that the overparameterized regime can be reached with fewer parameters than previously suggested by the theory. The theoretical explanation of this phenomenon has been a focal point of interest. It was shown in Jacot et al. (2020); Du et al.] (2019b); Li & Liang] (2018) that aandomly-intiallin networks, trained with gradient descent, achieve global optimality without moving away from their random initialization, which is known as the lazy training kernel regime. The main inspiration for this work is Chizat et al. (2020), where lazy training is first analyzed as a consequence of an artificial scaling factor. The existing works in the lazy training regime predominantly analyze overparameterized neural networks trained with (full-batch) gradient descent, and they do not extend to stochastic gradient descent, which is the core focus of our work. Additionally, it is essential to point out that, coutrary to other works, our primary goal is not the analysis of linearized dynamics defined by h(w) := M(wn) + Dh(wo) - wo)- It can be shown that gradient flow stays arbitrarily close the linearized dynamies for sufficiently large a. We will use very similar techniques to prove an exponential convergence rate of the empirical error for a sufficiently large a without looking into the case where a
overparameterized neural
Convergence of stochastic gradient descent for neural networks. Dynamics of stochastic gradient descent have been investigated in Lugosi & Nualart (0224) for deep lincar neural networks, which shows that a global optimum is achieved under some regularity assumptions on the loss landscape with a convergence rate that has an implicit dependency on the problem paramcters. In our work, inspired by the analysis in [ugosi & Nualart (2024), we analyze the stochastic gradient dynamics for training non-linear neural networks, and establish exponential convergence to the global optimum with an explicit and interpretable convergence rate under casily verifiable regularity assumptions.
Modeling of discrete training algorithms by SDEs: Continuous-time modeling of discrele-time stochastic gradient descent by using SDEs has recently been investigated in many works, as the SDE model provides essential insights into the discrete stochastic gradient descent training algorithm. In this work, we consider an SDE-model to approximate stochastic gradient descent dynamics, inspired by Li et al. (2017). which proved the viability of using stochastic gradient Langevin dynamics to model mini-batch stochastic
2

---

Published in Transactions on Machine Learning Research (02/2026)
gradient descent. Under certain assumptions, Li et al. (2017) shows that mini-batch stochastic gradient descent, is a weak order-one approximation of the stochastic gradient Langevin dynamics equation (see
Remark [p.
Our work is closely related to Chizat et al. 2020 and Lugosi & Nualart 2024).. However, there are many
substantial differences between prior work and our contributions. For instance, Chizat al. (2020) does not
consider the stochastic term needed to model stochastic learning algorithms. Compared to Lugosi & Nualart (2024), we investigate the dynamics in the lazy training regime, which allows us to formulate less restrictive and easily verifiable assumptions in machine learning context. An in-depth discussion on the differences of our contributions and prior works is presented in Appendix A
We limit our results to continuous models (as defined in the following section), although training typically employs stochastic gradient descent or other iterative methods in most practical applications. In addition to the multiple technical advantages arising from the continuous approximation of real (discrete) dynamics, such as the possibility of utilizing many tools from Itô calculus and the independence of the step size of the algorithm, this choice has yielded signifieant insights that apply to discrete models. One such instance is presented in Malladi et al.] (2024), where the authors derive contimuous model for the adaptive gradient algorithms, which is then used to motivate the optimality of the square root scaling rule for adjusting hyperparameters as batch sizes increase.
1.2 Notation
Id denotes the identity matrix. 3 indicates that B - A positive-definite. ||>||| = Vx x denotes the 2-norm of x E TR. We use Amia( (A) for the smallest eigenvalue of a Matrix A. For a positive-defintte matrix A, we define the vector norm || 13 == x" Ax. Additionally, || is the Frobenius norm of A and B, (xo) is the open ball of radius around the 2o. in the Euclidean norm. We denote the modulus of Lipschitz continuity of function f as Lip(J). The notation Dh denotes the Jacobian of h. Similarly, Vh represents the gradient of h and 4 represents its Hessian. For a function f R" - R, it holds that v"f= DDf We add an index to V to indicate, that the gradient is only taken with respect to one of the arguments. Рд can be seen as the gradient in the Hilbert space F. However, this can be replaced by the standard gradient with respect to the values that h E F takes on the finite set {= Lastly, l denotes the characteristic function, which is equal to 1 if the logical expression A holds true and 0 otherwise.
2 Problem Setup and Stochastic Gradient Langevin Dynamics
2.1 Problem Setting
In this work, we consider a smooth parametric predictor h : IP + F, where the parameter space is R2 and the output space is a Hilbert space IF (equipped with the norm || ||z). In addition, we assume that the empirical risk R. IF + R4 and the loss function é Rd F + I are smooth. Both functions are related by Ez (1x, Mw))] = R(h(w)), where E, denotes the expected value with rogards to some probability distribution P from which the data points x ~ P are sampled independently. The main objective in this paper to solve the following problem:
Find wf s.t. R(((w")) <E
for some c > O. Empirical risk minimization in supervised learning is special case of the formulation above, where P is a discrete probability measure on a given set of data points (( p)) and an unknown ground truth h+ e F, that connects z, and ys by 1² = yi- In empirical risk minimization, a popular choice for the loss function is given by [(x, h) =(h(x) - y)?, where y = h 1z).). For further details about the supervised learning formulation, see Section where we apply the results to train neural networks with smooth activations in a supervised learning setting.
3

---

Published in Transactions on Machine Learning Research (02/2026)
2.2
Lazy Training Stochastic Langevin Dynamics
We consider the scaled stochastic gradient Langevin dynamics defined by the Itō SDE
dur = - h(we) (in + yo (Ealar))
(1)
and Ta E Rt. We adopt a random initialization wo such that l(w)) 0 holds almost surely. For neural networks, this can be achieved by using the symmetric initialization schome (Bai & Lec 2020 Chizat et al.
where W, is a standard Brownian motion in RP, [a(w) is the covariance matrix of Val(; ah(w)) Val(;ah(w))
2020 Cayci Eryilmaz 2025).
Here, Ta (w) : (w)) denotes the matrix square root (w) (w), which is well-defined since (w) is symmetrri positive definite. The noise scaling factor The may depend on the output scaling factor in certain cases (see Appendix[] By.
Remark 1. For the empirical risk minimization problem where the risk function is defived by R(h(w)) is continuous-timg approzimation of the stochastic gradient
E l(T,, h(w)), the SDE in equation descenl (SGD)
na D (a + Via Vk.
Wk+1 = Wk
(2)
a
DJ M(wx) ah(wx))) and x ~ Unif(21.....,ail is sample from the distribution on the inpul space. In that case, Vilar is a random veclor with mean O and covariance matric na (xx)) In Li et al. (2017) it is shown that the stochastic continuuus-lime model that we consider in equation [ an order-one weak opprorimation of the SDE equation [ There are other stochastic models approzimate SGD in continuous time [L et al. 2017: Simsekli et al., 2019: Gurbuzbalabon al. 2021
where V& == (T (olu)))
3 Convergence of Stochastic Gradient Langevin Dynamics in the Lazy Training Regime
In this section, show the convergence in expectation of the loss function under suitable regularity as- sumptions. Contrary to other SDE models, equation [ contains a noise term that vanishes at a minimizer. Due to the introduced regularity of the loss é in the following Assumption and curvature constraint from Assumption 2 we can effectively keep the parameters from escaping a minimum.
Assumption 1. We make the following assumplions.
Regularity of the loss function: 8:R" x F R is Lipschitz-smooth
() Tne(xx,12)l)) Lip(V halle:
and m-strongly conver
(V - Val(a, h2))
(hr h22 2 milhi h211F.
Regularity of the model: The derivative of the model h R - F is Lip( continuous, and the newral langent kernel at initialization is striclly posilive definite Amin (Dh(wp) (T((o)) > 0.
Assumption holds for many loss functions, with the most common example being the squared error (see Appendix The smoothness assumption h holds for shallow neural networks (see Corollary [.. For deep neural networks, we will present an alternative that circumvents the smoothness of h (sce Corollary 2 and Appendix C The empirical neural (angent kernel is strictly positive definite for overparanneterized neural networks with high probability over the random initialization under reasonable data regularity conditions (Du al. 2019b/a Banerjoc al. 2023)..
The following assumption on the Hessian of the risk function is critical for the analysis in the stochastic setting.
4

---

Published in Transactions on Machine Learning Research (02/2026)
Assumption 2.
We assume that Tr ((a(w)) R(oh(()))a)
2a for all w E (,(w)), where r is the
R(ah(w))
radius that will be discussed later in this section (before Theorem
Assumption 2 antomatically holds for sufficiently overparametrized shallow and deep neural networks under
realistic assumptions. In Appendix B and C we provide the proof for shallow and deep neural networks with smooth activation functions.
Assumption 3. Assume that the gradient and the expectation of the loss interchange, ie.., V,,((h) = [VV.((x,h)].
For the empirical risk minimization in supervised learning, Assumption automatically holds due to the uniform sampling distribution over finite training set.
Lemma 1 (Regularity of R). Under Assumption 1 and Assumption 3 R is (ip(V.)) (VnL)-smooth and m-strongly COnVER.
Proof of Lemma
1 Using the convexity from Assumption
we have
[V V R(ha))
|E (.V..(( h1) 7xl(x,h2))l
< 1 h1)
< [Ei(Li (v - -hallel
= Lip( (Vp())) -h2lls-
Similarly, we can use the m-strong convexity from Assumption to obtain
(VnR(h)) Vn R((h))) (h1 h2) E. (Vvn(1r h2)( [h1 h2)
E (V h1) Vnl(I, (2)) h2)]
2 E2 [m|] h1 hallel
mmlhi h2 b2|
which concludes the proof.
Remark 2. At this point, it is essential to point out that Assumption 1 is already sufficient for the existence of a global minimizer of R in the supervised learning setting with a finite set of data points. Since we have shown in Lemma 7 that h R(h) is m-strong connex as funccion on the prediction space. Since the prediction space
be characterized by finite dimensional set if R only depends on a finite number of
values of h, we already know that there erists globol minimizer ht* E F. However, we do not assume the existence of w" E IP such that h(w²)) = h² (which is not needed for the theoretical findings of this poper). The eristence of wn is generally not guaranteed, since the mapping - R(h(w)) does not have to be strongly convez. Later in the numerical experiments as well as the proofs of Corollaries | and g we assumed the existence of w since it holds for the teacher-student setting, and simplify the colculations.
As key feature of analysis in the lany training regime, the weights wt do not deviate too far from their initial value, which implies that Ald = Dh(w:) ((((:))""(())))) h(wz) for some À >>> 0. Many publications have focused establishing strictly positive lower bounds for this value for deterministic (full-batch) gradient flow (sce, for instance, Du el al. (2019b) Chizat et al. (2020)); ); however, we are not aware of any results that guarantee for SGD or SGLD, We will solve this problem by introducing a
the positivity of À = Ami((D((u)) (D(w,,)D" stopping time, which the smallest time at which this condition no longer holds, and prove in Corollary 3 that this stopping time will be oc with high probability for a sufficiently large. The minimum eigenvalue of the NTK is strictly positive, at least long | -woll2 rr:- Lip(DaJ Therefore, will introduce the random stopping time
X
=== 2 ||ا, woll2 > r),
denoting the first time Ald 2 Dh(w,)D"h(a)) is degenerate. In addition, the inequality in Assumption 2 also holds at least for t S r. Working with this quantity and proving T oo with high probability is a major challenge in our analysis. In addition, we define min(TT r} =; T, where T 0 is some arbitrary but fixed runtime of the algorithm. Let < T in the following.
5

---

Published in Transactions on Machine Learning Research (02/2026)
Theorem (Convergence of stochastic gradient Langevin dynamics). Let ht be the global minimizer of R. Define the oplimality gop as R(h)- = R(hh) R(h*) > 0. Then, under Assumplions 3 il holds that
Bu(R(ac(())))] - R. [R(a((w))) exp(-mX21),
(3)
for any T.
Proof of Theorem 1 In the first part, this proof follows Lugosi & Nualart (2024) by using Itô's lemma find an exponential upper bound of the risk function. In the second part, we use the scaling factor a control the boundedness of the sccond-order (Hossian) term by using the interchangeability of the gradient and the expectation of the loss function. This directly connects the boundedness of the Hessian term to the curvature of the loss in parameter space.
As defined, Wt follows the SDE in equation 1 In the following, we apply ItO's formula to compute dg(w)), where g I² + R is defined by g(w) = log( ((ah(w))) for all with R(ah(wr)) O. The gradient and the Hessian of g are as follows:
Vug(w) R(ah(w)) aD (W"")WV
Vig(w) M(h((w)) a² D" (WW.A//HHHHHHH))))HHH))).....
M(hh(w) VER(ah(w))
This implies that g(wr) the solution of the following SDE
dg(we) (DDT)ar) Va dt
(4)
R(ah(we))
R(ah (a)
+ na Tr
(5)
2
R(ah(we))2
[a(we)) (aahhaa (a dt
Ia Tr
+
(6)
2a
[(ah(w.))
nna VJ, Rlah(i) dW.
+
(7)
R(ah(wi))
Similar to Lugosi & Nualart
2024 we observe that the quadratic variation the process
V R((hhws)) (M()))) dW.,
Mc = Th
R(ah(ws))
defined in equation given by
(M), No ( Tr V ((u ds,
R(ah(w.)))
which appears in equation 5
We can use the multi-dimennsonal Itô formula again on the process & == (M)), to prove that is a nonnegative martingale, which implies that.
E.[E, 1.
Returning to equation Lequaiio can rewrite this SDE in the integral form and take exponentials both sides, which gives
(-)-- (DTh(w)) V ds
R(ah(w))) = =[(ah(wo)) exp
R(ah(wa))
Ta
/" Tr
||a (wia) (33 ds Mi (M7)))
1
exp
2a
Jo
R(ach(ws))

---

Published in Transactions on Machine Learning Research (02/2026)
Next, will bound each of the integrals separately. It holds that
(i (n do [V Binte, DDT ds. R(ah(ws))
R(ah(w,,)
Since t < T, the cigenvalues of Dh are lower bounded by X2. Therefore, it follows that
Wil ds 2 X2 [/" [VR(ah(w>))] 2 22pm S ds = 2X2mt, R((h(wws)
M(ah(wa)) using the Folyak-Lojasiewin inequality, which holds due to the strong convexity of R. By Assumption we get
/ Tr ( ds S S x2m ds
na
2a
R(ah(ws))
Consequently, we obtain
1 (Mi)
R(ah(w))) S R(ah(wo)) exp
12X2ml X2 + Mr -
R(ah(wo)) exp
-X2mt + Mr 2
Taking expectations on both sides and using equation
the desired result follows.
Remark 3. It is crucial to point out that Theorem
does not provide the result
Bu ((((((u)))) 1o O,
but instead
t-000 0. E. [
We will return to this problem in Theorem to control the error stemming from E., (aa((((())))
Until now, we have only provided general result that be applied to different settings. The following results justify the applicability of Theorem to shallow as well as deep neural networks in supervised learning, Corollary 1 (Risk convergence for shallow neural networks). Consider a shallow neural network of the form
6(0, Σ coo(W))))
where only train the weights and randomly choose at initialization. Here, is the non-linear activation function tanh. This neural network fulfills the assumplions of Theorem 7 This implies that, for appropriate loss functions that fulfill Assumption (such as the mean squared loss function), equation holds for shallow neural networks.
Proof of Corollary The proof is presented in Appendix B
In the following, we extend the results for the shallow feedforward neural networks to deep neural networks of depth H based on Du al. (2019b).
Corollary 2 (Risk convergence for deep neural networks). Recursively define H-layer deep neural network
(()
Cr - (o(w(k)) 2(1-11),, for 1 1sks
V
f(rw)) (H)
for input data (0) E Rd and W(1) E Rxd as well as (Kk)) E Rmxm for 2 < < H. The parameter w contains both the hidden layer weights (Kk), as well as the output layer a. Tf a I--Lipschitz, [[(() S Gu.o m and (1x(00 (0)]l2 I 62.0, then equation holds for a Lip(Vh)-smooth and m-strongly conver loss function.

---

Published in Transactions on Machine Learning Research (02/2020)
Proof. The proof of Corollary can be found in Appendix
Remark 4 (First-exit time for deep neural networks). The theorem cannot be applied directly to deep neural networks, since they do not have Lipschita continuous derivative. Therefore we connot define the first-erit time
T :: inf
(t: inf(t: Ost
l1we woll > X/LAp(Dii))
However, using the Lemmas B1-B4 in Du et al. (2019a), we can define a different first-enit time that guarantees the positive definiteness of the NTK. Instead of using the radius r == L/Lip(Dh) we will use different radius, for which the above proof be conducted idendically. This circumvends the need for Dh. to be Lipschitz-smooth Further discussion is provided in Appendiz G
Our immediate next goal is to find an upper bound for the optimality gap (either in predictor or parameter space) that is independent a. This will allow us to increase a without changing the bound itself. it holds that
Corollary 3 (Convergence to the global empirical risk minimizer). Under Assumptions
Lip(VR) m B Bulllah(wo) - h-" 111l exp(eexXX
Ew Ewllaa((we) h*1
By choosing initialization for the neural network, such that h(wo) = 0, the right-hand side becomes independent of a.
Proof of Corollary 3 By the strong m-convexity of R follows that m [ah(wr) - 10 s R(ah(w:))-
We then combine this, with Theorem [ and the Lipschitz contimuity of R get
m
Re [[aa(ww)) - h h-]le] s E. (R((a(()))) exp(-m>>>
2
Lip(VR) R.g [lahh(w)) - h || mm)m)))))).
2
where we used the smoothness inequality R(ah(wo)) R(hh) /Lip ah (h(ww))))))))))) 11- 11 by Lemma 3.4 in Bubeck al. (2055). This leads to the desired equation
B[llah(we) - ht" Lip(VR)
E. [llah(ww) - 111 exp(-mx²
m
We use the above results to show that by choosing a sufficiently large, we can ensure that T = oo with high probability since AId < Dh(w:) D" holds. In the following result and throughout the rest of the paper, we denole the conaitional probability and expectation given the random initialization wi as Pa and E...
Theorem 2. Under Assumptions
3 it holds that
P[[[w: woll2 > r) < [[Dhaa))lsTap((
Lip(Ve) mX
1 a
(9)
Proof of Theorem To prove the above statement, we will first use equation and rewrite it to get an upper bound on the expected distance of wi from initialization. The resulting terms can then be further bounded by using the regularity of the loss function from Assumption 1 and the results from Theorem [ We connect the expected distance to the probability of exiting ball around the initialization using Markov's inequality E. -wlll Using equation we can rewrite the distance of wa to the origin
by P[[| woll e) <
|| woll2 111/
a 1 DDTM(. + Vila S' calwi)dBB
1
/
"110 [I l2dt yTa ff' oa (ur)dBil11
< a
8

---

Published in Transactions on Machine Learning Research (02/2026)
Using Cauchy-Schwand and Theorem 5.21 in Mao 2007
subsequently, we obtain
Eww ||we woll| - Eo [D"h(w))] VV
Va So (a.
a
L. [
B ISo Ir(Za(w))dde
x
We will first apply Markov's inequality to obtain
P[[ - woll2 2) Boll ("
E (w(Eaw
(10)
ra
ra
We will continue by first bounding the first part of the right-hand side Equation
10
Eallf Mll R(ah((;) lladt)
or [ / [lah(wr) ht 122d
Ta
/.' (ulllah(((r))) - 10- 112]3ett
[Di(wn))l=Lpp( /-' (ulllach((r))) 11 113]d
1 [D((w))llLip((())
Lip(VR)
EN [[adk(wo)) - h0 [2 exp( 1 22m22
or
m
Lip V
E [llaa(ww)) - 1-111] /' esp(-mA?tjdt
ar
772
Lip((V/ R)E [(((wo)) 11
5 ar
m= X2
where follows from Jensen's incquality and follows from Corollary term from the right-hand side of T0. We start by bounding
Next, we will bound the second
S (r(Ea(wa))ds
/ Tr (T (R.
R(ah(w))) (odh(wss)) ds
$
Tr (B., ( ds
E. (
(9)
Lip(Ve) ht" 11
Bes [lack(ww) Lip(V))) -1 [1:] / exp(-mX2s)ds,
m
where used the smoothness of in O and Corollary 3 in Jy exp((mA?ssdss MAz, which gives us the following upper bound
We can calculate that the value of
E
So Tr (ws))
Li)
Lip((VRJE//lah(uu)) 1+]]
ra
ramA
By choosing an initialization such that h(w)) = 0, we then get
P([|u - woll2 > r) [l(ww))llLipp(ee Lip(Ve)
1
Lip(V
or
m A2
mX
9

---

g does not Published in Transactions on Machine Learning Research (02/2026)
In the following corollary, we will show that by choosing appropriately, PP(r < T) is arbitrarily small. This implies that can effectively use our non-asymplolie bound that we established in Theorem | and Corollary to get an upper bound on the training error with arbitrary large probability (depending on a). Corollary 4. Under the some assumptions as Theorem 2
it follows for any T > 0
P(( < 7) - ()
(11)
Proof of Covollary 4 In the following let F ;= (in T). Since the right-hand side of equation depend on we get
P(r < T) PP < T)
S P([ woll rr)
O(a
where we used Theorem
The result provided by Corollary is the last ingredient needed to return to the problem int roduced in Remark and solve it.
Theorem (C'vvvneeee of SGLD). Under the same assumptions as Theorem it holds that
E [x([(a((())]] S exp(->²mtt) O(a
for any q >>>> and for any > 0.
Proof of Theorem To prove the above result, we will decompose the left-hand side of the result as follows [[[[H(((())))) (((((((((a))) + (e][((a
Theorem already proved that
En (M(aa( S exp(->2mt),
which implies that we only need show
E ((][(a(())) S O(a $).
Using Hölder's inequality with the Hölder conjugate pair (p, q), we get
Ey ( S Eu(R((h(wr)) )d E [1 ( 1
(12)
[u(aaaa p Re.l (r<T)]
(13)
EulR(ahh(wr) )rl P(r < T)
(14)
[u(a((h ) O(a -a)
(15)
where we used that the image of the characteristic function [(<300) is {0, 1} and Corollary in Additionally, it holds that
(TTh( | VR(ah(w>))1 ds M(ch(ws)
R(ah(wr)) =((ah(wo))) exp -DD
[oa(ww))lozg R(ah(ww)) ds c.E.
pMa
exp
J
Tr
2x <R(ah(wo))''''E
R(ah(wss)
10

---

Published in Transactions on Machine Learning Research (02/2026)
where used exp (-p.l" | hew, Viaahllll ds) which we upper bounded by 1. Taking the expectation on both sides leads to
identical
calculation to Theorem
except for the term
R(a)(w)))
B (R(ah())))) & R(ah(ua))" R(ah(wo))Pee
Plugging this result into equation 15 concludes the proof.
Since we have proven in Corollary that, with arbitrarily large probability, we stays in (y(w)), we can use this argument to motivate the proximity of the linearized dynamics Wa to wr As mentioned before, Chizat et al. (020)) uses the scaling factor a to prove the proximity of h to the linearized dynamics defined by h(w) ;= (wo) + (u(wo))(w - wo), which gives rise to the following linear stochastic gradient Langevin dynamics:
dw D" ((ao))((( + Ia (E(ar))
First, we demonstrate that the above results also apply the linearized neural network w.
Corollary 5. Assuming that Assumptions hold for the non-linearized neural network. In addition, assume that h(wo) = O. Then it follows that
Lip(VR) m |1 (xx()=m?t),
Emlllah((s)) h1
(16)
for all ST. In addition, it holds that
P( woll2 r) [DN(w))|FFLpp(( Lip(V!) Proof of Corollary 5
1
[ (a
or
m²x²
md
The proof is application of Corollary 3 and Theorem
From this, then the triangle inequality get the following results
Corollary 6. By using the bounds from Corollary 3 we can show that
Lip VR)
Exw Eullahh((st) - cc((w))lF) s 2
|/ || exp(-mX?
Proof of Proposition 0
It holds that
Eullaa((),) - ah(w))lle) < E. (llh((())))) - h" [|]] E. [[(a)(w)) h²1le]
(lllh()( h + Rew [Ilah(ww))) h 11
Lip R) | exp(-mX2t),
m
where we used Jensen's 's inequality in and Corollary and Corollary
5 in
This suggests that, within the lazy training regime, the error between the standard and linearized dynamics decays exponentially over time. For the lincarized models, following similar steps, we can obtain equivalent results Theorem 2 prove that P([l,, Welle > e) s O(2)) as long both we and W, are in the lazy training regime.
11

---

Published in Transactions on Machine Learning Research (02/2026)
4 Numerical Results
In this section, apply the results from the last section to shallow neural network in the teacher-student setting and deep neural network application. Consider the function defined by
((x==)) = E og tanh(w)a))
where w* = [w] ...w.ll E Rdxm' c Rm x E Rd. The training data generated as
y: o(znw") t L, n,
where w" E pdxmmm predetermined optimum parameter, 6, ~ NN(0, Id) and as ~ NN(0,, the empirical mean-squared training error
We consider
R()) Σ (1s f(x)))²
We consider single-hidden layer neural network of width m m' and tanh activations. In order to ensure that o(x;w)) wo) 0 for any RRd use the centered model ø(:w) -(:::)((((()))))) (((ww),, following Chizat et al. (2020). We simulate the stochastic gradient Langevin dynamics using the Euler-Maruyama scheme with a stepsize of Q1 = 10-2, 2 a noise factor 7/ = 10-2, ; and we repeat the experiment for a E (1/8, 8, 32, 256]. The input dimension is d = 16, and we chose m = 600 > m' = as well as a training set of size n = 800. We sample the teacher neural network weights from w ~ Unif([O, 1j26) and choose c* = 1. The student neural network is initialised by w/ ~ NO Id). We only train the hidden layer weights Wj and freeze the output layer weights ej initialization. The goal of this simulation is to compare the training loss, the minimum eigenvalue of the NTK, and the distance of the weights from initialization (Ez= ||اا w/ll22) for different values of a. The proofs, which show that this example fulfills all the needed assumptions, were moved to the appendix. The results are shown in Figure ! where the training loss rate exhibits an exponential decrease as the value of a increases. Here, we can also see that, even though we have achieved an exponential bound on the training error, this rate is still relatively slow, since the eigenvalues of the NTK are also small. Figure
/ (left), we can see that the values of W/ remain in the vicinity of the initial values for larger a. accordance with this, the minimum eigenvalue of the NTK (right) does decrease faster for smaller values of a. For 1/8, the minimum cigenvalue seems to decrease at (slow) exponential rate. For the larger values of a, this value seems to converge to value strictly larger than 0, realfirming our finding in the previous section.
Empiriaal risk time for different
10-1
20
40
62
80
200
120
140
Time C11
Figure Test-error evolution during training using Euler- Maruyama scheme with stepsize of öt 0.01 for E 8, 32, 256). We additionally include curve f(t) with the decay rate cxp(-X²t) (dotted line) reference, where 0.01122 is the minimum eigenvalue of the NTK at initialization.
12

---

Published in Transactions on Machine Learning Research (02/2026)
digenvalue
3.4
1.325 x
Dudante inbaitation
1.9
x.221x10"
2.1205 10-
00
102
1070
14395
0.125
0.125
255.000
255.200
10
- 0.030
32 02000
11385
- 8.002
22.0020 22.000
1.2
Umz
Figure 2: Distance from initialization (left) and value the minimal eigenvalue of the NTK (right) in time for different values a E {1/8, 8, 256}.
5 Conclusion
This work analyzed the lazy regime for stochastie gradient Langevin dynamics in the overparameterized setting. First, proved an exponential convergence rate for the expected optimality gap. Next, we es- tablished probabilistic conditions describing the probability of the system leaving this regime. Additionally, we empirically validated these insights in our numerical results. Interesting future rescarch directions in- clude the analysis of SGLD in the underparameterized regime, which was analyzed Chizat et al. (2020) for the deterministic gradient flow. Another interesting open question is the convergence study of different continuous- models for SGD in deep learning, such as Li et al. (2017) based on higher-order approxima- tions, and Gurbuzbalaaan et al. (001)) that models SGD with heavy-tailed Levy processes.
13

---

Published in Transactions on Machine Learning Research (02/2026)
References
Yu Bai and Jason D. Lee. Beyond linearization: On quadratic and higher-order approximation of wide neural networks, 2020. URL https://arxiv. http///awiii..rr//a///////9/0/ 1001001199900001199999
Arindam Banerjee, Pedro Cisneros- Libin Zhu, and Mikhail Belkin. Neural tangent kernel at imitial-
ization: linear width suffices. In Uncertainty in Artificial Intelligence, pp- 110-118. PMLR, 2023.
Peter I, Bartlett, Philip M Long, Gábor Lugosi, and Alexander Tsigler. Benign overfitting in linear regression. Proceedings of the National Academy of Sciences, :37003)330063 30070, 2020.
Mikhail Belkin, Daniel Hsu, Siyuan Ma, and Soumik Mandal. Reconciling modern machine-learning practice and the classical bias variance trade-off. Proceedings the National Academy of Sciences, 116(32): 15849-15854, 2019.
Sébastien Bubeck al. Convex optimization: Algorithms and complexity. Foundations and Trends@ Machine Learning, 8(3-4):231 357, 2015.
Semih Cayci and Atilla Eryilmaz. Convergence of gradient descent for recurrent neural networks: A nonasymptotic analysis. SIAM Journal on Mathematics of Data Science, 7(2):826 854, 2025.
Lenaic Chizat, Edouard Oyallon, and Francis Bach. On lazy training in differentiable programming, 2020.
URL https://arxiv..rrs http///.s//./..... 07956
Simon S. Du, Jason D. Lee, Haochuan Li, Liwei Wang, and Xiyu Zhai. Gradient descent finds global minima of deep neural networks, 2019a. URL https::/arxx..rrrggg httpgg/as//8111.... 03804
Simon S. Du, Xiyu Zhai, Barnabas Poczos, and Aarti Singh. Gradient descent provably optimizes over- parameterized neural networks. In International Conference Learning Representations, 2019b. URL https::/openreevie new/ffrummmtnti wtw.ooorunnindieen.i.. S1eK3199Y
Mert Gurbuzbalaban, Umut Şimgekli, and Lingjiong Zhu. The heavy-tail phenomenon in sgd, 2021. URL
sttps:///axxiv.or/aate 2006. 207600470000
Arthur Jacot, Franck Gabriel, and Clément Hongler. Neural tangent kernel: Convergence and generalization in neural networks, 2020. URL http:://aariiv.orrra///000...0.0
Qianxiao Li, Cheng Tai, and Weinan E. Stochastic modified equations and adaptive stochastic gradient algorithms, 2017. URL
Yuanzhi Li and Yingyu Liang. Learning overparameterized neural networks via stochastic gradient descent on structured data. Advances in noural information processing systems, 31, 2018.
Gabor Lugosi and Fulalia Nualart. Convergence of contimuous-time stochastic gradient descent with appli- cations to linear deep neural networks, 2024. URL http:://arriv..org/bb//2000 07401v1
Gabor Lugosi and Eulalia Nualart. Convergence of continuous-time stochassiic gradient descent with appli- cations to deep neural networks, 2025. URL http:pe//aaaxi.oogg//////2..////
Sadhika Malladi, Kaifeng Lyu, Abhishek Panigrahi, and Sanjeev Arora. On the sdes and scaling rules for adaptive gradient algorithms, 2024. URL
wtw.:orgg..........o.. 10287
X. Mao. Stochastic Differential Equations and Applications. Woodhead Publishing, 2007. ISBN 9780857099402. URL
Samet Oymak and Mahdi Soltanolotabi. Towards moderate overparameeteriiaationn global convergence guarantees for training shallow neural networks, 2019. URL httpss//aaaxv..or////////000 04674
Unut Simsekli, Levent Sagun, and Mert Gurbuzbalaban. A tail-index analysis of stochastic gradient noise in deep neural networks, 2019. URL httpps//aardd..orrr/a////O....1
Chiyuan Zhang, Samy Bengio, Moritz Hardt, Benjamin Rocht, and Oriol Vinyals. Understanding deep learning (still) requires rethinking generalization. Communications of the ACM, 64(3):107-115, 2021.
14

---

Published in Transactions on Machine Learning Research (02/2026)
A Comparison to Prior Work
Our work is closely related to Lugosi & Nualart (0024) and Chizat et al. (2020) In this scction, we explain the main differences and technical challenges in our work compared to these works. The original works in the lazy training regime, such as Chizat et al. 2020 Du et al. (2019b), consider deterministic dynamics, which corresponds to full-batch gradient descent. On the other hand, stochastic gradient descent (SGD) is widely used in practice because of its computational efficiency in large-scale problems. Motivated by this, we analyze continuous approximations of SGD by utilizing Itō stochastic differential equations (SDEs), which makes the analysis substantially different and more challenging compared to the ODEs studied Chizal al. (2020). To be more concrete, the dynamies presented in this previous work be rewritten
du-- - -- V
The dynamics that are treated in this publication
dui - V Raalaaa)) + Via (En (Erl)))
include the additional stochastic term that models subsampling in stochastic learning algorithms such stochastic gradient descent. The analysis in Chizat et al. (2020) is not directly applicable to the stochastic setting due to the diffusion term that accounts for stochasticity, which requires different and more sophis- ticated analysis tools to control the diffusion term. On the other hand, Lugosi & Nualart (0224) considers stochastic dynamics similar lo our setting. The key difference in our work that we analyze the Training dynamiis of deep neural networks the lazy training regime, characterized by an output scaling factor a that controls the parameter movement. This scaling factor and the analysis of the stopping time т (which is specific case of T, in Lugosi & Nualart (2024)) describing the exit time of the lazy training regime, gives us strictly positive lower bound on Amin (T, wo) which loosens the restriction on 0(r, 1)) in Lugosi & Nualart (2024). Additionally, by formulating the problem immediately as supervised learning problem, by considering w + R(h(ww) instead of w + ('w)) combined with the lazy training formulation, we make directly verifiable assumptions on the loss function and neural network model in a supervised learning setting, while the assumptions in Lugosi Nualart (2024) depend on more complicated problem parameters. Overall, this slight change to the setting yields more relaxed, easier-to-verify assumptions that achieve similar results.
A new version of Lugosi & Nualart (2024) has since been published, which includes similar results for non- lincar deep neural networks as well. There are two additional technical differences that we want to point out. First, for most of our results, we do not have to assume the existence of global minimizer in the parameter space wt (see Remark 2 Secondly, the comparison between the stochastic Langevin dynamics and the linearized dynamics in Corollary and as well as separated analysis in the function space (Theorem has not been conducted in the stochastic setting.
Chizat et al. 2020
Lugosi & Nualart (2024) Lugosi & Nualart
2025
Our work
stochastic dynamics lazy training non-linear DNN lincar comparisson
X
X
X
XXX
X
Table 1: Illustration of the key differences between our work and previous works.
B Applicability Proofs for Shallow Neural Networks
In this section, we will provide the proofs neoded show that the example from Scction ffulills Assumption Assumption follows immediately, since we are in the supervised learning setting, where the sample measure discrete.
and 2
15

---

Published in Transactions on Machine Learning Research (02/2026)
Proposition 1 (Smoothness and strong convexity of E). The loss- foss-function
l(x, = (y h(x))
is strongly convez and its gradient (wrt. h) is Lipschitz continuous with respect to
Proof. The gradient the loss function is given by
Val(x, = Vallb(x) v112 =2((2))x)))))))))))))))()))))- y).
We first prove the strong convexity
(Val(z,hi) - ((a h2(2)) -2(h (x) (zxx)) h2(x))
=2|1/1 h (x)
The Lipschitz continuity of the gradient follows immediately from the calculation of the gradient |V 12) - Vul(z, b2) = 2| (x) (2(2))l2
Next, we will prove the Lipschitz continuity of Dwg.
Proposition (Lipschitz continuity Dw The derivative Dwo of
((; = Σ Gollwa))
where a = tanh is Lipschitz continuous with respect to w.
Proof. The derivative of & with respect to wi given by
DR, o(x;w,c) == G Vm
- (wuf x)x²
We know that the derivative of tanh is Lipschitz continuous. It follows that
||Du, w) Dw, p(y a c)ll <
V Icalllo' x) o' d'(07 x)||[=
Vm"
Lip(o')
cc ||: |1 x w. x:ll
Vm
Lip(a')
cc |||
[w
2xll.
Vm
which proves the Tipschitz continuity.
The only requirement, that still needs be proven is Assumption 2 proof in two parts. First we will prove the following helpful lemma. Lemma 2. [Curvature of peremeter-los-mmppihe There exist parameters a, Ta such that (R((a((())) Id.
In order to do this, we will split the
Proof. To prove this property, we need to calculate the Hessian of the mapping + R(ah(w)) and then determine the maximum eigenvalue of the Hessian. To do this, we will first decompose the Hessian into different parts, and then write the eigenvalues the components using block-submatrices. The gradient of the above function is given by R( a n (.(w),)(w),
16

---

Published in Transactions on Machine Learning Research (02/2026)
where
ri(w): : aQ(2; c) y/ J.(w) Vub(za;
Using this notation, we can write the Hessian as
[)) a² J.(W)J/(W) Σ V(l
which implies that
a [[(W)]|+++ + a rr [
Amar (V ((ob(aww)))
In other words, we only need to bound ([(w))l((a(()) and / e)]l order to show the proposition.
We start off by bounding J (w)
|[((W)]
o' [w] m
m
| 2|
m
m Σ (w
| ||c||2
where we used Vz R: [x'(z) S 1.
The bound of r; follows immediately from
(n((w) (a0(z,w,)) Vsl S al
(wol x) + +lul s ||c|| + lps/-
Lastly, we still need to bound ( c)ll- We will do this by considering the blocks of the Hessian 22 M(.sw,))) First, we observe for j o k
v2 uyan baggw c) = Dwn,
( m))r)
5 0.
which means that this Hessian is block-diagonal matrix. In addition, it holds for j k
wy my (0(z;; =
vm -e,0 x.)aaa.
Taking the norm this equation, we receive
V Viyeen (k(a)w)))) $
4 3/3m e,l1z,ll
Amax
where we used Vz E R: [x)]l()) < 3/3
Putting all of the above logether, we get the upper bound
Amaz R(aolz;,, c))) < a² 1 4a Acllo
(allellı + Ilu|loo mm
3nv
which is in O(a²) for given dataset and predictor class. This implics that for na sufficiently small, the above proposition holds.
17

---

Published in Transactions on Machine Learning Research (02/2026)
The following proposition will rely the above Lemma to show that Assumption holds for shallow neural networks trained on finite training set with the mean squared error the loss function. Proposition 3. In the above shallow neural network setting holds that
Tr (allws) T R(,ah((w))) R(ah((w,)
2a X2
This result effectively justifying the application of Theorem to shallow neural networks.
Proof. In the following, we will assume that the global minimizer does have zero residual risk, ie.. R(h*) This implies that R(ah(ws)) R(ah(w.)). By using the basic property of the Trace operator, we get
Tr aa V ( (4s) = =T Tr
Ea (al ((ah)))
R(ah(w,))
R(ah(wss)
Dalws)
Tr
R(ah(ww))
From Lemma follows:
we know how to bound Tr(V2 For the other trace operator, we can proceed as
Ea(ws)
Ew VV ah [ (Maall/
Tr
= Tr
R(ah(wo))
R(ah(w..)
E. (Vlll, V ah),),))h()))))
< Tr
R(ahlwss)
Tr(V (/ (2
E
R(ah(w,))
VVl((s, ah(w.))
Ew
R(ah(ws))
Next we use, that E(x, ah (wa)) = (y - ah(wa))? to get
V Vax ah(w)))l1
[l ah(w.))ll
R(ah(ws)) 4.
En
4E.
[(ah(w.))
R(ahl(w))
R(ah(w)))
In total, we get which means that Proposition
Tr aa (ws) R(Ma(uus)))aaus < <0(a3)
R(ah(ws))
3 holds for appropriate TJa-
C Applicability Proofs for Deep Neural Networks
The goal of this section is to provide evidence that the results from this publication can also be used for deep neural networks under reasonable assumptions. As already explained in Remark | we can circumvent the need Lipschitz continuity of Dh by using a different stopping time that allows an identical proof to Theorem by guaranteeing the positive definiteness of Amin (Dh((n)))))))) D((w.)). Since we heavily rely on the results from Du et al.) (2019a), we will only consider the same class of deep neural networks as those presented in the original publication. Let z(() E Rd be the input, W(1) Rmx be the first weight matrix, W()) E Rmxm the weight at the k-th layer, for layers 2 < k < H. R is the output layer weights and o()) is an activation function. Additionally, define the normalizing scaling factor == (BewN(o.)lo(x))))
Then we recursively define
(k)
Cr (WW(H)))( for s k <
(17)
f(z,w) = at (H)
(18)
18

---

Published in Transactions on Machine Learning Research (02/2026)
and z(() = x is the imput of the neural network. We will present the results from Du et al. (2019a) in the following propositions. Although the original publication uses discrete time-steps, contrary to the continuous dynamics presented in this text, the proofs of Lemmas 1-4 in Appendix B never explicitly use the discrete time steps. Therefore, the proofs hold analogously in the continuous setting. We will only present Lemmas and in the following. The geometric series function (a(n) = Ei=o af will be used extensively.
Proposition 4 (Bound on output difference; Lemma B.3 in Du et al. (2019a)). Suppose for every k E {1,..., H (0)(2) I Cw,oym, (0)ll) I Cr,0 and (t) - (1) I /mR for constants Cw,0s Cr,0 > 0 and S Cw.0- (.) is L-Lipschilz, we have
((00())x((00)))) =0000 S Veobea,ognn (k) (1)0(1))
where c₂ LC
Proof.
See
Du et al. (019a) Appendix B.
The condition that we still need prove Amin (D (D"h(u))D(uu)) 2 A. To do this, we will look at the Gram matrix G(k) E [Rax, defined by
@f(ws:::))
Gs
Owa
Ow.
H
Of (ws) (f(Ws,
@5 (f(, 0f(w.. 2z)
+
!k! (K)
!k!
Ou,
Sa,
kml
G (s).
Each (s) is positive semi-definite matrix. Therefore, can bound the minimal Eigenvalue by only considering (s).
Proposition 5 (Lemma B.4 in Du et al. (019a)). ) Suppose a(-) is L-Lipschitz and B-smooth. Suppose for I m2.vmm,lla(0)) I 2/40m2/4 1 I [(xx() (0) I Cr,u; if c
k E (1...., H), (0) l2 cu.oymm.laa0
(s) ((k)) (0) |r, [a(s) - a(0)ll2 z VmR where R S cge, - and R S (H) -1 for some small constant e and = Lcw, we have
Proof. See Du ct al. (019a.) Appendix B.
Under the assumption that Amin (M((o) 2 X, this gives the needed result. to define new stopping time on the neural network weights. Instead of considering
inf 2 0: |lwe woll A/Lip(Dh)),
we need to consider the stopping time = inf (0..... (b) with
in ( 0: [[(())()) W()) (k) (0)l VmR),
if ifk if k= 0.
(k)
inf 2 0: |aa(t)a(00)l
- a0))]2 > VmR),
By using this stopping time, we can derive a similar result Theorem for deep neural networks. The following lemma will show that Assumption realistic for deep neural networks, by proving its validity for supervised learning with the mean squared loss function.
19

---

Published in Transactions on Machine Learning Research (02/2026)
Lemma 3. We consider the deep neural network defined in Equations 17 and activalion function. Addilionally, we consider the loss function
18 with taan its
((x, ( -y--(xx))
Then Assumption 2 holds, i.e,,
Tr Ta V
R(oh()))aaa) 2a² Xm
R(ah(w))
77
for all (wo))
Proof. In Appendix we have already shown that this loss function and the expected loss R(h) fulfill Assumption 1 (Proposition D. Assumption 3 follows again, since consider discrcte probability measure the finite set of training data. Similar to Proposition 3 to gel
Tr Ta
Tr Ea(w) Tr (Va((ah(w))) [((h(ww))
M(h(ww)
We can the same argument to bound the first trace operator we did in Proposition Additionally, we will bound the trace of the hessian as follows:
v H((h(ww) = V {xx, ah(w))
v lax ah(w)),
where we used that we only consider the expectation over a finite set of data points that
holds
(y
D(2(yv D124--ai((lo)))-aMMMM))))) 24Ck -
Taking the trace operator, we get
2 (V(()(()()))) < Σ (a -a(aa
Next, we can easily show that this trace in Oaa2), since
(a((w(x)) vl s (x)( + |o| S max [+(w)()]|+ max lg]. y=
3=
and w (z) for given æ, is a C² mapping by construction, it takes a maximum value on (,w)). By the same argument, Vwh(w)(z) and (z)()(()) are bounded on (ww for given æ. This implies that
2 Σ - - -
R((h)R(a <
which proofs Lemma 3 for sufficiently large Tl.-
Remark 5. In the proof of Lemma S we used crude upper bounds to show the validity of Assumption For a given urchilecture, loss function and dala distribution one can easily compule the exact hessian and find erplicit upper bounds that result in direct bounds for minimal choices of a² Th
20


Reviews:
[{"review_id": "k7docf27oI", "reviewer": "Reviewer_nqeY", "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": "", "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 analyses stochastic gradient Langevin dynamics (SGLD) in the lazy training (also referred to as \"kernel\") regime, where the model parameters remain close to their random initialisation. Authors use an Itô SDE model of SGLD and, under some ~reasonable assumptions (regularity of loss and Hessian), they show that it can provide some convergence guarantees. In Th. 1 (\"Convergence of SGLD) they show that the expected optimality gap decays exponentially fast (Eq. 3), while in Th. 2 they show that the probability that the training process leaves the lazy regime (i.e. it exits the sphere around the random initialisation point) can be made arbitrarily small by carefully selecting the scaling factor $\\alpha$ (Eq. 9). Authors also provide specialisations of their results (Corollaries 1, \"Risk convergence for shallow networks\", and 2, \"Risk convergence for deep networks\"), where they show that the results in Th. 1 can be applied to deep networks under specific regularity conditions. Then authors experimentally validate their findings by simulating SGLD via Euler-Maruyama (a method for finding approximating solutions to SDEs).", "claims_and_evidence_yn": "Yes", "claims_explanation": "Both theoretical findings and experimental results seem convincing.", "audience_yn": "Yes", "audience_explanation": "This is out of my domain of expertise but I am fairly sure many members of the community may find those findings interesting.", "requested_changes": "- Can you please compare your assumptions to e.g. the ones Chizat et al.? Having a table showing the assumptions and corresponding findings in related works would help immensely for those not familiar with the related literature.", "broader_impact_concerns": "I do not think this is necessary", "additional_comments": "I am not familiar with this domain -- I can't find anywhere in the TMLR review form to declare my level of expertise, hence I'm mentioning that here. I found this work interesting and novel (for me), but I am not deeply aware of the related literature on this topic. Please take this review with a grain of salt.", "extra_scores": {}, "extra_text": {}}, {"review_id": "qJBjKoTxBQ", "reviewer": "Reviewer_Gbbb", "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": "", "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 considers the question of performing a non-asymptotic convergence analysis of stochastic gradient Langevin dynamics in the lazy training regime. Three main results are given: 1) an inequality that establishes an exponential decay of the expected gap between the loss of the solution obtained by SGLD and the optimal loss, 2) a concentration inequality which shows that, with high probability, the model will not deviate significantly from the initial state, thus staying in the lazy training regime (which also translates to an upper bound on the expected optimality gap), and 3) applications to obtain upper bounds on the training loss of neural networks with finite depth and width.\n\nI believe this is generally a decent paper. The results are conceptually interesting, with the assumptions involved in them not appearing to be overly restrictive in the context of the overall literature. I guess one weak point could be the fact that there appears to be strong reliance to both Chizat et. al. (2020), and Lugosi and Nualart (2024). In that sense, the paper may not be the strongest in terms of technical novelty. For that reason, I'm interested in reading the comments by the authors where they discuss this aspect of the work.\n\nLenaic Chizat, Edouard Oyallon, and Francis Bach. On lazy training in differentiable programming, 2020. URL https://arxiv.org/abs/1812.07956.\n\nGabor Lugosi and Eulalia Nualart. Convergence of continuous-time stochastic gradient descent with applications to linear deep neural networks, 2024. URL https://arxiv.org/abs/2409.07401.", "claims_and_evidence_yn": "Yes", "claims_explanation": "I didn't notice any major issues in the proofs. That said, I do not actively work in this area.", "audience_yn": "Yes", "audience_explanation": "It's a paper on stochastic dynamics in neural networks, which is a core topic for TMLR's audience.", "requested_changes": "I would like to see a more detailed discussion of the technical challenges the authors encountered compared to prior work.", "broader_impact_concerns": "No concerns", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "ybOj7rZdyN", "reviewer": "Reviewer_1uFt", "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": "", "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 provides a rigorous, non-asymptotic convergence analysis of stochastic gradient Langevin dynamics (SGLD) viewed as a continuous-time Itô stochastic differential equation approximation of stochastic gradient descent. Focusing on the lazy training regime, the authors prove that under suitable regularity assumptions on the Hessian of the loss function, SGLD with multiplicative, state-dependent noise:\n\n- Maintains a non-degenerate neural tangent kernel throughout training with high probability, ensuring stable dynamics;\n\n- Achieves exponential convergence in expectation to the empirical risk minimizer, with explicit finite-time and finite-width bounds on the optimality gap.\n\n- The theoretical results are supported by numerical experiments in a regression setting, illustrating the practical relevance of the analysis.", "claims_and_evidence_yn": "Yes", "claims_explanation": "I thank the authors for their work. The paper is overall well written. The results are new and interesting for the community. The theoretical results are supported by numerical simulations.\n\nIn my opinion, the work merits publication. However, I have identified some possible errors in the proofs, so I recommend that the authors address these issues before the paper is accepted. Please refer to my comment in the \"Requested Changes\" section.", "audience_yn": "Yes", "audience_explanation": "The TMLR community would be interested in this work because it provides rigorous, non-asymptotic convergence guarantees for SGLD in a realistic lazy training regime, incorporating multiplicative, state-dependent noise. By connecting the dynamics to the neural tangent kernel and establishing finite-time, finite-width bounds, the paper advances theoretical understanding of optimization in deep learning. \n\nContrary to previous works, that focused on establishing strictly positive lower bounds on the eigenvalues of the Neural Tangent Kernel for full-batch gradient flow, the paper extends those results to SGD.", "requested_changes": "Here are my major concerns:\n \n- I might have found a wrong statement in the proof of Thoerem 2. The authors claim on the top of page 9 that $\\int_0^t \\mathrm{Tr}(\\Sigma_{\\alpha}(\\omega_s))ds \\leq 0$, but this term appear with a square root on the line above. This is problematic in case the integral is negative. Could the authors correct/clarify this part of the proof ?\n\n- Could the authors clarify why it is licit to write $\\mathbb{P}(\\tau <\\infty) \\leq \\mathbb{P}(\\|\\omega_{\\tau}-\\omega_0\\|>r)$ in the proof of Corollary 4 ? More precisely, why the object $\\omega_{\\tau}$ is well-defined when $\\tau=\\infty$ ? \n\n\nHere are some minor comments:\n\n- It would be helpful to ensure consistent notation. For example, the expectation operator is sometimes written with a subscript $x$ and sometimes not (see Assumption 3 and the proof of Lemma 1). Similarly, it would help to specify the variable with respect to which the gradient is taken in the nabla notation. This is done in Assumption 2, but not in Assumptions 1 and 3.\n\n- In Section 2.1, should $R$ satisfy some properties to ensure that $\\omega^*$ is well defined (i.e., that a global minimizer exists)?\n\n- page 5: \"stray\" should be stay.\n\n- I don't understand the use of Dh: it seems the authors use this notation as if it referred to the gradient associated (while based page 2, it refers to the differential). Why not maintain consistency and use the ∇ (nabla) notation instead?\n\n- There are some inconsistencies in the notation: sometimes D^{\\top}h is used (e.G. in section 2.2), sometimes Dh^{\\top} is used (e.g. in assumption 2)\n\n- In Corollary 1, the notation \\sigma is not defined\n\n- In Eq.9, there is typo regarding the style of \"F\" in the subscript of the norm. Same in the statement and the proof of Theorem 2.\n\n- I think there are some inconsistencies in the paper regarding the use of either $\\| \\cdot \\|_2$ or $\\| \\cdot \\|_{\\mathbb {F}}$. For example, in $h(\\omega)$ should be in $\\mathbb {F}$, but the norm $\\| \\cdot \\|_2$  is often used for those objects such as in the statement and proof of Corollary 3.\n\n- There is a typo at the end of the proof of Theorem 3 with \"proofn\".\n\n- In Section 4, $c^{\\star}$ should be in $\\mathbb{R}^{m'}$ and not $\\mathbb R$. Also, is a transpose missing in $\\omega_j^* x$ ?", "broader_impact_concerns": "I don't see any concerns on the ethical implications of the work.", "additional_comments": "", "extra_scores": {}, "extra_text": {}}]

Author rebuttal:

Decision: Accept (poster)