Mean-field analysis for heavy ball methods: Dropout-stability, connectivity, and global convergence
Abstract
The stochastic heavy ball method (SHB), also known as stochastic gradient descent (SGD) with
Polyak’s momentum, is widely used in training neural networks.
However, despite the remarkable success of such algorithm in practice, its theoretical characterization remains limited.
In this paper, we focus on neural networks with two and three layers and provide a rigorous understanding of the properties of the solutions found by SHB: (i) stability after dropping out part of the neurons, (ii) connectivity along a low-loss path, and (iii) convergence to the global optimum.
To achieve this goal, we take a mean-field view and relate the SHB dynamics to a certain partial differential equation in the limit of large network widths.
This mean-field perspective has inspired a recent line of work focusing on SGD while, in contrast, our paper considers an algorithm with momentum.
More specifically, after proving existence and uniqueness of the limit differential equations, we show convergence to the global optimum and give a quantitative bound between the mean-field limit and the SHB dynamics of a finite-width network.
Armed with this last bound, we are able to establish the dropout-stability and connectivity of SHB solutions.
1. Introduction
Despite the exceptional empirical success of gradient-based algorithms in the training of overparameterized neural networks, their convergence properties are still not well understood, given that the optimization landscape is known to be highly non-convex and to contain spurious local minima [5, 31, 45].
A popular line of work starting from [7, 25, 30, 36] has proposed a new methodology to analyze the behavior of stochastic gradient descent (SGD), namely, the mean-field regime.
The idea is that, as the number of neurons of the network grows, the SGD training dynamics converges to the solution of a certain Wasserstein gradient flow.
This perspective has facilitated the study of architectures with multiple layers [4, 13, 23, 27], and it has given a rigorous justification to a number of properties displayed by SGD solutions, including convergence towards a global optimum [6, 7, 17, 25, 28], dropout-stability and connectivity [34], and implicit bias [8, 33, 42].
Optimization with momentum, e.g., the heavy ball method [29] or Adam [18], is widely used in practice [38].
However, all the aforementioned works consider the vanilla SGD algorithm and, in general, the theoretical understanding of algorithms with momentum has lagged behind.
To address this gap, [20] defines a mean-field limit for the stochastic heavy ball (SHB) method – also known as SGD with Polyak’s momentum – for a two-layer model.
In particular, the convergence to the mean-field limit is proved, as well as that the solution of the mean-field equation approaches a global optimum.
However, [20] leaves as an open problem finding a quantitative bound between the infinite-width limit and the finite-width network, and the analysis is restricted to two layers.
In this paper, we define a mean-field limit for the heavy ball method in two-layer and three-layer networks.
We show global convergence in the three-layer setting, and give quantitative bounds for networks with finite widths.
This last result opens the way to providing a rigorous understanding of effects commonly observed in practice, such as the connectivity of solutions via low-loss paths [10, 12, 14, 15].
Furthermore, we highlight that the explicit characterization of the SHB dynamics via the mean-field limit could prove useful to analyze the robustness and reliability of the solution found by the optimization algorithm.
More specifically, our key technical contributions can be summarized as follows:
We show existence and uniqueness of the mean-field equations capturing the SHB training.
We give non-asymptotic convergence results of the SHB dynamics of a finite-width neural network towards the mean-field limit.
Our bounds are tight in terms of the network width, and they exhibit a mild dependence on the input dimension.
As a consequence of these bounds, we discuss how SHB solutions can be connected via a simple piece-wise linear path, along which the increase in loss vanishes as the width of the network grows.
Finally, we prove a global convergence result for three-layer networks, under certain assumptions on the mode of convergence of the dynamics.
Related work
For two-layer networks, our approach extends the line of work [25, 26] to the heavy ball method.
Because of the presence of momentum, in this case the mean-field limit is described by a second-order differential equation (instead of the first-order one capturing the SGD dynamics).
The mean-field limit for heavy ball methods was first considered in [20], which deals with a setting regularized by noise and does not provide quantitative bounds.
Various approaches have been proposed to define a mean-field limit for neural networks with more than two layers [4, 13, 36], and here we follow the “neuronal embedding” framework [27, 28].
Global optimality for three-layer networks was shown in [28], under a convergence assumption in the same spirit of [7].
Results for networks with more than three layers exploit a special initialization [27] or skip connections [13, 23].
Another recent line of work has considered training neural networks in the neural tangent kernel (NTK) regime (or lazy regime)
[1, 9, 11, 16], and this type of analysis has been adapted to the heavy ball method by [40] and to other adaptive methods by [43].
However, we remark that, unlike in the mean-field regime, neural networks are unable to perform feature learning in the NTK regime [44].
Beyond the training of neural networks, the continuous limit of momentum-based stochastic gradient descent algorithms has been studied in [35, 37, 41], and such dynamics are known to be closely related to sampling methods such as MCMC [24].
2. Problem setup
For space reasons, in the paper we focus on three-layer networks, and the results for two-layer networks are deferred to Appendix B.
The network has n 1 and n 2 neurons in the first and second hidden layer, respectively:
[Equation]
Here x , w 1 ( j 1 ) ∈ R D , w 2 ( j 1 , j 2 ) , w 3 ( j 2 ) ∈ R , and we use W ∈ R n 1 D + n 2 n 1 + n 2 to denote the collection of all parameters.
The training data z = ( x , y ) is generated i.i.d.
from a distribution D , and the neural network (1) is trained to minimize the population risk function R ( W ) =
E z [ R ( y, f ( x ; W ))] via the following one-pass stochastic heavy ball (SHB) method 1 : W SHB [Equation] k
[Equation]
=
[Equation]
SHB
[Equation]
+
[Equation]
Here (cid:98) ∇ W SHB R ( y ( k ) , f
( x ( k ); W SHB ) denotes the scaled gradient, and the choice of the scaling factors ensures that each component of the gradients is of order 1 (i.e., independent of the layer widths n 1 , n 2 ), see Appendix A for the details.
We make the following assumptions:
(A1)
There exists a universal constant K > 0 such that ∥ σ 1 ∥ ∞ , ∥ σ ′ 1 ∥ ∞ , ∥ σ ′′ 1 ∥ ∞ , ∥ σ 2 ∥ ∞ , ∥ σ ′ 2 ∥ ∞ , ∥ σ ′′ 2 ∥ ∞ ≤ K .
The data distribution D is such that | y | , ∥ x ∥ 2 ≤ K almost surely.
Furthermore, σ ′ 2 ( x ) ̸ = 0 for all x , | ∂ 2 R ( y, f ( x ; W )) | is K -Lipschitz continuous with respect to the second argument and K -bounded for any W .
(A2)
w 1 (0 , j 1 ) , w 2 (0 , j 1 , j 2 ) , w 3 (0 , j 2 ) have an i.i
d. initialization from ρ 10 × ρ 20 × ρ 30 , i.e. w 1 (0 , j 1 ) i.i
d. ∼ ρ 10 , w 2 (0 , j 1 , j 2 ) i.i.d.
∼ ρ 20 , w 3 (0 , j 2 ) i.i.d.
∼ ρ 30 for j 1 ∈ [ n 1 ] , j 2 ∈ [ n 2 ] .
Furthermore, w 1 (0 , j 1 ) is K 2 -sub-Gaussian, | w 2 (0 , j 1 , j 2 ) | , | w 3 (0 , j 2 ) | ≤ K almost surely, and r (0)
= 0 .
The assumptions above hold e.g. for tanh or sigmoid activation functions, and logistic or Huber loss.
Even if ∂ 2 R ( y, f ( x ; W )) is not K -bounded for the square loss, we still expect similar results to hold, since it suffices that ∂ 2 R ( y, f ( x ; W )) is bounded with high probability.
Finally, we remark that the boundedness of the initialization w 2 (0 , j 1 , j 2 ) , w 3 (0 , j 2 ) is purely to simplify the proof, which could be generalized to a K 2 -sub-Gaussian initialization.
3. Derivation of the mean-field limit
We start by defining the neuronal embedding .
In [28, Proposition 7], it is proved that, for the i.i.d. initialization (A2), there exists a product probability space (Ω 1 × Ω 2 , F 1 ×F 2 , P 1 × P 2 ) and functions w 1 (0 , · ) : Ω 1 −→ R D , w 2 (0 , · , · ) :
Ω 1 × Ω 2 −→ R , w 3 (0 , · ) :
Ω 2 −→ R s.t.
for any n 1 , n 2 > 0 , { w 1 (0 , C 1 ( j 1 )) , w 2 (0 , C 1 ( j 1 ) , C 2 ( j 2 )) , w 3 (0 , C 2 ( j 2 )) , for j 1 ∈ [ n 1 ] , j 2 ∈ [ n 2 ] } d = { w 1 (0 , j 1 ) , w 2 (0 , j 1 , j 2 ) , w 3 (0 , j 2 ) , for j 1 ∈ [ n 1 ] , j 2 ∈ [ n 2 ] } , where d = denotes equality in distribution, C 1 ( j 1 ) i.i.d.
∼ P 1 , C 2 ( j 2 ) i.i.d.
∼ P 2 , for j 1 ∈ [ n 1 ] , j 2 ∈ [ n 2 ] and we use the short-hand [ n i ] := { 1 , . .
. , n i } .
The neuronal embedding is (cid:8) (Ω 1 × Ω 2 , F 1 × F 2 ,
P 1 × P 2 ) , w 1 (0 , · ) , w 2 (0 , · , · ) , w 3 (0 , · ) (cid:9) , and the mean-field limit is the ODE tracks the functions
[Equation]
where c 1 ∈ Ω 1 , c 2 ∈ Ω 2 are dummy variables.
The output of the neural network under the meanfield limit (3) is described via the following forward pass:
[Equation]
Furthermore
[Equation]
the quantities ∆ W 3 , ∆ W 2 , ∆ W 1 appearing in (3) are described via the backward pass:
[Equation]
By analyzing a Picard type of iteration [39]
[Equation]
we can show the existence and uniqueness of the mean-field limit (see Appendix D.2 for the proof)
[Equation]
Theorem 1 Under Assumptions (A1)-(A2), there exists a unique solution of the mean-field ODE (3).
4. Convergence to the mean-field limit
Let W SHB ( k ) =
(cid:0) [Equation]
w
[Equation] ] (cid:1) be obtained via the SHB iteration (2).
Before stating our result, let us discuss how to couple this SHB dynamics to the mean-field ODE (3).
First, sample a finite neural network w.r.t.
the neuronal embedding, i.e., C 1 ( j 1 ) i.i.d.
∼ P 1 , C 2 ( j 2 ) i.i.d.
∼ P 2 , for j 1 ∈ [ n 1 ] , j 2 ∈ [ n 2 ] and w 1 (0 , · ) , w 2 (0 , · , · ) , w 3 (0 , · ) .
Given w 1 (0 , · ) , w 2 (0 , · , · ) , w 3 (0 , · ) , let the mean-field ODE (3) evolve up to some t , thus obtaining w 1 ( t, · ) , w 2 ( t, · , · ) , w 3 ( t, · ) .
Next, initialize the weights corresponding to the SHB evolution according to the initialization of the mean-field ODE, i.e., w SHB(0 , j 1 ) =
w 1 (0 , C 1 ( j 1 )) , w SHB(0 , j 1 , j 2 ) =
w 2 (0 , C 1 ( j 1 ) , C 2 ( j 2 )) and w SHB(0 , j 2 ) =
w 3 (0 , C 2 ( j 2 )) , and let them evolve according to SHB dynamics (2), thus obtaining w SHB( k, j 1 ) , w SHB( k, j 1 , j 2 ) , w SHB( k, j 2 ) for k = ⌈ t/ϵ ⌉ .
Finally, define the following distance metric that measures the difference between the mean-field and the SHB dynamics:
[Equation]
Theorem 2 Let Assumptions (A1)-(A2) hold.
Consider the coupled SHB dynamics (2) and meanfield ODE (3), and the distance metric (6).
Then, with probability at least 1 − exp( − δ 2 ) ,
[Equation]
where n max = max { n 1 , n 2 } , n min = min { n 1 , n 2 } , and K ( γ, T ) is a constant depending only on γ, T .
The proof of Theorem 2 is deferred to Appendix F.
This result shows that the approximation error between the SHB dynamics and the mean-field limit vanishes as n 1 , n 2 grow large, with n max
=
o ( e n min ) and ε
= o (1 / √ D + log n 1 n 2 ) .
The bound (7) is dimension-free , in the sense that n 1 , n 2 do not need to scale with the input dimension D .
The constant K ( γ, T ) scales rather poorly in T , i.e., K ( γ, T )
=
O ( e e T ) , as common in existing mean-field results [25–28].
An interesting open problem is to improve such dependence, e.g., by using ideas from [32].
An application of Theorem 2 consists in characterizing the dropout-stability and connectivity of the solutions found by SHB.
Given two non-empty sets A 1 ⊆ [ n 1 ] , A 2 ⊆ [ n 2 ] , we say that W is ϵ D -dropout stable if
[Equation]
where R drop ( W ; A 1 , A 2 ) is obtained by replacing the two-layer network f ( x ; W ) defined in (1) with the dropout network
[Equation]
Furthermore, two solutions W and W ′ are ϵ C -connected if there exists a continuous path in parameter space that starts at W , ends at W ′ and along which the risk R ( · ) is upper bounded by max { R ( W ) , R ( W ′ ) } + ϵ C .
The connectivity of solutions obtained via gradient descent methods has been empirically observed in [10, 14], and it has been related to dropout-stability in [21].
The fact that SGD solutions enjoy dropout-stability and connectivity properties has been proved in [34] and, by combining this analysis with Theorem 2, similarly strong guarantees can be obtained for heavy ball methods.
In particular, after k ≤ ⌊ T/ε ⌋ steps of the iteration (2), the resulting parameters are ϵ D -dropout stable and ϵ C -connected, where
[Equation]
with probability at least 1 − exp( − δ 2 ) .
Here, A max = max {| A 1 | , | A 2 |} , A min = min {| A 1 | , | A 2 |} , n max = max { n 1 , n 2 } and n min = min { n 1 , n 2 } .
The path connecting the two solutions is piecewise linear, and it can be explicitly constructed as in [21, 34].
Finally, let us highlight that our mean-field perspective can shed light on the thought-provoking conjecture of [12], where it is empirically observed that, after a suitable permutation, the solutions of the optimization algorithm enjoy linear connectivity.
In fact, Theorem 2 shows that, by running the SHB training algorithm (2) multiple times, all the resulting solutions satisfy (7).
This readily implies that, after a permutation of the neurons, the distance between such solutions can also be upper bounded by the RHS of (7).
5. Global convergence of mean-field dynamics for three-layer network
In order to show the global convergence result, we first need to make some extra assumptions.
(B1)
The activation σ 1 exhibits a universal approximation property, i.e., { σ 1 ( ⟨ w , ·⟩ ) : w ∈ R D } has dense span in L 2 ( D x ) , where D x denotes the x -marginal of the data distribution D .
(B2) ρ 10 has full support.
(B3)
The mean-field ODE (3) converges to the limit (cid:0) w 1 ( ∞ , c 1 ) , w 2 ( ∞ , c 1 , c 2 ) ,
w 3 ( ∞ , c 2 ) (cid:1) s.t.
Pr [ w 3 ( ∞
C
̸ = 0] > 0 .
Formally, we have that, as t → ∞ ,
[Equation]
The universal approximation property is the key assumption to obtain a global convergence result.
This requirement is mild, since most activation functions used in practice are universal approximators.
The assumption on full support is also mild, since widely used initialization schemes (e.g.,
He’s or LeCun’s initialization) employ a Gaussian distribution, which indeed has full support.
The assumption on the mode of convergence is purely technical, and it is an open question whether it can be relaxed.
We remark that these requirements also appear in [28], with the exception of
Pr [ w 3 ( ∞ , C 2 ) ̸ = 0] > 0 , which is needed to handle the heavy ball dynamics.
Theorem 3 Let Assumptions (A1)-(A2) and (B1)-(B3) hold, and assume further that R ( y, f ( x ; W )) is convex in f ( x ; W ) .
Let W ( t ) be the solution of the mean-field ODE (3).
Then, we have that
[Equation]
The detailed proof is deferred to Appendix G and we provide here a sketch.
First, we show a degenerate property for the mean-field ODE, i.e., there exist deterministic functions w ∗ 1 ( · , · ) :
R [-≥ 0 × R D −→ R D , w ∗ 2 ( · , · , · , · ) : R ≥ 0 × R D × R × R −→ R , w ∗ 3 ( · , · ) : R ≥ 
-][Equation]
R such that
[Equation]
Next, we show that, for any finite t , w ∗ 1 ( · , · ) is continuous in both arguments and w 1 ( t, C 1 ) is full support.
Finally, the convergence to the global minimum is obtained by combining the argument that w 1 ( t, C 1 ) is full support for all finite t with the mode of convergence assumption.
Theorem 3 is rather different from the global convergence result for the heavy ball method presented in [20].
In fact, [20] consider noisy dynamics (i.e., with additive isotropic noise), and show the convergence of the mean-field ODE to the global minimum of a certain free energy, which represents an entropic regularization of the loss function.
In this setup, the convergence is guaranteed by the noise term in the dynamics and by the regularization term in the free energy functional.
In contrast, we consider noiseless dynamics and do not prove its convergence.
Instead, we show that, when the mean-field ODE converges, it must do so towards the global minimum of the unregularized loss function.
At the technical level, our proof strategy is an adaptation to the heavy ball case of the argument for SGD in [28], which also crucially relies on the universal approximation property of the activation function.
A similar idea was first proposed in [23], and it also appears in [13].
However, our contribution is the first to tackle the case of optimization with momentum.
Appendix A. Details on training dynamics
We recall that the training data z
= ( x , y ) is generated i.i.d.
from a distribution D .
The neural network is trained to minimize the population risk function R ( W ) =
E z
[ R ( y, f ( x ; W ))] via the following one-pass stochastic heavy ball (SHB) method:
[Equation]
where we use (cid:98) ∇ W R
( y ( k ) , f ( x ( k ); W ) to denote the scaled gradient, and the scaling factors for each parameters are specified below.
This is a one-pass method in the sense that, at each step, we sample a new data point z ( k ) independent from the previous ones.
In order to define a continuous-time ODE for the heavy ball method, we pick β = (1 − γε ) and η = ε 2 , so the one-pass SHB method can be equivalently written as follows:
[Equation]
which is exactly the form in (2).
The formulation in [35] is similar: in [35, Eq. 1.2], β = 1 − γϵ 1+ γϵ , while here we let β = 1 − γε ; hence, the two choices are basically the same when ε is small.
The corresponding continuous ODE, also studied in [20, Eq. 6], is given by
[Equation]
We remark that there are different ways to derive a continuous dynamics from (12), and (2) is obtained by applying the Euler scheme based on the second-order Taylor expansion.
The corresponding ODE (14) is denoted as the low-resolution ODE in [35].
It is an interesting and challenging task to analyze other types of ODEs associated to the SHB method, for example the high-resolution
ODE proposed in [35].
We leave this to future works.
We also remark that similar formulation of the continuous counterpart of heavy ball methods with fixed momentum are studied in [19, 22].
We conclude this part by discussing the scaling factors for the gradient in (2):
[Equation] [Equation]
where
[Equation]
In words
[Equation]
the scaling factor is n 2 for the third layer
[Equation]
n 1 × n 2 for the second layer
[Equation]
and n 1 for the first layer.
Hence
[Equation]
the SHB dynamics can be expressed in the following more explicit form:
[Equation]
Notation.
In the following sections, we will use w 1 ( t, j ) or w 1 ( k, j ) to represent the weights at time t or time step k .
The same notations also applies to w 2 , w 3 .
For convenience, we will also use the lighter notation
[Equation]
to denote the quantities defined in (5).
Appendix B. Results for two-layer networks B
1. Derivation of the mean-field limit
We consider a two-layer neural network with n neurons and input x ∈ R D :
[Equation]
Here, we use the short-hand [ n ] := { 1 , . . . , n } and, for j ∈ [ n ] , the parameters of the j -th neuron are denoted by θ ( j )
= ( w 1 ( j ) , w 2 ( j )) , with w 1 ( j ) ∈ R D and w 2 ( j ) ∈ R .
The parameters are updated according to (2), with
[Equation]
where
[Equation]
We make the following assumptions:
(C1)
There exists a universal constant K > 0 such that ∥ σ ∥ ∞ , ∥ σ ′ ∥ ∞ , ∥ σ ′′ ∥ ∞ ≤ K .
The data distribution D is such that, almost surely, | y | , ∥ x ∥ 2 ≤ K .
Furthermore, | ∂ 2 R ( y, f ( x ; W )) | is K -Lipschitz continuous in f ( x ; W ) and K -bounded for any W .
(C2)
At initialization, w 1 (0 , j ) , w 2 (0 , j ) i.i.d.
∼ ρ 0 , where ρ 0 is such that w 1 (0 , j ) is K 2 -sub-Gaussian, and | w 2 (0 , j ) | ≤ K almost surely.
Furthermore, r (0) = 0 .
The idea of defining the mean-field limit in two-layer case is that the output of the network can be viewed as an expectation over the empirical distribution of the weights, that is:
[Equation]
where σ ⋆ ( x ; θ ( j ))
= w 2 ( j ) σ ( w 1 ( j ) T x ) and (cid:98) ρ θ = 1 n (cid:80) nj
=
δ
θ ( j ) .
Thus, the evolution of the parameters θ ( t ) according to (14) can be viewed as the evolution of (cid:98) ρ θ ( t ) according to a certain distributional dynamics induced by (14).
Since we assume i.i.d. initialization, as the number of neurons n −→ ∞ , we expect that (cid:98) ρ θ (0) −→ ρ 0 .
In this limit, the distributional dynamics induced by (14), can be described by a certain PDE, with initial condition ρ 0 .
Let
[Equation]
Then, we define the mean-field PDE associated to the heavy ball method as
[Equation]
The existence and uniqueness of the solution of (22) is given by the following result, which is proved in Appendix D.1.
Theorem 4 Under Assumptions (C1)-(C2), there exists a unique solution of the mean-field PDE (22).
B.2. Convergence to the mean-field limit
We recall that the mean-field PDE is defined in (22), and the SHB dynamics can be expressed as
[Equation]
[Equation]
where θ SHB ( k, j ) denotes the parameter associated to the j -th neuron at step k , (cid:98) Ψ is defined in (21), and ρ θ SHB ( k ) =
1 n (cid:80) nj =1 δ θ SHB ( k,j ) denotes the empirical distribution of the parameters { θ SHB ( k, j ) } j ∈ [ n ] .
We couple the mean-field PDE (22) and the SHB dynamics (23), in the sense that they share the same initialization: θ (0) ∼ ρ θ (0) and θ SHB (0 , j ) i.i.d.
∼ ρ θ (0) .
Let us define the following distance metric that measures the difference between the mean-field dynamics and the SHB dynamics:
[Equation]
Theorem 5 Let Assumptions (C1)-(C2) holds.
Consider the mean-field PDE (22), the SHB dynamics (23) and the distance metric (24).
Then, with probability at least 1 − exp( − δ 2 ) ,
[Equation]
where K ( γ, T ) is a constant depending only on γ, T .
We remark that the RHS of (25) is also an upper bound on sup t ∈ [0 ,T ] W 2 ( ρ θ ( t ) , ρ θ SHB ( ⌊ t/ε ⌋ )) , which follows directly from the definition of the Wasserstein W 2 distance.
The guarantees of Theorem 5 are similar to those of Theorem 2, and considerations analogous to those at end of Section 4 can be done as concerns the dropout-stability and connectivity of the solutions found by SHB for two-layer networks.
Appendix C. A-priori estimates C.1. Two-layer networks
Lemma 6 Assume that (C1)-(C2) hold, and let f ( x ; ρ ) , Ψ( θ ; ρ ) , ∇ θ Ψ( θ ; ρ ) be defined in (21).
Then, for any fixed T
[Equation]
there exist universal constants K, K 2 ( γ, T ) , where the latter depends only on γ, T , such that the following results hold
(Boundedness)
We have that, for any θ , ρ ,
[Equation]
(Boundedness for mean-field ODE)
We have that, for any t ≤ T , w 2 ( t ) as governed by (22)
satisfies
[Equation]
(Lipschitz continuity):
[Equation]
Proof 1.
By the definition and assumption (C1), we have that
[Equation]
By writing down the integral form of the ODE, we have
[Equation]
By setting K 2 ( γ, T ) :
= ( K + KT + KT 2 ) e γT , the proof of (27) is complete.
For the Lipschitz continuity argument, we have
[Equation]
Thus
[Equation]
We define the Bounded Lipschitz (BL) divergence as follows:
[Equation]
We have the following relationship between the BL-divergence and the Wasserstein distance (see for example [7, Appendix A.1] for more details):
[Equation]
Hence
[Equation]
which implies that the RHS of (30) is upper bounded by
[Equation]
This concludes the proof of (28).
The Lipschitz continuity of ∇ θ Ψ( θ ; ρ ) follows from the same argument
[Equation]
C.2. Three-layer networks
Lemma 7 Assume that (A1)-(A2) hold and let H 2 , f, ∆ W 1 , ∆ W 2 , ∆ W 3 , ∆ H 1 , ∆ H 2 be defined in (4) and (5).
Then, for any fixed T and given a neuronal embedding { (Ω 1 × Ω 2 , F 1 × F 2 , P 1 × P 2 ) , w 1 (0 , · ) , w 2 (0 , · , · ) , w 3 (0 , · ) } , there exists a universal constant K and universal constants K 3 , 2 ( γ, T ) , K 3 , 3 ( γ, T ) only depending on γ, T such that the following results hold.
(Boundedness)
We have that, for any W , z , for any t ∈ [0 , T ] and for any c 1 ∈ Ω 1 , c 2 ∈ Ω 2 ,
• | f
( x ; W ( t )) | ≤ K ess sup C| w 3 ( t, C 2 ) | • H 2 ( x , c 2 ; W ( t ))
| ≤ K ess sup C 1 ,C| w 2 ( t, C 1 , C 2 ) | • | ∆ W 3 ( z , c 2 ; W ( t ))
| ≤ K • | ∆ H 2 ( z , c 2 ; W ( t )) | ≤ K ess sup C | w 3 ( t, C 2 ) | • | ∆ W 2 ( z , c 1 , c 2 ; W ( t ))
| ≤ K (cid:32) ess sup C 1 ,C| w 2 ( t, C 1 , C 2 ) | (cid:33)
ess sup C| w 3 ( t, C 2 ) | • | ∆ H 1 ( x , c 1 ; W ( t ))
| ≤ K ess sup C | w 3 ( t, C 2 ) | ess sup C 1 ,C| w 2 ( t, C 1 , C 2 ) | • ∥ ∆ W 1 ( x , c 1 ; W (
[Equation]
(Boundedness for mean-field ODE)
We have that, for any t ≤ T ,
[Equation]
(Lipschitz continuity)
We have that, for any t ≤ T , •
[Equation]
t ))
[Equation] [Equation]
[Equation]
• | ∆ H 2 ( z , c 2 ; W ( t ))
− ∆ H 2 ( z , c 2 ; ˜ W ( t ))
[Equation] [Equation]
Proof
By the definition and assumption (A1), we have • | f
( x ; W ( t )) | = | E C 2 w 3 ( t, C 2 ) σ 2 ( H 2 ( t, x , C 2 ))
|
[Equation] [Equation]
• | H 2 ( x , c 2 ; W ( t ))
| =
| E C 1 w 2 ( t, C 1 , c 2 ) σ 1 ( H 1 ( t, x , C 1 )) | ≤
[Equation]
sup
[Equation] [Equation]
t
[Equation]
C
[Equation]
C
[Equation]
) σ 1 ( H 1 ( x ,
[Equation]
W (
[Equation]
|
[Equation]
• | ∆ W 3 ( x , c 2 ; W ( t ))
| = | ∂ 2 R ( y ; f ( x ; W ( t )))
σ 2 ( H 2 ( x , c 2 ; W ( t )))
| ≤ K • | ∆ H 2 ( x , c 2 ; W ( t ))
| = | ∂ 2 R ( y ; f ( x ; W ( t )))
w 3 ( t, c 2 ) σ ′ 2 ( H 2 ( x , c 2 ; W ( t )))
| ≤ ess sup C| w 3 ( t, C 2 ) σ ′ 2 ( H 2 ( x , C 2 ; W ( t )))
| ≤ K ess sup C
[Equation]
w 3 ( t, C 2 ) | • | ∆ W 2 ( x , c 1 , c 2 ; W ( t ))
| = | ∆ H 2 ( x , c 2 ; W ( t ))
σ 2 ( H 1 ( x , c 1 ; W ( t )))
|
[Equation] [Equation]
• | ∆ H 1 ( x , c 1 ; W ( t ))
| =
| E C 2 ∆ H 2 ( x , C 2 ; W ( t ))
w 2 ( t, c 1 , C 2 ) σ ′ 1 ( H 1 ( x , c 1 ; W ( t )))
| ≤ ess sup C| ∆
H 2 ( x , C 2 ; W ( t ))
| ess sup C 1 ,C| w 2 ( t, C 1 , C 2 )
| ess sup C | σ ′ 1 ( H 1 ( x , C 1 ; W ( t )))
| ≤ K ess sup C| w 3 ( t, C 2 ) | ess sup C 1 ,C| w 2 ( t, C 1 , C 2 ) | • ∥ ∆ W 1 ( x , c 1 ; W ( t ))
∥ 2 = ∥
∆ H 1 ( x , c 1 ; W ( t ))
x ∥ 2 ≤ ess sup
[Equation]
[Equation]
which readily gives the first claim.
Next, we write
[Equation]
which by Gronwall’s lemma
[Equation]
implies that
[Equation]
For the Lipschitz continuity argument, we have
[Equation]
• | f
( x ; W ( t )) − f ( x ; ˜ W ( t ))
[Equation]
• | ∆ H
[Equation] [Equation]
, c
[Equation] [Equation]
t )
[Equation]
• | ∆ W 2 ( z , c 1 , c 2 ; W ( t ))
[Equation]
∆ W 2 ( z , c 1 , c 2 ; ˜ W ( t )) |
[Equation]
[Equation]
Appendix D. Existence and uniqueness of the mean-field limit D.1. Two-layer networks
In this section
[Equation]
we prove the existence and uniqueness of the mean-field limit for two-layer networks.
We recall the mean-field ODE again here:
[Equation]
The proof follows from constructing a Picard type of iteration, similarly to [36, Section 4], [17,
Theorem C
[Equation]
Below is an adaptation of the strategy in [39, Theorem 1.1].
We first write the integral form of the mean-field ODE:
[Equation]
where ρ ( t ) is the law of ( θ ( t ) , r ( t )) , and we use ρ θ ( t ) , ρ r ( t ) to denote the θ and r marginals, respectively.
We define the space P 2 ( R D × R D ) to be the space of probability measures on R D × R D equipped with Wasserstein metric W 2 , and we have ρ ( t ) ∈
P 2 ( R D × R D ).
We define the space
C (cid:0) [0 , T ] , P 2 ( R D × R D ) (cid:1) to be the space of continuous maps ρ ( · ; T ) : [0 , T ] → P 2 ( R D × R D ) .
We omit T when there’s no confusion.
The space is equipped with the following metric: d T ( ρ 1 , ρ 2 ) =
sup t ∈ [0 ,T ] W 2 ( ρ 1 ( t ) , ρ 2 ( t )) .
Note that the space (cid:0) P 2 ( R D × R D ) , W 2 (cid:1) is a complete space [2, Theorem 8.7].
Thus for any fixed 0 < T < ∞ , the space (cid:0) C (cid:0) [0 , T ] , P 2 ( R D × R D ) (cid:1)
× d T (cid:1) is also complete.
Next, we define the operator H T ( · , θ (0)) :
C (cid:0) [0 , T ] , P 2 ( R D × R D ) (cid:1)
→ C (cid:0) [0 , T ] , P 2 ( R D × R D ) (cid:1) as follows:
[Equation]
where θ (0) denotes the parameters of the mean-field ODE (33) at initialization, which means that the stochastic process we defined in (35) is coupled with the mean-field ODE.
Note that the ρ θ 1 ( t ) in (35) is no longer the law of (cid:101) θ ( t ) , but the input distribution.
We use
H T ( ρ ( t )) to denote H T ( ρ ; θ (0))
( t ) , that is the distribution of the solution (35) at time t .
We omit θ (0) when there is no confusion.
The definition of H T can be interpreted as follows: it maps ρ ∈ C (cid:0) [0 , T ] , P 2 ( R D × R D ) (cid:1) as input to output the law of ( θ ( t ) , r ( t )) which evolves according to the stochastic process induced by the probability measure ρ ( t ) .
It is easy to see that the fixed point of H T is the solution of the non-linear dynamics (33).
Thus, our goal is to show that there exist a T 0 such that H T 0 has unique fixed point, or equivalently that
H T 0 is a strict contraction.
Proposition 8 Under Assumptions (C1)-(C2), there exists a T 0 only depending on K, γ and a
C ( T 0 ) ∈ (0 , 1) such that, for all ρ 1 , ρ 2 ∈ C (cid:0) [0 , T 0 ] , P 2 ( R D × R D ) (cid:1) with the same initialization ( θ 1 (0) , r 1 (0))
= ( θ 2 (0) , r 2 (0)) , we have:
[Equation]
Proof We first fix any 0 < T < ∞ , and the space C (cid:0) [0 , T ] , P 2 ( R D × R D ) (cid:1) .
Given ρ 1 , ρ 2 ∈
C (cid:0) [0 , T ] , P 2 ( R D × R D ) (cid:1) , we define two dynamics as follows:
[Equation]
where θ 1 ( t ) = (cid:16) w (1)1 , w (1)2 (cid:17) and θ 2 ( t ) = (cid:16) w (2)1 , w (2)2 (cid:17) .
We want to upper bound the difference between these two dynamics, which will give us an upper bound on
[Equation]
For all t ∈ [0 , T ] , we have
[Equation]
Now, by Lemma 6, we have that ∥∇
[Equation]
where θ i ( t ) =
( w ( i ) 1 ( t ) , w ( i ) 2 ( t )) , i ∈ 1 ,
Thus we have that:
[Equation]
Similarly for ∥ r 1 ( t ) − r 2 ( t ) ∥
[Equation]
we have that:
[Equation]
Putting these two results together we have:
[Equation]
t ∈ [0 ,T ] W 2 ( ρ θ 1 ( t )
[Equation]
ρ θ 2 ( t ))
+ 2 γ (cid:90) t (cid:13)(cid:13)(cid:13)(cid:13)(cid:18) θ
[Equation]
( s )
[Equation]
( s ) (cid:19) − (cid:18)
[Equation]
( s ) r
[Equation]
( s ) (cid:19)(cid:13)(cid:13)(cid:13)(cid:13) 2 ds + 4 K (1 + K 2 ( γ, T )) (cid:90) t (cid:90)
[Equation]
(cid:13)(cid:13)(cid:13)(cid:13)(cid:18)
[Equation]
(cid:19)
[Equation]
(cid:19)(cid:13)(cid:13)(cid:13)(cid:13)
[Equation]
du ds
By Corollary 27, we have that:
[Equation]
≤ 4 K
[Equation]
1 + K
[Equation] [Equation]
(cid:18)
[Equation]
+ exp
[Equation]
2 γ (cid:19)(cid:19) max t
[Equation] [Equation]
] W
[Equation]
ρ θ 1 ( t )
[Equation]
ρ θ
[Equation]
( t ))
Thus, we have that (cid:13)(cid:13) (cid:13)(cid:13) (cid:18) θ 1 ( t ) r 1 ( t ) (cid:19) − (cid:18) θ 2 ( t ) r 2 ( t )
[Equation]
which implies that
[Equation]
where we set K ( γ, T ) = 4 K (1 + K 2 ( γ, T )) (cid:16) 1 + exp 4 γ 2 +4 K
(1+ K 2 ( γ,T ))
T 4 γ (cid:17) .
Let C ( T )
=
T 2 K ( γ, T ) .
Then, we could always find a T 0 such that C ( T 0 ) < 1 since C (0) =
and C ( T ) is continuous in T , which finishes our proof.
By Banach’s fixed point theorem, there exist a T 0 > 0 such that the mean-field ODE has a unique solution in time interval [0 , T 0 ] .
Now, we show the existence and uniqueness of the solution of the mean-field ODE for any time period [0 , T ] .
Theorem 9 Under Assumptions (C1)-(C2), for any T > 0 , there exists a unique solution for the mean-field ODE (22) in the interval [0 , T ] .
Proof The idea is to separate the time interval [0 , T ] into subintervals of length T 0 , that is, we consider the intervals [0 , T 0 ] , [ T 0 , 2 T 0 ] , ..., [ ⌊ TT 0 ⌋ T 0 , T ] .
Note that the contraction property we proved in
Proposition 8 only depends on the length of the time interval, so the proof can be done recursively.
That is:
In the interval [0 , T 0 ] , (22) with initialization ( θ (0) , r (0)) has a unique solution { ρ ( t ) } t ∈ [0 ,T 0 ] .
In the interval [ T 0 , 2 T 0 ] , we consider (22) with initial distribution ρ ( T 0 ) , and it has a unique solution { ρ ( t ) } t ∈
[ T 0 , 2 T 0 ] .
Recursively do the above steps until the interval [ ⌊ TT 0 ⌋ T 0 , T ] .
Thus we have that, for any T > 0 , there exists a unique solution for (22) in the interval [0 , T ] .
D.2. Three-layer networks
In this section, we prove the existence and the uniqueness of the mean-field ODE (3).
The integral form of the mean-field ODE is given by
[Equation]
In order to prove the existence and the uniqueness
[Equation]
we follow the same Picard’s iteration arguments as for the two-layers case.
Given a neuronal embedding
[Equation]
, we first define the following norm:
[Equation]
where C 1 i
[Equation]
i.d ∼ P 1 , C 2 i.i.d ∼ P
Next
[Equation]
we define the following metric for two sets of mean-field parameters:
[Equation]
[Equation]
Note that the metric we define above is not the metric induced by the norm, since in the definition of the norm we only require the boundedness of w 2 and w 3 .
We define the following functional space of the mean-field parameters:
[Equation]
which means that all the (cid:102) W ∈ W T ( W (0)) have the same initialization W (0) .
By Lemma 7, we know that:
[Equation]
It is easy to see that the space W T ( W (0)) is complete w.r.t.
the metric D T ( W , (cid:102) W ) .
Let us define the operator: H T : W T ( W (0))
−→ W T ( W (0))
as follows:
Input
[Equation]
Output: { (cid:102) W ( t ) } t ∈ [0 ,T ] , such that:
[Equation]
We aim to show the following proposition.
Proposition
Under Assumptions (A1)-(A2), there exists a T 0 only depending on K, γ and
C ( T 0 ) ∈ (0 , 1) , such that, for all W 1 , W 2 ∈ W T ( W (0)) , we have:
[Equation]
Proof For simplicity of notation, we denote the output of H T ( W 1 ) to be (cid:102) W 1 , which is composed of (cid:101) w 13 , (cid:101) w 12 , (cid:101) w 11 .
The output of H T ( W 2 ) is denoted similarly.
By the definition of the mean-field ODE, we have that, for any t ≤ T ,
[Equation]
By Lemma
[Equation]
we have that:
[Equation]
Thus, we have:
[Equation]
which implies that
[Equation]
Since ( γT + T 2 ) K ( γ, T ) = 0 when T = 0 , and ( γT + T 2 ) K ( γ, T ) is continuous in T , we can pick a T 0 such that ( γT 0 + T 20 ) K ( γ, T 0 ) < 1 , which finishes the proof.
Since W T ( W (0)) is complete, by the Banach fixed point Theorem, there exists a unique fixed point for the operator H T 0 , which implies that the mean-field ODE (3) has a unique solution in [0 , T 0 ] .
By following the same argument of the proof of Theorem 9 (separate the interval [0 , T ] into sub-intervals of length T 0 and successively apply Proposition 10 to each of them), we readily obtain our main result concerning the existence and uniqueness of (3) in [0 , T ] .
Appendix E. Convergence to the mean-field limit – Two-layer networks
In this section, we prove the convergence to the mean-field limit for two-layer neural networks (Theorem 5).
Our proof’s structure is inspired from [26].
Before going into the arguments, we first recall the definition of the mean-field ODE and the stochastic heavy ball method (SHB) for twolayer networks.
Then, we define two auxiliary dynamics: the particle dynamics (PD) and the heavy ball dynamics (HB).
First, recall the mean-field ODE as follows:
[Equation]
) and the corresponding integral form
[Equation]
The SHB dynamics is as follows:
[Equation]
where ρ θ SHB ( k ) =
1 n (cid:80) nj =1 δ θ ( k,j ) is the empirical measure.
In order to describe the convergence to mean-field limit, we define the following particle dynamics (PD):
[Equation] [Equation]
where ρ θ PD ( t )
= 1 n (cid:80) nj =1 δ θ PD ( t,j ) is the empirical distribution at time t .
Furthermore, the heavy ball (HB) dynamics is defined as
[Equation]
We remark that (52), (53) and (54) have the same initialization
[Equation]
that is :
[Equation]
Define the following distance metrics:
[Equation]
E.1. Bound between mean-field ODE and particle dynamics
In this section, we bound the difference between the mean-field ODE defined in (51) and the particle dynamics defined in (53), whose integral form is as follows:
[Equation]
Proposition 11 Under Assumptions (C1) - (C2), we have that, with probability at least 1 − exp( − δ 2 ) ,
[Equation]
where K ( γ, T ) is a constant depending only on γ, T .
Before proving Proposition 11, we first prove the following lemma, which characterizes the
Lipschitz continuity of the mean-field ODE and the particle dynamics.
Lemma 12 Under Assumptions (C1) - (C2), there exists a universal constant K ( γ, T ) depending only on γ, T such that, for any t, τ > 0 such that t, t + τ < T ,
[Equation]
The same holds for the particle dynamics θ PD ( t, j ) , ∀ j ∈ [ n ] .
Proof We only prove the results for the mean-field ODE, and the proof for the particle dynamics follows from the same arguments.
We first try to bound the increments ∥ θ ( t ) − θ (0) ∥ 2 .
By the definition of the mean-field dynamics, we have that:
[Equation]
where in the last step we use that ∥∇ θ Ψ( θ ( u ); ρ θ ( u ))
∥ 2 ≤ K 1 ( γ, T ) , which follows from Lemma 6.
By Gronwall’s lemma, this implies that, for any t ≤ T ,
[Equation]
Next, by definition of the mean-field ODE, we have that:
[Equation]
Thus,
[Equation]
where we use that fact that
[Equation]
For the Lipschitz continuity of ρ θ , we just note that by definition of the W 2 distance, we have:
[Equation]
Now we are ready to prove Proposition 11.
Proof In order to bound the difference, we first define n i.i.d mean-field dynamics:
[Equation]
where ρ θ ( t, j ) is the law of θ ( t, j ) , and we coupled the n i.i.d dynamics with the particle dynamics at initialization, that is, we let:
[Equation]
We also define the empirical distribution of θ ( t, j ) , that is: (cid:98) ρ θ ( t ) = 1 n (cid:80) nj =1 δ θ ( t,j ) .
Since the n mean-field dynamics are i.i.d, we have that ρ θ ( t, j )
=
ρ θ ( t ) , ∀ j ∈ [ n ] , thus we use the notation of ρ θ ( t ) to denote the the law of θ ( t, j ) for each j ∈ [ n ] .
By Lemma 6 and Lemma 12, we know that:
[Equation]
We have that
[Equation]
and our goal is to bound:
[Equation]
Now we aim to bound the quantity
[Equation]
An application of the triangle inequality gives
[Equation]
Recall by definition that
[Equation]
Similarly
[Equation]
∇
[Equation]
For the term in (60)
[Equation]
we use concentration inequalities to give an upper bound.
By the Lipschitz continuity of ∂ 2 R in Assumption (C1), we have
[Equation]
For the term | f ( x ; ρ ( u )) − f ( x ; (cid:98) ρ θ ( u )) | , we have |
[Equation]
Note that, by Lemma 6, we know that
[Equation]
By Lemma 24, we have that, with probability at least 1 − exp( − n ( δ ′ ) 2 ) ,
[Equation]
n n (cid:88)
[Equation]
By Lemma
[Equation]
we know that
[Equation]
is K ( γ, T ) -Lipschitz continuous in u .
Thus, by taking a union bound over j ∈ [ n ] and t ∈ (cid:110) 0 , η, ..., ⌊ Tη ⌋ η (cid:111) , we have that, with probability at least 1 − nTη exp( − n ( δ ′ ) 2 ) ,
[Equation]
Take η = 1 √ n , δ ′
=
(cid:114)
δ 2 +log (cid:16) n 32 T (cid:17) n .
Then, with probability at least 1 − exp( − δ 2 ) , max j ∈ [ n ] sup t ∈
[0 ,T ] (cid:12)(cid:12)(cid:12)(cid:12) (cid:12) n n (cid:88) i =
[Equation]
which implies that, for term (60), max j
[Equation]
with probability 1 − exp( − δ 2 ) .
For the term in (61), we use the Lipschitz continuity of ∇ θ Ψ .
By Lemma 6, we have that, for each j ∈ [ n ] , ∥∇
[Equation]
Note that (cid:98) ρ θ ( t ) , ρ θ PD ( t ) are discrete measures, thus we have:
[Equation]
Hence,
[Equation]
Combining the above results
[Equation]
we have that, with probability 1 − exp( − δ 2 ) ,
[Equation]
An application of Corollary 27 concludes the proof
[Equation]
E.2. Bound between particle dynamics and heavy ball dynamics
In this section, we bound the difference between the particle dynamics defined in (53) and the heavy ball dynamics defined in (54).
We recall that the distance we aim to bound is defined in (56).
Note that the heavy ball dynamics is a discretization of the particle dynamics.
Thus we aim to bound the difference at time point kε .
Proposition 13 Under Assumptions (C1)-(C2), there exist a universal constant K ( γ, T ) depending only on γ, T , such that
[Equation]
Proof By the Taylor expansion
[Equation]
we have the following approximation for the particle dynamics:
[Equation]
Also by Taylor expansion we have:
[Equation]
By plugging (61) into (63)
[Equation]
we have that
[Equation] [Equation]
Now we get a discrete iteration equation for the particle dynamics, with an approximation error of at most O ( ε 3 ) .
By accumulating the ∇ θ Ψ( θ PD ( lε, j ); ρ θ PD ( lε )) term from l = 1 , ..., k , we have
[Equation]
65) where c ( k ) l
[Equation]
The heavy ball dynamics can be written in a similar fashion:
[Equation]
Thus
[Equation]
we have that
[Equation]
By Lemma 6, we have that
[Equation]
Since ρ θ PD ( lε )
[Equation]
ρ θ HB ( lε )) are discrete distributions
[Equation]
we have that
[Equation] [Equation]
[Equation]
which implies that
[Equation]
As a result
[Equation]
we have
[Equation]
Finally
[Equation]
an application of the discrete Gronwall’s lemma concludes the proof
[Equation]
E.3. Bound between heavy ball dynamics and stochastic heavy ball dynamics
In this section, we bound the difference between the heavy ball dynamics defined in (54) and the stochastic heavy ball dynamics defined in (52).
We recall that the distance we aim to bound is defined in (57).
The manipulations of the previous section imply that the heavy ball dynamics can be written as
[Equation]
Similarly
[Equation]
the stochastic heavy ball dynamics can be written as
[Equation]
Proposition 14 Under Assumptions (C1)-(C2), there exists a universal constant K ( γ, T ) depending only on γ, T , such that, with probability 1 − exp( − δ 2 ) ,
[Equation]
Proof By using (67) and (68)
[Equation]
we have
[Equation]
By triangle inequality, we have that
[Equation] [Equation]
For the term in (70)
[Equation]
by the Lipschitz continuity of ∇ θ (cid:98) Ψ
[Equation]
we obtain
[Equation]
Since ρ θ HB , ρ θ SHB are discrete distributions, we have that
[Equation]
Thus,
[Equation]
For the term in (71)
[Equation]
note that, since the z ( l ) ’s are sampled i.i.d.
at each step by definition
[Equation]
we have
[Equation]
Thus,
[Equation]
is a martingale difference.
By Lemma 6, we have that, for all l ∈ { 1 , ...,
⌊ T/ε ⌋} ,
[Equation]
Hence
[Equation]
an application of Lemma 25 gives that
[Equation]
with probability at least 1 − exp( − δ 2 ) , max
[Equation]
∈
[Equation]
., ⌊ T/
[Equation]
By taking a union bounds on j ∈ [ n ] , we have that, with probability at least 1 − exp( − δ 2 ) ,
[Equation]
By combining the above result, we conclude that
[Equation]
Finally, an application of the discrete Gronwall’s lemma concludes the proof.
E.4. Proof of Theorem
Proof The proof follows from combining Proposition 11, 13, 14, and the fact that:
[Equation]
Appendix F. Convergence to the mean-field limit – Three-layer networks
In this section, we prove the convergence of the training dynamics to the mean-field limit for a three-layer neural network.
Our proof’s structure is inspired from [28].
Before going into the proofs, let’s first recall the definition of the mean-field ODE and the
SHB dynamics, and then define two auxiliary dynamics, namely the HB dynamics and the particle dynamics.
For the convenience of further computation, we define these continuous dynamics in integral form.
We define the random variable corresponding to the stochastic heavy ball dynamics, the heavy ball dynamics, the particle dynamics, and the mean-field ODE as W SHB , W HB , W PD , W respectively.
The mean-field ODE (3) in integral form is the following:
[Equation]
The SHB dynamics is as follows:
[Equation]
where z ( k ) is the data point sampled at time step k .
We define the particle dynamics as a continuous dynamics without mean-field interaction:
[Equation]
We define the HB dynamics by replacing the stochastic gradient in the SHB dynamics by the true gradient.
That is:
[Equation]
Note that the HB dynamics can be viewed as the discrete version of the PD dynamics.
In order to present our theoretical results
[Equation]
we define the following distance metrics to quantify the level of correspondence of these dynamics:
[Equation]
[Equation]
We acknowledge the abuse in notation from reusing D T for multiple metrics, which is done to avoid proliferation of notation.
It is clear that:
[Equation]
where D T ( W , W SHB ) is defined in (24).
In the following subsections, we bound the terms D T ( W , W PD ) , D T ( W PD , W HB ) and finally
D T,ε ( W HB , W SHB ) .
F.1. Bound between mean-field ODE and particle dynamics
In this section, we bound the difference between the mean-field ODE defined in (73) and the particle dynamics defined in (75).
We recall that the distance we aim to bound is defined in (77).
Proposition 15 Under Assumptions (A1)-(A2), we have that, with probability at least 1 − exp( − δ 2 ) ,
[Equation]
where K ( γ, T ) is a universal constant depending only on γ, T , n min = min { n 1 , n 2 } and n max = max { n 1 , n 2 } .
In order to prove Proposition 15, we need the following auxiliary lemma, which characterizes the Lipschitz continuity of the mean-field ODE and of the particle dynamics.
Lemma 16 Under Assumptions (A1)-(A2), there exists a universal constant K ( γ, T ) depending only on γ, T such that, for any t, τ > 0 with t, t + τ ≤ T ,
[Equation]
The same holds for the particle dynamics w PD 1 ( t, j
[Equation]
, w PD 2 ( t,
[Equation]
1 , j
[Equation]
, w PD 3 ( t, j 2 ) .
Proof We do the proof for w 1 ( t, c 1 ) , and the same argument applies to w 2 ( t, c 1 , c 2 ) , w 3 ( t, c 2 ) .
First, we derive a bound on the increments up to time t , ∥ w 1 ( t, c 1 ) − w 1 (0 , c 1 ) ∥ 2 .
By the definition of the mean-field ODE (73), we have that
[Equation]
which implies that
[Equation]
where in the last step we use that
[Equation]
for some constant K 1 ( γ, T ) depending only on γ, T ,
[Equation]
which holds by Lemma
[Equation]
Thus, by Gronwall’s lemma
[Equation]
we have that
[Equation]
Now
[Equation]
by using again the definition of the mean-field ODE (73)
[Equation]
we have that:
[Equation]
where in the second line we use that ∥
E z ess sup c 1 sup u ∈ [0 ,T ] ∆ W 1 ( z , c 1 ; W ( u ))
∥ 2 ≤ K 1 ( γ, T ) by
Lemma 7, and in the last passage we use (83).
By setting K ( γ, T ) =
K 2 ( γ, T ) + K 1 ( γ, T ) T , we obtain the desired result and the proof is complete.
By the above Lemma 16 and Lemma 7, we immediately get the following corollary.
Corollary 17 Under Assumptions (A1)-(A2), there exists a universal constant K ( γ, T ) depending only on γ, T such that, for any t, τ > 0 with t, t + τ ≤ T , the following functions
[Equation]
are K ( γ, T ) -Lipschitz continuous in t .
The same holds for the particle dynamics, i.e., the functions
[Equation]
are K ( γ, T ) -Lipschitz continuous in t .
Proof We do the proof for f ( x ; W ( t )) , and the same argument applies to the other cases.
By
Lemma 7, we have that
[Equation]
By Lemma
[Equation]
we have that
[Equation]
Furthermore
[Equation]
by Lemma
[Equation]
we have that:
[Equation]
Thus
[Equation]
we conclude
[Equation]
which gives the desired result.
Now we are ready to prove Proposition 15.
Proof Let us recall that the quantity ∆ W 3 is defined in (5).
We start with computing the difference in the term ∆ W 3 :
[Equation]
where in the last inequality we use the boundedness and Lipschitz continuity of ∂ 2 R and σ 2 obtained from Lemma
[Equation]
Similarly, for ∆ W 1 , ∆ W 2 , we have that
[Equation]
t, j 1 ) ∥ 2 + K | w 3 ( t, C 2 ( j 2 )) − w PD 3 ( t, j 2 ) | ≤ E z K 3 , 3 ( γ, T ) (cid:0) | f
( x ; W ( t )) − f ( x ; W PD ( t ))
| + | H 2 ( x , C 2 ( j 2 ); W ( t )) − H 2 ( x , j 2 ; W PD ( t )) | (cid:1) + K 3 , 3 ( γ, T ) ∥
w 1 ( t, C 1 ( j 1 )) − w PD 1 ( t, j 1 ) ∥
+ K | w 3 ( t, C 2 ( j 2 )) − w PD 3 ( t, j 2 ) | , (85) and that
[Equation]
Here, we remark that the expectation E C 2 [∆ H 2 ( · , C 2 ; · ) w 2 ( · , · , C 2 )] for the mean-field ODE corresponds to the average 1 n 2 (cid:80) n 2 j 2 =1 ∆ H 2 ( · , j 2 ; · ) w 2 ( · , · , j 2 ) for the particle dynamics.
Now, our goal is to upper bound the following quantities:
[Equation]
To do so
[Equation]
we follow [
[Equation]
Appendix C.2, Proof of Theorem
[Equation]
Claim 2].
Then, we have that, for any δ 1 , δ 2 , δ 3 > 0 , max
[Equation]
with probability at least
[Equation]
By Corollary 17, we know that f
[Equation]
x ; W (
[Equation] [Equation]
H
[Equation]
c
[Equation]
W
[Equation]
E
[Equation]
∆
[Equation]
C
[Equation]
W
[Equation]
t ))
w
[Equation]
c
[Equation]
C
[Equation]
) (cid:3) are K ( γ, T ) -Lipschitz continuous, and the corresponding quantities for the particle dynamics are also K ( γ, T ) -Lipschitz continuous.
Thus, by taking a union bound on t ∈ { 0 , η, ..., ⌊ T/η ⌋}
[Equation]
we have max sup
[Equation]
with probability at least
[Equation]
In particular
[Equation]
we pick
[Equation]
(cid:16)(cid:112) log n
[Equation]
Then, max sup
[Equation]
with probability at least 1 − exp( − δ 2 ) .
Next, we combine (87) with (84), (85) and (86) to provide high-probability bounds on ∆ W 3 , ∆ W 2 , ∆ W 1 .
By recalling the definition of the mean-field ODE (73) and the analogous definition of the particle dynamics (75), we finally obtain that, for all t ≤ T , with probability at least 1 − exp( − δ 2 ) ,
[Equation]
An application of Corollary 27 gives the desired result (81) and concludes the proof.
F.2. Bound between particle dynamics and heavy ball dynamics
In this part we bound the difference between the particle dynamics defined in (75) and the heavy ball dynamics defined in (76).
We recall that the distance we aim to bound is defined in (78).
Proposition 18 Under Assumptions (A1)-(A2), there exists a universal constant K ( γ, T ) depending only on γ, T such that
[Equation]
Proof Note that the heavy ball dynamics is just a discretization of the particle dynamics, so we first bound the difference at each time point kε .
By a second order Taylor expansion, we have the following approximation for the particle dynamics.
We do the computation for w 3 as a representative, and the proofs for w 1 , w 2 are the same.
We have that
[Equation]
Also by Taylor expansion, we have that
[Equation]
By plugging (90) into (89), we obtain w PD 3 (( k + 1) ε
[Equation]
j 2 )
= w PD
[Equation]
( kε
[Equation]
j 2 ) + w PD
[Equation]
( kε
[Equation]
[Equation] [Equation]
+
[Equation] [Equation]
w PD
[Equation]
kε
[Equation] [Equation]
+ O (
[Equation]
) =
[Equation] [Equation]
+
[Equation]
( kε
[Equation] [Equation]
PD
[Equation]
(( k
[Equation]
1) ε
[Equation]
j 2 ) + ( − γ∂ t w PD 3 (
[Equation]
j 2 )
[Equation]
where in the last step we use again (90).
By unrolling the recursion
[Equation]
we can write the particle dynamics in the following form:
[Equation]
where
[Equation]
Similarly for w
[Equation]
w
[Equation]
we have:
[Equation]
We can write analogous expressions for the heavy ball dynamics:
[Equation]
Let us define the following notation, for k ∈ (cid:8) 1 , ...,
⌊ Tε ⌋ (cid:9) ,
[Equation]
Recall that, by construction, w PD 3 (0 , j 2 )
=
w HB 3 (0 , j 2 ) , w PD 2 (0 , j 1 , j 2 ) =
w HB 2 (0 , j 1 , j 2 ) and w PD 1 (0 , j 1 ) =
w HB 1 (0 , j 1 ) for all j 1 , j 2 .
Thus, by computing the difference between w PD 1 , w PD 2 , w PDand w HB 1 , w HB 2 , w HB 3 , we have that D T ( W PD , W HB ; k ) satisfies the following induction inequality:
[Equation]
where we have used the Lipschitz continuity of ∆ W 3 , ∆ W 2 and ∆ W 1 obtained via Lemma 7.
Thus, by the discrete Gronwall’s lemma, we obtain that, for any k ∈ (cid:8) 1 , ...,
⌊ Tε ⌋ (cid:9) ,
[Equation]
Finally, an application of Lemma 16 gives that w PD 1 , w PD 2 , w PDare K 3 ( γ, T ) -Lipschitz continuous in time.
Thus, for any t ≤ T ,
[Equation]
Similar results hold also for | w PD 2 ( t, j 1 , j 2 ) − w HB 2 ( ⌊ t/ε ⌋ , j 1 , j 2 ) | and | w PD 1 ( t, j 1 ) − w HB 1 ( ⌊ t/ε ⌋ , j 1 ) | .
As a result, we conclude that
[Equation]
which gives the desired result.
F.3. Bound between heavy ball dynamics and stochastic heavy ball dynamics
In this part we bound the difference between the heavy ball dynamics defined in (76) and the stochastic heavy ball dynamics defined in (74).
We recall that the distance we aim to bound is defined in (79).
Proposition 19 Under Assumptions (A1)-(A2), we have that, with probability at least 1 − exp( − δ 2 ) ,
[Equation]
where K ( γ, T ) is a universal constant depending only on γ, T .
Before proving Proposition 19, we state and prove a result concerning the boundedness of the
SHB dynamics.
Lemma 20 Under Assumptions (A1)-(A2), we have that, for any k ∈ { 1 , . . .
, ⌊ T ε ⌋} ,
[Equation]
where K is a universal constant.
Proof By following passages analogous to those leading to (91), we have that the SHB dynamics can be written as
[Equation]
Recall that
[Equation]
which implies that
[Equation]
Thus
[Equation]
we have
[Equation]
where in the last step we use that kε ≤ T .
For | w SHB( k, j 1 , j 2 ) | , we recall that
[Equation]
Thus, we have
[Equation]
which gives the desired result.
At this point, we are ready to prove Proposition 19.
Proof Throughout this argument, we fix ε and consider k ∈
⌊ Tε ⌋ .
Recall that the HB and SHB dynamics can be written as in (91) and (94), respectively.
Furthermore,
[Equation] [Equation]
Recall that, by construction, w HB 3 (0 , j 2 ) =
w SHB(
[Equation]
j 2 ) , w HB 2 (0 , j 1 , j 2 ) =
w SHB(0 , j 1 , j 2 ) and w HB 1 (0 , j 1 ) =
w SHB(0 , j 1 ) for all j 1 , j 2 .
Thus, by computing the difference between the expressions in (91) and (94)-(95)
[Equation]
we have
[Equation] [Equation]
To bound (96)
[Equation]
(99) and (101)
[Equation]
we use the Lipschitz continuity of ∆ W 3 , ∆ W 2 and ∆ W 1 , together with the fact that c ( k ) l ≤ ε/γ .
In particular,
[Equation]
Similarly, we have (cid:12)(cid:12)(cid:12)(cid:12)(cid:12) k − 1 (cid:88) l =
[Equation]
(cid:1)(cid:12)(cid:12)(cid:12)(cid:12)(cid:12)
[Equation]
(cid:13)(cid:13)(cid:13)(cid:13) (cid:13) k − 1 (cid:88) l
=
[Equation]
To bound (102),(100) and (97), we first define the filtration F 3 ( k ) as the sigma-algebra generated by ( { w 3 (0 , j 2 ) } j 2 ∈ [ n 2 ] , z (0) , ..., z ( k )) .
We define the filtration F 2 ( k ) , F 1 ( k ) in the same way.
Recall that, in a one-pass algorithm, we take i.i.d.
samples at each step and, hence, we can write, for all l ∈ (cid:8) 1 , ...,
⌊ Tε ⌋ (cid:9) ,
[Equation]
Clearly, we have that (cid:8) ∆ W 3 ( z ( l ) , j
[Equation]
W
[Equation]
( l )) , l ∈
[Equation]
..., ⌊ Tε ⌋ (cid:9)(cid:9) are mutually independent, which implies that
[Equation]
is a martingale difference with respect to the filtration F 3 ( l ) .
Thus,
[Equation]
is a martingale (same for ∆ W 2 and ∆ W 1 ).
Next, we show that the martingale differences are bounded, so that we can use martingale convergence results to bound these terms
[Equation]
Combining Lemma 20 with the same strategy of the a-priori estimations of Lemma
[Equation]
we have the following upper bounds:
[Equation]
Thus
[Equation]
an application of Lemma 25 gives
[Equation]
[Table]
[Equation]
[Table]
By taking a union bound over j 1 , j 2 , we have that, with probability at least 1 − exp( − δ 2 ) ,
[Equation]
Combining the results
[Equation]
we conclude that, with probability at least 1 − exp( − δ 2 ) ,
[Equation]
T,ε
[Equation]
An application of the discrete Gronwall’s lemma gives the desired result (93) and concludes the proof.
F.4. Proof of Theorem
Proof The proof follows from combining Proposition 15, 18, 19 and the fact that:
[Equation]
Appendix G. Global convergence of the mean-field ODE
In this section, we aim to prove the global convergence result through the recipe below: 1.
We show the following degenerate property for the mean-field ODE: there exist deterministic functions w ∗ 1 ( · , · ) : R ≥ 0 × R D −→ R D , w ∗ 2 ( · , · , · , · ) : R ≥ 0 × R D × R × R −→ R , w ∗ 3 ( · , · ) :
R ≥ 0 × R −→ R such that
[Equation]
We show that (i) w ∗ 1 ( · , · ) is continuous in both arguments for any finite t , and that (ii) if w 1 (0 , C 1 ) is full support, then w 1 ( t, C 1 ) is full support for any finite t . 3.
Combining the argument that w 1 ( t, C 1 ) is full support for all finite t and the mode of convergence assumption, we show that the mean-field ODE must convergence to the global minimum.
We first show the degenerate property of the mean-field ODE in the following lemma:
Lemma 21 Under Assumptions (A1)-(A2), there exist deterministic functions w ∗ 1 ( · , · ) : R ≥ 0 ×
R D −→ R D , w ∗ 2 ( · , · , · , · ) : R ≥ 0 × R D × R × R −→ R , w ∗ 3 ( · , · ) : R ≥
× R −→ R such that
[Equation]
Proof We follow the proof in [28, Appendix D.2].
To shorten the notations, we make the following definition: we define the sigma-algebras generated by
[Equation]
as S 1 , S 123 , S 3 respectively.
The lemma is equivalent to prove that w 1 ( t, C 1 ) , w 2 ( t, C 1 , C 2 ) , w
[Equation]
are S 1 , S 123 , S 3 -measurable, respectively.
In order to prove the measurability result, we define a reduced dynamics as follows:
[Equation]
Note that the reduced dynamics w RD 1 , w RD 2 , w RDis clearly S 1 , S 123 , S 3 -measurable.
Furthermore, the reduced dynamics is not self-contained, namely, the gradient terms E (cid:2) ∆ W 3 ( z , C 2 ; W ( t )) | S 3 (cid:3) ,
E (cid:2) ∆ W 2 ( z , C 1 , C 2 ; W ( t )) | S 123 (cid:3) and E (cid:2) ∆ W 1 ( z , C 1 ; W ( t )) | S 1 (cid:3) are induced by the mean-field ODE
W ( t ) .
In order to state the next result, we define the following metric :
[Equation]
Next, we aim to show that the reduced dynamics is equivalent to the mean-field ODE, i.e., for any
T > 0 ,
[Equation]
The key step is to prove that
[Equation]
where K ( γ, T ) is a universal constant depending only on T, γ .
Here, | E (cid:2) ∆ W 3 ( z , C 2 ; W ( t )) | S
[Equation]
−
E z ∆ W 3 ( z , C 2 ; W ( t )) | is a random variable, and the ess sup in (106) is taken with respect to it.
The same remark applies to the ess sup in (107) and in (108), which are intended to be taken with respect to the corresponding random variables.
We now prove that (106) holds.
Note that
[Equation]
Using the Lipschitz-continuity of ∆ W 3 , we have that:
[Equation]
By following the argument in [28, Appendix D.2] (which does not depend on the dynamics, but only on the structure of the gradient), we have that E z ∆ W 3 ( z , C 2 ; W RD ( t )) is S 3 -measurable, i.e.,
[Equation]
By combining (109), (110) and (111), we obtain that (106) holds.
The arguments giving (107) and (108) are analogous.
From this, we can compute the difference between the reduced dynamics and the mean-field
ODE as
[Equation]
which, after applying Corollary 27, gives that D T ( W , W RD )
= 0 .
This implies that W = W RD and, hence, w 1 ( t, C 1 ) , w 2 ( t, C 1 , C 2 ) , w 3 ( t, C 2 ) are S 1 , S 123 , S 3 -measurable, respectively.
Next, we show the continuity of the function w ∗ 1 ( · , · ) : R ≥ 0 × R D −→ R D in both arguments.
Lemma 22 Under Assumptions (A1)-(A2), we have that, for all t ∈ [0 , T ] and for all u 1 , u ′ 1 ∈ R D ,
[Equation]
Proof In order to prove the lemma
[Equation]
we first need to derive the dynamics that characterize the evolution of the functions w ∗ 1 ( t, u 1 ) , w ∗ 2 ( t, u 1 , u 2 , u 3 ) , w ∗ 3 ( t, u 3 ) .
This dynamics is induced by the mean-field ODE, whose form we recall below:
[Equation]
Recall also that w 3 ( t, c 2 ) =
w ∗ 3 ( t, w 3 (0 , c 2 )) .
Thus, in order to get the dynamics of w ∗ 3 ( t, u 3 ) , we replace w 3 (0 , c 2 ) by u 3 , w 2 (0 , c 1 , c 2 ) by u 2 , and w 1 (0 , c 1 ) by u 1 into (114).
By doing the same replacements into (115) and (116) for w 2 ( t, c 1 , c 2 ) and w 1 ( t, c 1 ) , respectively, we obtain
[Equation]
where we have the following modified forward and backward paths:
[Equation]
, u 3 ) , ∆ W 3 ( t, z , u 3 ) =
∂ 2 R ( y, f ( x ; W ( t )))
σ 2 ( H 2 ( t, x , u 3 )) , ∆ W 2 ( t, z , u 1 , u 2 , u 3 ) =
∂ 2 R ( y, f ( x ; W ( t )))
w 3 ( t, u 3 ) σ ′ 2 ( H 2 ( t, x , u 3 ))
σ 1 ( H 1 ( t, x , u 1 )) , ∆ W 1 ( t, z , u 1 ) =
E u 2 ,u 3 (cid:2) ∂ 2 R ( y, f ( x ; W ( t )))
w 3 ( t, u 3 ) σ ′ 2 ( H 2 ( t, x , u 3 )) · w 2 ( t, u 1 , u 2 , u 3 ) σ ′
( H 1 ( t, x , u 1 ))
x (cid:3) .
Thus, we have that
[Equation]
An application of Lemma 7 gives that
[Equation]
Similarly for w ∗ 2 , we have that
[Equation]
and another application of Lemma 7 gives that
[Equation]
By combining (117), (118), (119) and (120), we obtain | w
[Equation]
Thus, by Corollary 27, we have that:
[Equation]
−
[Equation]
which implies that ∥ w ∗ 1 ( t, u 1 ) − w ∗ 1 ( t, u ′ 1 ) ∥ 2 ≤ K 4 ( γ, T ) | u
1 − u ′ 1 | , and concludes the proof of (113).
The Lipschitz continuity (112) of w ∗ 1 ( t, u 1 ) is already proved in Lemma 16.
At this point, we show that, if w 1 (0 , c 1 ) :
Ω −→ R D has full support, then w ∗ 1 ( t, u 1 ) has full support.
Lemma 23 Under Assumptions (A1)-(A2) and (B1)-(B3), we have that w ∗ 1 ( t, u 1 ) has full support for any t < ∞ .
Proof By the continuity argument in Lemma 22, we have that
[Equation]
We want to show that, for any x ∈ R D , there exist a v such that w ∗ 1 ( t, v ) = x .
For any x ∈ R D , define a map g x ( t, v ) = x − ( w ∗ 1 ( t, v ) − v ) .
It is easy to see that if v a fixed point of g x ( t, · ) , then w ∗ 1 ( t, v ) =
x as
[Equation]
By (121), we have that g x ( t, · ) : R D −→ B ( x , K ( γ, T ) t ) , where B ( x , K ( γ, T ) t ) is the closed ball centered at x with radius K ( γ, T ) t .
Now, if we restrict g x ( t, v ) on B ( x , K ( γ, T ) t ) , we have that it is a map from B ( x , K ( γ, T ) t ) , which is a compact set, to itself.
Furthermore, g x ( t, v ) is continuous in v , since w ∗ 1 ( t, v ) is continuous in v by (122).
Thus, by the Brouwer fixed point theorem, we have that there exist a fixed point v ∈ B ( x , K ( γ, T ) t ) , which finishes the argument.
Finally, we are ready to prove Theorem 3.
Our proof follows similar steps as that of [28, Proof of Theorem 8].
Proof By Assumption (B3), we have that
[Equation]
[ | E z
[Equation]
By the definition of ∆ W 2 ( t, z , C 1 , C 2 )
[Equation]
we have
[Equation]
Recall from Lemma 23 that
[Equation]
for all finite t
[Equation]
w 1 ( t, C 1 ) has full support.
Hence
[Equation]
we have that, for uin a dense subset of R D ,
[Equation]
Our aim is to conclude that, for almost all x , we have that E z [ ∂ 2 R ( y, f ( x ; W ( ∞ )))
| x ] = 0 .
By definition of the backward path
[Equation]
we have that
[Equation] [Equation]
By Assumption (B3)
[Equation]
the RHS of (123) converges to 0 as t → ∞ .
Hence, by taking the limit on both sides
[Equation]
we have that, for u 1 in a dense subset of R D ,
[Equation]
which implies that, for almost all c 2 ,
[Equation]
By definition of ∆ H 2 ( z , C 2 ; W ( ∞ ))
we have that, for almost all c 2 ,
[Equation]
Note that Assumption (A1) gives that σ ′ 2 ̸ = 0 , and Assumption (B3) that w 3 ( ∞ , c 2 ) ̸ = 0 with probability > 0 (where the probability is intended over c 2 ).
Hence, we have that, with probability > 0 (over c 2 ),
[Equation]
Recall that σ 1 ( u T 1 x ) is a function of x
[Equation]
but ∂ 2 R ( y, f ( x ; W ( ∞ ))) depends on both y and x .
Thus
[Equation]
we can re-write (125) as
[Equation]
Now
[Equation]
we want to use the universal approximation property of σ 1 to conclude that, for almost every x ,
[Equation]
The idea is that linear combinations of σ 1 ( u T 1 x ) can approximate any function in L 2 ( D x ) .
Thus, if E y (cid:2) ∂ 2 R ( y, f ( x ; W ( ∞ ))) | x (cid:3) σ ′ 2 ( H 2 ( x , c 2 ; W ( ∞ ))) is in L 2 ( D x ) , we have that there exist a sequence of index sets { I k } k ∈ N , such that:
[Equation]
To simplify the notation
[Equation]
we define:
[Equation]
From (126) and by linearity of expectation, we have that, for all k ,
[Equation]
Thus we have
[Equation]
2 (cid:105) , which implies that
[Equation]
Hence, we have that
[Equation]
which implies that (127) holds.
Furthermore, to see that E y (cid:2) ∂ 2 R ( y, f ( x ; W ( ∞ ))) | x (cid:3) σ ′ 2 ( H 2 ( x , c 2 ; W ( ∞ ))) is indeed in L 2 ( D x ) , it suffices to note that, by Assumption (A1),
[Equation]
By Assumption (A1), we also have that σ ′ 2 ( x ) ̸ = 0 for all x .
Hence, (127) implies that, for almost every x ,
[Equation]
Since the loss is convex in f ( x ; W ( ∞ )) , we have
[Equation]
where the last passage follows from (128).
Thus, we conclude that
[Equation]
Finally
[Equation]
we want to show that
[Equation]
To see this
[Equation]
we write
[Equation]
and use again Assumption (B3).
By combining (129) and (130)
[Equation]
we obtain the desired result
[Equation]
Appendix H. Technical lemmas
Lemma 24 (Corollary of McDiarmid inequality) [26, Lemma 30]
Let { X i } i ∈ [ n ] ∈ R
d be a sequence of i
[Equation]
i.d random variable, with ∥ X i ∥ 2 ≤ K ans E
[ X i ] = 0 , then we have:
[Equation]
Lemma 25 (Azuma-Hoeffding bound)
[26, Lemma 31] Let ( X k ) k ≥ 0 is a martingale taking value in R D with respect to the filtration ( F k ) , with X 0 = 0 .
Assume that the martingale difference at time k is L k -subgaussian, which means the following holds almost surely for all λ ∈ R D :
[Equation]
Then, we have
Pr
[Equation]
Note that, if L k ≤ L for all k , then
[Equation]
Lemma 26 (Pachpatte’s inequality) [3, Chapter 1, Theorem 1.7.1]
Let u , f
and g be non-negative continuous functions defined on [0 , T ] , for which the inequality u ( t ) ≤
[Equation]
+
[Equation]
g
[Equation]
u (
[Equation]
dr (cid:19) ds holds, where u 0 is a non-negative constant.
Then we have:
[Equation]
Corollary 27 (Pachpatte’s inequality for constants) Let u be a non-negative continuous function defined on [0 , T ] , and γ, K be positive real numbers.
Assume the following inequality holds:
[Equation]
Then
[Equation]
we have
[Equation]