Time-independent Generalization Bounds for SGLD in Non-convex Settings
Year: 2021

Abstract:
We establish generalization error bounds for stochastic gradient Langevin dynamics (SGLD) with constant learning rate under the assumptions of dissipativity and smoothness, a setting that has received increased attention in the sampling/optimization literature. Unlike existing bounds for SGLD in non-convex settings, ours are time-independent and decay to zero as the sample size increases. Using the framework of uniform stability, we establish time-independent bounds by exploiting the Wasserstein contraction property of the Langevin diffusion, which also allows us to circumvent the need to bound gradients using Lipschitz-like assumptions. Our analysis also supports variants of SGLD that use different discretization methods, incorporate Euclidean projections, or use non-isotropic noise.

heading: 1 Introduction
text: Investigating the generalization error of a learning algorithm is a fundamental problem in machine learning that has motivated the development of a rich theory connecting notions of model complexity and sensitivity to the ability of a learning algorithm to generalize well to unseen data. Recently, there has been increased interest in studying the generalization capabilities of stochastic gradient descent (SGD) and its variants, e.g. [13, 1, 6, 17, 22, 19]. Despite them being some of the most important methods of modern machine learning and statistics, their learning capabilities have not been fully explored, particularly in non-convex settings.
One variant that has attracted a lot of attention is stochastic gradient Langevin dynamics (SGLD) [29]. With the addition of independent Gaussian noise to each iteration, SGLD combines the SGD framework with the Langevin diffusion, a stochastic process that converges to the Gibbs distribution of a given objective function. By tuning the scale of the noise applied, SGLD has been shown to work well as both a sampling scheme and a statistical learning algorithm in a variety of settings, e.g. [29, 25, 4, 14]. Using its relationship with the Langevin equation, various tools from stochastic analysis have been adopted to develop a rich theory analyzing SGLD in settings where no comparable results seem to be available for SGD, e.g. [23, 32, 5, 31]. In particular, a recent work from Raginsky et al. [23] provides non-asymptotic excess risk bounds for SGLD for a broad class of non-convex learning problems and these bounds were derived using the exponential ergodicity of Langevin diffusions. Following this work, there has been ongoing progress in the development of bounds for both optimization error and mixing times, e.g. [32, 5, 31, 7, 4]. Many of these works are motivated by recent results showing that Langevin diffusions contract in particular Wasserstein distances [8, 9].
In parallel to this literature, there has been a growing body of work that considers generalization error bounds for SGLD. Using uniform stability, Mou et al. [17] showed that under the assumption of bounded gradient updates and bounded loss functions, SGLD attains generalization bounds that grow slower with time than those that are known for SGD in identical settings [13, 6]. By bounding
35th Conference on Neural Information Processing Systems (NeurIPS 2021).
ar X
iv :2
11 1.
12 87
6v 1
[ st
at .M
L ]
2 5
N ov
the generalization error using the mutual information between its input and output [30], Pensia et al. [22] established similar results while replacing the assumption of bounded loss functions with the weaker assumption that the loss function applied to the data is subgaussian. This technique has also been used to develop data-dependent generalization bounds that depend on the gradient along the trajectory of the iterates [18, 12, 17, 19].
Although the assumptions made in the aforementioned bounds cover several settings of interest, we are not aware of any convergence analysis that considers Lipschitz, bounded, or subgaussian objective functions. In general, the Gibbs distribution, which is an object that is fundamental to the design and known convergence analyses for SGLD [29, 25, 23, 5, 32], cannot be defined in the case that the objective is bounded over the Euclidean space, as the distribution would not integrate to one. Similarly, we are not aware of any generalization bounds for SGLD that use the assumptions of dissipativity and smoothness that are consistently applied in the non-convex sampling/optimization literature, e.g. [23, 5, 32].
Another commonality in existing generalization bounds for SGLD is that they grow indefinitely with time. Sampling from the Gibbs distribution, the limiting behavior of the Langevin diffusion, is known to be uniformly stable in multiple settings [30, 23] and it has also been shown that the continuoustime Langevin dynamics are uniformly stable when applied to bounded Lipschitz functions with weight decay [16]. However, convergence in generalization error for SGLD seems to have only been observed in preliminary empirical evidence that uses data-dependent bounds [18, 12]. Timeindependent bounds have been obtained for SGD, for example, under the assumption of strong convexity in [13] and in non-convex settings, Lei and Ying [15] obtain bounds that decrease with time but do not decay to zero as the sample size increases.

heading: 1.1 Our contributions
text: In this paper, we obtain expected generalization error bounds for SGLD for learning problems under dissipativity and smoothness assumptions. We analyze both the discrete-time algorithm as well as a continuous-time version and in both cases, we obtain bounds that converge with time. Taking the supremum with respect to time yields time-independent bounds that, with the appropriate scaling of the learning rate, decay to zero as the sample size increases.
At first, we focus on the special case of Lipschitz loss functions with weight decay, which is the primary example of a dissipative objective given in Raginsky et al. [23]. For the continuous-time algorithm, we obtain a bound with rate O(n−1) and for the discrete-time algorithm, we obtain a bound with the slower rate O(n−1 + η1/2), where η is the learning rate. Then we extend the result to the full dissipative case without Lipschitz requirements. In this setting, we obtain bounds with rate O(n−1η−1/2) for the continuous-time algorithm and with rate O(n−1η−1/2 + η1/2) in the discretetime case. We also discuss how our method allows for the consideration of different discretization techniques and how it can be used to obtain generalization error bounds for modifications of SGLD that incorporate Euclidean projections or non-isotropic noise. From our bounds, a scheme for choosing η follows: in the Lipschitz setting, we obtain dimension-free O(n−1) bounds by setting η ∝ n−2d−1 where d is the model dimension, and in the dissipative setting we find that a scaling of η ∝ n−1 leads to O(n−1/2) bounds that, in general, scale exponentially in dimension. To derive time-independent bounds, our proof technique depends fundamentally on the two sources of noise that occur in the algorithm: the random mini-batches and the injected Gaussian noise. Using the framework of uniform stability, in the sense defined in [10], we derive generalization bounds by estimating how much SGLD diverges in Wasserstein distance when an element of the data set is changed. We exploit recent results that use reflection couplings to show that under dissipativitytype assumptions, Langevin diffusions contract in Wasserstein distance. Though this property has been adopted extensively in the sampling/optimization literature, we are not aware of any results for generalization bounds based on this property. Using the convexity of the Wasserstein distance, we combine this with the stability induced by using stochastic gradients to obtain bounds that are time-independent.
A peculiarity of our results that arises from the methodology we use is that some of our bounds diverge as η → 0. This contrasts with the usual approach of stability-based generalization bounds based on non-expansivity that often require the learning rate to decay sufficiently fast or the time-horizon to be sufficiently small to guarantee bounds that are non-vacuous [13, 17]. Furthermore, for our bounds to
converge with time, we require the mini-batch size to be less than n. While our bounds support the case of full-batch gradient descent, we find that our bounds grow indefinitely with time.

heading: 2 Background and notation

heading: 2.1 Stability and generalization
text: In this paper, we consider a loss function f : Rd ×Z → R where the Euclidean space Rd represents the set of possible model parameters and Z represents the data instance space. A common objective in learning theory is to minimize the population risk which, given a data distribution P , is defined by
FP (x) := EZ∼P f(x, Z).
In practice, we often cannot compute FP (x) so we instead collect independent samples from P to form a data set and we compute the average loss, or empirical risk, over the data set. We will use the notation S = (z1, ..., zn) to denote the data set and define the empirical risk as
FS(x) := 1
n n∑ i=1 f(x, zi).
When the parameter is chosen by a random algorithm that depends on the data set, denoted A(S), we define the object central to this paper, the generalization error, as follows:
gen(A) := FP (A(S))− FS(A(S)).
To bound this quantity we employ the following notion of uniform stability.
Definition 2.1 ([13], Definition 2.1). An algorithm A is ε-uniformly stable if
εstab(A) := sup S∼=Ŝ sup z∈Z
E [ f ( A ( S ) , z ) − f ( A ( Ŝ ) , z )] ≤ ε,
where the first supremum is over data sets S, Ŝ ∈ Zn that differ by one element, denoted by S ∼= Ŝ.
The connection between generalization and stability under changes in the data set has received increased attention since the paper of Bousquet and Elisseeff [2], and these results have been extended to account for random algorithms by Elisseeff et al. [10]. The precise notion of stability that we consider in this paper is given in Hardt et al. [13].
Proposition 2.2 ([13], Theorem 2.2). SupposeA is an ε-uniformly stable algorithm, then the expected generalization error is bounded by ∣∣EA,S gen(A)∣∣ ≤ ε. Though we will present the results of this paper as bounds on the expected generalization error, uniform stability bounds can also give high probability bounds for the generalization error [11, 3]. Outside of generalization bounds, the concept of uniform stability has also been shown to be fundamentally related to differential privacy and learnability [28, 24].

heading: 2.2 Stochastic gradient Langevin dynamics
text: We will now define the algorithm of interest in this paper. Given a mini-batch B ⊂ [n] := {1, ..., n}, let the mini-batch average be defined by
FS(x,B) := 1 |B| ∑ i∈B f(x, zi).
In our analysis we will consider uniformly sampled random mini-batches of fixed size k ≤ n. We now define SGLD, which given an initial distribution µ0, is characterized by the update
xt+1 = xt − η∇FS(xt, Bt+1) + √ 2β−1ηξt+1, x0 ∼ µ0, (1)
where (Bt)∞t=1 is an i.i.d. sequence of random variables distributed uniformly on the set {B ⊂ [n] : |B| = k} and (ξt)∞t=1 is an i.i.d. sequence of N(0, Id) standard Gaussian random variables. The parameters η, β > 0 are tunable and are referred to as the learning rate and inverse temperature respectively. Here and throughout the paper we use ∇ to refer to the gradient with respect to the model parameter.
If k = n then (1) describes the Euler-Maruyama discretization of the (overdamped) Langevin equation with potential FS(·). This is a stochastic differential equation of the form
dXt = −∇FS(x)dt+ √ 2β−1dWt, X0 ∼ µ0,
where Wt is a d-dimensional Wiener process [21]. In the more common case that k < n we can still define a stochastic differential equation that approximates SGLD, given by
dXt = −∇FS(Xt, Bdt/ηe)dt+ √ 2β−1dWt, X0 ∼ µ0. (2)
We will refer to the solution Xt as the continuous-time SGLD algorithm. Note that, under the smoothness assumptions imposed throughout this paper, all stochastic differential equations considered have strong solutions (see Theorem 3.1 of [21]).
Since xt is a Markov process, we can define its Markov kernel Rx. Though we will not frequently reference this fact, we will use the notation µRsx to denote the law of xt+s under the condition xt ∼ µ. Because of its dependence on Bk, Xt is not necessarily a Markov process in continuous-time, but the discrete-time process (Xtη)∞t=0 does in fact satisfy the Markov property and so we denote its kernel by RX .
Given a deterministic set B ⊂ [n] we use the notation µPBt to denote the law of θt, the solution to the SDE, dθt = −∇FS(θt, B)dt+ √
2β−1dWt, θ0 ∼ µ, (3) which is a Markov (diffusion) process. Thus µRX can be computed by integrating µPBη over B with respect to the mini-batch distribution.
Since this work is concerned with stability under changes in the data set, we will also be interested in SGLD when S is replaced with some data set Ŝ that differs by a single element. We will use the notation x̂t and X̂t to denote the respective counterparts of xt of Xt when trained using Ŝ instead of S and similarly we will use R̂x, R̂X and P̂Bt .

heading: 2.3 Wasserstein distance
text: For probability measures µ, ν on Rd with finite pth moment, we define the Wasserstein distance
Wp(µ, ν) := (
inf π∈C(µ,ν)
∫ ‖x− y‖p π(dx, dy) )1/p ,
where ‖ · ‖ is the Euclidean norm. Here the infimum is over all couplings of µ and ν, that is, the set of all probability measures with marginals µ and ν. In our analysis, we will also consider Wasserstein distances based on any semimetric ρ, which under certain integrability assumptions is defined by
Wρ(µ, ν) := inf π∈C(µ,ν)
∫ ρ(x, y)π(dx, dy).
Note that throughout this paper all probability measures considered outside of the data distribution, are Borel probability measures on Rd. A property of the Wasserstein distance that will prove to be of central importance in our results is its convexity. Lemma 2.3 (Convexity of the Wasserstein distance). Suppose that ρ is a semimetric and µ1, µ2, ν1, ν2 are probability measures. Then, for any r ∈ [0, 1],
Wρ(µ, ν) ≤ rWρ(µ1, ν1) + (1− r) Wρ(µ2, ν2), where we define µ(dx) = rµ1(dx) + (1− r)µ2(dx) and ν(dx) = rν1(dx) + (1− r)ν2(dx).
The proof of this result, as well as a more general statement, can be found in [26], Theorem 4.8.

heading: 3 The Lipschitz setting with weight decay regularization
text: To begin our analysis, we focus on a special case of the dissipativity and smoothness assumptions. We will assume that the loss function f(x, z) is Lipschitz in x for any fixed z, and when performing gradient updates we will use the empirical risk with an added weight decay term:
F̃S(x,B) = FS(x,B) + λ
2 ‖x‖2.
This setting has been considered in the data-dependent literature for SGLD, for example by Mou et al. [17], and under the additional assumption of boundedness, this setting has also been considered in the stability-based analysis of the continuous Langevin dynamics by Li et al. [16].
The following assumptions are imposed: Assumption 3.1. For each z ∈ Z , f(·, z) is L-Lipschitz: for all x1, x2 ∈ Rd and z ∈ Z , |f(x1, z)− f(x2, z)| ≤ L‖x1 − x2‖. Assumption 3.2. For each z ∈ Z , f(·, z) is differentiable and M -smooth: for all x1, x2 ∈ Rd and z ∈ Z , ‖∇f(x1, z)−∇f(x2, z)‖ ≤M‖x1 − x2‖. Assumption 3.3. The initial condition µ0 has finite first moment σ1 := µ0(‖ · ‖) <∞.
When using weight decay, Assumption 3.1 guarantees that F̃S is dissipative which allows us to obtain Wasserstein contractions. Assumption 3.2 guarantees the existence of a solution to the SDE in (2) and allows us to control the error that is brought about when using it to approximate the discrete-time algorithm. Assumption 3.3 is used to guarantee that SGLD has finite first moment so that we can perform our analysis using the 1-Wasserstein distance. Note that the weight decay term only affects the gradient updates of SGLD and is not used in the computation of the generalization error, which is still defined in terms of FS and FP and not F̃S . Theorem 3.1. Suppose Assumptions 3.1–3.3 hold and η < 1, then for any t ∈ N the continuous-time algorithm attains the generalization bound
|E gen(Xηt)| < C2 min { ηt, (C1 + 1)n
n− k
} k
n .
Furthermore, if η < λ−1 then the discrete-time algorithm attains the generalization bound |E gen(xt)| < C3 min { ηt, (C1 + 1)n
n− k
}( k
n + (ηd)1/2
) .
The positive constants C1 ≡ C1(M,L, λ, β), C2 ≡ C2(M,L, λ, β), C3 ≡ C3(M,L, λ, β, σ1) are given in equations (5), (7) and (8) respectively. Remark 3.2. This result suggests that by using the step-size scaling η = O(1/n2d) we obtain dimension-free time-independent generalization bounds for SGLD that scale with O(k/n).
The bound for the continuous-time algorithm has no explicit dependence on dimension and in the discrete-time case the bound grows with the rate d1/2. In general, the constants C1, C2 and C3 depend exponentially on M,L, λ and β. However, in the case of λ ≥ M , the constants become polynomial in these parameters and C2 becomes independent of β. A similar dependence was found in the time-independent bounds for the Langevin equation [16] and it is also found in convergence analyses in non-convex settings, leading to the time horizon required to attain a certain guaranteed accuracy having to scale exponentially with parameters including d and β, e.g. [23, 5, 32].

heading: 3.1 Methodology
text: In this section, we present the method used to obtain the bounds in Theorem 3.1, in particular we proceed with the stability framework presented in Section 2.1. Without loss of generality, we require that the two data sets S and Ŝ differ only in the nth coordinate.
Under Assumption 3.1, to obtain uniform stability, it is sufficient to control the quantities W1(µ0R t X , µ0R̂ t X) and W1(µ0R t x, µ0R̂ t x). It follows directly from the definition of εstab that
εstab(A) ≤ L sup S∼=Ŝ W1
( law ( A ( S )) , law ( A ( Ŝ ))) . (4)
A property similar to this has been utilized in [23], where it was used to prove the stability of the Gibbs algorithm.
To prove Theorem 3.1, we begin with a weak estimate for the divergence in 1-Wasserstein distance using synchronous couplings between Xt and X̂t, that is, we couple the processes by having them share the same Brownian motion. Lemma 3.3. Suppose Assumption 3.1 holds, then for any two probability measures µ, ν on Rd,
W1 ( µPBt , νP̂ B t ) ≤W1(µ, ν) + 2Lt.
Additionally, via a similar technique, we obtain discretization error bounds. Lemma 3.4. Suppose Assumptions 3.1-3.3 hold and η < λ−1, then with µ = µ0Rtx for any t ∈ N,
W1(µRx, µRX) ≤ η(λ+M) [ η(λσ1 + 2L) + 2 √ 2dβ−1η ] exp(M + 1).
We refer to Appendix B for the proofs of these two lemmas.
In the case P̂B = PB , we can obtain sharper bounds than those given in Lemma 3.3. In a recent result from Eberle [8], reflection couplings were used to show that under conditions not dissimilar from dissipativity, diffusion processes contract in Wasserstein distance. The Wasserstein distance considered in this result does not use a standard metric, it uses the metric ρg(x, y) := g(‖x − y‖) where g is a strictly-increasing concave function that is constructed depending on the objective function (or drift term) used. In the lemma that follows we give a special case of such a result in the setting that we consider. Throughout this section we use the notation Wg(µ, ν) := Wρg (µ, ν). We refer to Appendix C for the proof and a broader discussion on this result. Lemma 3.5. Suppose Assumptions 3.1 and 3.2 hold, then there exists a strictly-increasing concave function g : R+ ∪ {0} → R+ ∪ {0} such that for any two probability measures µ, ν on Rd, any B ⊂ [n] and t ≥ 0, we have
Wg ( µPBt , νP B t ) ≤ e−t/C1 Wg(µ, ν), C1 := { c1c2, if λ < M, c3, if λ ≥M,
(5)
with constants c1 = exp[2βL2(M − λ)/λ2], c2 = 8(L2β/λ + 1)/λ, c3 = max(16L2β/λ2, 2/λ). Furthermore, Wg is equivalent to the 1-Wasserstein distance in the sense that
1
2 max(c1, 1) W1(µ, ν) ≤Wg(µ, ν) ≤W1(µ, ν).
In the case that n is not in the random mini-batch, i.e. n 6∈ B, which occurs with probability 1− k/n, it follows that P̂B = PB and so the processes Xtη and X̂tη contract with respect to the Wasserstein distance Wg. In the case that n is in the random mini-batch, n ∈ B, Lemma 3.3 gives uniform bounds for how much Xtη and X̂tη can diverge. The remainder of the proof works to combine these facts and shows that if n is sufficiently large, Xtη and X̂tη can only diverge by some fixed amount. We combine the two cases using the convexity of the Wasserstein distance and we find that given a sufficiently large n, the resulting bound does indeed converge over time.
Proof of Theorem 3.1. As was pointed out in (4), it is sufficient to bound the 1-Wasserstein distance between the processes to obtain stability bounds, so this is how we will proceed. Since the probability
of a random mini-batch containing the element n is k/n, Lemma 2.3 is used to bound the Wasserstein distance with the decomposition
Wg ( µRX , νR̂X ) ≤ k n sup B:n∈B Wg ( µPBη , νP̂ B η ) + ( 1− k n ) sup B:n 6∈B Wg ( µPBη , νP̂ B η ) .
Bounding the first term using Lemma 3.3 and the second term using Lemma 3.5, it follows that Wg(µRX , νR̂X) ≤ c̃1 Wg(µ, ν) + c̃2 where c̃1 := kn + (1− k n )e −η/C1 and c̃2 = 2Lη kn . It follows by induction that
Wg ( µ0R t X , µ0R̂ t X ) ≤ t∑ s=1 c̃t−s1 c̃2 + c̃ t 1 Wg(µ0, µ0) = 1− c̃t1 1− c̃1 c̃2, (6)
since c̃1 < 1. Combining this with (4), it follows that εstab(Xηt) ≤ 2L(c1 ∨ 1) Wg ( µRtX , νR̂ t X ) ≤ C2η ( 1− c̃t1 1− c̃1 ) k n ,
where, C2(M,L, λ, β) := 4L
2(c1 ∧ 1). (7) This bound is simplified with the approximation 1 − c̃t1 ≤ 1 ∧ (1 − c̃1)t and, using the bound ex ≥ 1 + x, it follows that (1− e−x)−1 ≤ 1 + 1/x for x > 0 and so
1
1− c̃1 =
1 (1− k/n)(1− e−η/c) ≤ n(c/η + 1) n− k .
This result is extended to the discrete-time case by applying the triangle inequality, Wg ( µRx, νR̂x ) = Wg ( µRx, µRX ) + Wg ( µRX , νR̂X ) + Wg ( νR̂X , νR̂x ) .
The first and third terms are bounded using the discretization error bound in Lemma 3.4 and the same argument is applied to obtain the bound,
εstab(xt) ≤ 2L(c1 ∧ 1) Wg ( µ0R t x, µ0R̂ t x ) ≤ C3η
1− c̃t1 1− c̃1
( k
n + (dη)1/2
) .
where C3(M,L, λ, β, σ1) := 4L(c1 ∨ 1) ( L+ (λ+M) ( λσ1 + 2L+ 2 √ 2β−1 )) . (8)
Remark 3.6. For the discrete-time algorithm, the bound is obtained by appending the one-step discretization error to the bound in (6) for the continuous-time algorithm. Thus, the result can easily be extended to a broader range of discretizations of (Xt)t≥0. Suppose Kx is a Markov kernel on Rd such that for any probability measure µ on Rd, W1(µKx, µRX) ≤ δ(η), then we immediately obtain the bound
W1 ( µ0K t x, µ0K̂ t x ) ≤ O
( 1
n + δ(η) η
) .
In the case of the standard discrete-time algorithm we obtain δ(η) = O(η3/2). A simple example of a more accurate discretization follows by instead of approximating Xη with a single Euler-Maruyama step we use T = b1/ηc steps with step-size γ = η/T , then we can show that δ(η) ≤ O(√γ) ≤ O(η2). Remark 3.7. With our methodology, it is also possible to derive bounds for projected versions of SGLD. If Ω ⊂ Rd is a compact set, we can define the algorithm,
xΠt+1 = ΠΩ [ xΠt − η∇FS(xt, Bt+1) + √ 2β−1ηζt+1 ] , (9)
where ΠΩ is the Euclidean projection onto the set Ω. A well-known property that is often used in the optimization literature is that for any x1, x2 ∈ Rd, ‖ΠΩ(x1) − ΠΩ(x2)‖ ≤ ‖x1 − x2‖ and, given two measures µ, ν on Rd, this naturally extends to
W1(µ ∗ΠΩ, ν ∗ΠΩ) ≤W1(µ, ν), where µ∗ΠΩ denotes the push-forward of µ with respect to the mapping ΠΩ. If we let RΠx be the kernel of the process given in (9), it follows immediately that,
W1 ( µRΠx , νR̂ Π x ) ≤W1 ( µRx, νR̂x ) ,
and hence the Wasserstein distance bound given in the proof also applies to the projected algorithm and therefore, the generalization bound in Theorem 4.1 also holds for the projected algorithm.
Remark 3.8. Our technique supports the case where, instead of assuming ξt is a standard Gaussian random vector, we set ξt ∼ N(0,Σ) for some symmetric positive semidefinite matrix Σ ∈ Rd×d. If the operator norm of Σ is 1, the conclusions of Lemma 3.3 and Lemma 3.5 still hold true and, as a result, the conclusion of Theorem 3.1 would also hold.

heading: 4 The dissipative setting
text: Now we extend the results of the previous section to the full dissipative smooth case. We refer to Section 4 of [23] for a detailed discussion on dissipativity. Assumption 4.1. For each z ∈ Z, f(·, z) is (m, b)-dissipative: for all x ∈ Rd and z ∈ Z ,
〈∇f(x, z), x〉 ≥ m‖x‖2 − b. Assumption 4.2. Same as Assumption 3.2. Assumption 4.3. The initial condition µ0 has finite fourth moment σ4 := µ0(‖ · ‖) <∞.
Without the Lipschitz assumption, the analysis is more challenging as it is no longer sufficient to control the uniform argument stability or 1-Wasserstein distance. Furthermore, we can no longer guarantee that the difference between FS and FŜ is upper bounded. Theorem 4.1. Suppose Assumptions 4.1–4.3 hold. If η ∈ (0, 1) then for any t ∈ N, the continuoustime algorithm attains the generalization bound
|E gen(Xηt)| < C5 min { ηt, (C4 + 1)n
n− k
} k
nη1/2 .
Furthermore, if η ≤ 1/2m, then the discrete-time algorithm attains the generalization bound |E gen(xt)| < C6 min { ηt, (C4 + 1)n
n− k
}( k
nη1/2 + η1/2
) .
The constants C4 ≡ C4(M,m, b, d, β), C5 ≡ C5(M,m, b, d, β, σ4), C6 ≡ C6(M,m, b, d, β, σ4) are given in equations (11), (12) and (13). Remark 4.2. Choosing η = O(n−1/2) leads to O(n−1/2) generalization bounds.
In general, the constants C4, C5, and C6 depend exponentially on parameters M,m, b, d, β. Improving the dependence on these parameters in specific settings would require improvements on estimates of the contraction rate and therefore, improvements on the coupling arguments given in [9].
A peculiarity of this result is that the bounds explode as η → 0. This is due to the fact that when η is smaller, there is less time for the processes to contract at each iteration. In the setting of Section 3, this is combated by the fact that the divergence bounds scale with rate O(η), but in the present setting we only obtain divergence bounds with rate O(η1/2). This is a result of the metric that we use in our analysis that is designed to suit the coupling arguments given in [9].

heading: 4.1 Methodology
text: Since this result relies on similar techniques to the proof of Theorem 4.1, we will postpone the proof to the appendix and in this section, we will focus on the challenges faced when extending to the dissipative setting. Without the Lipschitz assumption it is no longer sufficient to control the 1-Wasserstein distance and so we must turn to different metrics. Raginsky et al. [23] used continuity with respect to the 2-Wasserstein metric to derive stability bounds for the Gibbs sampler but, as has been noted in [8] and [27], contractions in 2-Wasserstein distance are notably more difficult to obtain than in 1-Wasserstein distance.
We use a semimetric of the form,
ρ(x, y) = g(‖x− y‖)(1 + 2ε+ ε‖x‖2 + ε‖y‖2), where g : R+ ∪ {0} → R+ ∪ {0} is a non-decreasing concave function and ε ∈ (0, 1). As will be exhibited in the appendix, Eberle et al. [9] show that there exists g such that for ε sufficiently small,
Wρ ( µPBt , νP B t ) ≤ e−t/C4 Wρ ( µ, ν ) ,
for any two probability measures µ, ν on Rd. Additionally, this function is constant for r > R and has the property ϕr ≤ g(r) ≤ r for some ϕ,R ∈ R+. We show that this result can be used to derive stability bounds in the following lemma.
Lemma 4.3. Suppose Assumptions 4.1 and 4.2 hold and let A be a random algorithm, then εstab(A) ≤ M ( b/m+ 1 ) ϕε(R ∨ 1)
sup S∼=Ŝ Wρ
( law ( A ( S )) , law ( A ( Ŝ ))) .
The proof of Theorem 4.1 proceeds by establishing similar results to Lemma 3.3 and Lemma 3.5 using synchronous and reflection couplings. However, the argument is made markedly more difficult due to the fact that Wρ is not a metric, we do not have access to the triangle inequality and it is not possible to show Wρ is directly comparable to any standard Wasserstein distances.

heading: 5 Conclusion
text: In this paper, we derive time-independent generalization error bounds for SGLD under the assumptions of dissipativity and smoothness. We obtain bounds scaling as O(n−1/2η−1/2) for the continuous-time algorithm and scaling as O(n−1/2η−1/2 + η1/2) in the discrete-time case. In the special case of Lipschitz loss functions with weight decay regularization, we obtain faster rates of O(n−1) for the continuous-time algorithm and O(n−1 + η1/2) for the discrete-time case. In the latter case, we show that by having the step-size scaling η = O(d−1n−2), we obtain dimension-free O(n−1)-generalization bounds. Within the framework of uniform stability, we use a combination of synchronous and reflection couplings to control the Wasserstein distance between versions of the algorithm with perturbations on the data set. Using the convexity of the Wasserstein distance, we show that the two versions of the algorithm can only diverge by a fixed amount that, with the appropriate scaling of η, decays to zero as n increases.
The methodology used in our analysis allows for the extension to a broader class of discretizations. We are also able to explore modifications of the algorithm that incorporate Euclidean projections or non-isotropic Gaussian noise.
The fact that our proof relies heavily on properties of the Langevin diffusion does introduce some shortcomings. For example, the discretization error introduced when considering the discrete-time algorithm leads to an explicit dependence on the dimension and, furthermore, the dependence on model parameters aside from d is often exponential. In addition, Gaussian noise is fundamental in the construction and analysis of the Langevin diffusion, so our analysis does not directly extend to algorithms where different types of noise are used.

heading: Acknowledgements
text: Tyler Farghly was supported by the Engineering and Physical Sciences Research Council (EP/T517811/1). Patrick Rebeschini was supported in part by the Alan Turing Institute under the EPSRC grant EP/N510129/1.

Review summaries:
The paper proves generalization bounds for stochastic gradient Langevin dynamics (SGLD) for nonconvex functions. Previous bounds for SGLD in similar settings typically goes to infinity as time or the number of steps increases. This work gives bounds that are independent of time, and also decreases to 0 as the sample size $n$ grows.  Similar to previous bounds, the authors prove that SGLD is uniformly stable, which is known to imply good generalization. The proofs of stability proceed by starting with the continuous version of the algorithm, and then apply discretization techniques to get guarantees for the discrete-time SGLD. This discretization leads to a slower rate: the bounds are $O(1/n)$ vs $O(\sqrt{d/n})$ for Lipschitz loss functions plus weight decay, and $O(1/n^{1/3})$ vs $O(1/n^{1/6})$ for the full dissipative case without the Lipschitz-ness assumption.  For Lipschitz loss functions, it is sufficient to control the 1-Wasserstein distance ($W_1$) between the distributions induced by two neighbouring datasets. The nontrivial step is Lemma 3.4, which shows that conditioning that the differing data point, $n$, is not in the minibatch, some proxy for $W_1$, $W_g$, actually gets contracted due to the regularization term. Then, as long as the minibatch size is relatively small, the overall distance also gets shrinked. This leads to the time-independent bound on stability.  The proof becomes more challenging without the Lipschitz-ness assumption, since controlling $W_1$ alone does not guarantee stability. The authors show that controlling $W_2$ instead would work, though the contraction proof becomes more involved.

The paper derives time-independent generalization bounds for (discrete- and continuous- time) SGLD under dissipativity and smoothness assumptions. The main novelty of the new bounds appears to be the fact that they decay to zero as the sample size increases. The authors claim that the theoretical approach is different to prior related works on SGLD as it relies on the framework of uniform stability.

In this paper, the authors investigate the stability of stochastic gradient Langevin dynamics (SGLD) with varying data sets. They prove a time-independent stability bound for SGLD in non-convex settings. And their analysis can also be used to analyze other variants of SGLD.  More specifically, the reviewer summarizes their contribution as follows: In this paper, they bound the generalization error to show the stability of SGLD. As claimed by the authors, unlike existing bounds for SGLD in non-convex setting, their bound is time-independent and decays to zero as data size increases.