Second-Order Path Kernel Interpolation Formulas in Machine Learning Jin Guo1 , Roy Y. He1 , and Jean-Michel Morel2,
arXiv:2606.07495v1 [cs.LG] 5 Jun 2026
1
B
Department of Mathematics, City University of Hong Kong, Hong Kong 2,B Division of Industrial Data Science, Lingnan University, Hong Kong
Abstract Understanding how training data shape neural network predictions is a central problem in modern learning theory. In 2020, Pedro Domingos proposed an interpolation formula valid for every model learned by deterministic gradient descent. It expresses the model’s prediction as an integral, along the optimization path, of a data-dependent kernel that aligns the model’s gradients at the test and training data. Such a first-order characterization remains valid for models trained with batchbased stochastic optimization. In this paper, we develop second-order forms of these interpolation formulas. We show that the leading path-kernel interpolation is supplemented by a curvatureweighted interpolation term. For stochastic gradient descent, an additional sampling-induced component appears, coupling the curvature of the prediction with the covariance of mini-batch gradient noise. We also extend the representation to stochastic gradient descent with momentum, where the interpolation structure is preserved but with the weights modified by a memory-related factor. Moreover, we establish a concentration estimate for the terminal prediction, identifying the fluctuation scale around the expected second-order representation. Together, these results provide a refinement of the path-kernel interpretation of neural network prediction.
1
Introduction
Deep learning has achieved remarkable empirical success across a wide range of tasks [16, 33]. This success has motivated growing interest in characterizing models’ predictions under gradient-descent training and explaining their behavior on unseen data. Kernel perspectives provide an important approach to this problem. The neural tangent kernel (NTK) [12, 27] interprets the learning process as a linearized model around initialization in the infinite-width regime [17]. However, fixed-kernel or purely linearized descriptions are incomplete for finite-width networks obtained through optimization, where the effective kernel may evolve along the training trajectory and depend on the data. An alternative kernel perspective was developed by Domingos [6], who showed that models trained by gradient descent, including neural networks, induce path kernels that determine the model’s prediction. Since path kernels accumulate gradient information from all training points throughout the training process and match it against the gradient at the query in a Hilbert space, this theory yields a clear picture of the influence of training data. More specifically, let f (·, Θ) : Rp → Rm be a differentiable PN p m model trained on {(xn , yn∗ )}N with differentiable loss L(Θ) = N1 n=1 ℓ(f (xn , Θ), yn∗ ) n=1 ⊂ R × R and learning rate η > 0 by gradient descent; then for any input x ∈ Rp , the model output at time T satisfies N Z 1 X T ∂ℓ lim f (x, ΘT ) = − Kt (x, xn ) f (xn , Θt ), yn∗ dt + f (x, Θ0 ), (1) η→0 N n=1 0 ∂f where
Kt (x, x′ ) := ∇Θ f (x, Θt ), ∇Θ f (x′ , Θt ) ,
x, x′ ∈ Rp ,
is the gradient kernel. In [10], Guo et al. extended (1) to a non-asymptotic, stochastic setting to characterize the predictions of models learned through mini-batch training. In particular, they show that, under a first-order weak approximation [19] of stochastic batch gradient descent, the expected network output admits a representation analogous to (1) with optimizer-specific weights. 1
The above extension is closely related to the approximation of stochastic iterative algorithms. Classical theory studies their continuous-time limits [15, 21], and diffusion approximations have been used to analyze SGD through weak convergence and the evolution of its probability distribution [11]. More recently, stochastic modified equations have provided a systematic weak-approximation framework for capturing finite-step-size and mini-batch effects beyond the leading-order continuous-time limit [19]. This viewpoint has also been extended to adaptive methods such as RMSprop and Adam [24], and to distributed and local variants of SGD [9]. These works show that higher-order terms are not merely discretization remainders but can encode meaningful properties of the training dynamics. Meanwhile, Domingos’ formula (1), as well as the representations established in [10], captures only the leading tangent geometry of the training path, leaving the influence of higher-order geometry unaddressed. For instance, the curvature of the loss function in parameter space has long been recognized as a central element in the optimization of neural networks. Empirical Hessian analyses show that neural loss landscapes have highly structured spectra, often with a large near-zero bulk and a small number of outlier directions associated with data-dependent structure [32]. Loss-landscape visualization further reveals that the local geometry around trained solutions carries important information about the stability of neural network minima and their generalization behavior [18]. Related work based on local entropy similarly highlights the role of wide valleys and flat local geometry in the solutions found through neural network training [2], while large-scale Hessian spectrum studies show that curvature evolves throughout training and shapes optimization behavior [8]. These findings also provide a foundation for deep learning privacy [30]. Hence, investigating the higher-order influences on a model’s predictions can yield valuable insights, yet remains under-explored. Motivated by this gap, we develop the second-order counterpart of the first-order path-kernel representation. First-order path kernels capture the leading tangent-feature alignment, but they do not describe how second-order geometry modifies the contribution of training samples to test predictions. We show that curvature plays a direct, prediction-level role: it reweights the pathwise contribution of each training sample to test outputs. Moreover, we find that this curvature effect is a common mechanism across gradient-based methods, including gradient descent (GD), stochastic gradient descent (SGD) [31], and SGD with momentum (SGDM) [34], all of which contain a curvature-weighted correction to the first-order path-kernel term. Our theory thus lifts this higher-order viewpoint to the prediction level, yielding explicit second-order corrections to the pathwise formula for test outputs. More interestingly, once we move beyond first order in the stochastic setting, we find an additional mechanism that is absent in deterministic training. In SGD and SGDM, mini-batch sampling produces terms that couple sampling fluctuations with the curvature of the network in parameter space. This connects our result to a line of work showing that SGD noise, gradient covariance, and curvature jointly influence the minima reached by stochastic training [13, 20, 39]. Unlike these parameter-space analyses, which typically study stability, escape behavior, or Hessian-based sharpness [5, 37–39], our result offers a prediction-level perspective on the influence of the flatness or sharpness of local minima. Curvature matters not only because it affects the motion of the parameters near a minimum, but also because it directly reshapes the pathwise contribution of training samples to predictions. The main contributions of this paper are summarized as follows. • We identify a curvature-weighted interpolation term shared by GD, SGD, and SGDM, which reweights the training residuals through the Hessian of the loss. This extends the first-order path-kernel formula from a leading tangent-feature description to a refined output expansion. • For stochastic methods, we identify sampling-induced corrections that are absent in deterministic GD. These terms couple mini-batch fluctuations with the curvature of the network in parameter space, thereby making the dependence of the model’s predictions on batch size and curvature explicit. • For SGDM, we further show that momentum preserves the interpolation structure while reshaping the temporal weights through a memory kernel. • We support the sampling-induced representation with a localized concentration estimate for the terminal prediction, together with numerical experiments that verify the predicted remainder scaling and sensitivity effects. This paper is organized as follows. Section 2 introduces the basic notation and preliminaries, including path kernels and the SDE approximation of SGD. Section 3 presents the main second-order output 2
representations for GD, SGD, and SGDM. Section 4 discusses several implications of the theory, including curvature-weighted interpolation, the concentration of model predictions, and the role of sampling-induced variation and batch size. Section 5 provides numerical experiments that verify predicted scaling orders and illustrate the effect of batch size and curvature-based batch selection on prediction sensitivity. Section 6 concludes the paper. The appendix contains the extension to learning-rate schedules and the auxiliary lemmas used in the proofs.
2
Preliminaries
In this section, we define the key mathematical objects and provide the necessary background.
2.1
Parameterized learning models
Consider a learning model f : Rp ×Rd → Rm parameterized by Θ ∈ Rd . This includes classical methods such as support vector machines (SVMs) as well as modern neural networks of any architecture. Let p m the training data {(xn , yn∗ )}N n=1 ⊆ R ×R be drawn independently from an unknown distribution, and denote by P(x) the marginal distribution of x. For an input x, the model output is y(x) = f (x, Θ) ∈ Rm , and the loss function is ℓ(·, ·) : Rm × Rm → R. Training the model corresponds to finding an optimal parameter Θ that minimizes the expected risk Z ℓ f (x, Θ), y ∗ (x) dP(x), where y ∗ (x) denotes the true response at x. Since the data distribution is unknown, the expected risk cannot be minimized directly. Instead, one minimizes the empirical risk N
L(Θ) =
1 X ℓ f (xn , Θ), yn∗ . N n=1
Standard gradient descent on the parameters Θ yields the iterative update Θk = Θk−1 − η ∇L(Θk−1 ), where k is the iteration index, η > 0 is the learning rate, and Θ0 is the vector of initial parameters. For large datasets, computing the full gradient at each iteration can be expensive. A standard alternative is stochastic gradient descent (SGD), which approximates the full-batch loss using randomly sampled mini-batches. Let |B| ≤ N denote the batch size, Γ the set of all subsets of {1, . . . , N } of size |B|, and B a uniformly random element of Γ. The mini-batch loss is LB (Θ) =
1 X ℓ f (xi , Θ), yi∗ . B i∈B
Since Γ is finite and B is uniformly distributed, the expectation EB [LB (Θ)] is well defined and coincides with the empirical risk, EB [LB (Θ)] = L(Θ). Interchanging the gradient with the finite sum then gives EB ∇Θ LB (Θ) = ∇Θ EB LB (Θ) = ∇Θ L(Θ), so the mini-batch gradient is an unbiased estimator of the full-batch gradient. Let {Bk : k = 0, 1, 2, . . . } be a sequence of independent and identically distributed (i.i.d.) uniform random batches. Starting from an initial point Θ0 ∈ Rd , the SGD update at step k is Θk+1 = Θk − η ∇LBk (Θk ), where η > 0 is the learning rate.
3
(2)
2.2
Gradient kernel and path kernel
Domingos [6] established a first-order kernel representation (1) for gradient descent. The key components of this interpolation theory are the gradient kernel and the optimization-induced path kernel. Definition 2.1 (Gradient kernel and path kernel). Let Θt , t ∈ [0, T ], denote the parameter trajectory of the gradient flow (5). The gradient kernel between two inputs (x, x′ ) ∈ Rp × Rp at time t is Kt (x, x′ ) := ∇Θ f (x, Θt ), ∇Θ f (x′ , Θt ) . The path kernel is its time integral over the training trajectory: Z T ′ Kpath (x, x ) := Kt (x, x′ ) dt.
(3)
(4)
0
Intuitively, the gradient kernel (3) measures the instantaneous similarity between the model’s gradients at two inputs, and the path kernel (4) accumulates this similarity over the entire training trajectory.
2.3
SDE approximation of SGD
Building on the first-order stochastic extension of Domingos’ theorem, we now focus on identifying and characterizing the second-order structure of the output formula. Rather than re-establishing the stochastic representation itself, we investigate how finite-learning-rate effects, loss curvature, and minibatch sampling noise contribute additional second-order corrections. To this end, we use the stochastic differential equation (SDE) approximation of SGD as a continuous-time framework for deriving and interpreting these higher-order terms. Assumption 2.1. For every integer n = 1, . . . , N , the per-sample loss ℓ(f (xn , Θ), yn∗ ) is continuously differentiable in Θ, and its gradient satisfies a uniform local bound: for each R > 0 there exists a constant MR > 0 such that max ∥∇Θ ℓ(f (xn , Θ), yn∗ )∥ ≤ MR
∥Θ∥≤R
for all n = 1, . . . , N.
In addition, for each fixed test point x ∈ Rp , the parameter gradient ∇Θ f (x, Θ) is locally bounded in Θ. This assumption holds for any smooth network and loss on a bounded domain. It does, however, exclude networks with nonsmooth activations such as ReLU and nonsmooth losses such as the ℓ1 loss. With the SGD iterates (2), the random mini-batch sampling is inherently discrete. In contrast, Domingos’ theorem (Theorem 3.2) relies on a continuous gradient flow (5) and does not directly extend to the stochastic, finite-learning-rate setting. In [10], a first-order continuous-time SDE approximation of stochastic gradient algorithms [11,19,25] was used to extend the representation to data-driven models trained with stochastic gradient descent (SGD). We briefly review the relevant concepts in their general form. Let T > 0, η ∈ (0, 1 ∧ T ) (where a ∧ b := min{a, b}), and let α ≥ 1 be an integer. Set K = ⌊T /η⌋ (where ⌊·⌋ denotes the floor function). Let G be the set of continuous functions g : Rd → R satisfying |g(z)| ≤ κ1 1 + |z|2κ2 , z ∈ Rd , for some κ1 , κ2 > 0. For α ≥ 1, define G α := g ∈ C α (Rd ) : ∂ β g ∈ G for all |β| ≤ α , where β is a multi-index. Thus G α is a subset of C α , the space of α-times continuously differentiable functions. Definition 2.2 (α-th order weak approximation [19]). We say that a continuous-time stochastic process {Zt : t ∈ [0, T ]} in Rd is an α-th order weak approximation of a discrete stochastic process {Θk : k = 0, . . . , K} in Rd if for every g ∈ G α+1 , there exists a positive constant C, independent of η, such that max E[g(Θk )] − E[g(Zkη )] ≤ Cη α . k=0,...,K
4
In the context of Definition 2.2, functions from G α+1 serve as test functions. This space collects sufficiently smooth functions with at most polynomial growth, which is a condition mild enough to include the observables of interest, yet strong enough to ensure integrability and to justify the Itô– Taylor expansion rigorously. This framework is widely employed to analyze the discretization errors of stochastic differential equations (SDEs) when modeling SGD via continuous-time approximations [7, 14,19]. Crucially, the resulting error bounds hold uniformly over all SGD iterations up to ⌊T /η⌋ steps, not merely at the final iteration.
3
Main Results
In this section, we present our main results, which characterize the role of second-order geometry in the predictions of parameterized learning models. We summarize our findings as follows. • We show that, in addition to the path-kernel-induced gradient alignment in (1), an analogous alignment takes place in a different Hilbert space, one modified by the local curvature of the loss landscape through the Hessian (Theorem 3.2). This constitutes a second-order extension of [6], valid for deterministic gradient descent. • Passing to the more realistic stochastic descent framework, we find that, under mild conditions, this second-order extension carries over to the mean predictions of models trained via SGD, but with a new sampling-induced perturbation (Theorem 3.3). • We show that the momentum in SGDM induces additional tangent alignments in Hilbert spaces generated by time-transported kernels (15)–(16), which vanish when the network gradient and the loss Hessian remain bounded (Theorem 3.2). We focus primarily on the role and implications of the second-order terms for the model’s predictions, and refer the reader to [6, 10] for discussion of the first-order terms.
3.1
Second-order Domingos Theorem
To capture the effects of finite step size more accurately, we replace the plain gradient flow with the second-order modified equation [1, 4] of gradient descent, as is standard in backward error analysis: η Θ̇t = −∇L(Θt ) − ∇2 L(Θt )∇L(Θt ) + O(η 2 ). 2
(5)
Definition 3.1 (Curvature-induced kernel). The curvature-induced kernel between two data points x, x′ is the Hessian-weighted inner product of the model gradients at the two points along the parameter trajectory induced by gradient descent; that is, Ktcur (x, x′ ) := ⟨∇Θ f (x, Θt ), ∇Θ f (x′ , Θt )⟩∇2 L(Θt ) Θ
(6)
= ∇Θ f (x, Θt )⊤ ∇2Θ L(Θt )∇Θ f (x′ , Θt ). The kernel Ktcur (x, x′ ) measures how the alignment between the two tangent features is reweighted by the local curvature of the loss landscape during training. In contrast to the standard gradient inner product, this quantity assigns greater weight to directions associated with larger curvature of L at Θt . The following theorem shows that the output of any model trained by gradient descent can be expressed as a weighted sum of path-kernel and curvature-kernel contributions evaluated between the test input and each training point. Theorem 3.2. Consider a differentiable model f (·, Θ) trained on a set {(xn , yn∗ )}N n=1 via gradient PN descent on the empirical loss L(Θ) = N1 n=1 ℓ(f (xn , Θ), yn∗ ) with learning rate η > 0. The model
5
output satisfies N
f (x, ΘK ) = f (x, Θ0 ) −
1 X N n=1
Z T 0
∂ℓ Kt (x, xn ) (f (xn , Θt ), yn∗ ) dt ∂f | {z } path-kernel term
−
N Z T X
η 2N n=1
0
∂ℓ Ktcur (x, xn ) (f (xn , Θt ), yn∗ ) dt + O(η 2 ). ∂f | {z }
(7)
curvature-kernel term
Proof. Let Θt ∈ R denote the parameter trajectory of the model f (·, Θ). Following the standard modified equation (5) for gradient descent, we approximate the discrete update Θk+1 = Θk −η∇L(Θk ) by continuous-time dynamics η (8) Θ̇t = −∇Θ L(Θt ) − ∇2Θ L(Θt )∇Θ L(Θt ) + O(η 2 ). 2 d
By the chain rule and (8), we obtain df (x, Θt ) = ∇Θ f (x, Θt )⊤ Θ̇t dt
η = ∇Θ f (x, Θt )⊤ −∇Θ L(Θt ) − ∇2Θ L(Θt )∇Θ L(Θt ) + O(η 2 ) 2 N 1 X ∂ℓ =− ⟨∇Θ f (x, Θt ), ∇Θ f (xn , Θt )⟩ (f (xn , Θt ), yn∗ ) N n=1 ∂f N
−
η X ∂ℓ ⟨∇Θ f (x, Θt ), ∇Θ f (xn , Θt )⟩∇2Θ L(Θt ) (f (xn , Θt ), yn∗ ) + O(η 2 ), 2N n=1 ∂f
where ⟨∇Θ f (x, Θt ), ∇Θ f (xn , Θt )⟩∇2Θ L(Θt ) = ∇Θ f (x, Θt )⊤ ∇2Θ L(Θt )∇Θ f (xn , Θt ). Integrating both sides from 0 to T yields the result. The curvature term also admits an interpolation interpretation. Here, Ktcur (x, xn ) acts as a ∂ℓ ∗ curvature-weighted coefficient applied to ∂f f (xn , Θt ), yn , which is interpretable as a training residual . Since this term carries a factor of η, it represents a second-order, curvature-dependent interpolation correction rather than the contribution of the leading path-kernel. A natural question is whether Domingos’ theorem continues to hold under a learning-rate schedule. The answer is yes, and we prove in Appendix A that the schedule reweights both the path-kernel term and its curvature-dependent correction along the optimization trajectory.
3.2
Sampling-induced kernel from random batches
To derive a refined representation beyond the first-order path-kernel expansion, we employ a secondorder weak approximation of SGD. For a given learning rate η > 0, the discrete-time iterates are weakly approximated by a continuous process that retains the first-order drift-diffusion structure while incorporating second-order corrections generated by loss curvature and mini-batch fluctuations. The following lemma provides this approximation and yields an O(η 2 ) bound on the weak error between the discrete iterates and the continuous process. Lemma 3.1 (Second-order weak approximation of SGD [19]). Let T > 0, η ∈ (0, 1 ∧ T ), and set K = ⌊T /η⌋. Let {Θk : k ≥ 0} be the SGD iterates defined in (2). Assume that L is twice continuously differentiable, ∇|∇L|2 and ∇LB are Lipschitz, L ∈ G 4 , and the test function g ∈ G 3 . Define {Zt : t ∈ [0, T ]} as the stochastic process satisfying the SDE √ 1/2 dZt = −∇ L(Zt ) + 14 η |∇L(Zt )|2 dt + η Σ|B| (Zt ) dWt , Z0 = Θ0 , (9) where Σ|B| (Z) = EB (∇LB (Z) − ∇L(Z))(∇LB (Z) − ∇L(Z))⊤ . Then {Zt : t ∈ [0, T ]} is a secondorder weak approximation of SGD; that is, there exists a constant C > 0, independent of η, such that max E[g(Θk )] − E[g(Zkη )] ≤ Cη 2 . k=0,...,K
6
Lemma 3.1 ensures that, under the regularity assumptions stated above, the continuous process (9) captures the expected behavior of SGD uniformly over all K iterates, to second order in the learning rate. This refined continuous-time representation serves as the foundation for deriving our main results. For a mini-batch gradient estimator LB (Z) with batch size |B|, the gradient-noise covariance scales as O(|B|−1 ) at a fixed parameter value Θ [26]. Under i.i.d. sampling with replacement, one has the exact identity 1 Σsingle (Zs ), Σ|B| (Zs ) = |B| ⊤ PN where Σsingle (Zs ) := N1 n=1 ∇ℓ(xn , Zs ) − ∇L(Zs ) ∇ℓ(xn , Zs ) − ∇L(Zs ) is the single-sample gradient covariance [36]. Under uniform sampling without replacement from a dataset of size N , the exact formula becomes N − |B| Σsingle (Zs ), Σ|B| (Zs ) = |B|(N − 1) which reduces to the same |B|−1 scaling, up to the finite-population correction factor. In particular, when |B| ≪ N , the covariance is well approximated by a term proportional to 1/|B|. The following theorem identifies how this batch-size-dependent covariance contributes to the expected prediction through the second-order sampling-induced term. Theorem 3.3 (Stochastic Domingos’ Theorem (SGD)). Consider a learning model y = f (x, Θ), with f (x, ·) ∈ G 3 . The parameter Θ of the model is learned from the dataset {(xn , yn∗ )}N n=1 by SGD with learning rate η and Assumption 2.1 holds. In addition, assume that E[LB (Θ)] ∈ G 4 . Then we have "Z # N T ∂ℓ 1 X ∗ E f (xn , Zs ), yn ds Ks (x, xn ) E[f (x, ΘK )] = f (x, Θ0 ) − N n=1 ∂f 0 | {z } path-kernel term
"Z
N X
η − E 2N n=1 "Z
+
η E 2
T
0
T
∂ℓ Kscur (x, xn )
# f (xn , Zs ), yn∗
∂f {z
0
|
Curvature–induced term
ds
}
# 2 Tr ∇Θ f (x, Zs ) Σ|B| (Zs ) ds + O(η 2 ). | {z }
(10)
Sampling-induced term
Proof. For an input x, the expectation of the output after k iterations can be written as E[f (x, ΘK )] = E[f (x, ZT )] + E[f (x, ΘK )] − E[f (x, ZT )] (11)
= E[f (x, ZT )] + O(η 2 ).
The last equation comes from Lemma 3.1. Since ZT is a continuous Itô process, the Itô formula gives d Z T X ∂f f (x, ZT ) =f (x, Z0 ) + (x, Zs ) dZsj j ∂Θ 0 j=1 +
d
Z T
j,l=1
0
1 X 2
∂2f (x, Zs ) d⟨Z j , Z l ⟩s , ∂Θj ∂Θl
1/2 where Zsj denotes the j-th coordinate of Zs . Denote by ajl (Z) j,l=1 the (j, l)-entry of Σ|B| (Z). From Eq.(9), we have Z T ⊤ f (x, ZT ) = f (x, Θ0 ) − (∇Θ f (x, Zs )) ∇Θ L(Zs )ds 0
η − 2 +
Z T
⊤
(∇Θ f (x, Zs )) ∇2Θ L(Zs )∇L(Zs )ds +
Z T η
0
0
d Z η X T
2
√
j,l=1
0
2
∂ f (x, Zs ) ∂Θj ∂Θl
d X
ajh (Zs )alh (Zs )ds
h=1
7
⊤
1
2 (∇Θ f (x, Zs )) Σ|B| (Zs )dWs
Z T
⊤
(∇Θ f (x, Zs )) ∇Θ L(Zs )ds
=f (x, Θ0 ) − 0
η − 2 +
√
Z T
⊤
(∇Θ f (x, Zs )) ∇2Θ L(Zs )∇L(Zs )ds
0
Z T η 0
1
⊤
2 (∇Θ f (x, Zs )) Σ|B| (Zs )dWs +
η 2
Z T
Tr ∇2Θ f (x, Zs )Σ|B| (Zs ) ds.
0
Expressing ∇Θ L(Zs ) in terms of the empirical average over samples {(xn , yn∗ )}N n=1 , we can rewrite the above result as N
f (x, ZT ) = f (x, Θ0 ) − Z T
1 X N n=1
Z T ∇Θ f (x, Zs ), ∇Θ f (xn , Zs ) 0
∂ℓ (f (xn , Zs ), yn∗ )ds ∂f
N
1 X ∂ℓ (∇Θ f (xn , Zs )) (f (xn , Zs ), yn∗ ))ds N n=1 ∂f 0 Z T Z 1 η T √ ⊤ 2 (Zs )dWs + + η (∇Θ f (x, Zs )) Σ|B| Tr ∇2Θ f (x, Zs )Σ|B| (Zs ) ds 2 0 0 Z N T 1 X ∂ℓ = f (x, Θ0 ) − Ks (x, xn ) (f (xn , Zs ), yn∗ )ds N n=1 0 ∂f Z Z T 1 η T √ ⊤ 2 + η Tr ∇2Θ f (x, Zs )Σ|B| (Zs ) ds (Zs )dWs + (∇Θ f (x, Zs )) Σ|B| 2 0 0 Z N T η X ∂ℓ − Kscur (x, xn ) (f (xn , Zs ), yn∗ )ds. 2N n=1 0 ∂f
−
η 2
⊤
(∇Θ f (x, Zs )) ∇2Θ L(Zs )
Plugging the above equation into (11) and noting that "Z T
E 0
⊤
1 2
#
(∇Θ f (x, Zs )) Σ|B| (Zs )dWs = 0,
we obtain the desired result. Equation (10) decomposes the expected prediction into three parts. • Path-kernel term. Along the training path (Zs )s≥0 , the initial prediction f (x, Θ0 ) is modified by ∂ℓ an accumulated contribution from all training samples: each sample’s residual ∂f (f (xn , Zs ), yn∗ ) is weighted by the path kernel. • Curvature-induced term. Geometrically, the curvature-weighted path kernel extends the standard path kernel by incorporating the local geometry of the loss landscape. This term retains an interpolation structure, but now along random training trajectories. Moreover, since it carries a factor of η, it represents a second-order correction to the leading path-kernel term rather than a dominant contribution. A large magnitude of Kscur (x, xn ) indicates strong alignment between the tangent features of the test and training points through the local Hessian geometry of the loss; the corresponding training residual then exerts a larger second-order effect on the prediction. • Sampling-induced term. This noise-induced term is specific to SGD and depends on the minibatch sampling noise. From a geometric perspective, Σ|B| (Zs ) is the covariance matrix of the mini-batch gradient noise at the parameter Zs , and ∇2Θ f (x, Zs ) quantifies the local curvature of the predictor at x. The term thus captures the interaction between mini-batch randomness and the curvature of the model output. Since smaller batch sizes yield a larger covariance, the effect of sampling-induced noise becomes stronger as the batch size decreases. This contribution is absent in the deterministic case. Theorem 3.3 extends Domingos’ theorem to the setting of a second-order weak approximation of SGD. Beyond the first-order characterization in terms of the weighted path kernels, the result introduces 8
additional curvature-dependent terms that capture the influence of the loss Hessian and the secondorder derivatives of the model. In particular, the sampling-induced term reflects the interaction between the noise covariance and the Hessian of the model output, highlighting the role of mini-batch sampling in shaping the optimization trajectory. Consequently, the trained model can be viewed as a correction of the initialization f (x, Θ0 ) by the training residuals, weighted by path kernels and curvature-induced kernels, where these weights encode the subtle effects of loss-landscape geometry and gradient noise in the SGD dynamics.
3.3
Second-order transported path kernel representation for SGDM
In the presence of momentum, the learning dynamics acquire a temporal memory effect, in which past gradients continue to influence the update with exponentially decaying weight. This mechanism reshapes how training examples contribute during optimization, leading to a modified first-order form of the model’s output under SGDM. To capture the finer structure of momentum dynamics under stochastic perturbations, we extend Domingos’ theorem via a second-order weak approximation of SGDM. The following theorem provides a precise formulation of this result. Lemma 3.2 (Second-order weak approximation of SGDM [19, 35]). Let T > 0, η ∈ (0, 1 ∧ T ), and set K = ⌊T /η⌋. Assume the same conditions as in Lemma 3.1 and the drift terms in the following SDEs be fixed, and define Mt , Zt : t ∈ [0, T ] as the stochastic process satisfying are Lipschitz. Let µ = 1−β η the SDEs h i η η dMt = − µI + [µ2 I − ∇2Θ L(Zt )] Mt + 1 + µ ∇Θ L(Zt ) dt (12) 2 2 1 √ 2 + η Σ|B| (Zt ) dWt , M0 = 0, h i η η dZt = 1 − µ Mt − ∇Θ L(Zt ) dt, Z0 = Θ0 , (13) 2 2 where I is the identity matrix and Σ|B| (Zt ) is defined in Lemma 3.1. Then {Zt : t ∈ [0, kη]} is a second-order weak approximation of SGDM. Lemma 3.2 provides the continuous-time second-order approximation of the SGDM iterates. We next use this representation to derive the corresponding expansion for the expected model output. Theorem 3.4 (Stochastic Domingos’ theorem (SGDM)). Consider a learning model y = f (x, Θ) with f (x, ·) ∈ G 3 . Suppose the parameter Θ is learned from a dataset {(xn , yn∗ )}N n=1 using stochastic gradient descent with momentum (SGDM), and ∇2Θ L(Zt ) is bounded on [0, T ]. Assume the same conditions as in Theorem 3.3, and the boundedness conditions in Lemma B.2 hold along the trajectory. Then, in the sense of a second-order weak approximation of SGDM, we have N
1 X E E[f (x, ΘK )] = f (x, Θ0 ) − N n=1
Z T Z t 0
∂ℓ ∗ f (xn , Zs ), yn ds dt Kt,s (x, xn ) ∂f |0 {z } memory-weighted path kernel
N
η X − E 2N n=1
−
N X
Z T Z t
η E 2N n=1
0
Z T 0
|0
∂ℓ cur Kt,s (x, xn )
∂f {z
f (xn , Zs ), yn∗
accumulated curvature-induced term
ds dt }
∂ℓ Kt (x, xn ) f (xn , Zt ), yn∗ dt ∂f {z } | path-kernel term
+
√
Z T Z t ηE 0
1/2 e−µ(t−s) Tr (Ds Zt )⊤ ∇2Θ f (x, Zt ) Σ|B| (Zs ) ds dt |0 {z } transported sampling-induced term
2
+ O(η ),
9
(14)
where µ = 1−β η is fixed, and µ2 Kt,s (x, xn ) := ∇Θ f (x, Zt ), e 1− (t − s)η ∇Θ f (xn , Zs ) , 2 D E cur Kt,s , (x, xn ) := ∇Θ f (x, Zt ), e−µ(t−s) ∇Θ f (xn , Zs ) R t 2
−µ(t−s)
s
∇Θ L(Zr ) dr
(15) (16) (17)
Kt (x, xn ) := ⟨∇Θ f (x, Zt ), ∇Θ f (xn , Zt )⟩ .
Here, Ds denotes the Malliavin derivative [28, Definition 1.2.1] with respect to the driving Brownian motion at time s. Proof. For any input x, the expectation of the output in the k-th iteration is E[f (x, ΘK )] = E[f (x, ZT )] + E[f (x, ΘK )] − E[f (x, ZT )] = E[f (x, ZT )] + O(η 2 ).
(18)
Denote At := µI + η2 µ2 I − ∇2Θ L(Zt ) , then the SDE (12) for Mt in Lemma 3.2 can be written as i h η √ 1/2 (19) dMt = − At Mt + 1 + µ ∇Θ L(Zt ) dt + η Σ|B| (Zt ) dWt . 2 For any t, s with t ≥ s ≥ 0, let Φt,s be the fundamental matrix associated with the linear system ∂t Φt,s = −At Φt,s ,
Φs,s = I.
(20)
Since At is bounded along the trajectory, the standard theory of nonautonomous linear systems ensures the existence of the fundamental matrix Φt,s , see [22, Corollary 2.3]. −1 −1 From Φ−1 t,s Φt,s = I, we obtain ∂t Φt,s = Φt,s At . Moreover, we note that the fundamental matrix satisfies Φt,s = Φt,0 Φ−1 (21) s,0 , 0 ≤ s ≤ t. Taking the differential of Φ−1 t,0 Mt and using Eq.(19) we have η −1 √ −1 1/2 d(Φ−1 η Φt,0 Σ|B| (Zt ) dWt . t,0 Mt ) = − 1 + µ Φt,0 ∇Θ L(Zt )dt + 2
(22)
Integrating both sides of Eq.(22), we have Z t
d(Φ−1 s,0 Ms ) = −
0
Z Z t η t −1 √ 1/2 Φs,0 ∇Θ L(Zs )ds + η 1+ µ Φ−1 s,0 Σ|B| (Zs ) dWs . 2 0 0
(23)
Multiplying both sides of Eq.(23) by Φt,0 , and combining with Eq.(21), we obtain, since M0 = 0, Z t η √ 1/2 Φt,s Σ|B| (Zs ) dWs . Φt,s 1 + µ ∇Θ L(Zs ) ds + η 2 0 0
Z t Mt = − By Eq.(12) in Lemma 3.2,
dZt =
(24)
h i η η 1 − µ Mt − ∇Θ L(Zt ) dt. 2 2
Since Zt is a continuous finite-variation process, we have d⟨Zj , Zℓ ⟩t = 0. Hence, by chain rule we obtain d dZt f (x, Zt ) = ∇Θ f (x, Zt )⊤ dt dt (25) η η = ∇Θ f (x, Zt )⊤ 1 − µ Mt − ∇Θ L(Zt ) . 2 2 Substituting Eq.(24) into Eq.(25) and reorganizing terms gives d (1) (2) (3) f (x, Zt ) = It (x) + It (x) + It (x) + O(η 2 ), dt
10
(26)
where (1) It (x) := −
Z t η η 1− µ 1+ µ ∇Θ f (x, Zt ), Φt,s ∇Θ L(Zs ) ds , 2 2 0
η (2) It (x) := − ⟨∇Θ f (x, Zt ), ∇Θ L(Zt )⟩ , 2 Z t η √ 1/2 (3) η ∇Θ f (x, Zt ), Φt,s Σ|B| (Zs ) dWs . It (x) := 1 − µ 2 0
(27) (28) (29)
From Lemma B.3, we have Z η t 2 µ2 Φt,s = e I+ ∇ L(Zr )dr − (t − s)η I + O(η 2 ). 2 s Θ 2 Plugging Eq.(30) into Eq.(27), since 1 − η2 µ 1 + η2 µ = 1 + O(η 2 ), we obtain Z t (1) It (x) = − ∇Θ f (x, Zt ), Φt,s ∇Θ L(Zs ) ds + O(η 2 ) −µ(t−s)
(30)
0
=−
N Z 1 X t
N n=1
e−µ(t−s) K t,s (x, xn )
0
N Z t X
η − 2N n=1
e
−µ(t−s)
∂ℓ f (xn , Zs ), yn∗ ds ∂f
∂ℓ cur K t,s (x, xn ) ∂f
0
(31) f (xn , Zs ), yn∗
ds
N Z µ2 X t ∂ℓ + η (t − s)e−µ(t−s) K t,s (x, xn ) f (xn , Zs ), yn∗ ds + O(η 2 ). 2N n=1 0 ∂f
where K t,s (x, xn ) := ⟨∇Θ f (x, Zt ), ∇Θ f (xn , Zs )⟩ , cur
K t,s (x, xn ) := ⟨∇Θ f (x, Zt ), ∇Θ f (xn , Zs )⟩R t ∇2 L(Zr )dr . s
Θ
With the definition of Kt (x, x′ ) (Eq.(3)), we can also rewrite N
(2)
It (x) = −
η X ∂ℓ Kt (x, xn ) f (xn , Zt ), yn∗ . 2N n=1 ∂f
(32)
(3)
Now we expand Φt,s inside It . Invoking Lemma B.3 again, we obtain Z t η √ µ2 3/2 (3) 1/2 It (x) = e−µ(t−s) 1 − µ η− η (t − s) ∇Θ f (x, Zt )⊤ Σ|B| (Zs ) dWs 2 2 0 Z t √ Z η η t −µ(t−s) 1/2 + e ∇Θ f (x, Zt )⊤ ∇2Θ L(Zr ) dr Σ|B| (Zs ) dWs + O(η 2 ). 2 0 s
(33)
By the same argument as in Lemma B.2, after taking expectation, each stochastic integral in (33) is √ of order O( η), hence the two terms carrying the prefactor η 3/2 can be absorbed into the remainder. Thereby, "Z # "Z # Z t T T √ (3) ⊤ −µ(t−s) 1/2 ∇Θ f (x, Zt ) e Σ|B| (Zs ) dWs dt + O(η 2 ). (34) E It (x)dt = ηE 0
0
0
The expectation in Eq.(34) does not vanish in general, since its integrand contains the future quantity ∇Θ f (x, Zt ) and is therefore not adapted to Fs . To identify its contribution, we use the Malliavin duality formula [28, Definition 1.3.1]; see also [3, Lemma 2.1]. It yields that # "Z Z t T 1/2 E ∇Θ f (x, Zt )⊤ e−µ(t−s) Σ|B| (Zs ) dWs dt 0
0
"Z
T Z t
e
=E 0
0
−µ(t−s)
# ⊤ 1/2 Tr Ds (∇Θ f (x, Zt )) Σ|B| (Zs ) dsdt ,
11
where Ds denotes the Malliavin derivative with respect to the Brownian motion at time s. Using the chain rule [28, Proposition 1.2.3], Ds (∇Θ f (x, Zt )) = ∇2Θ f (x, Zt )Ds Zt , Eq.(34) then becomes "Z T
E
#
(3) It (x)dt
# "Z Z t T √ 1/2 −µ(t−s) ⊤ 2 e Tr (Ds Zt ) ∇Θ f (x, Zt ) Σ|B| (Zs ) dsdt = ηE 0
0
0
(35)
2
+ O(η ). Integrating Eq.(26) over [0, T ] and taking expectations, we thus have "Z # T (1) (2) (3) E[f (x, ZT )] − f (x, Z0 ) = E It (x) + It (x) + It (x) dt + O(η 2 ).
(36)
0
Plugging (31), (32) and (35) into the above equation, and using Eq.(18) yields the desired result. Equation (14) decomposes the expected prediction under SGDM into four contributions. • Memory-weighted path kernel. This is the leading deterministic memory contribution in the SGDM representation. It shows that SGDM replaces the instantaneous path kernel with a timetransported kernel: a residual signal generated by sample xn at time s is propagated to a later 2 time t with weight e−µ(t−s) 1 − µ2 (t − s)η , and contributes to the prediction at the test point through the alignment of tangent features. • Accumulated curvature-induced term. This term describes how loss curvature modifies the transport of past residual signals. A training residual generated at time s is first propagated forward by the exponential memory factor e−µ(t−s) and then reweighted by the accumulated Hessian along the path from s to t. Curvature thus changes how past training information is retained and transported before it affects the tangent-feature comparison at the test point. • Path-kernel term. This term retains the ordinary path-kernel structure, but with a coefficient proportional to η. It measures the instantaneous alignment between the tangent feature of the test point and that of each training sample at the same time t. • Transported sampling-induced term. This term shows that mini-batch sampling noise injected at time s is carried forward to time t through the momentum dynamics with exponential weight e−µ(t−s) . Its contribution to the prediction is then measured by the tangent feature at the test point.
4
Implications of the Second-order Theory
In this section, we discuss some implications of our second-order theory for the model’s predictions.
4.1
Curvature-weighted interpolation
For the common case of a quadratic loss ℓ(f, y ∗ ) = (f −y ∗ )2 , the derivative ∂ℓ/∂f equals 2 f (xn , Θt )− yn∗ . Hence Eq. (7) reduces to an interpolation formula: N
2 X y(T ) − y(0) = − N n=1 −
Z T 0
N Z T X
η N n=1
∇Θ f (x, Θt ), ∇Θ f (xn , Θt ) f (xn , Θt ) − yn∗ dt
0
⟨∇Θ f (x, Θt ), ∇Θ f (xn , Θt )⟩∇2Θ L(Θt ) f (xn , Θt ) − yn∗ dt
+ O(η 2 ),
12
(37)
for any x, where y(T ) := f (x, ΘT ) and y(0) := f (x, Θ0 ). Eq. (37) shows that y(T ) − y(0) is a weighted sum of all training-point prediction corrections f (xn , Θt ) − yn∗ , with weights given by the gradient kernel integrated over the training path. In the leading term, these corrections are weighted by the path kernel. The second term has the same interpolation structure, but with a curvature-weighted kernel. Consequently, the effect of a training sample on the output depends not only on tangent-feature alignment, but also on how this alignment interacts with the local geometry of the loss landscape.
4.2
Concentration of model predictions
Theorem 3.3 describes the mean prediction of models trained with SGD. In fact, we can say more about the variability of the prediction: it is governed jointly by the learning rate, the covariance of the sampling-induced noise, and the local geometry of the loss landscape. Proposition 4.1 (Concentration of the terminal prediction). Let T = Kη, and let Zt be the stochastic process defined in Lemma 3.1. Let U ⊂ Rd be an open region, and define the exit time τ := inf{t ≥ 0 : Zt ∈ / U}. Assume that, for a fixed test point x, there exist constants GL , M2 , M3 , KΣ > 0 and batch-size-dependent constants S|B| > 0 and Gf (x) > 0 such that, for every z ∈ U , ∥∇L(z)∥ ≤ GL ,
∥∇2 L(z)∥op ≤ M2 ,
∥∇3 L(z)∥op ≤ M3 ,
2 ∥Σ|B| (z)∥op ≤ S|B| , 1/2
and ∥∇Θ f (x, z)∥ ≤ Gf (x). Assume further that the diffusion coefficient Σ|B| satisfies d X
1/2
D(Σ|B| )q (z) v
2
≤ KΣ2 ∥v∥2 ,
∀z ∈ U, ∀v ∈ Rd ,
q=1 1/2 1/2 e η := M2 + η (M3 GL + M 2 ) + η K 2 . Then, where (Σ|B| )q (z) denotes the q-th column of Σ|B| (z). Set K 2 2 2 Σ for every r > 0, 2 e K r η . P (|f (x, ZT ∧τ ) − E[f (x, ZT ∧τ )]| ≥ r) ≤ 2 exp − (38) eη T 2 2 2 K ηS|B| Gf (x) e −1
Proof. Since f (x, Zt ) contains a nonzero drift term, it is not a martingale. To construct a martingale whose terminal value is the stopped terminal prediction f (x, ZT ∧τ ), we consider the same stochastic dynamics with initial condition specified at an arbitrary pair (t, z). For 0 ≤ t ≤ s ≤ T and z ∈ U , let Zt,z denote the process evolving according to the same dynamics as Zs , but initialized at the s intermediate time t from the state z, i.e. Zt,z t = z: h i η 2 √ 1/2 t,z t,z t,z dZt,z = −∇ L(Z ) − ∇ L(Z )∇ L(Z ) ds + η Σ|B| (Zt,z Θ Θ s s Θ s s s ) dWs . 2 Let τ t,z := inf{r ≥ t : Zt,z / U} be the corresponding exit time from U, and define the stopped value r ∈ function ux (t, z) := E f x, Zt,z . T ∧τ t,z By the Markov property, ux (t ∧ τ, Zt∧τ ) = E[f (x, ZT ∧τ ) | Ft∧τ ], so the process ux (t ∧ τ, Zt∧τ ) is a martingale with terminal value f (x, ZT ∧τ ). Moreover, ux (0, Θ0 ) = E[f (x, ZT ∧τ )] , ux T ∧ τ, ZT ∧τ = f (x, ZT ∧τ ). Hence Mt := ux (t ∧ τ, Zt∧τ ) − ux (0, Z0 )
(39)
is a martingale satisfying M0 = 0 and MT = f (x, ZT ∧τ ) − E[f (x, ZT ∧τ )]. Before the exit time τ , the stopped value function ux satisfies the backward Kolmogorov equation [29, Theorem 8.1.1] on U , ∂t ux (t, z) + Lη ux (t, z) = 0, (40) where, for all φ ∈ C 2 (Rd ), ⊤ η 2 η Lη φ(z) = −∇Θ L(z) − ∇Θ L(z)∇Θ L(z) ∇φ(z) + Tr Σ|B| (z)∇2 φ(z) . 2 2 13
Applying Itô’s formula to ux (t, Zt ) up to the bounded stopping time τ ∧ T , we obtain 1 dux (t, Zt ) = ∂t ux (t, Zt ) dt + ∇ux (t, Zt )⊤ dZt + Tr ∇2 ux (t, Zt ) d⟨Z⟩t 2 h i √ 1/2 = ∂t ux (t, Zt ) + Lη ux (t, Zt ) dt + η ∇ux (t, Zt )⊤ Σ|B| (Zt ) dWt .
(41)
Substituting the backward Kolmogorov equation (40) into (41), we have Z t∧τ √ 1/2 ux (t ∧ τ, Zt∧τ ) − ux (0, Z0 ) = η ∇ux (s, Zs )⊤ Σ|B| (Zs ) dWs . 0
Recall Eq. (39), we get MT = ux (T ∧ τ, ZT ∧τ ) − ux (0, Z0 ) =
√
Z T ∧τ η 0
1/2
∇ux (s, Zs )⊤ Σ|B| (Zs ) dWs .
Then the quadratic variation of M is Z T ∧τ ⟨M ⟩T = η ∇ux (s, Zs )⊤ Σ|B| (Zs ) ∇ux (s, Zs ) ds.
(42)
(43)
0
Since the process is stopped upon leaving U , all the bounds in the assumptions hold on the integration interval; hence the gradient estimate in Lemma B.1 also applies to the stopped value function ux . Then we obtain ∇ux (s, Zs )⊤ Σ|B| (Zs ) ∇ux (s, Zs ) ≤ ∥Σ|B| (Zs )∥op ∥∇ux (s, Zs )∥2 2 ≤ S|B| G2f (x) e2Kη (T −s) , e
and so Eq.(43) satisfies 2 ⟨M ⟩T ≤ ηS|B| G2f (x)
Z T
2 e2Kη (T −s) ds = ηS|B| G2f (x) e
0
e2Kη T − 1 . eη 2K e
By the exponential martingale inequality, P(|MT | ≥ r) ≤ 2 exp −
eη r2 K eη T 2 G2 (x)(e2K ηS|B| − 1) f
! .
Since MT = f (x, ZT ∧τ ) − E[f (x, ZT ∧τ )], the desired inequality follows. Proposition 4.1 provides a locally stability bound for the terminal prediction around its mean. Combined with the stochastic Domingos representation, it shows that the stopped prediction f (x, ZT ∧τ ) concentrates around its expected kernel-type representation. On the event {τ > T }, this stopped prediction coincides with the original terminal prediction f (x, ZT ). The bound indicates that prediction stability is governed by the learning rate, the sampling-induced covariance, the curvature of the loss landscape, and the parameter sensitivity of the test prediction. The dependence on batch size en2 ters through the covariance bound S|B| : larger batches strengthen concentration around the expected representation by reducing the sampling-induced covariance in the bound. The loss-curvature bounds e η . The factor Gf (x) further shows that M2 and M3 affect the concentration probability through K test points whose predictions are more sensitive to parameter changes are less stable under stochastic training. Denote by A2 (x, T ) the second-order approximation in Theorem 3.3, excluding the O(η 2 ) remainder, so that E[f (x, ZT )] = A2 (x, T ) + Cη 2 . For any δ ∈ (0, 1), define v u u e2Ke η T − 1 2 log . Rδ (x, T ) := S|B| Gf (x)tη e δ Kη 14
Figure 1: Scaling verification for the output expansion under gradient descent. Left: the mean firstorder residual Rη over interpolation points scales as O(η) after subtracting the leading path-kernel term. Right: the mean second-order remainder Rη (x) − Curvη (x) over interpolation points and extrapolation points scales as O(η 2 ) after further subtracting the curvature-induced term. Although the extrapolation error has a larger magnitude, it follows the same second-order scaling. From Proposition 4.1, we obtain the following estimate: P |f (x, ZT ∧τ ) − A2 (x, T )| ≤ |C|η 2 + Rδ (x, T ) ≥ 1 − δ.
(44)
The parameter δ only enters through the logarithmic factor log(2/δ), reflecting the usual trade-off between the size of the error radius and the probability of the event.
4.3
Sampling-induced variation and batch size effects
When the observable is the loss itself, then the model loss at T is "Z # kη 2 E[L(Θk )] = L(Θ0 ) − E ∥∇Θ L(Zs )∥ ds 0
"Z
η − E 2
#
kη
∇Θ L(Zs )
⊤
∇2Θ L(Zs )∇Θ L(Zs )ds
0
"Z # kη η 2 + E Tr ∇Θ L(Zs )Σ|B| (Zs ) ds + O(η 2 ). 2 {z } | 0 Sampling-induced term
The representation above makes the role of mini-batch noise explicit: the sampling-induced term couples the local Hessian of the loss with the covariance of the mini-batch gradient noise. The batchsize dependence follows directly from the scaling of Σ|B| . Larger batches exert weaker pressure against high-curvature regions and can therefore tolerate sharper minima. Conversely, small-batch SGD is strongly penalized in sharp regions and is thus naturally biased toward flatter ones.
5
Experiments
In this section, we validate our second-order theory from different aspects via a series of numerical experiments.
5.1
Scaling order verification
We illustrate Theorem 3.2 and Theorem 3.3 on a one-dimensional regression problem with target function f∗ (x) = sin(2x) + 0.4 sin(x). 15
Figure 2: Empirical scaling of the first- and second-order remainders in the SGD output expansion at the test point x = 1.2. For each learning rate η, we run SGD for 20 independent rounds with the same initialization and dataset, and display the distribution of the absolute remainder across rounds (2) by boxplots. Left: the first-order remainder |Rη (x)|. Right: the second-order remainder |Rη (x)|. A two-hidden-layer tanh network: f (x, Θ) = a(2)
M X
(1) (2) tanh(a(1) , m x + bm ) + b
m=1 (1)
(1)
with learnable parameter Θ = {am , bm , m = 1, . . . , M, a(2) , b(2) } and width M = 8 is trained by gradient descent on N = 20 samples from [−2, 2], with terminal time T = 2.0 and learning rates η ∈ {10−1 , 5 × 10−2 , 2 × 10−2 , 10−2 , 5 × 10−3 }. We evaluate the expansion Eq.(10) on interpolation points {−1.22, −0.67, 0.24, 1.53} and extrapolation points {−3.5, 2.5, 2.7, 3.0}. For each test point x, we compute the first-order residual N
Rη (x) = f (x, ΘT ) − f (x, Θ0 ) +
1 X N n=1
Z T Kt (x, xn ) 0
∂ℓ f (xn , Θt ), yn∗ dt, ∂f
and the second-order remainder N
Rη − Curη :=Rη (x) +
η X 2N n=1
Z T
Ktcur (x, xn )
0
∂ℓ f (xn , Zs ), yn∗ dt. ∂f
The time integrals are evaluated numerically along the discrete GD trajectory using trapezoidal quadrature. The left panel of Figure 1 plots the mean interpolation residual and shows a slope 1 on a log-log scale, which confirms that the error after subtracting the leading path-kernel term is of order O(η). The right panel plots the mean second-order remainder for both interpolation and extrapolation points. In both cases, the curves exhibit slope 2, showing that after subtracting the curvature-induced term, the remaining error is of order O(η 2 ). While the extrapolation points have a larger error constant, the scaling law remains unchanged. This provides numerical evidence that the curvature-induced term is indeed the leading second-order contribution in the Theorem 3.2. Figure 2 reports the empirical scaling of the SGD reminder at the test point x = 1.2. We consider the learning rates η ∈ {1, 0.5, 0.2, 0.1, 0.05}. For each learning rate η, we perform 20 independent experiments using the same dataset and initialization with independently sampled mini-batch sequences, and visualize the resulting distribution of the remainder by boxplots. We define the second-order
16
(a)
(b)
(c)
Figure 3: Effect of batch size on output stability for SGD on the two-moons classification task. A two-hidden-layer tanh network of width 32 trained by SGD, while the batch size varies over {4, 8, 16, 32, 64, 128, 800}. (a) Mean standard deviation of the final logit across runs decreases monotonically as the batch size increases. (b) Centered final-logit distribution at the representative point (2.5, 1.2). (c) Heatmaps of the output standard deviation for batch sizes 4, 32, and 800. remainder by "Z # N T η X ∂ℓ cur ∗ E Kt (x, xn ) f (xn , Zs ), yn dt 2N n=1 ∂f 0 "Z # T η − E Tr ∇2Θ f (x, Zt ) Σ|B| (Zt ) dt . 2 0
Rη(2) :=Rη (x) +
In the left panel of Figure 2, the median of |Rη (x)| increases approximately linearly with η, which is broadly consistent with the predicted first-order scaling. In the right panel, the second-order remainder shows a qualitatively similar trend to the reference slope 2, although the spread across runs remains visible. Overall, this provides numerical evidence that, under SGD training, the first and second-order remainder terms follow the scaling behavior predicted by Theorem 3.3.
5.2
Batch size and sensitivity of predictions
We next study how the mini-batch size affects output stability in stochastic training on a nonlinear classification task. We use the two-moons dataset with 1000 samples and noise level 0.12, of which 800 samples are used for training. The classifier is a two-hidden-layer tanh network of width 32, trained by SGD with learning rate 0.05. To align the comparison across batch sizes with the discrete-time training index, we fix the total number of parameter updates to 3000 and vary the batch size in |B| ∈ {4, 8, 16, 32, 64, 128, 800}. 17
Figure 4: Effect of loss-Hessian-based batch selection on output sensitivity. (a) One-step experiment: from a fixed checkpoint, candidate mini-batches are used to perform one SGD update, and the resulting St+1 (x) is measured. The sensitivity is positively correlated with the largest eigenvalue of the full loss Hessian. (b) Multi-step experiment: we compare random batch selection with strategies that repeatedly choose the batch with the largest or smallest loss Hessian. The plotted trajectories of St (x) show that batch selection changes the stability of the prediction along training. For each batch size, we repeat 20 training rounds from the same initialization, changing only the mini-batch sampling order. For a test point x, we define output stability through the variability of the final logit f (x, Θk ) between runs, where k = 3000 is the total number of SGD iterations. Figure 3 summarizes the results. Panel (a) shows that the mean output standard deviation decreases monotonically as the batch size increases. This trend is particularly pronounced for test points near the decision boundary, indicating that the stability is sensitive to test data. Panel (b) displays the centered final-logit distribution at the representative test point (2.5, 1.2), where the spread between runs is much larger for small batch sizes and collapses as the batch size approaches full-batch training. Panel (c) visualizes the spatial distribution of the output standard deviation in the input domain for batch sizes 4, 32, and 800. The instability is strongest for small batch sizes and is progressively reduced as the batch size grows, becoming nearly negligible in the full-batch regime. These results provide direct numerical evidence that mini-batch sampling noise significantly affects prediction stability and that increasing the batch size stabilizes the network output.
5.3
Sensitivity control via curvature
We conduct a controlled one dimensional regression experiment to examine how curvature guided minibatch selection affects prediction stability. Training data are generated from the target function (5.1) with additive Gaussian noise of standard deviation 0.05. We use N = 100 training samples uniformly distributed in the interval [−2, 2]. The network is (5.1) with width M = 24. Training is performed for 100 iterations with a mini-batch size 40 and a learning rate 10−2 . For each strategy, we run the experiment five times. In the sensitivity evaluation, parameter perturbations are sampled from an isotropic Gaussian with standard deviation 10−3 , and the sensitivity statistics are estimated from 20 perturbation samples. To compare different geometric biases induced by stochastic training, we consider the following batch selection rules. At every iteration, instead of directly drawing one mini batch and updating the |B| network, we first sample |B| = 20 candidate mini batches of size 40. For each candidate batch {Bi }i=1 , we compute the one step SGD update (B )
Θt+1i = Θt − η∇Θ LBi (Θt ), where LBi is the empirical loss on batch Bi . We then evaluate a curvature score at the one step look (B ) ahead parameter Θt+1i , and choose the candidate batch that maximizes or minimizes this score. We
18
Figure 5: Curvature-based batch selection and loss sensitivity on a reduced MNIST task. We train on MNIST with 50 training images per class. Left: selected-curvature trajectory. Right: loss-sensitivity trajectory, computed from Gaussian parameter perturbations with σsens = 10−3 . use two different curvature scores. The first is the largest eigenvalue of the full loss Hessian, (B ) λmax ∇2Θ L(Θt+1i ) , (B ) the second is the smallest eigenvalue of the full loss Hessian λmin ∇2Θ L(Θt+1i ) . For a fixed test point x, we define the pointwise sensitivity at iteration t by perturbing the trained parameter Θt with 2 I), where σsens controls the perturbation scale, and measuring isotropic Gaussian noise ξ ∼ N (0, σsens the magnitude of the prediction change,
St (x) := Eξ f (x, Θt + ξ) − f (x, Θt ) . Figure 4 reports the pointwise sensitivity St (x) defined above. The one-step experiment shows that batches leading to larger loss-Hessian eigenvalues tend to produce larger post-update sensitivity St+1 (x). For a fixed test point x = 0.51, the maximum-eigenvalue strategy typically yields the largest sensitivity trajectory, while minimum-eigenvalue strategy stays lower and closer to the random baseline. This indicates that forcing the optimization path toward larger curvature of the full training objective leads to more fragile predictions under parameter perturbations, whereas selecting batches that reduce the curvature of the full loss leads to more stable predictors. Overall, it supports the claim that flatter minima, as measured through the geometry of the full loss, are associated with improved prediction stability. We next study whether the curvature-based batch selection affects the sensitivity of the loss to parameter perturbations in a small-scale MNIST experiment. We use 50 training images per class. At each SGD iteration, we sample 10 candidate mini-batches of size 5, perform a one-step look-ahead update for each candidate, estimate the largest eigenvalue of the loss Hessian and select the batch with either maximal or minimal curvature. We repeat each experiment for 5 runs and record the loss sensitivity every 20 steps. Figure 5 summarizes the resulting training dynamics. The left panel shows the trajectory of the selected curvature. The maximal-curvature rule consistently chooses batches with substantially larger curvature values than the minimal-curvature rule throughout training. The right panel shows the corresponding trajectory of the loss sensitivity. We observe that the maximum-eigenvalue strategy also leads to a consistently higher loss-sensitivity trajectory than both the minimal-eigenvalue strategy rule and the standard SGD baseline. This suggests that the curvature bias introduced by the batch selection rule has a direct impact on the sensitivity of the loss to parameter perturbations.
6
Conclusion
In this work, we developed second-order output representations for models trained by gradient-based optimization. Starting from the path-kernel interpretation of gradient descent, we show that the 19
second-order expansion can still be viewed as an interpolation of training residuals along the optimization trajectory. The interpolation weights are refined by an additional curvature-induced component, which records how the loss geometry modifies the propagation of training residuals to a test prediction. Under stochastic training, a further sampling-induced contribution appears, showing that mini-batch sampling affects the expected prediction through the parameter curvature of the prediction map along the directions of the sampling covariance. For SGDM, these mechanisms are propagated further through a momentum-induced memory effect. The concentration result complements the second-order representation by showing how the learning rate, the batch-size-dependent covariance, and the loss Hessian control the size of prediction fluctuations. Our numerical experiments support these conclusions: the observed scaling of the first- and second-order remainders is consistent with the theory. Overall, the network stores information about the training data through the accumulated residuals along the optimization path, and the final prediction is obtained by interpolating these residuals through datadependent kernel weights. The second-order viewpoint yields a refined interpolation framework, in which a prediction is shaped not only by path kernels but also by curvature- and sampling-induced contributions.
A
Second-order representation with scheduler
A widely used learning-rate schedule in deep learning is the cosine scheduler [23], which smoothly anneals the step size from a large initial value to a small final value according to 1 k π , k = 0, . . . , K, ηk = ηmin + (ηmax − ηmin ) 1 + cos K 2 where ηmax > 0 and ηmin ≥ 0 are the maximum and minimum learning rates, T is the total number of training iterations, and t ∈ {0, 1, . . . , T } is the iteration index. Unlike a constant step size or a piecewise decay schedule, cosine annealing produces a smooth and progressively decreasing learning rate. Thus, the optimization initially takes relatively large steps, while the update magnitude gradually decreases during the later stages of training, reducing the possibility of oscillatory behavior near a minimizer. To apply the second-order interpolation formula to this setting, we first identify the modified equation associated with gradient descent under a schedule learning-rate weight. k Lemma A.1. Let L ∈ C 3 (Rd ). For scheduler learning rate ηk = ηmin + 21 (ηmax − ηmin ) 1 + cos K π , define η = ηmin , let tk = kη, k = 0, . . . , K, with T = Kη. Consider the scheduled gradient-descent iteration Θk+1 = Θk − ηk ∇Θ L(Θk ), k = 0, · · · , K − 1. (45) Define scalar weight w(t) = 1 + 12 (κ − 1) 1 + cos( πt T ) , κ = ηmax /ηmin , t ∈ [0, T ]. Equivalently, ηk = ηw(tk ). Then the corresponding modified equation is η η Θ̇t = −w(t)∇L(Θt ) + ẇ(t)∇L(Θt ) − w(t)2 ∇2 L(Θt )∇L(Θt ) + O(η 2 ). 2 2
(46)
Proof. Denote b(t, Θt ) = −w(t)∇Θ L(Θt ), following the standard backward error analysis, we determine a(t, Θt ) to get the truncated modified equation Θ̇t = b(t, Θt ) + ηa(t, Θt ).
(47)
Higher-order corrections in the modified vector field start at order η 2 . Our method is to determine a(t, Θ) by matching the second-order Taylor expansion of this modified equation with the scheduled gradient-descent update. Starting from tk = kη, Taylor expansion over one step gives Θ(tk + η) = Θ(tk ) + η Θ̇(tk ) +
η2 Θ̈(tk ) + O(η 3 ). 2
(48)
Differentiating Eq. (47) with respect to t, we obtain Θ̈t =
d [b(t, Θt ) + η a(t, Θt )] dt
h i =∂t b(t, Θt ) + DΘ b(t, Θt )Θ̇t + η ∂t a(t, Θt ) + DΘ (a(t, Θt ))Θ̇t . 20
(49)
Since Θ̇t = b(t, Θt ) + O(η), DΘ b(t, Θt ) = −w(t)∇2Θ L(Θt ), we have DΘ b(t, Θt )Θ̇t =DΘ b(t, Θt )b(t, Θt ) + O(η) = −w(t)∇2Θ L(Θt )b(t, Θt ) + O(η) = w(t)2 ∇2Θ L(Θt )∇Θ L(Θt ) + O(η).
(50)
Plugging Eq.(50) into Eq.(49), we have Θ̈t = ∂t b(t, Θt ) + w(t)2 ∇2Θ L(Θt )∇Θ L(Θt ) + O(η).
(51)
1 Θ(tk + η) =Θ(tk ) + ηb(tk , Θtk ) + η 2 a(tk , Θtk ) + ∂t b(tk , Θtk ) 2 1 + (DΘ b(tk , Θtk ))b(tk , Θtk ) + O(η 3 ). 2
(52)
Then Eq.(48) becomes
Then we choose a(t, Θt ) so that the coefficient of η 2 vanishes, that is 1 1 a(t, Θt ) + ∂t b(t, Θt ) + (DΘ b(t, Θt ))b(t, Θt ) = 0. 2 2 Substituting the expressions for ∂t b(t, Θt ) and DΘ b(t, Θt ), we obtain a(t, Θt ) =
1 1 ẇ(t)∇Θ L(Θt ) − w(t)2 ∇2Θ L(Θt )∇Θ L(Θt ). 2 2
(53)
Plugging Eq.(53) into Eq.(47), the truncated modified equation is η η Θ̇t = −w(t)∇L(Θt ) + ẇ(t)∇L(Θt ) − w(t)2 ∇2 L(Θt )∇L(Θt ). 2 2 Since T is fixed, the O(η 3 ) local error accumulated over K = T /η steps gives an O(η 2 ) error at the grid points tk = kη, under the usual local Lipschitz bounds. Applying Theorem 3.2 to this time-rescaled gradient flow (Eq.(46)) yields the following schedulerweighted path-kernel evolution. Corollary A.1 (Scheduler-weighted path-kernel evolution). Assume the same regularity conditions as in Theorem 3.2, and consider the modified equation of the cosine learning-rate schedule with learning rate w(t). Then for any model output f (x, ΘT ), the terminal prediction satisfies N Z 1 X T η ∂ℓ f (x, ΘT ) =f (x, Θ0 ) − w(t) − ẇ(t) Kt (x, xn ) f (xn , Θt ), yn∗ dt N n=1 0 2 ∂f Z N ∂ℓ η X T − w(t)2 Ktcur (x, xn ) f (xn , Θt ), yn∗ dt + O(η 2 ), 2N n=1 0 ∂f
(54) (55)
with the scalar weight w(t) = 1 + 12 (κ − 1) 1 + cos( πt T ) , t ∈ [0, T ]. This expression shows that cosine annealing does not alter the path-kernel structure, while introducing time-dependent weights into both the leading and curvature-dependent contributions. The leading path-kernel term is weighted by w(t) − η2 ẇ(t), the second-order curvature-weighted term is scaled by w(t)2 .
B
Lemmas
We collect some technical lemmas used for proving our main results.
21
Lemma B.1. Let T > 0. For 0 ≤ t ≤ T and z ∈ U, let Zt,z s , s ∈ [t, T ], be the solution of t,z dZt,z s = bη (Zs ) ds +
√
1/2
η Σ|B| (Zt,z s ) dWs ,
Zt,z t = z,
where bη (z) = −∇L(z) − η2 ∇2 L(z)∇L(z). For a fixed test point x, define ux (t, z) = E f (x, Zt,z T ) . Assume that the restarted trajectory Zt,z s remains in U, and that the following bounds hold on U: ∥∇2 L(z)∥op ≤ M2 ,
∥∇L(z)∥ ≤ GL ,
∥∇3 L(z)∥op ≤ M3 , 1/2
and ∥∇Θ f (x, z)∥ ≤ Gf (x) for the fixed test point x. Assume further that the diffusion coefficient Σ|B| satisfies d X 2 1/2 D(Σ|B| )q (z)v ≤ KΣ2 ∥v∥2 , ∀z ∈ U, ∀v ∈ Rd , q=1
1/2 1/2 e η = M2 + η (M3 GL + M 2 ) + η K 2 . Then where (Σ|B| )q (z) denotes the q-th column of Σ|B| (z). Set K 2 2 2 Σ
∥∇z ux (t, z)∥ ≤ Gf (x)eKη (T −t) . e
Proof. First, we estimate the derivative of the drift. Since η bη (z) = −∇L(z) − ∇2 L(z)∇L(z), 2 we have
η ∇bη (z) = −∇2 L(z) − ∇ ∇2 L(z)∇L(z) . 2
For any direction v, ∇ ∇2 L(z)∇L(z) v = (∇v ∇2 L(z))∇L(z) + ∇2 L(z)∇2 L(z)v. Therefore,
∇ ∇2 L(z)∇L(z) op ≤ ∥∇3 L(z)∥op ∥∇L(z)∥ + ∥∇2 L(z)∥2op .
Using the bounds on ∇L, ∇2 L, and ∇3 L, we obtain η ∥∇bη (z)∥op ≤ M2 + (M3 GL + M22 ) =: Kb . 2 be the derivative of the stochastic flow with respect to the initial point. Let Js,t (z) = ∇z Zt,z s Differentiating the SDE with respect to z, we obtain d
dJs,t (z) = ∇bη (Zt,z s )Js,t (z) ds +
√ X 1/2 q D(Σ|B| )q (Zt,z η s )Js,t (z) dWs ,
Jt,t (z) = I.
q=1
Fix a unit vector v ∈ Rd , and set Ys = Js,t (z)v. Then d
dYs = ∇bη (Zt,z s )Ys ds +
√ X 1/2 q η D(Σ|B| )q (Zt,z s )Ys dWs . q=1
Applying Itô’s formula to ∥Ys ∥2 gives d∥Ys ∥2 =2Ys⊤ ∇bη (Zt,z s )Ys ds + η
d X
1/2
D(Σ|B| )q (Zt,z s )Ys
2
ds
q=1 d
√ X ⊤ 1/2 q Ys D(Σ|B| )q (Zt,z +2 η s )Ys dWs . q=1
Taking expectations, the stochastic integral has zero expectation. Hence " # d X 2 d 1/2 2 ⊤ t,z t,z E∥Ys ∥ = E 2Ys ∇bη (Zs )Ys + η D(Σ|B| )q (Zs )Ys . ds q=1 22
Using ∥∇bη (z)∥op ≤ Kb and
Pd
2
1/2
q=1
D(Σ|B| )q (z)v
≤ KΣ2 ∥v∥2 , we get
d E∥Ys ∥2 ≤ (2Kb + ηKΣ2 )E∥Ys ∥2 . ds Since Yt = v and ∥v∥ = 1, Gronwall’s inequality yields 2
E∥Js,t (z)v∥2 ≤ e(2Kb +ηKΣ )(s−t) = e2Kη (s−t) , e
e η = Kb + η K 2 . Next, by differentiating ux (t, z) = E[f (x, Zt,z )] with respect to z, we obtain where K T 2 Σ ∇z ux (t, z) = E JT,t (z)⊤ ∇Θ f (x, Zt,z T ) . Thus, for every unit vector v, ⊤ |v ⊤ ∇z ux (t, z)| = E ∇Θ f (x, Zt,z T ) JT,t (z)v ≤ Gf (x)E [∥JT,t (z)v∥] ≤ Gf (x) E∥JT,t (z)v∥2
1/2
≤ Gf (x)eKη (T −t) . e
Taking the supremum over all unit vectors v, we obtain ∥∇z ux (t, z)∥ ≤ Gf (x)eKη (T −t) . e
This completes the proof. We next eastimate the stochastic integral that appears in the SGDM expansion. Lemma B.2. Assume that there exists CT > 0, independent of η, such that for every z ∈ {Zr : 0 ≤ r ≤ T }, 1/2
∥Σ|B| (z)∥F ≤ CT ,
1/2
∥D(Σ|B| (z))[v]∥F ≤ CT ∥v∥, ∀v ∈ Rd ,
∥∇2Θ L(z)∥op + ∥∇3Θ L(z)∥op ≤ CT ,
∥∇2Θ f (x, z)∥op ≤ CT ,
sup ∥Mr ∥ ≤ CT . 0≤r≤T
Here and below, CT denotes a generic constant, independent of η, whose value may change from line to line. Then for 0 ≤ t ≤ T , Z t √ 1/2 ⊤ E ∇Θ f (x, Zt ) Σ|B| (Zs ) dWs ≤ CT η. 0
Proof. Since the Brownian noise enters only through the momentum M, we have Ds Ms =
√
1 2 ηΣ|B| (Zs ), Ds Zs = 0.
Taking the Malliavin derivative of Mr gives Z r η √ 1/2 Ds Mr = η Σ|B| (Zs ) − µI + (µ2 I − ∇2Θ L(Zζ )) Ds Mζ dζ 2 s Z r Z η η r 2 + ∇3Θ L(Zζ )[Ds Zζ ] Mζ dζ − 1 + µ ∇Θ L(Zζ )Ds Zζ dζ 2 s 2 s Z r √ 1/2 + η DΣ|B| (Zζ )[Ds Zζ ] dWζ . s
By Itó isometry we obtain Z r Z r 2 2 √ 1/2 1/2 E η DΣ|B| (Zζ )[Ds Zζ ] dWζ = η E DΣ|B| (Zζ )[Ds Zζ ] dζ F s s Z Z r r E∥Ds Zζ ∥2 dζ. ≤ CT η E∥Ds Zζ ∥2 dζ ≤ CT s
23
s
By Cauchy–Schwarz inequality, we have Z r
2
E[∥Ds Mr ∥ ] ≤ CT η + CT
E ∥Ds Mζ ∥2 + ∥Ds Zζ ∥2 dζ.
s
Taking the Malliavin derivative of Z gives Z r h i η η Ds Zr = 1 − µ Ds Mζ − ∇2Θ L(Zζ )Ds Zζ dζ. 2 2 s By Cauchy-Schwarz inequality, E∥Ds Zr ∥2 ≤ CT
Z r
E∥Ds Mζ ∥2 + E∥Ds Zζ ∥2 dζ.
s
Define Yr = E[∥Ds Mr ∥ ] + E[∥Ds Zr ∥ ] we obtain 2
2
Z r Yr ≤ CT η + CT
Yζ dζ. s
By Gronwall’s inequality, (56)
E[∥Ds Zr ∥2 ] ≤ Yr ≤ CT η exp CT (r − s) ≤ CT η exp CT T ≤ CT η. By Malliavin duality, Z t Z t 1/2 1/2 E ∇Θ f (x, Zt )⊤ Σ|B| (Zs )dWs = E Tr Ds (∇Θ f (x, Zt ))⊤ Σ|B| (Zs ) ds 0 0 Z t 1/2 =E Tr (∇2Θ f (x, Zt )Ds Zt )⊤ Σ|B| (Zs ) ds . 0
Therefore, by Eq.(56) we obtain Z t Z t 1/2 ⊤ 2 ⊤ 1/2 E ∇Θ f (x, Zt ) Σ|B| (Zs )dWs ≤ E Tr (∇Θ f (x, Zt )Ds Zt ) Σ|B| (Zs ) ds 0
0
Z t ≤ CT
Z t E∥Ds Zt ∥ds ≤ CT
0
1
(E∥Ds Zt ∥2 ) 2 ds
0
√ ≤ CT η.
Lemma B.3. For 0 ≤ s ≤ t ≤ T := Kη, let Φt,s be the fundamental matrix associated with the linear system ∂t Φt,s = −At Φt,s , Φs,s = I, where At = µI+ η2 [µ2 I−∇2Θ L(Zt )]. Assume that sup0≤t≤T ∥∇2Θ L(Zt )∥ ≤ CT , where CT is independent of η. Then Z µ2 η t 2 −µ(t−s) ∇ L(Zr )dr − (t − s)η I + O(η 2 ), (57) Φt,s = e I+ 2 s Θ 2 where the remainder is understood in operator norm. Proof. We now derive an explicit approximation of Φt,s . Defining η
η
2
2
Ψt,s := e(µ+ 2 µ )(t−s) Φt,s , i.e. Φt,s = e−(µ+ 2 µ )(t−s) Ψt,s , and differentiating with respect to t, we have η 2 η 2 η ∂t Φt,s = − µ + µ2 e−(µ+ 2 µ )(t−s) Ψt,s + e−(µ+ 2 µ )(t−s) ∂t Ψt,s . 2 24
(58)
(59)
On the other hand, plugging Eq.(58) into Eq.(20) we have η 2 η 2 η η ∂t Φt,s = − µ + µ2 e−(µ+ 2 µ )(t−s) Ψt,s + ∇2Θ L(Zt )e−(µ+ 2 µ )(t−s) Ψt,s . 2 2
(60)
Comparing Eq.(59) and Eq.(60) we deduce that ∂t Ψt,s =
η 2 ∇ L(Zt )Ψt,s , Ψs,s = I. 2 Θ
(61)
Integrating Eq.(61) between s and t, we have Ψt,s − I =
η 2
Z t
∇2Θ L(Zr )Ψr,s dr.
(62)
s
Define Rt,s := Ψt,s − I −
η 2
Z t
∇2Θ L(Zr )dr,
(63)
s
combining with this definition (62) we obtain Z η t 2 ∇ L(Zr )(Ψr,s − I)dr. Rt,s = 2 s Θ
(64)
Using Eq.(61) and Grönwall inequality, we have Z Z η r ηCT r ∥Ψr,s ∥ ≤ 1 + ∥∇2Θ L(Zq )∥∥Ψq,s ∥dq ≤ 1 + ∥Ψq,s ∥dq 2 s 2 s ηCT ηCT ≤ exp (r − s) ≤ exp Kη . 2 2 Noticing that Ψs,s = I, we have ∥Ψr,s − I∥ ≤
ηCT 2
Z r
exp
s
ηCT ηCT ηCT Kη dq = exp Kη (r − s). 2 2 2
Plugging Eq.(65) into Eq.(64), we obtain Z t Z η CT t ηCT ηCT η 2 CT2 ηCT ∥Rt,s ∥ ≤ exp Kη (r − s)dr ≤ exp Kη (r − s)dr 2 2 2 4 2 s s = O(η 2 ). Going back to the definition (63) of Rt,s , this yields Z η t 2 Ψt,s = I + ∇ L(Zr )dr + O(η 2 ). 2 s Θ Combining this result with (58), we obtain Z η 2 η t 2 ∇Θ L(Zr )dr + O(η 2 ) Φt,s = e−(µ+ 2 µ )(t−s) I + 2 s Z η t 2 −µ(t−s) − η2 µ2 (t−s) 2 =e e ∇ L(Zr )dr + O(η ) . I+ 2 s Θ η
2
Taking Taylor expansion to e− 2 µ (t−s) , we get Z µ2 η t 2 −µ(t−s) 2 2 Φt,s = e 1− (t − s)η + O(η ) I+ ∇ L(Zr )dr + O(η ) 2 2 s Θ Z µ2 η t 2 = e−µ(t−s) I − (t − s)η I + ∇ L(Zr )dr + O(η 2 ). 2 2 s Θ
25
(65)
References [1] D. G. Barrett and B. Dherin, arXiv:2009.11162, (2020).
Implicit gradient regularization,
arXiv preprint
[2] P. Chaudhari, A. Choromanska, S. Soatto, Y. LeCun, C. Baldassi, C. Borgs, J. Chayes, L. Sagun, and R. Zecchina, Entropy-SGD: Biasing gradient descent into wide valleys, Journal of Statistical Mechanics: Theory and Experiment, 2019 (2019), p. 124018. [3] A. Debussche, Weak approximation of stochastic partial differential equations: the nonlinear case, Mathematics of Computation, 80 (2011), pp. 89–117. [4] S. Di Giovacchino, D. J. Higham, and K. Zygalakis, Backward error analysis and the qualitative behaviour of stochastic optimization algorithms: Application to stochastic coordinate descent, arXiv preprint arXiv:2309.02082, (2023). [5] L. Dinh, R. Pascanu, S. Bengio, and Y. Bengio, Sharp minima can generalize for deep nets, in International Conference on Machine Learning, PMLR, 2017, pp. 1019–1028. [6] P. Domingos, Every model learned by gradient descent is approximately a kernel machine, arXiv preprint arXiv:2012.00152, (2020). [7] Y. Feng, T. Gao, L. Li, J.-G. Liu, and Y. Lu, Uniform-in-time weak error analysis for stochastic gradient descent algorithms via diffusion approximation, arXiv preprint arXiv:1902.00635, (2019). [8] B. Ghorbani, S. Krishnan, and Y. Xiao, An investigation into neural net optimization via hessian eigenvalue density, in International Conference on Machine Learning, PMLR, 2019, pp. 2232–2241. [9] X. Gu, K. Lyu, L. Huang, and S. Arora, Why (and when) does local SGD generalize better than SGD?, arXiv preprint arXiv:2303.01215, (2023). [10] J. Guo, R. Y. He, and J.-M. Morel, On interpolation formulas describing neural network generalization, arXiv preprint arXiv:2603.13872, (2026). [11] W. Hu, C. J. Li, L. Li, and J.-G. Liu, On the diffusion approximation of nonconvex stochastic gradient descent, arXiv preprint arXiv:1705.07562, (2017). [12] A. Jacot, F. Gabriel, and C. Hongler, Neural tangent kernel: Convergence and generalization in neural networks, Advances in neural information processing systems, 31 (2018). [13] S. Jastrzębski, Z. Kenton, D. Arpit, N. Ballas, A. Fischer, Y. Bengio, and A. Storkey, Three factors influencing minima in SGD, arXiv preprint arXiv:1711.04623, (2017). [14] P. Kloeden and E. Platen, Numerical Solution of Stochastic Differential Equations, Stochastic Modelling and Applied Probability, Springer Berlin Heidelberg, 2013. [15] H. J. Kushner and G. Yin, Stochastic approximation algorithms for parallel and distributed processing, Stochastics: An International Journal of Probability and Stochastic Processes, 22 (1987), pp. 219–250. [16] Y. LeCun, Y. Bengio, and G. Hinton, Deep learning, nature, 521 (2015), pp. 436–444. [17] J. Lee, L. Xiao, S. Schoenholz, Y. Bahri, R. Novak, J. Sohl-Dickstein, and J. Pennington, Wide neural networks of any depth evolve as linear models under gradient descent, Advances in neural information processing systems, 32 (2019). [18] H. Li, Z. Xu, G. Taylor, C. Studer, and T. Goldstein, Visualizing the loss landscape of neural nets, Advances in neural information processing systems, 31 (2018). [19] Q. Li, C. Tai, and E. Weinan, Stochastic modified equations and dynamics of stochastic gradient algorithms i: Mathematical foundations, Journal of Machine Learning Research, 20 (2019), pp. 1–47. 26
[20] X. Li, Z. Shen, L. Zhang, and N. He, A hessian-aware stochastic differential equation for modelling SGD, Mathematical Programming, (2026), pp. 1–80. [21] L. Ljung, G. Pflug, and H. Walk, Stochastic approximation and optimization of random systems, vol. 17, Birkhäuser, 2012. [22] H. Logemann, E. P. Ryan, et al., Ordinary differential equations: analysis, qualitative theory and control, Springer, 2014. [23] I. Loshchilov and F. Hutter, SGDR: Stochastic gradient descent with warm restarts, arXiv preprint arXiv:1608.03983, (2016). [24] S. Malladi, K. Lyu, A. Panigrahi, and S. Arora, On the SDEs and scaling rules for adaptive gradient algorithms, Advances in Neural Information Processing Systems, 35 (2022), pp. 7697–7711. [25] S. Mandt, M. Hoffman, and D. Blei, A variational analysis of stochastic gradient algorithms, in International conference on machine learning, PMLR, 2016, pp. 354–363. [26] S. Mandt, M. D. Hoffman, and D. M. Blei, Stochastic gradient descent as approximate bayesian inference, Journal of Machine Learning Research, 18 (2017), pp. 1–35. [27] R. Novak, L. Xiao, J. Hron, J. Lee, A. A. Alemi, J. Sohl-Dickstein, and S. S. Schoenholz, NEURAL TANGENTS: Fast and easy infinite neural networks in Python, arXiv preprint arXiv:1912.02803, (2019). [28] D. Nualart, The Malliavin calculus and related topics, Springer, 2006. [29] B. Oksendal, Stochastic differential equations: an introduction with applications, Springer Science & Business Media, 2013. [30] D. Ravikumar, E. Soufleri, and K. Roy, Curvature clues: Decoding deep learning privacy with input loss curvature, Advances in Neural Information Processing Systems, 37 (2024), pp. 20003–20030. [31] H. Robbins and S. Monro, A stochastic approximation method, The Annals of Mathematical Statistics, 22 (1951), pp. 400–407. [32] L. Sagun, U. Evci, V. U. Guney, Y. Dauphin, and L. Bottou, Empirical analysis of the hessian of over-parametrized neural networks, arXiv preprint arXiv:1706.04454, (2017). [33] J. Schmidhuber, Deep learning in neural networks: An overview, Neural networks, 61 (2015), pp. 85–117. [34] I. Sutskever, J. Martens, G. Dahl, and G. Hinton, On the importance of initialization and momentum in deep learning, in International conference on machine learning, PMLR, 2013, pp. 1139–1147. [35] R. Wang, S. Malladi, T. Wang, K. Lyu, and Z. Li, The marginal value of momentum for small learning rate SGD, arXiv preprint arXiv:2307.15196, (2023). [36] J. Wu, W. Hu, H. Xiong, J. Huan, V. Braverman, and Z. Zhu, On the noisy gradient descent that generalizes as SGD, in International Conference on Machine Learning, PMLR, 2020, pp. 10367–10376. [37] L. Wu, M. Wang, and W. Su, The alignment property of sgd noise and how it helps select flat minima: A stability analysis, in Advances in Neural Information Processing Systems, S. Koyejo, S. Mohamed, A. Agarwal, D. Belgrave, K. Cho, and A. Oh, eds., vol. 35, Curran Associates, Inc., 2022, pp. 4680–4693. [38] Z. Xie, I. Sato, and M. Sugiyama, A diffusion theory for deep learning dynamics: Stochastic gradient descent exponentially favors flat minima, arXiv preprint arXiv:2002.03495, (2020).
27
[39] Z. Zhu, J. Wu, B. Yu, L. Wu, and J. Ma, The anisotropic noise in stochastic gradient descent: Its behavior of escaping from sharp minima and regularization effects, arXiv preprint arXiv:1803.00195, (2018).
28