Time-independent Generalization Bounds for SGLD in Non-convex Settings

1 Introduction . 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 ] , can not 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 . 1.1 Our contributions . 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 . 2 Background and notation . 2.1 Stability and generalization . 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 can not 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 ] . 2.2 Stochastic gradient Langevin dynamics . 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 .

Summary: 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.