Title: ON GENERALIZATION ERROR BOUNDS OF NOISY GRADIENT METHODS FOR NON-CONVEX LEARNING

Abstract: Generalization error (also known as the out-of-sample error) measures how well the hypothesis learned from training data generalizes to previously unseen data. Proving tight generalization error bounds is a central question in statistical learning theory. In this paper, we obtain generalization error bounds for learning general non-convex objectives, which has attracted significant attention in recent years. We develop a new framework, termed Bayes-Stability, for proving algorithm-dependent generalization error bounds. The new framework combines ideas from both the PAC-Bayesian theory and the notion of algorithmic stability. Applying the Bayes-Stability method, we obtain new data-dependent generalization bounds for stochastic gradient Langevin dynamics (SGLD) and several other noisy gradient methods (e.g., with momentum, mini-batch and acceleration, Entropy-SGD). Our result recovers (and is typically tighter than) a recent result in Mou et al. ( 2018) and improves upon the results in Pensia et al. (2018) . Our experiments demonstrate that our data-dependent bounds can distinguish randomly labelled data from normal data, which provides an explanation to the intriguing phenomena observed in Zhang et al. (2017a). We also study the setting where the total loss is the sum of a bounded loss and an additional 2 regularization term. We obtain new generalization bounds for the continuous Langevin dynamic in this setting by developing a new Log-Sobolev inequality for the parameter distribution at any time. Our new bounds are more desirable when the noise level of the process is not very small, and do not become vacuous even when T tends to infinity.

1. INTRODUCTION:

Non-convex stochastic optimization is the major workhorse of modern machine learning. For instance, the standard supervised learning on a model class parametrized by R d can be formulated as the following optimization problem:

min w∈R d E z∼D [F (w, z)] ,

where w denotes the model parameter, D is an unknown data distribution over the instance space Z, and F : R d × Z → R is a given objective function which may be non-convex. A learning algorithm takes as input a sequence S = (z 1 , z 2 , . . . , z n ) of n data points sampled i.i.d. from D, and outputs a (possibly randomized) parameter configuration ŵ ∈ R d .

A fundamental problem in learning theory is to understand the generalization performance of learning algorithms-is the algorithm guaranteed to output a model that generalizes well to the data distribution D? Specifically, we aim to prove upper bounds on the generalization error err gen (S) = L( ŵ, D) -L( ŵ, S), where L( ŵ, D) = Ez∼D[L( ŵ, z)] and L( ŵ, S) = 1 n n i=1 L( ŵ, z i ) are the population and empirical losses, respectively. We note that the loss function L (e.g., the 0/1 loss) could be different from the objective function F (e.g., the cross-entropy loss) used in the training process (which serves as a surrogate for the loss L).

Classical learning theory relates the generalization error to various complexity measures (e.g., the VC-dimension and Rademacher complexity) of the model class. Directly applying these classical complexity measures, however, often fails to explain the recent success of over-parametrized neural networks, where the model complexity significantly exceeds the amount of available training data (see e.g., Zhang et al. (2017a)) . By incorporating certain data-dependent quantities such as margin and compressibility into the classical framework, some recent work (e.g., Bartlett et al. (2017) ; Arora et al. (2018) ; Wei & Ma (2019) ) obtains more meaningful generalization bounds in the deep learning context.

An alternative approach to generalization is to prove algorithm-dependent bounds. One celebrated example along this line is the algorithmic stability framework initiated by Bousquet & Elisseeff (2002) . Roughly speaking, the generalization error can be bounded by the stability of the algorithm (see Section 2 for the details). Using this framework, Hardt et al. (2016) study the stability (hence the generalization) of stochastic gradient descent (SGD) for both convex and non-convex functions. Their work motivates recent study of the generalization performance of several other gradient-based optimization methods: Kuzborskij & Lampert (2018) ; London (2016) ; Chaudhari et al. (2017) ; Raginsky et al. (2017) ; Mou et al. (2018) ; Pensia et al. (2018) ; Chen et al. (2018) .

In this paper, we study the algorithmic stability and generalization performance of various iterative gradient-based method, with certain continuous noise injected in each iteration, in a non-convex setting. As a concrete example, we consider the stochastic gradient Langevin dynamics (SGLD) (see Raginsky et al. (2017) ; Mou et al. (2018) ; Pensia et al. (2018) ). Viewed as a variant of SGD, SGLD adds an isotropic Gaussian noise at every update step:

EQUATION

where g t (W t-1 ) denotes either the full gradient or the gradient over a mini-batch sampled from training dataset. We also study a continuous version of (1), which is the dynamic defined by the following stochastic differential equation (SDE):

EQUATION

where B t is the standard Brownian motion.

1.1. RELATED WORK:

Most related to our work is the study of algorithm-dependent generalization bounds of stochastic gradient methods. Hardt et al. (2016) first study the generalization performance of SGD via algorithmic stability. They prove a generalization bound that scales linearly with T , the number of iterations, when the loss function is convex, but their results for general non-convex optimization are more restricted. London (2017) and Rivasplata et al. (2018) also combine ideas from both PAC-Bayesian and algorithm stability. However, these works are essentially different from ours. In London (2017) , the prior and posterior are distributions on the hyperparameter space instead of distributions on the hypothesis space. Rivasplata et al. (2018) study the hypothesis stability measured by the distance on the hypothesis space in a setting where the returned hypothesis (model parameter) is perturbed by a Gaussian noise. Our work is a follow-up of the recent work by Mou et al. (2018) , in which they provide generalization bounds for SGLD from both stability and PAC-Bayesian perspectives.

Another closely related work by Pensia et al. (2018) derives similar bounds for noisy stochastic gradient methods, based on the information theoretic framework of Xu & Raginsky (2017) . However, their bounds scale as O( T /n) (n is the size of the training dataset) and are sub-optimal even for SGLD.

We acknowledge that besides the algorithm-dependent approach that we follow, recent advances in learning theory aim to explain the generalization performance of neural networks from many other perspectives. Chizat et al. (2019) . Most of these results are stated in the context of neural networks (some are tailored to networks with specific architecture), whereas our work addresses generalization in non-convex stochastic optimization in general. We also note that some recent work provides explanations for the phenomenon reported in Zhang et al. (2017a) from a variety of different perspectives (e.g., Bartlett et al. (2017) ; Arora et al. (2018; 2019) ).

Welling & Teh (2011) first consider stochastic gradient Langevin dynamics (SGLD) as a sampling algorithm in the Bayesian inference context. Raginsky et al. (2017) give a non-asymptotic analysis and establish the finite-time convergence guarantee of SGLD to an approximate global minimum. Zhang et al. (2017b) analyze the hitting time of SGLD and prove that SGLD converges to an approximate local minimum. These results are further improved and generalized to a family of Langevin dynamics based algorithms by the subsequent work of Xu et al. (2018) .

1.2. OVERVIEW OF OUR RESULTS:

In this paper, we provide generalization guarantees for the noisy variants of several popular stochastic gradient methods.

The Bayes-Stability method and data-dependent generalization bounds. We develop a new method for proving generalization bounds, termed as Bayes-Stability, by incorporating ideas from the PAC-Bayesian theory into the stability framework. In particular, assuming the loss takes value in [0, C], our method shows that the generalization error is bounded by both 2C Ez[ 2KL(P,

Q z )] and 2C Ez[ 2KL(Q z , P )],

where P is a prior distribution independent of the training set S, and Q z is the expected posterior distribution conditioned on z n = z (i.e., the last training data is z). The formal definition and the results can be found in Definition 5 and Theorem 7.

Inspired by Lever et al. (2013) , instead of using a fixed prior distribution, we bound the KLdivergence from the posterior to a distribution-dependent prior. This enables us to derive the following generalization error bound that depends on the expected norm of the gradient along the optimization path:

EQUATION

)

Here S is the dataset and g e (t)

= EW t-1 [ 1 n n i=1 ∇F (W t-1 , z i ) 2 ]

is the expected empirical squared gradient norm at step t; see Theorem 11 for the details. (Mou et al., 2018, Theorem 1) , where L is the global Lipschitz constant of the loss, our new bound (3) depends on the data distribution and is typically tighter (as the gradient norm is at most L). In modern deep neural networks, the worstcase Lipschitz constant L can be quite large, and typically much larger than the expected empirical gradient norm along the optimization trajectory. Specifically, in the later stage of the training, the expected empirical gradient is small (see Figure 1 (d) for the details). Hence, our generalization bound does not grow much even if we train longer at this stage.

Compared with the previous O:

LC n t γ 2 t σ 2 t bound in

Our new bound also offers an explanation to the difference between training on correct and random labels observed by Zhang et al. (2017a) . In particular, we show empirically that the sum of expected squared gradient norm (along the optimization path) is significantly higher when the training labels are replaced with random labels (Section 3.1, Figure 1 , Appendix C.2).

We would also like to mention the PAC-Bayesian bound (for SGLD with 2 -regularization) proposed by Mou et al. (2018) . (This bound is different from what we mentioned before; see Theorem 2 in their paper.) Their bound scales as O(1/ √ n) and the numerator of their bound has a similar sum of gradient norms (with a decaying weight if the regularization coefficient λ > 0). Their bound is based on the PAC-Bayesian approach and holds with high probability, while our bound only holds in expectation.

terov's accelerated gradient method in Nesterov (1983) (Theorem 26), and Entropy-SGD proposed by Chaudhari et al. (2017) (Theorem 27). 2. Extension to general noises. The proof of the generalization bound in Mou et al. (2018) relies heavily on that the noise is Gaussian1 , which makes it difficult to generalize to other noise distributions such as the Laplace distribution. In contrast, our analysis easily carries over to the class of log-Lipschitz noises (i.e., noises drawn from distributions with Lipschitz log densities). 3. Pathwise stability. In practice, it is also natural to output a certain function of the entire optimization path, e.g., the one with the smallest empirical risk or a weighted average. We show that the same generalization bound holds for all such variants (Remark 12). We note that the analysis in an independent work of Pensia et al. (2018) also satisfies this property, Generalization bounds with 2 regularization via Log-Sobolev inequalities. We also study the setting where the total objective function F is the sum of a C-bounded differentiable objective F 0 and an additional 2 regularization term λ 2 w 2 2 . In this case, F can be treated as a perturbation of a quadratic function, and the continuous Langevin dynamics (CLD) is well understood for quadratic functions. We obtain two generalization bounds for CLD, both via the technique of Log-Sobolev inequalities, a powerful tool for proving the convergence rate of CLD. One of our bounds is as follows (Theorem 15):

yet their bound is O C 2 L 2 n -1 T t=1 η 2 t /σ 2 t (see

EQUATION

The above bound has the following advantages:

1. Applying e -x ≥ 1 -x, one can see that our bound is at most O( √ T /n), which matches the previous bound in (Mou et al., 2018, Proposition 8) 3 .

2. As time T grows, the bound is upper bounded by and approaches to 2e 4βC CLn -1 β/λ (unlike the previous O( √ T /n) bound that goes to infinity as T → +∞). 3. If the noise level is not so small (i.e., β is not very large), the generalization bound is quite desirable.

Our analysis is based on a Log-Sobolev inequality (LSI) for the parameter distribution at time t, whereas most known LSIs only hold for the stationary distribution of the Markov process. We prove the new LSI by exploiting the variational formulation of the entropy formula.

2. PRELIMINARIES:

Notations. We use D to denote the data distribution. The training dataset S = (z 1 , . . . , z n ) is a sequence of n independent samples drawn from D. S, S ∈ Z n are called neighboring datasets if and only if they differ at exactly one data point (we could assume without loss of generality that z n = z n ). Let F (w, z) and L(w, z) be the objective and the loss functions, respectively, where w ∈ R d denotes a model parameter and z ∈ Z is a data point. Define F (w, S) = 1 n n i=1 F (w, z i ) and F (w, D) = Ez∼D[F (w, z)]; L(w, S) and L(w, D) are defined similarly. A learning algorithm A takes as input a dataset S, and outputs a parameter w ∈ R d randomly. Let G be the set of all possible mini-batches. G n = {B ∈ G : n ∈ B} denotes the collection of mini-batches that contain the n-th data point, while

G n = G \ G n . Let diam(A) = sup x,y∈A x -y 2 denote the diameter of a set A. Definition 1 (L-lipschitz). A function F : R d × Z → R is L-lipschitz if and only if |F (w 1 , z) - F (w 2 , z)| ≤ L w 1 -w 2 2 holds for any w 1 , w 2 ∈ R d and z ∈ Z.

Definition 2 (Expected generalization error). The expected generalization error of a learning algorithm A is defined as

err gen := E S∼D n [err gen (S)] = E S∼D n ,A [L(A(S), D) -L(A(S), S)].

Algorithmic Stability. Intuitively, a learning algorithm that is stable (i.e., a small perturbation of the training data does not affect its output too much) can generalize well. In the seminal work of Bousquet & Elisseeff (2002) (see also Hardt et al. (2016) ), the authors formally defined algorithmic stability and established a close connection between the stability of a learning algorithm and its generalization performance. 

sup z∈Z |E A [L(w S , z)] -E A [L(w S , z)]| ≤ n ,

where w S and w S denote the outputs of A on S and S respectively. Lemma 4 (Generalization in expectation). (Hardt et al. (2016) ) Suppose a randomized algorithm A is n -uniformly stable. Then, |err gen | ≤ n .

3. BAYES-STABILITY METHOD:

In this section, we incorporate ideas from the PAC-Bayesian theory (see e.g., Lever et al. (2013) ) into the algorithmic stability framework. Combined with the technical tools introduced in previous sections, the new framework enables us to prove tighter data-dependent generalization bounds.

First, we define the posterior of a dataset and the posterior of a single data point. Definition 5 (Single-point posterior). Let Q S be the posterior distribution of the parameter for a given training dataset S = (z 1 , . . . , z n ). In other words, it is the probability distribution of the output of the learning algorithm on dataset S (e.g., for T iterations of SGLD in (1), Q S is the pdf of W T ). The single-point posterior Q (i,z) is defined as ,...,zi-1,z,zi+1,...,zn) .

Q (i,z) = E (z1,...,zi-1,zi+1,...zn) Q (z1

For convenience, we make the following natural assumption on the learning algorithm: Assumption 6 (Order-independent). For any fixed dataset S = (z 1 , . . . , z n ) and any permutation p, Q S is the same as Q S p , where S p = (z p1 , . . . , z pn ).

Assumption 6 implies Q (1,z) = • • • = Q (n,z)

, so we use Q z as a shorthand for Q (i,z) in the following. Note that this assumption can be easily satisfied by letting the learning algorithm randomly permute the training data at the beginning. It is also easy to verify that both SGD and SGLD satisfy the order-independent assumption. Now, we state our new Bayes-stability framework, which holds for any prior distribution P over the parameter space that is independent of the training dataset S.

Theorem 7 (Bayes-Stability). Suppose the loss function L(w, z) is C-bounded and the learning algorithm is order-independent (Assumption 6). Then for any prior distribution P not depending on S, the generalization error is bounded by both

2C E z 2KL(P, Q z ) and 2C E z 2KL(Q z , P ) .

Remark 8. Our Bayes-Stability framework originates from the algorithmic stability framework, and hence is similar to the notions of uniform stability and leave-one-out error (see Elisseeff et al. (2003) ). However, there are important differences. Uniform stability is a distribution-independent property, while Bayes-Stability can incorporate the information of the data distribution (through the prior P ). Leave-one-out error measures the loss of a learned model on an unseen data point, yet Bayes-Stability focuses on the extent to which a single data point affects the outcome of the learning algorithm (compared to the prior).

To establish an intuition, we first apply this framework to obtain an expectation generalization bound for (full) gradient Langevin dynamics (GLD), which is a special case of SGLD in (1) (i.e., GLD uses the full gradient ∇ w F (W t-1 , S) as g t (W t-1 )). Theorem 9. Suppose that the loss function L is C-bounded. Then we have the following expected generalization bound for T iterations of GLD:

err gen ≤ 2 √ 2C n E S∼D n T t=1 γ 2 t σ 2 t g e (t) ,

where

g e (t) = Ew∼W t-1 [ 1 n n i=1 ∇F (w, z i ) 2 2

] is the empirical squared gradient norm, and W t is the parameter at step t of GLD.

Proof The proof builds upon the following technical lemma, which we prove in Appendix A.2. Lemma 10. Let (W 0 , . . . , W T ) and (W 0 , . . . , W T ) be two independent sequences of random variables such that for each t ∈ {0, . . . , T }, W t and W t have the same support. Suppose W 0 and W 0 follow the same distribution. Then,

KL(W ≤T , W ≤T ) = T t=1 E w<t∼W<t [KL(W t |W <t = w <t , W t |W <t = w <t )],

where W ≤t denotes (W 0 , . . . , W t ) and W <t denotes W ≤t-1 .

Define P = E S∼D n-1 [Q (S,0) ],

where 0 denotes the zero data point (i.e., F (w, 0) = 0 for any w). Theorem 7 shows that

err gen ≤ 2C E z 2KL(Q z , P ).

(5)

By the convexity of KL-divergence, for a fixed z ∈ Z, we have

EQUATION

Let (W t ) t≥0 and(W t ) t≥0 be the training process of GLD for S = (S, z) and S = (S, 0), respectively. Note that for a fixed w <t , both

W t |W <t = w <t and W t |W <t = w <t are Gaussian distributions. Since KL(N (µ 1 , σ 2 I), N (µ 2 , σ 2 I)) = µ1-µ2 2 2 2σ 2 (see Lemma 18 in Appendix A.2). KL(W t |W <t = w <t , W t |W <t = w <t ) = γ 2 t ∇F (w t-1 , z) 2 2 σ 2 t n 2 .

Applying Lemma 10 and

KL(W T , W T ) ≤ KL(W ≤T , W ≤T ) gives KL(Q S , Q S ) ≤ 1 n 2 T t=1 γ 2 t σ 2 t E w∼Wt-1 ∇F (w, z) 2 2 .

Recall that W t-1 is the parameter at step t -1 using S = (S, z) as dataset. In this case, we can rewrite z as z n since it is the n-th data point of S. Note that SGLD satisfies the order-independent assumption, we can rewrite z as z i for all i ∈ [n]. Together with ( 5), (6), and using

1 n n i=1 √ x i ≤ 1 n n i=1

x i , we can prove this theorem.

More generally, we give the following bound for SGLD. The proof is similar to that of Theorem 9; the difference is that we need to bound the KL-divergence between two Gaussian mixtures instead of two Gaussians. This proof is more technical and deferred to Appendix A.3. Theorem 11. Suppose that the loss function L is C-bounded and the objective function f is Llipschitz. Assume that the following conditions hold:

1. Batch size b ≤ n/2.

2. Learning rate γ t ≤ σ t /(20L).

Then, the following expected generalization error bound holds for T iterations of SGLD (1):

err gen ≤ 8.12C n E S∼D n T t=1 γ 2 t σ 2 t g e (t) , (empirical norm)

where

g e (t) = Ew∼W t-1 [ 1 n n i=1 ∇F (w, z i ) 2 2 ]

is the empirical squared gradient norm, and W t is the parameter at step t of SGLD.

Furthermore, based on essentially the same proof, we can obtain the following bound that depends on the population gradient norm:

err gen ≤ 8.12C n E S T t=1 γ 2 t σ 2 t E w∼W t-1 E z∼D ∇F (w, z) 2 2

.

The full proofs of the above results are postponed to Appendix A, and we provide some remarks about the new bounds.

Remark 12. In fact, our proof establishes that the above upper bound holds for the two sequences W ≤T and W ≤T :

KL(W ≤T , W ≤T ) ≤ 8.12 n 2 T t=1 γ 2 t σ 2

t g e (t). Hence, our bound holds for any sufficiently regular function over the parameter sequences:

KL(f (W ≤T ), f (W ≤T )) ≤ 8.12 n 2 T t=1 γ 2 t σ 2

t g e (t). In particular, our generalization error bound automatically extends to several variants of SGLD, such as outputting the average of the trajectory, the average of the suffix of certain length, or the exponential moving average. Remark 13 (High-probability bounds). By relaxing the expected squared gradient norm term to L 2 and using the uniform stability framework, our proof can be adapted to recover the O(LC √ T /n) bound in (Mou et al., 2018, Theorem 1) . Then, we can apply the recent results of Feldman & Vondrak (2019) to provide a generalization error bound of Õ(LC √ T /n + 1/ √ n) that holds with high probability. (Here Õ hides poly-logarithmic factors.) When T is at least linear in n, the additional 1/ √ n term is not dominating.

3.1. EXPERIMENT:

Distinguish random from normal. Inspired by Zhang et al. (2017a) , we run both GLD (Figure 1 ) and SGLD (Appendix C.2) to fit both normal data and randomly labeled data (see Appendix C for more experiment details). As shown in Figure 1 and Figure 3 in Appendix C.2, a larger random label portion p leads to both much higher generalization error and much larger generalization error bound. Moreover, the shapes of the curves of our bounds look quite similar to those of the generalization error curves.

Note that in (b) and (c) of Figure 1 , the scales in the y-axis are different. We list some possible reasons that may explain why our bound is larger than the actual generalization error.

(1) as we explained in Remark 12, our bounds (Theorem 9 and 11) hold for any trajectory-based output, and are much stronger than upper bounds for the last point on the trajectory.

(2) The constant we can prove in Lemma 21 may not be very tight.

(3) The variance of Gaussian noise σ 2 t is not large enough in our experiment. However, if we choose a larger variance, fitting the random labeled training data becomes quite slow. Hence, we use a small data size (n = 10000) for the above reason. We also run an extra experiment for GLD on the full MNIST dataset (n = 60000) without label corruption (see Figure 2 in the Appendix C). We can see that our bound is non-vacuous (since GLD-which computes the full gradients-took a long time to converge, we stopped when we achieved 90% training accuracy). 4Published as a conference paper at ICLR 2020 Relax the step size constraint. The condition on the step size in Theorem 11 may seem restrictive in the practical use. 5 We provide several ways to relax this constraint:

1. The proof of Theorem 11 still goes through if we replace L with

max i∈[n] ∇F (W t-1 , z i ) 2 in the constraint.

2. The maximum gradient norm can be controlled by gradient clipping, i.e., multiplying

min(C L , ∇F (Wt-1,zi) 2 ) ∇F (Wt-1,zi) 2 to each ∇F (W t-1 , z i ).

3. Replacing the constant 20 with 2 in this constraint will only increase the constant of our bound from 8.12 to 84.4.

We also provide an experiment combining the above ideas to make our Theorem 11 applicable in the practical use (see Figure 4 in Appendix C).

4. GENERALIZATION OF CLD AND GLD WITH 2 REGULARIZATION:

In this section, we study the generalization error of Continuous Langevin Dynamics (CLD) with 2 regularization. Throughout this section, we assume that the objective function over training set S is defined as F (w, S) = F 0 (w, S) + λ 2 w 2 2 , and moreover, the following assumption holds. Assumption 14. The loss function L and the original objective F 0 are C-bounded. Moreover, F 0 is differentiable and L-lipschitz.

The Continuous Langevin Dynamics is defined by the following SDE:

dW t = -∇F (W t , S) dt + 2β -1 dB t , W 0 ∼ µ 0 , (CLD)

where (B t ) t≥0 is the standard Brownian motion on R d and the initial distribution µ 0 is the centered Gaussian distribution in R d with covariance 1 λβ I d . We show that the generalization error of CLD is upper bounded by O e 4βC n -1 β/λ , which is independent of the training time T (Theorem 15). Furthermore, as T goes to infinity, we have a tighter generalization error bound O βC 2 n -1 (Theorem 39 in Appendix B). We also study the generalization of Gradient Langevin Dynamics (GLD), which is the discretization of CLD:

W k+1 = W k -η∇F (W k , S) + 2ηβ -1 ξ k , (GLD)

where ξ k is the standard Gaussian random vector in R d . By leveraging a result developed in Raginsky et al. ( 2017), we show that, as Kη 2 tends to zero, GLD has the same generalization as CLD (see Theorems 15 and 39). We first formally state our first main result in this section.

Theorem 15. Under Assumption 14, CLD (with initial probability measure dµ 0 = 1 Z e -λβ w 2 2 dw) has the following expected generalization error bound:

EQUATION

)

In addition, if L is M -smooth and non-negative, by setting λβ > 2, λ > 0 and η ∈ 0, 1 ∧ λ 8M 2 , GLD (running K iterations with the same µ 0 as CLD) has the expected generalization error bound:

EQUATION

)

where C 1 is a constant that only depends on M , λ, β, b, L and d.

The following lemma is crucial for establishing the above generalization bound for CLD. In particular, we need to establish a Log-Sobolev inequality for µ t , the parameter distribution at time t, for every time step t > 0. In contrast, most known LSIs only characterize the stationary distribution of the Markov process. The proof of the lemma can be found in Appendix B. Lemma 16. Under Assumption 14, let µ t be the probability measure of W t in CLD (with

dµ 0 = 1 Z e -λβ w 2 2 dw).

Let ν be a probability measure that is absolutely continuous with respect to µ t . Suppose dµ t = π t (w) dw and dν = γ(w) dw. Then, it holds that

KL(γ, π t ) ≤ exp(8βC) 2λβ R d ∇ log γ(w) π t (w) 2 2 γ(w) dw.

We sketch the proof of Theorem 15, and the complete proof is relegated to Appendix B.

Proof Sketch of Theorem 15 Suppose S and S are two neighboring datasets. Let (W t ) t≥0 and (W t ) t≥0 be the process of CLD running on S and S , respectively. Let γ t and π t be the pdf of W t and W t . Let F S (w) denote F (w, S). We have

d dt KL(γ t , π t ) = -1 β R d γ t ∇ log γ t π t 2 2 dw + R d γ t ∇ log γ t π t , ∇F S -∇F S dw ≤ -1 2β R d γ t ∇ log γ t π t 2 2 dw + β 2 R d γ t ∇F S -∇F S 2 2 dw. ≤ -λ e 8βC KL(γ t , π t ) + 2βL 2 n 2 (Lemma 16)

Solving this inequality gives KL(γ t , π t ) ≤ 1 n 2 λ 2βL 2 e 8βC (1 -e -λt/e 8βC ). Hence the generalization error of CLD can be bounded by 2C 1 2 KL(γ T , π T ), which proves the first part. The second part of the theorem follows from Lemma 36 in Appendix B.

Our second generalization bound for CLD (Theorem 39 in Appendix B) is

err gen ≤ 8βC 2 n + 4C exp -λT e 4βC βC.

One future direction is to study the discrete noise introduced in SGD (in this case the KL-divergence may not be well defined). For either SGLD or CLD, if the noise level is small (i.e., β is large), it may take a long time for the diffusion process to reach the stable distribution. Hence, another interesting future direction is to consider the local behavior and generalization of the diffusion process in finite time through the techniques developed in the studies of metastability (see e.g., Bovier et al. (2005) ; Bovier & den Hollander (2006) ; Tzen et al. (2018) ). In particular, the technique may be helpful for further improving the bounds in Theorems 15 and 39 (when T is not very large).

A PROOFS IN SECTION 3:

A.1 BAYES-STABILITY FRAMEWORK Lemma 17. Under Assumption 6, for any prior distribution P not depending on the dataset S = (z 1 , . . . , z n ), the generalization error is upper bounded by

E z E w∼P L(w, z) -E w∼Qz L(w, z) + E z E w∼P L(w) -E w∼Qz L(w) ,

where L(w) denotes the population loss L(w, D).

Proof of Lemma 17:

Let err train = ES Ew∼Q S L(w, S) and err test = ES Ew∼Q S L(w). We can rewrite generalization error as err gen = err test -err train , where

err test = E z E w∼Q (1,z) L(w) = E z E w∼Qz L(w) (Assumption 6) = E z R d (Q z (w) -P (w))L(w) dw + R d P (w)L(w) dw.

and

err train = 1 n n i=1 E S E w∼Q S L(w, z i ) = 1 n n i=1 E z E w∼Q (i,z) L(w, z) = E z E w∼Qz L(w, z) (Assumption 6) = E z R d (Q z (w) -P (w))L(w, z) dw + R d P (w) E z L(w, z) dw (P is a prior) = E z R d (Q z (w) -P (w))L(w, z) dw + R d P (w)L(w) dw.

(definition of f (w))

Thus, we have

|err gen | = |err test -err train | = E z R d (Q z (w) -P (w))L(w) dw -E z R d (Q z (w) -P (w))L(w, z) dw ≤ E z E w∼Qz L(w, z) -E w∼P L(w, z) + E z E w∼Qz L(w) -E w∼P L(w) .

Now we are ready to prove Theorem 7, which we restate in the following.

Theorem 7 (Bayes-Stability). Suppose the loss function L(w, z) is C-bounded and the learning algorithm is order-independent (Assumption 6), then for any prior distribution P not depending on S, the generalization error is bounded by both 2C E z 2KL(P, Q z ) and 2C E z 2KL(Q z , P ) .

Proof By Lemma 17,

err gen ≤ E z E w∼P L(w, z) -E w∼Qz L(w, z) + E z E w∼P L(w) -E w∼Qz L(w) ≤ E z [2C • TV(P, Q z ) + 2C • TV(P, Q z )] (C-boundedness) ≤ 4C E z 1 2 KL(P, Q z ) (Pinsker's inequality)

The other bound follows from a similar argument.

Published as a conference paper at ICLR 2020 A.2 TECHNICAL LEMMAS Now we turn to the proof of Theorem 11. The following lemma allows us to reduce the proof of algorithmic stability to the analysis of a single update step. Lemma 10. Let (W 0 , . . . , W T ) and (W 0 , . . . , W T ) be two independent sequences of random variables such that for each t ∈ {0, . . . , T }, W t and W t have the same support. Suppose W 0 and W 0 follow the same distribution. Then,

KL(W ≤T , W ≤T ) = T t=1 E w<t∼W<t [KL(W t |W <t = w <t , W t |W <t = w <t )],

where W ≤t denotes (W 0 , . . . , W t ) and W <t denotes W ≤t-1 .

Proof By the chain rule of the KL-divergence,

KL(W ≤t , W ≤t ) = KL(W <t , W <t ) + E w<t∼W<t [KL(W t |W <t = w <t , W t |W <t = w <t )].

The lemma follows from a summation over t = 1, . . . , T .

The following lemma (see e.g., (Duchi, 2007, Section 9 )) gives a closed-form formula for the KLdivergence between two Gaussian distributions. Lemma 18. Suppose that P = N (µ 1 , Σ 1 ) and Q = N (µ 2 , Σ 2 ) are two Gaussian distributions on R d . Then,

KL(P, Q) = 1 2 tr(Σ -1 2 Σ 1 ) + (µ 2 -µ 1 ) Σ -1 2 (µ 2 -µ 1 ) -d + ln det(Σ 2 ) det(Σ 1

) .

The following lemma (Topsoe, 2000, Theorem 3 ) helps us to upper bound the KL-divergence. Definition 19. Let P and Q be two probability distributions on R d . The directional triangular discrimination from P to Q is defined as

∆ * (P, Q) = +∞ k=0 2 k • ∆ 2 -k P + (1 -2 -k )Q, Q ,

where

∆ (P, Q) = R d (P (w) -Q(w)) 2 P (w) + Q(w) dw.

Lemma 20. For any two probability distributions P and Q on R d ,

KL(P, Q) ≤ ln 2 • ∆ * (P, Q) .

Recall that G is the set of all possible mini-batches. 

G n = {B ∈ G : n ∈ B} denotes the collection of mini-batches that contain n, while G n = G \ G n . diam(A) = sup x,

1. µ B -µ B ≤ β for B ∈ G n and µ B = µ B for B ∈ G n . 2. diam({µ B : B ∈ G} ∪ {µ B : B ∈ G}) ≤ σ/

2 k • ∆ 2 -k P + (1 -2 -k )P , P = ln 2 • +∞ k=0 2 k • R d 4 -k (P (w) -P (w)) 2 2 -k P (w) + (2 -2 -k )P (w)

dw.

The numerator of the above integrand is upper bounded by

4 -k (P -P ) 2 = 4 -k 1 |G| B∈G (p µ B ,σ -p µ B ,σ ) 2 = 4 -k |G n | 2 |G| 2 1 |G n | B∈Gn (p µ B ,σ -p µ B ,σ ) 2 ≤ 4 -k b 2 n 2 • 1 |G n | B∈Gn (p µ B ,σ -p µ B ,σ ) 2 , (9)

while the denominator can be lower bounded as follows:

2 -k P + (2 -2 -k )P ≥ 2 -k |G| B∈Gn p µ B ,σ + 2 -2 -k |G| B∈Gn p µ B ,σ = 2 |G| B∈Gn p µ B ,σ (µ B = µ B for B ∈ G n ) = 1 |G n | • 2(n -b) n B∈Gn p µ B ,σ ≥ 1 |G n | B∈Gn p µ B ,σ , (b ≤ n/2)

which implies, by the convexity of 1/x, that

EQUATION

Inequalities ( 9) and ( 10) together imply

EQUATION

(11)

Now we bound the right-hand side of (11) for fixed A and B. By applying a translation and a rotation, we can assume without loss of generality that µ A = 0, and the last d -2 coordinates of µ B and µ B are all zero. Note that the integral is unchanged when we project the space to the twodimensional subspace corresponding to the first two coordinates. Thus, it suffices to prove a bound for d = 2. We rewrite (11) as

EQUATION

Let I be the integral in the right-hand side of (12). Note that µ B σ ,

µ B σ ≤ 0.1 and µ B -µ B σ ≤ β σ . Let δ = β σ and r = w σ . Our goal is to bound max δ∈[0,0.1] (Iδ -2 ). Let (x) + = max(x, 0). Since I ≤ σ 2 ∞ 0 max y∈[(r-0.1) + ,r+0.1] (e -y 2 -e -(y+δ) 2 ) 2

e -r 2 2πr dr.

We have

max δ∈[0,0.1] I δ 2 ≤ σ 2 ∞ 0 e r 2 2πr max y∈[(r-0.1) + ,r+0.1] max δ∈[0,0.1] (e -y 2 -e -(y+δ) 2 ) 2 δ 2 dr.

Let φ(y, δ) = ( e -y 2 -e -(y+δ) 2 δ

) 2 , we make two claims which we will prove later:

1. For all y, δ ≥ 0, φ(y, δ) ≤ 2 e . 2. For all y ≥ 1 √ 2 , φ(y, δ) is non-increasing in δ.

The above claims imply that:

1. For any r ∈ 0, 1 √ 2 + 0.1 , max y∈[(r-0.1) + ,r+0.1],δ∈[0,0.1] [φ(y, δ)] ≤ 2 e .

2. For any r ∈ 1

√ 2 + 0.1, +∞ , we have max y∈[(r-0.1) + ,r+0.1],δ∈[0,0.1] φ(y, δ) ≤ max y∈[(r-0.1) + ,r+0.1] lim δ→0 [φ(y, δ)] = max y∈[(r-0.1) + ,r+0.1]

4y 2 e -2y 2 = 4(r -0.1) 2 e -2(r-0.1) 2 .

The last step holds since y → y 2 e -2y 2 is decreasing on 1 √ 2 , +∞ .

Thus we have

max δ∈[0,0.1] I δ 2 ≤ σ 2 1 √ 2 +0.1 0 e r 2 2πr • 2 e dr + σ 2 +∞ 1 √ 2 +0

.1 e r 2 2πr • 4(r -0.1) 2 e -2(r-0.1) 2 dr ≤ 18.6487σ 2 . Plugging the above into (12) gives

∆(2 -k P + (1 -2 -k )P , P ) ≤ 4 -k b 2 n 2 σ 2 π δ 2 max δ∈[0,0.1] (I/δ 2 ) ≤ 4 -k b 2 n 2 π • 18.6487δ 2 .

We conclude that

KL(P, P ) ≤ ln 2 +∞ k=0 2 k • ∆(2 -k P + (1 -2 -k )P , P ) ≤ 37.2974b 2 δ 2 ln 2 n 2 π ≤ 8.23b 2 β 2 n 2 σ 2 .

Finally, we prove the two claims used above:

1. For all y, δ ≥ 0, let h(x) = e -x 2 , we have e -y 2 -e -(y+δ

) 2 = y y+δ h (t) dt. Since |h (t)| = |e -t 2 (-2t)| = e -t 2 (2t). Let ∂ ∂t |h (t)| = 0, we have t = 1/ √ 2. Thus, |h (t)| ≤ e 1/2 √

2 and e -y 2 -e -(y+δ) 2 ≤ δ 2/e.

2.. Suppose y:

≥ 1/ √ 2, we have ∂ ∂δ e -y 2 -e -(y+δ) 2 δ = e -(y+δ) 2 (2yδ + 2δ 2 + 1) -e -y 2 δ 2 = e -y 2 δ 2 [e -2yδ-δ 2 (2yδ + 2δ 2 + 1) -1].

Let g(y, δ) = e -2yδ-δ 2 (2yδ + 2δ 2 + 1) -1. Note that Published as a conference paper at ICLR 2020

(a) lim δ→0 ∂ ∂δ ( e -y 2 -e -(y+δ) 2 δ ) = 0. (b) ∂ ∂δ g(1/ √ 2, δ) = -4e -δ(δ+ √ 2) δ 2 (δ + √ 2) < 0 and g(1/ √ 2, 0) = 0. It implies that ∂ ∂δ ( e -(1/ √ 2)) 2 -e -(1/ √ 2+δ) 2 δ

) ≤ 0 for δ > 0. Since ∂ ∂y g(y, δ) = -4δ 2 e -δ(δ+2y) (δ + y) < 0 for y ≥ 0, we conclude that for any y ≥ 1/ √ 2:

∂ ∂δ e -y 2 -e -(y+δ) 2 δ ≤ ∂ ∂δ e -(1/ √ 2) 2 -e -(1/ √ 2+δ) 2 δ ≤ 0.

A.3 MAIN THEOREM:

Recall that SGLD on dataset S is defined as

W t ← W t-1 -γ t ∇ w F (W t-1 , S Bt ) + σ t √ 2 N (0, I d ).

Here γ t is the step size. B t = {i 1 , . . . , i b } is a subset of {1, . . . , n} of size b, and

S Bt = (z i1 , . . . , z i b ) is the mini-batch indexed by B t . Recall that F (w, S) denotes 1 |S| |S| i=1 F (w, z i ).

We restate and prove Theorem 11 in the following. Theorem 11. Suppose that the loss function L is C-bounded and the objective function F is Llipschitz. Assume that the following conditions hold:

1. Batch size b ≤ n/2. 2. Learning rate γ t ≤ σ t /(20L).

Then, the following expected generalization error bound holds for T iterations of SGLD (1):

err gen ≤ 8.12C n E S∼D n T t=1 γ 2 t σ 2 t g e (t) , (empirical norm)

where

g e (t) = Ew∼W t-1 [ 1 n n i=1 ∇F (w, z i ) 2 2

] is the empirical squared gradient norm, and W t is the parameter at step t of SGLD.

Proof By Theorem 7, we have

EQUATION

for any prior distribution P . In particular, we define the prior as P (w) = E S∼D n-1 [P S (w)], where P S (w) = Q (S,0) . By the convexity of KL-divergence,

EQUATION

)

Fix a data point z ∈ Z. Let (W t ) t≥0 and (W t ) t≥0 be the training process of SGLD for S = (S, z) and S = (S, 0), respectively. Fix a time step t and w <t = (w 0 , . . . , w t-1 ). Let P t and P t denote the distribution of W t and W t conditioned on W <t = w <t and W <t = w <t , respectively. By the definition of SGLD, we have

P t = 1 |G| B∈G p µ B and P t = 1 |G| B∈G p µ B , where µ B = w t-1 - γ t ∇ w F (w t-1 , S B ), µ B = w t-1 -γ t ∇ w F (w t-1 , S B )

, and p µ denotes the Gaussian distribution 

N (µ, σ 2 t 2 I d ). We note that: 1. µ B -µ B ≤ γt ∇F (wt-1,z) 2 b for B ∈ G n and µ B = µ B for B ∈ G n . 2. diam({µ B : B ∈ G} ∪ {µ B : B ∈ G}) ≤ 2γ t L ≤ σ t /

≤ T t=1 E w∼Wt-1 8.23γ 2 t ∇F (w, z) 2 2 σ 2 t n 2

, which implies that

KL(Q S , Q S ) = KL(W T , W T ) ≤ KL(W ≤T , W ≤T ) ≤ 8.23 n 2 T t=1 γ 2 t σ 2 t E w∼Wt-1 ∇F (w, z) 2 2 .

Together with ( 13) and ( 14), we have

err gen ≤ 2C E z 2 E S 8.23 n 2 T t=1 γ 2 t σ 2 t E w∼Wt-1 ∇F (w, z) 2 2 ≤ 2C 2 E S 8.23 n 2 T t=1 γ 2 t σ 2 t E w∼Wt-1 ∇F (w, z n ) 2 2 . (concavity of √ x)

Since SGLD is order-independent, we can replace ∇F (w, z n ) with ∇F (w, z i ) for any i ∈ [n] in the right-hand side of the above bound. Our theorem then follows from the concavity of √ x. Furthermore, if we bound KL(P, Q z ) instead of KL(Q z , P ) in the above proof, we obtain the following bound that depends on the population squared gradient norm:

err gen ≤ 8.12C n E S T t=1 γ 2 t σ 2 t E w∼W t-1 E z∼D ∇F (w, z) 2 2 .

A.4 EXTENSION TO GENERAL NOISES:

We can extend the generalization bounds in previous sections, which require the noise to be Gaussian, to other general noises, namely the family of log-lipschitz noises.

Definition 22 (Log-Lipschitz Noises). A probability distribution on R d with density p is L-loglipschitz if and only if ∇ ln p(w) ≤ L holds for any w ∈ R d . A random variable ζ is called an L-log-lipschitz noise if and only if it is drawn from an L-log-lipschitz distribution.

The analog of SGLD, noisy momentum method (Definition 24), and noisy NAG (Definition 25) can be naturally defined by replacing the Gaussian noise ζ t at each iteration with an independent L-log-lipschitz noise in the definition.

The following lemma is an analog of Lemma 21 under L-log-lipschitz noises. Recall that G denotes a collection of mini-batches of size b. Lemma 23 readily implies the analogs of Theorems 11, 26 and 27 under more general noise distributions.

Lemma 23. Suppose that batch size b ≤ n/2 and N is an L noise -log-lipschitz distribution on R d . {µ B : B ∈ G} and {µ B : B ∈ G} are two collections of points in R d that satisfy the following conditions for some constant β ∈ 0, 1 Lnoise :

1. µ B -µ B ≤ β for B ∈ G n and µ B = µ B for B ∈ G n . 2. diam({µ B : B ∈ G} ∪ {µ B : B ∈ G}) ≤ 1.

For µ ∈ R d , let p µ denote the distribution of ζ +µ when ζ is drawn from N . Let P = 1 |G| B∈G p µ B and P = 1 |G| B∈G p µ B be mixture distributions over all mini-batches. Then,

KL(P, P ) ≤ C 0 b 2 β 2 n 2

for some constant C 0 that only depends on L noise .

Proof of Lemma 23 Following the same argument as in the proof of Lemma 21, we have

EQUATION

where

EQUATION

Fixed A ∈ G n and B ∈ G n . Let p noise denote the density of the noise distribution N . Since µ A -µ B ≤ 1 and p noise is L noise -log-lipschitz, we have

p µ B (w) = p noise (w -µ B ) ≤ p noise (w -µ A ) • e Lnoise µ A -µ B ≤ e Lnoise p µ A (w).

Similarly, since µ B -µ B ≤ β, we have e -βLnoise p µ B (w) ≤ p µ B (w) ≤ e βLnoise p µ B (w).

Then, it follows from βL noise ≤ 1 that

(p µ B (w) -p µ B (w)) 2 ≤ (e βLnoise -1) 2 p µ B (w) 2 ≤ β 2 L 2 noise p µ B (w) 2 .

Therefore, the integral on the righthand side of ( 16) can be upper bounded as follows:

R d (p µ B (w) -p µ B (w)) 2 p µ A (w) dw ≤β 2 L 2 noise R d p µ B (w) 2 p µ A (w) dw ≤β 2 L 2 noise R d p µ B (w) • e Lnoise dw =β 2 L 2 noise e Lnoise .

Plugging the above inequality into (15) and ( 16) gives

∆(2 -k P + (1 -2 -k )P , P ) ≤ 4 -k b 2 n 2 |G n ||G n | A∈Gn B∈Gn β 2 L 2 noise e Lnoise = L 2 noise e Lnoise • 4 -k b 2 β 2 n 2 .

and

KL(P, P ) ≤ ln 2 • +∞ k=0 2 k L 2 noise e Lnoise • 4 -k b 2 β 2 n 2 = 2 ln 2L 2 noise e Lnoise • b 2 β 2 n 2 .

Algorithm 1: Entropy-SGD Input: Training set S = (z 1 , .., z n ) and loss function g(w, z).

Hyper-parameters: Scope γ, SGD learning rate η, SGLD step size η and batch size b.

1 for t = 1 to T do 2 //SGD iteration 3 W t,0 , µ t,0 ← W t-1,K+1 ; 4 for k = 0 to K -1 do 5 //SGLD iteration 6 B t,k ← mini-batch with size b; 7 W t,k+1 ← W t,k -η ∇ w g(W t,k , S B t,k ) + η γ(W t-1,K+1 -W t,k ) + √ η εN (0, 1 2 I d ); 8 µ t,k+1 ← (1 -α)µ t,k + αW t,k ; 9 end 10 W t,K+1 ← W t,K -ηγ(W t,K -µ t,K ); 11 end 12 return W T,K+1 ;

Theorem 27. Suppose that the loss function L is C-bounded and the objective function F is Llipschitz. If batch size b ≤ n/2 and √ η ≤ ε/(20L), the following expected generalization error bound holds for Entropy-SGD:

EQUATION

where

g e (t, k) = Ew∼W t,k [ 1 n n i=1 ∇F (w, z i ) 2 2 ]

is the empirical squared gradient norm, and W t,k denotes the training process with respect to S.

Since g e (t, k) is at most L 2 , it further implies the generalization error of Entropy-SGD is bounded by

O C √ η L εn √ T K .

Proof of Theorem 27 Define the history before time step (t, k) as follows:

W ≤(t,k) = (W 0,0 , ..., W 0,K+1 , ..., W t-1,0 , ..., W t-1,K+1 , W t,0 , ..., W t,k ).

Since µ is only determined by W , we only need to focus on W . This proof is similar to the proof of Theorem 11. By setting ,0) ]. Suppose S = (S, z) and S = (S, 0) are fixed, let W and W denote their training process, respectively. Considering the following 3 cases:

P = E S [Q (S

1. W t,0 ← W t-1,K+1 : In this case, for a fixed w ≤(t-1,K+1) , we have

KL W t,0 |W ≤(t-1,K+1) = w ≤(t-1,K+1) , W t,0 |W ≤(t-1,K+1) = w ≤(t-1,K+1) = 0. 2. W t,k+1 ← W t,k -η ∇ w g(W t,k , S B t,k ) + η γ(W t-1,K+1 -W t,k ) + √ η εN (0, 1 2 I d ):

In this case, fix a w ≤(t,k) , applying Lemma 21 gives

KL W t,k+1 |W ≤(t,k) = w ≤(t,k) , W t,k+1 |W ≤(t,k) = w ≤(t,k) ≤ 8.23η ∇F (w t,k , z) 2 2 ε 2 n 2 .

3. W t,K+1 ← W t,K -ηγ(W t , K -µ t,K ): In this case, for a fixed w ≤(t,K) , we have

KL W t,K+1 |W ≤(t,K) = w ≤(t,K) , W t,K+1 |W ≤(t,K) = w ≤(t,K) = 0.

By applying Lemma 10, we have

KL(W T,K+1 , W T,K+1 ) ≤ 8.23η ε 2 n 2 T t=1 K-1 k=0 g e (t, k),

and Where g e (t, k) is the empirical squared gradient norm of the k-th SGLD iteration in the t-th SGD iteration, respectively. The rest of the proof is the same as the proof of Theorem 11.

Lemma 31 (Bounded perturbation). Assume that the probability measure ν satisfies LS(α) (with respect to Γ). Let µ be a probability measure such that 1/b ≤ dµ/dν ≤ b for some constant b > 1.

Then µ satisfies LS(b 2 α) (with respect to Γ).

In fact, Lemma 31 is a simple consequence of the following variational formula in the special case that φ(x) = x log x, which we will also need in our proof:

Lemma 32 (Variational formula). (see .g., (Bakry et al., 2013, p. 240 )) Let φ : I → R on some open interval I ⊂ R be convex of class C 2 . For every (bounded or suitably integrable) measurable function f : R d → R with values in I,

EQUATION

)

It is worth noting the integrand of the right-hand side is nonnegative due to the convexity of φ.

B.2 LOGARITHMIC SOBOLEV INEQUALITY FOR CLD:

Recall that F S (w) = F (w, S) := F 0 (w, S) + λ w 2 2 /2 is the sum of the empirical original objective F 0 (w, S) 2 regularization. Let dµ = 1 Zµ exp(-βF S (w)) dw be the invariant (Gibbs) measure of CLD, and ν is the centered Gaussian measure dν = 1

Zν exp(-βλ w 2 2 /2) dw. Invoking Lemma 30 with σ 2 = 1 λβ shows that ν satisfies LS(1/λ) (with respect to Γ). Consider the density h(w) = dµ dν = Zν Zµ exp(-βF 0 (w, S)). If the original objective function F 0 is C-bounded, we have exp(-2βC) ≤ h(w) ≤ exp(2βC). By applying Lemma 31 with b = exp(2βC), we have the following lemma.

Lemma 33. Under Assumption 14, let Γ(f, g) = β -1 ∇f, ∇g be the carré du champ operator of the diffusion semigroup associated to CLD, and µ be the invariant measure of the SDE. Then, µ satisfies LS(e 4βC /λ) with respect to Γ. Let µ t be the probability measure of W t . By definition of P t , for any real-valued bounded measurable function f on R d and any s, t ≥ 0,

EQUATION

In particular, if the invariant measure µ = µ ∞ exists, we have

EQUATION

The following lemma is crucial for establishing the first generalization bound for CLD. In fact, we establish a Log-Sobolev inequality for µ t , the parameter distribution at time t, for any time t > 0.

Note that our choice of the initial distribution µ 0 is important for the proof.6 

Lemma 34. Under Assumption 14, let µ t be the probability measure of W t in (CLD) with initial probability measure dµ 0 = 1 Z e -λβ w 2 2

dw. Let Γ be the carré du champ operator of diffusion semigroup associated to (CLD). Then, for any f : R d → R + in D(E):

Ent µt (f ) ≤ e 8βC /λ 2 R d Γ(f ) f dµ t

Proof Let µ be the invariant measure of CLD. By Lemma 33 and Definition 29,

EQUATION

By applying Lemma 32 with φ(x) = x log x, we rewrite the left-hand side as

Ent µ (f ) := R d f log f dµ - R d f dµ log R d f dµ = inf r∈I R d [φ(f ) -φ(r) -φ (r)(f -r)] dµ = inf r∈I R d [P t (φ(f ) -φ(r) -φ (r)(f -r))] dµ.

where the last equation holds by the definition of invariant measure P t f dµ = f dµ. Thus, we have

EQUATION

Let µ t be the probability measure of W t . Lemma 32 and ( 22) together imply that

EQUATION

Since P t (φ(f ) -φ(r) -φ (r)(f -r)) ≥ 07 and dµ0 dµ ≤ exp(2βC), we have

EQUATION

Since dµ dµ0 ≤ exp(2βC) and µ is the invariant measure, we conclude that

EQUATION

Lemma 16. Under Assumption 14, let µ t be the probability measure of W t in CLD (with dµ 0 = 1 Z e -λβ w 2 2 dw). Let ν be a probability measure that is absolutely continuous with respect to µ t . Suppose dµ t = π t (w) dw and dν = γ(w) dw. Then it holds that:

EQUATION

Proof Let f (w) = γ(w)/π t (w), by Lemma 34 and R d f dµ t = 1, we have

EQUATION

We can see that the left-hand side is equal to KL(γ, π t ) 8 , and the right-hand side is equal to

e 8βC 2λβ R d ∇ γ(w) πt(w) 2 2 γ(w)/π t (w) π t (w) dw = e 8βC 2λβ R d ∇ log γ(w) π t (w) 2 2 γ(w) dw.

This concludes the proof. . We can rewrite F S (w) = 1 n n i=1 h(w, z i ). Define µ S,k and ν S,t as the probability measure of W k (in GLD) and W t (in CLD), respectively. Raginsky et al. (2017) provided a bound of KL(µ S,k , ν S,ηK ) under Assumption 35. This bound enables us to derive a generalization error bound for the discrete GLD from the bound for the continuous CLD. We use the assumption from Raginsky et al. (2017) . Their work considers the following SGLD:

W k+1 = W k -ηg S (W k ) + 2ηβ -1 ξ k .

Where g S (w k ) is a conditionally unbiased estimate of the gradient ∇F S (w k ). In our GLD setting,

g S (W k ) is equal to ∇F S (w k ). Assumption 35. Let F S (w) = 1 n n i=1 h(w, z i ) = F 0 (w, S) + λ 2 w 2 2 .

1. The function h takes non-negative real values, and there exist constants 4. There exists a constant δ ∈ [0, 1), such that, for each S ∈ Z n ,

A, B ≥ 0, such that |h(0, z)| ≤ A and ∇h(0, z) 2 ≤ B ∀z ∈ Z. 2. For each z ∈ Z, the function h(•, z) is M -smooth: for some M > 0, ∇h(w, z) -∇h(v, z) 2 ≤ M w -v 2 , ∀w, v ∈ R d .

E[ g S (w) -∇F S (w) 2 2 ] ≤ 2δ M 2 w 2 2 + B 2 , ∀w ∈ R d .

5. The probability law µ 0 of the initial hypothesis W 0 has a bounded and strictly positive density p 0 with respect to the Lebesgue measure on R d , and dw) has the following expected generalization error bound:

κ 0 := log R d e w 2 2 p 0 (w) dw < ∞.

err gen ≤ 2e 4βC CL n β λ 1 -exp - λT e 8βC (31) 8 Indeed, R d f log f dµt = R d γ π t log( γ π t )πtdw = KL(γ, πt)

In addition, if F 0 is also M -smooth and non-negative, by setting λβ > 2, λ > 0 and η ∈ [0, 1 ∧ λ 8M 2 ), the GLD (running K iterations with the same µ 0 as CLD) has the expected generalization error bound:

EQUATION

where C 1 is a constant that only depends on M , λ, β, b, L and d.

Proof of Theorem 15:

We apply the uniform stability framework. Suppose S and S are two neighboring datasets that differ on exactly one data point. Let (W t ) t≥0 and (W t ) t≥0 be the process of CLD running on S and S , respectively. Let γ t and π t be the pdf of W t and W t . We have

EQUATION

According to Fokker-Planck equation (see Risken (1996) ) for CLD, we know that

∂γ t ∂t = 1 β ∆γ t + ∇ • (γ t ∇F S ), ∂π t ∂t = 1 β ∆π t + ∇ • (π t ∇F S ).

It follows that

I := R d dγ t dt log γ t π t dw = R d 1 β ∆γ t + ∇ • (γ t ∇F S ) log γ t π t dw = -1 β R d ∇ log γ t π t , ∇γ t dw - R d

∇ log γ t π t , γ t ∇F S dw, (integration by parts)

and

J := R d γ t π t dπ t dt dw = R d γ t π t 1 β ∆π t + ∇ • (π t ∇F S ) dw = -1 β R d ∇ γ t π t , ∇π t dw - R d ∇ γ t π t , π t ∇F S dw.

(integration by parts)

Together with (33), we have

d dt KL(γ t , π t ) = I -J = -1 β R d ∇γ t γ t - ∇π t π t , ∇γ t - ∇γ t π t - γ t ∇π t π 2 t , ∇π t dw - R d ∇ log γ t π t , γ t ∇F S - γ t π t ∇ log γ t π t , π t ∇F S dw = -1 β R d γ t ∇ log γ t π t 2 2 dw + R d γ t ∇ log γ t π t , ∇F S -∇F S dw ≤ -1 2β R d γ t ∇ log γ t π t 2 2 dw + β 2 R d γ t ∇F S -∇F S 2 2 dw.

The following Lemma shows that P t ( dµ0 dµ ) = µ t in our diffusion process. Lemma 38. Let P denote the diffusion semigroup of CLD. Let µ denote the invariant measure of P and let µ t denote the probability measure of W t . Then P t ( dµ0 dµ ) = µ t .

Proof Let dµ = µ(x) dx and dµ t = µ t (x) dx. As shown in (Pavliotis, 2014, page 118) , our diffusion process (Smoluchowski dynamics) is reversible, which means µ(x)p t (x, y) = µ(y)p t (y, x).

Thus for any g(x), we have Since g is arbitrary, P t ( dµ0 dµ ) and µ t must be the same.

Theorem 39. Suppose that n > 8βC. Under Assumption 14, CLD (with initial distribution dµ 0 =

1 Z e -λβ w 2 2 dw) has the following expected generalization error bound:

err gen ≤ 8βC 2 n + 4C exp -λT e 4βC βC.

In addition, if F 0 is also M -smooth and non-negative, by setting λβ > 2, λ > 1 2 and η ∈ [0, 1 ∧ 2λ-1 8M 2 ), the GLD process (running K iterations with the same µ 0 as CLD) has the expected generalization error bound:

err gen ≤ 2C 2KC 1 η 2 + 8βC 2 n + 4C exp -ληK e 4βC βC,

where C 1 is a constant that only depends on M , λ, β, b, L and d.

Proof of Theorem 39 Suppose S and S are two datasets that differ on exactly one data point. Let (W t ) t≥0 and (W t ) t≥0 be their processes, respectively. Let dµ t = π t (w) dw and dµ t = π t (w) dw be the probability measure of W t and W t , respectively. The invariant measure of CLD for S and S are denoted as µ and µ , respectively. Recall that dµ = 1 Z µ e -βF S (w) dw, dµ = 1 Z µ e -βF S (w) dw.

The total variation distance of µ and µ is 2 as g e (t) to compute our bound in Figure 1 . The estimation of g e (t) at every step t is shown in Figure 1 (d). Since g e (t) is not very stable, in our figure, we plot its moving average over a window of size 100 to make the curve smoother (i.e., g avg (t) = 1 100 t+100 τ =t g e (τ )).

TV(µ, µ ) = 1 2 R d 1 - dµ dµ dµ = 1 2 R d 1 - Z µ Z µ exp(-β(F S (w) -F S (w))) 1 Z µ e -βF

C.2 EXPERIMENTAL RESULTS FOR SGLD:

In this subsection, we present some experiment results for running SGLD on both MNIST and CIFAR10 datasets, to demonstrate that our bound (see Theorem 11), in particular the sum of the empirical squared gradient norms along the training path, can distinguish normal dataset from dataset that contains random labels. As shown in Figure 3 , the curves of our bounds look quite similar to the generalization curves. Due to the sub-optimal constants in our bound, the bound is currently greater than 1, and hence we omit the numbers on the y-axis.

We note that in our experiments presented in Figure 3 , the learning rate that we choose is larger than that required by the second condition of Theorem 11. 

EQUATION

where C L is defined in Section 3.1. In order to show that our observation ("random > normal") still holds when the step size satisfies the requirement of our theory, we run an experiment that using gradient clipping trick with C L = 1. The model is trained on a small subset of MNIST as fitting the original data set with random labels under such a large Gaussian noise is extremely slow. As shown in Figure 4 , the experimental results remain unchanged when all the conditions of our bound are met.

These experiments indicate that the sum of squared empirical gradient norms is highly related to the generalization performance, and we believe by further optimizing the constants in our bound, it is possible to achieve a generalization bound that is much closer to the real generalization error. Learning rate γ t = 0.01 • 0.95 t/60 . Note that in the early stage, the testing accuracy is even higher than the training accuracy, thus we plot the absolute value of generalization error. As shown in this figure, even when the training accuracy approaches 90%, our bound is still relatively small.

:

In particular, their proof leverages the Fokker-Planck equation, which describes the time evolution of the density function associated with the Langevin dynamics and can only handle Gaussian noise.

They assume the loss is sub-Gaussian. By Hoeffding's lemma, C-bounded random variables are sub-Gaussian with parameter C.

3 The proof of their O( √ T /n) bound can be easily extended to our setting with 2 regularization.

We highlight another difficulty in proving non-vacuous generalization error bounds when the data are randomly labeled. Consider a 10-class classification setting where all the labels are random. For any sufficiently small data size, there is always a deep neural network that perfectly fits the dataset. Thus, the training error is zero while the population error is 90%. In this case, any valid generalization error bound should be larger than 0.9. Then, the theoretical bound would still be vacuous even if it is only loose by a factor of 2.

The condition γt = O(σt/L) is also required in(Mou et al., 2018, Theorem 1)

The high level idea to prove this bound is very similar to that inRaginsky et al. (2017). We first observe that the (stationary) Gibbs distribution µ has a small generalization error. Then, we bound the distance from µ t to µ. In our setting, we can use the Holley-Stroock perturbation lemma which allows us to bound the Logarithmic Sobolev constant, and we can thus bound the above distance easily.5 FUTURE DIRECTIONSIn this paper, we prove new generalization bounds for a variety of noisy gradient-based methods. Our current techniques can only handle continuous noises for which we can bound the KL-divergence.

ACKNOWLEDGEMENTWe would like to thank Liwei Wang for several helpful discussions during various stages of the work. The research is supported in part by the National Natural Science Foundation of China Grant 61822203, 61772297, 61632016, 61761146003, and the Zhongguancun Haihua Institute for Frontier Information Technology and Turing AI Institute of Nanjing.

For arbitrary initial distribution, it is impossible to prove similar inequality for any t ≥ 0 (unless the loss is strongly convex).

This is because φ is convex and Pt is a positive operator.

Figures and Tables:

Figure 1: Training MLP with GLD (σ t = 0.2 √ 2γ t ) on a smaller version of MNIST with different random label portion p. (a) shows the training accuracy. (b) shows the generalization error, i.e., the gap between the 0/1 loss L 01 on the training data and on the test data. (c) plots our bound in Theorem 9. (d) shows that for p = 0, the gradient norms become much smaller at later stages of training.

y∈A x -y denotes the diameter of set A. The following technical lemma upper bounds the KL-divergence between two Gaussian mixtures induced by sampling a mini-batch from neighbouring datasets. Lemma 21. Suppose that batch size b ≤ n/2. {µ B : B ∈ G} and {µ B : B ∈ G} are two collections of points in R d labeled by mini-batches of size b that satisfy the following conditions for some constant β ∈ [0, σ]:

THE DISCRETIZATION LEMMA FROM RAGINSKY ET AL. (2017) Let h(w, z) = F 0 (w, z) +

For each z ∈ Z, the function h(•, z) is (m, b)-dissipative: for some m > 0 and b ≥ 0, w, ∇h(w, z) ≥ m w 2 2 -b, ∀w ∈ R d .

(Raginsky et al., 2017, Lemma 7) Suppose that Assumption 35 holds and set µ S,0 = ν S,0 = µ 0 . Then, for any k ∈ N and any η ∈ (0, 1 ∧ m 4M 2 ), the following inequality holdsKL(µ S,k , ν S,ηk ) ≤ (C 0 βδ + C 1 η)kη,where C 0 and C 1 are constants that only depend on M , κ 0 , m, b, β, B and d.B.4 PROOFS FOR MAIN THEOREMSTheorem 15. Under Assumption 14, CLD (with initial probability measure dµ 0 = 1 Z e -λβ w 2 2

x)] = g(x)µ(x)(P t (dµ 0 /dµ))(x)dx = g(x)µ(x)dx µ 0 (y)p t (x, y)/µ(y)dy = g(x)µ(x)dx µ 0 (y)p t (y, x)/µ(x)dy= g(x)µ(x)µ t (x)/µ(x)dx = E x∼µt [g(x)].

Figure 2: Training MLP with GLD (σ t = 0.2γ t ) on the full MNIST dataset without label corruption.Learning rate γ t = 0.01 • 0.95 t/60 . Note that in the early stage, the testing accuracy is even higher than the training accuracy, thus we plot the absolute value of generalization error. As shown in this figure, even when the training accuracy approaches 90%, our bound is still relatively small.

10.

Table content: <table><tr><td>By applying Lemma 21 with β =</td><td colspan="3">γt ∇F (wt-1,z) 2 b</td><td>and σ = σ t ,</td></tr><tr><td colspan="3">KL(P t , P t ) ≤</td><td colspan="2">8.23γ 2 t ∇F (w t-1 , z) σ 2 t n 2</td><td>2 2</td><td>.</td></tr><tr><td>By Lemma 10,</td><td/><td/><td/></tr><tr><td/><td/><td>T</td><td/></tr><tr><td colspan="2">KL(W ≤T , W ≤T ) =</td><td>t=1</td><td colspan="2">E w&lt;t∼W&lt;t</td><td>[KL(P t , P t )]</td></tr></table>

Since it is time-consuming to compute the exact g e (t), in our experiment, we use an unbiased estimation instead. At every step, we randomly sample a mini-batch B with batch size 200 from the training data, and use 1 200 i∈B ∇f (W t-1 , z i )

Table content: <table><tr><td colspan="3">is EW t-1 [ 1 n</td><td>n i=1 ∇f (W t-1 , z)</td><td/><td/><td/><td/><td/><td/><td/><td/></tr><tr><td/><td/><td/><td/><td/><td/><td/><td/><td/><td/><td/><td/><td>(39)</td></tr><tr><td>Since</td><td>Zµ Z µ</td><td colspan="4">exp(-β(F S (w) -F S (w))) ∈ e -4βC n , e</td><td>4βC n</td><td/><td colspan="5">and 4βC n &lt; 1/2, we have</td></tr><tr><td/><td/><td/><td>TV(µ, µ ) ≤ max</td><td>1 2</td><td>1 -e -4βC n</td><td>,</td><td>1 2</td><td>e</td><td>4βC n</td><td>-1</td><td>≤</td><td>4βC n</td><td>.</td><td>(40)</td></tr><tr><td/><td/><td/><td/><td/><td/><td/><td/><td/><td/><td/><td/><td>(41)</td></tr></table>

This is because the global Lipschitz constant L is hard to estimate and the model is not able to fit training data under a very large noise. As discussed in Section 3.1, we can relax σ t ≥ (20L)γ t to σ t ≥ (2 max i∈[n] ∇F (W t-1 , z i ) )γ t . By applying gradient clipping trick, we can further relax this condition to σ t ≥ min{C L , (2 max

Table content: <table/>

Peer review (iclr2020):
This paper aims at developing a better understanding of generalization error for increasingly prevalent non-convex learning problems. For many such problems, the existing generalization bounds in the statistical learning theory literature are not very informative. To address these issues, the paper explores algorithm-specific generalization bounds, especially focusing on various types of noisy gradient methods. The paper employs a framework that combines uniform stability and PAC-Bayesian theory to obtain generalization bound for the noisy gradient methods. For gradient Langevin dynamic (GLD) and stochastic gradient Langevin dynamics (SGLD), using this Bayes-Stability framework, the paper obtains a generalization bound on the expected generalization error that scales with the expected empirical squared gradient norm. As argued in the paper, this provides an improvement over the existing bounds in the literature. Furthermore, this bound enables the treatment of the setting with noisy labels. For this setting the expected empirical squared gradient norm along the optimization path is higher, leading to worse generalization bound. The paper then extends their results to the setting where an regularization is added to the non-convex objective. By using a new Log-Sobolev inequality for the parameter distribution at time t, the paper obtains new generalization bounds for continuous Langevin dynamic (CLD). These bounds subsequently provide bounds for GLD as well. The paper demonstrates the utility of their generalization bound via empirical evaluation on MNIST and CIFAR dataset. The obtained generalization bounds are informative as they appear to capture the trend in the generalization error. Overall, the paper is very well written with a clear comparison with the existing generalization bounds. The results in the paper are interesting and novel. That said, the discussion in the introduction and abstract appears a bit misleading as it gives the impression that this is the first paper that combines the ideas from stability and PAC-Bayesian theory to obtain generalization bounds. This is not the case, e.g. see [1]. As noted by the authors, some of the bounds obtained in this paper share similarities with one of the bounds in Mou et al. as all these bounds contain the expected empirical squared gradient norm. The bound in Mou et al. holds with high probability and decays as , whereas the bounds in this paper are on expected generalization error and decay as . Could authors comment on extending their results to hold with high probability and how it would affect their bounds? [1] Rivasplata et al., PAC-Bayes bounds for stable algorithms with instance-dependent priors. ----------------------- Post author response ------------- Thank you for addressing my comments. I have decided to keep my original score unchanged.