Title: Time-Independent Information-Theoretic Generalization Bounds for SGLD
Venue: NeurIPS 2023
Accepted: yes
Keywords: SGLD

Langevin dynamics

Generalization

Information theoretic analysis
Abstract: We provide novel information-theoretic generalization bounds for stochastic gradient Langevin dynamics (SGLD) under the assumptions of smoothness and dissipativity, which are widely used in sampling and non-convex optimization studies.
Our bounds are time-independent and decay to zero as the sample size increases, regardless of the number of iterations and whether the step size is fixed.
Unlike previous studies, we derive the generalization error bounds by focusing on the time evolution of the Kullback--Leibler divergence, which is related to the stability of datasets and is the upper bound of the mutual information between output parameters and an input dataset.
Additionally, we establish the first information-theoretic generalization bound when the training and test loss are the same by showing that a loss function of SGLD is sub-exponential.
This bound is also time-independent and removes the problematic step size dependence in existing work, leading to an improved excess risk bound by combining our analysis with the existing non-convex optimization error bounds.

Time-Independent Information-Theoretic Generalization Bounds for SGLD

Abstract

We provide novel information-theoretic generalization bounds for stochastic gradient Langevin dynamics (SGLD) under the assumptions of smoothness and dissipativity, which are widely used in sampling and non-convex optimization studies. Our bounds are time-independent and decay to zero as the sample size increases, regardless of the number of iterations and whether the step size is fixed. Unlike previous studies, we derive the generalization error bounds by focusing on the time evolution of the Kullback-Leibler divergence, which is related to the stability of datasets and is the upper bound of the mutual information between output parameters and an input dataset. Additionally, we establish the first information-theoretic generalization bound when the training and test loss are the same by showing that a loss function of SGLD is sub-exponential. This bound is also time-independent and removes the problematic step size dependence in existing work, leading to an improved excess risk bound by combining our analysis with the existing non-convex optimization error bounds.

1 Introduction

Stochastic optimization, including stochastic gradient descent (SGD), is central to realizing practical large-scale or deep-learning models. There are currently considerable active discussions on accurately determining the generalization performance of models trained by SGD or its variants. In particular, stochastic gradient Langevin dynamics (SGLD) [13, 39, 29], a noisy variant of SGD, has garnered much attention in this type of study since it provides a useful theoretical framework for generalization error analysis based on the Langevin diffusion context [29]. Our study aims to contribute to a more accurate understanding and evaluation of the generalization performance for SGLD.

There are two main approaches to generalization analysis in SGLD. One is the information-theoretic analysis proposed by Russo and Zou [30] and Xu and Raginsky [40], by which a generalization error bound is derived using the mutual information (MI) between the learned parameters and the training dataset. Recently, some extensions using gradient information have been made to investigate the generalization properties of SGLD, for example, upper-bounding the MI with the norm of gradients [28] and the sum of gradient variances [25, 35, 37, 36]. Information-theoretic generalization bounds are applicable to a wide range of noisy iterative algorithms such as differentially private SGD [11] and stochastic gradient Hamiltonian Monte Carlo [5] modified to include a noisy momentum.

The other approach is stability analysis , by which the effects of changes in the learning algorithm due to the addition or removal of a single training data point on the generalization performance are investigated. Raginsky et al. [29] derived the non-asymptotic generalization and excess risk bound of SGLD via the exponential ergodicity of Langevin diffusion. Starting with the study by Raginsky

et al. [29], there have been many attempts to improve the generalization analysis in SGLD from the stability perspective, such as those by Zhang et al. [43], Mou et al. [23] and Li et al. [20].

Unfortunately, these existing generalization bounds are time-dependent ; namely, they diverge with increasing number of iterations unless the step size is adjusted so that the order of bound values is O ( n -1 ) or O ( n -1 / 2 ) , where n is the sample size (see Section 5 for details). Farghly and Rebeschini [10] attempted to avoid this problem by Wasserstein stability analysis through reflection coupling [9] under the smoothness and dissipativity [14] assumptions commonly used in sampling and non-convex optimization communities [29, 44]. Although their bounds bypass the divergence problem when taking a supremum over time, the geometry introduced for the reflection coupling yields an unnatural dependence on step size, resulting in a vacuous bound as the step size decreases (see Table 1).

In this paper, we provide novel generalization bounds for SGLD under smooth and dissipative loss functions obtained by the information-theoretic approach. We focus on the upper bound of the MI, namely, the Kullback-Leibler (KL) divergence between the distributions of parameters learned from different training datasets. We then analyze its time evolution caused by the update of the SGLD algorithm through the Fokker-Planck (FP) equation (Lemma 1). On the basis of this analysis, we obtain time-independent generalization error bounds that decay to zero as n →∞ regardless of the number of iterations or whether the step size is fixed (Theorem 4 and Corollary 1). Conventional information-theoretic generalization bounds [29, 28, 35, 36] are derived by bounding the MI between the parameters at all iterations and the training dataset. Therefore, these bounds grow linearly with the number of iterations, resulting in a time-dependent generalization error bound. Our analysis based on time evolution eliminates both this linearity issue and the unnatural dependence on step size (the inverse of step size) in the time-independent bound of Farghly and Rebeschini [10].

Another contribution is providing the first information-theoretic generalization bound and the excess risk bound when the same loss is used for training and the generalization performance evaluation. In the conventional information-theoretic approach, deriving generalization error bounds under this setting was challenging owing to the unknown tail behavior of a loss function of SGLD. We overcome this difficulty with our discovery that a smooth and dissipative loss function of SGLD is sub-exponential.

2 Preliminaries

2.1 Problem settings and stochastic gradient Langevin dynamics

We represent random variables in capital letters, such as X , and deterministic values in lowercase letters, such as x , and express the Euclidean inner product and distance as · and ∥ · ∥ . Let µ be an unknown generating distribution on the instance space Z and w ∈ W ⊆ R d be the d -dimensional parameters such as weights of neural networks, where W is the space of the parameters. We consider a loss function l : W×Z → R and the following optimization problem:

which cannot be computed since µ is unknown. Instead, we typically minimize the empirical risk estimated using the dataset S := { Z i } n i =1 :

where { Z i } n i =1 are independent and identically distributed (i.i.d.) samples from µ , i.e., Z i i . i . d . ∼ µ .

Stochastic gradient Langevin dynamics. In this paper, we use the SGLD algorithm [39] to solve the empirical risk minimization. SGLD utilizes the gradient information of the loss function; however, some loss functions, such as the 0 -1 loss, are not differentiable. In this case, it is common to use the differentiable surrogate loss function f : W×Z → R (e.g., the cross-entropy loss) and minimize the following empirical risk: F S ( w ) := 1 n ∑ n i =1 f ( w,Z i ) . Given a mini-batch B ⊂ [ n ] := { 1 , · · · , n } with k = | B | ≤ n , we define its mini-batch version as

The SGLD algorithm updates the parameters using the following recursion:

where P W 0 is a given initial distribution, ∇ F ( w,B ) is a stochastic gradient, t is the number of iterations, η t is the step size, β t is the inverse temperature, and ( B t ) ∞ t =0 is an i.i.d. sequence of random variables distributed uniformly on { B ⊂ [ n ] : | B | = k } . In addition, ( ξ t ) ∞ t =0 is an i.i.d. sequence of standard Gaussian random variables, i.e., ξ t ∼ N (0 , I d ) , where I d is the d -dimensional identity matrix. The output parameters W ∈ W obtained using SGLD can be seen as the samples from a conditional distribution P W | S : Z n →W . We express the t -th output of the SGLD algorithm as W t .

2.2 Expected generalization error and its bounds

The focus of this paper is the expected generalization error , defined as

where the expectation is taken over the joint distribution of ( S, W ) , i.e., µ n ⊗ P W | S .

Information-theoretic generalization bounds. Russo and Zou [30] and Xu and Raginsky [40] have shown that Eq. (1) can be bounded by the MI between the input dataset S and the output parameters W under the following sub-Gaussian assumption.

Assumption 1 (sub-Gaussian losses) . A loss function l ( w,Z ) is sub-Gaussian under Z ∼ µ for all w ∈ W , that is, there is a positive constant σ 2 g such that log E Z [exp( λ ( l ( w,Z ) -E l ( w,Z )))] ≤ λ 2 σ 2 g / 2 for all constant λ ∈ R .

For example, bounded or Lipschitz-continuous loss functions satisfy this assumption. Assumptions regarding the tail behavior of the loss function distributions as in the above are necessary for the information-theoretic generalization error analysis. Bu et al. [3] have investigated informationtheoretic generalization bounds with another tail-behavior assumption such as sub-exponential losses.

We introduce the following standard information-theoretic generalization bound.

Theorem 1 (Russo and Zou [30] and Xu and Raginsky [40]) . Suppose that Assumption 1 holds. Then, we have

under a training dataset S = { Z i } n i =1 and the algorithm's output W , where I ( W ; S ) is the MI between W and S .

In the SGLD context, I ( W,S ) of Eq. (2) can be upper-bounded in a form that incorporates the gradient variance [28, 25, 35, 36]. Given the output of the T -th iterate of the SGLD algorithm, W T , the following upper bound can be obtained.

Theorem 2 (Modified bound of Pensia et al. [28]) . Let f ( · , z ) be an L -Lipschitz continuous function, namely, there is a constant L > 0 such that ∥ f ( w,z ) -f ( ¯ w,z ) ∥ ≤ L ∥ w -¯ w ∥ holds for all w, ¯ w ∈ W and all z ∈ Z . Then, we obtain

where Var[ ∇ f ( W,B ) | W ] := EE B [ ∥∇ W f ( W,B ) -E B [ ∇ W f ( W,B )] ∥ 2 | W ] is the conditional variance.

Note that bound (i) can be obtained by the data-processing inequality [6].

The merit of such information-theoretic generalization bounds is that we can evaluate the bound value using the empirically estimated gradient variance per iteration. However, unfortunately, from Eq. (3), this bound is time-dependent ; namely, the bound value can diverge unless the gradient variance or β t η t approaches 0 as T → ∞ . This is due to the data-processing inequality when deriving upper bound (i) in Eq. (3). By the data-processing inequality, we obtain I ( W T ; S ) ≤ I ( W ( T ) ; S ) , where

W ( T ) := ( W 0 , W 1 , · · · , W T ) denotes the joint random variables appearing in all the iterations in the algorithm. Since W ( T ) is treated simultaneously, the bound is inevitably linear in T .

Another limitation of the information-theoretic approach appears in the setting where training losses ( f ) are also used for performance evaluation, which is often employed in sampling and non-convex optimization studies of SGLD [29, 41]. In this setting, the generalization error is defined as

where F µ := E Z [ f ( w,Z )] . We cannot conduct the information-theoretic analysis for Eq. (4) because the tail behavior of the distribution of the training loss is unclear.

Time-independent generalization bounds for Eq. (4) . To solve the above problems, Farghly and Rebeschini [10] provided the generalization error bounds of Eq. (4) from the stability perspective under the following assumptions widely used in the non-convex optimization analysis of SGLD [29, 41, 19].

Assumption 2 (Smoothness) . For each z ∈ Z , f ( · ,z ) is differentiable and M -smooth. That is, there is a positive constant M for all w, ¯ w ∈ W and all z ∈ Z such that

Assumption 3 (Dissipativity [14]) . For each z ∈ Z , f ( · , z ) is ( m,b ) -dissipative. 1 That is, there are positive constants m and b for all w ∈ W and z ∈ Z such that

The discussion regarding loss functions that satisfy Assumption 3 is presented in Appendix B.

Hereafter, we eliminate the time dependence of the step size and temperature by setting η t = η and β t = β . With this notation, Farghly and Rebeschini [10] derived the following generalization bound.

Theorem 3 (Farghly and Rebeschini [10]) . Suppose that Assumptions 2 and 3 hold. Assume that the initial law of W 0 has a finite fourth moment σ . Then, if η ≤ 1 / 2 m , for any T ∈ N , we have

where ( x ∧ y ) = min { x, y } , and C 1 and C 2 are the positive constant terms w.r.t. { M,m,b,d,β,σ } and { M,m,b,d,β } , respectively.

Farghly and Rebeschini [10] utilized the Wasserstein stability on the basis of the contraction property of Langevin diffusion under reflection coupling. The important technique to derive the above bound is that we only focus on W T differently from W ( T ) of the information-theoretic approach when deriving the contraction property. In this way, the resulting bounds do not suffer from divergence as T →∞ ; however, it still has a problem. That is, Eq. (5) depends on the factor η -1 / 2 , which implies that it becomes vacuous or even diverges with decreasing η (= η T ) as T →∞ .

In this paper, we propose new generalization bounds to address the drawbacks of the informationtheoretic and stability-based approaches. Specifically, the proposed bounds are established on the basis of the two expected generalization errors outlined in Eqs. (1) and (4), which remain timeindependent and do not diverge as the step size decreases.

3 Time-independent generalization error bound for SGLD

Here, we explain our time-independent bound of gen( µ, P W | S ; L ) for SGLD. We first introduce the main result (Section 3.1) and then summarize its proof outline (Sections 3.2 and 3.3). Finally, in Section 3.4, we provide a detailed discussion on our bound with concrete examples.

3.1 Main result

Our key idea is to derive the generalized error bound using the FP equation. To use the FP equation, we impose the following regularity condition for P W 0 .

Assumption 4 (Regularity of the initial distribution) . The initial distribution of W 0 : P W 0 is a Gaussian distribution 2 with a finite variance s 2 > 0 , which is independent of η and T .

Our analysis is also grounded in the time evolution of the FP equation using the logarithmic Sobolev inequality (LSI) [2] associated with π described as follows. We state that π satisfies the LSI with constant c LS , if for any ρ ≪ π , the following relation holds:

Raginsky et al. [29] showed the existence of c LS under Assumptions 2, 3, and β ≥ 2 /m . Note that c LS is expressed by the problem-dependent constant (see Appendix F.1 for details).

We now introduce our generalization error bound.

Theorem 4. Suppose that Assumptions 1, 2, 3, and 4 are satisfied. Then, for any β ≥ 2 /m and η ∈ (0 , 1 ∧ m 5 M 2 ∧ 4 βc LS ) and any T ∈ N , we have

where c 1 , c 2 , and V ∇ are the positive constant terms w.r.t. { M,m,b,d,β,s 2 } .

The above theorem shows | gen( µ, P W T | S ; L ) | = O ( √ ( ηT ∧ 1) /n ) , which implies time independence since it does not diverge even if T →∞ and thus converges as n →∞ .

In Eq. (6), the term V ∇ corresponds to stability , which is expressed as the upper bound of the difference of the expected conditional gradients with respect to changes in training datasets at each iteration. This shows a certain similarity to existing information-theoretic generalization bounds, such as Theorem 2, expressed by the variance of gradients with respect to the training datasets. This similarity is discussed in detail in Section 3.4. Additionally, detailed information on the explicit expression of c 1 , c 2 , and V ∇ can be found in Appendix F.

3.2 Proof outline of Theorem 4

In this section, we present how to derive our bound in Theorem 4. Our aim here is to share the ideas behind our analysis and an outline of the proof, providing the detailed proof in Appendix F.

We adopt the information-theoretic approach and focus on the MI in Eq. (2). By using the Jensen inequality, we have the following upper bound of the MI:

where S and S ′ are random variables drawn independently from µ n , and KL( P W T | S | P W T | S ′ ) is the KL divergence from P W T | S ′ to P W T | S . Note that this KL divergence indicates the stability of the learned parameter from two datasets, S and S ′ . We also note that P W T | S ′ can be regarded as the data-dependent prior. Thus, this KL divergence is tighter than that of the data-independent prior, which is often used in the probably approximately correct (PAC)-Bayes bound 3 .

The key idea is to analyze the time evolution of the KL divergence , which is summarized in the following lemma:

Lemma 1. Suppose that the same assumptions in Theorem 4 hold. Then, for any t ∈ N , we have

where V ∆ and c 3 is the constant term w.r.t. { M,m,b,d,β,s 2 } .

We will discuss the details of Lemma 1 in Section 3.3. By recursively applying Eq. (8) from t = 0 to T , we obtain

which is based on the fact that KL( P W 0 | S | P W 0 | S ′ ) = 0 . Note that bound (i) can be obtained from e -η 4 βc LS < 1 -η 4 βc LS + η 2 16 β 2 c 2 LS when η 4 βc LS ≤ 1 , and e -ηT 4 βc LS ≥ 1 -ηT 4 βc LS .

3.3 Proof outline of Lemma 1 under the continuous Langevin diffusion

Here, we organize our ideas for the proof of Lemma 1 that are important in the derivation of Theorem 4. For simplicity, we now provide an intuitive explanation and an outline of the proof under the continuous Langevin diffusion setting. Note that the results of Theorem 4 and Lemma 1 are based on the SGLD setting, and their proofs are shown in Appendix F.

The Langevin diffusion is defined as

where d H t is the standard Brownian motion in R d . Note that, in this context, t expresses the continuous time and the full-batch gradient ∇ F ( W t , S ) is used. The stationary distribution of Eq. (9) is given as the Gibbs distribution π (d w ) ∝ exp( -βF ( w,S )) .

With some abuse of notation, let us denote P W t | S as the conditional distribution obtained using Eq. (9) and express its density as ρ t . Then, the FP equation [2] for Eq. (9) can be obtained as

Similarly, we can define the Langevin diffusion when we use dataset S ′ and the conditional distribution using that diffusion as P W t | S ′ with the density γ t , obtaining the FP equation in the form of ρ t replaced by γ t in Eq. (10).

Now we analyze the time evolution of KL( P W t | S | P W t | S ′ ) = KL( ρ t | γ t ) at time t , i.e., ∂ KL( ρ t | γ t ) /∂t . By utilizing the FP equations of ρ t and γ t and the Cauchy-Schwartz inequality, we obtain the following upper bound:

The second term on the right-hand side of Eq. (11) represents the stability of the gradient with respect to the randomness of the training dataset S, S ′ ∼ µ n , which leads to V ∇ in Lemma 1 under the SGLD setting. Hereafter, we define E ∥∇ F ( W t , S ) -∇ F ( W t , S ′ ) ∥ 2 as ˜ V ∇ t .

By introducing ∇ log π ( w ) into E ∥∇ log ρ t -∇ log γ t ∥ 2 in Eq. (11), we obtain

where the first inequality is from the fact that -x 2 ≤ -∥ x -y ∥ 2 / 2 + y 2 for x, y ∈ R d and the second one is from the LSI. We introduced Ω( ρ t , γ t , π ) := E ρ t ∥∇ log π ∥ 2 +2 E ρ t ∇ log ρ t · ∇ log γ t to simplify the notation.

By integrating e t 4 βc LS ∂ KL( ρ t | γ t ) ∂t in Eq. (12) over t ∈ [0 , η ] and rearranging it, we obtain

In Appendix F, we show that the terms related to π in Eq. (13) can be bounded by using the techniques of Raginsky et al. [29] and Vempala and Wibisono [32].

We next derive an upper bound for the following terms in Eq. (13): E [log γ t π ] and Ω( ρ t , γ t , π ) by using the parametrix method for the FP equation [12, 27], which allows us to expand the FP equation's solution via the heat kernel. On the basis of this expansion, we can upper bound Eq. (13) as

By combining Eq. (14) with Eq. (13), we obtain the continuous version of Lemma 1.

The same procedure can be used for the SGLD setup. The difference from the continuous Langevin diffusion case is that the discretization errors and the effects of using a stochastic gradient are taken into account, resulting in the appearance of an additional constant 4 in the above bounds (see Appendix F for details).

3.4 Additional discussion on our bound in terms of stability

We conclude this section by presenting further discussion on our bound in terms of stability with a concrete example.

As shown in Eq. (7), the information-theoretic generalization bound is closely related to the stability in KL divergence under the different training datasets. However, our bound in Theorem 4 incorporates the constant term c 2 , which is irrelevant to stability, alongside the stability term V ∇ . If we can avoid the occurrence of c 2 , the resulting upper bound of KL( P W t | S | P W t | S ′ ) would be dominantly expressed by V ∇ , and as a result, we may obtain a bound where the relationship between generalization and stability is more directly represented.

The problematic constant term c 2 arises from c 3 η in Lemma 1 analyzing the time evolution of stability in KL divergence. Specifically, the term c 3 η is the byproduct of treating the general dissipative function using LSI. Actually, it is possible to avoid the problematic constant term c 3 η and derive bounds that are evaluated solely on the basis of stability-related metrics in specific examples, such as strongly convex or bounded (non-convex) losses with l 2 -regularization. For simplicity, we show this fact using the following theorem under the Langevin diffusion (LD) setting, where the probability induced by Eq. (9) is expressed as P W T | S .

Theorem 5. Suppose that Assumptions 1 and 2 are satisfied and that F ( w,z ) is R -strongly convex ( 0 < R < ∞ ). Then, for any T ∈ R + , we have

and

A similar bound in Eq. (16) (with R replaced by λ/e 8 βC ) can be obtained for bounded non-convex losses with l 2 -regularization, where F ( w,z ) = F 0 ( w,z ) + λ 2 ∥ w ∥ 2 ( 0 < λ < ∞ ) and F 0 ( w,z ) is C -bounded ( 0 ≤ C < ∞ ). The full proof is summarized in Appendix F.4.

When comparing with Lemma 1, we can see that, in Eq. (15), stability-unrelated constants do not appear in the time evolution of KL divergence at each time step. Therefore, the resulting generalization bound is also independent of such constants. Furthermore, when compared with Theorem 2, which adds up the stability terms at all time steps, our bound is dominated by the stability terms near the final time step, as those at earlier time steps decrease geometrically by e -R 4 . This indicates that the stability around the initial time steps is of lesser importance in evaluating the final generalization performance.

Note that our bounds are closely related to the bound indicated in Proposition 9 of Mou et al. [23], which was also derived by focusing on stability. The bound of Mou et al. [23] primarily assesses

generalization errors focusing on the gradient norm near the conclusion of training. In contrast, our bounds evaluate it through the norm of differences in gradients , emphasizing the state in the proximity of training completion. In other words, our bound allows for the evaluation of generalization errors using a stability measure that is more closely related to generalization performance than the gradient norm. This benefit originates from our approach, which tracks the time evolution of MI-related stability in Eq. (15) on the basis of information-theoretic generalization bounds, in contrast to the PAC-Bayes bounds derived from the direct analysis of stability measures as in Mou et al. [23].

4 Generalization analysis for SGLD directly using a training loss

In this section, we consider the setting that the generalization performance is measured by a training loss f directly as in Eq. (4). We show that this is possible by demonstrating that loss functions of SGLD are sub-exponential under smooth and dissipative assumptions (Section 4.1). On the basis of this fact, we obtain for the first time an information-theoretic generalization bound of SGLD that is similar to Theorem 4. Finally, combining these results with existing optimization error bounds provides an excess risk bound with improved convergence (Section 4.2).

4.1 Smooth and dissipative loss function of SGLD is sub-exponential

To perform an information-theoretic analysis for SGLD, it is necessary to know the tail behavior of f ( W,Z ) . Our contribution here is showing that a loss function of SGLD under smooth and dissipative assumptions is sub-exponential.

Theorem 6. Suppose that Assumptions 2, 3 and 4 are satisfied. Let P W T = E S [ P W T | S ] be the marginal distribution of the output obtained using the SGLD algorithm at the T -th iteration. Then, for any η ∈ (0 , 1 ∧ m 5 M 2 ) and T ∈ N , f ( W T , Z ) is sub-exponential under the distribution P W T ⊗ µ . That is, there exist positive constants σ 2 e and ν w.r.t. { m,β,M,b,d,s 2 } 5 such that

Proof sketch. First, note that under Assumptions 2 and 3, for any z ∈ Z , we obtain

where A is a positive constant (see Lemma 8 in Appendix G.1 for its explicit form). By employing Lemma B.2 from Farghly and Rebeschini [10], we can show the following fact: for any p ∈ N ,

where

This implies that W T is a sub-Gausssian random variable according to Proposition 2.5.2 in Vershynin [33]. To show the sub-exponential property, we directly upper-bound E W T ⊗ Z [ e λ ( f ( W T ,Z ) -E WT ⊗ Z [ f ( W T ,Z )]) ] by considering the Taylor expansion of the exponential moment and using Eqs. (17) and (18) (see Appendix G.2 for the complete proof).

Remark 1. In previous information-theoretic analysis studies [28, 25, 35], it is often assumed that a loss function l ( w,Z ) is sub-Gaussian under the distribution µ for all w ∈ W . In contrast, Theorem 6 holds under the distribution P W T ⊗ µ , not conditioned on w ∈ W .

We can interpret the sub-exponential property of SGLD intuitively as follows. Under Assumptions 2 and 3, the loss function grows at most as a quadratic function shown in Eq. (17). The conditional distribution of the parameters follows the Gaussian distribution, and the square of the Gaussian random variable is known as the chi-square ( χ 2 ) random variable [34]. According to these facts, we expect that the behavior of the loss function resembles that of the χ 2 -random variable; therefore, it is sub-exponential since the χ 2 -distribution is also sub-exponential [34]. Theorem 6 validates this intuition.

4.2 Generalization bounds for SGLD using the same loss for training and evaluation

On the basis of Theorem 6, we can derive the following information-theoretic generalization bound for SGLD even if a surrogate loss is not used. In contrast to Theorem 4, an assumption regarding the tail behavior of a loss function such as Assumption 1 is not necessary.

Corollary 1. Suppose that Assumptions 2, 3, and 4 are satisfied. Then, for any β ≥ 2 /m , η ∈ (0 , 1 ∧ m 5 M 2 ∧ 4 βc LS ) , and T ∈ N , we obtain

where

,

c 1 and c 2 are the same as in Theorem 4, and σ 2 e and ν are the same as in Theorem 6.

Proof sketch. This is the direct consequence of the sub-exponential property from Theorem 6 and the upper bound of MI in Eq. (7) (see Appendix G.3 for the complete proof).

Remark 2. Despite the assumptions of Corollary 1 being the same as those made by Farghly and Rebeschini [10] except for the initial distribution and step size, the resulting bound becomes 0 as n →∞ without being dependent on inverse stepsize.

We conclude this section by introducing our excess risk bound. Let us define the excess risk as follows: Excess( µ, P W | S ) := E W,S [ F µ ( W ) -F µ ( w ∗ )] , where w ∗ = argmin w ∈W F µ ( w ) . Under this definition, we derive the following upper bound for the excess risk by utilizing Corollary 1.

Corollary 2. Suppose that Assumptions 2, 3, and 4 are satisfied. Then, for any β ≥ 2 /m , η ∈ (0 , 1 ∧ m 5 M 2 ∧ 4 βc LS ) , and T ∈ N , we obtain

where c err is the positive constant w.r.t. { M,m,b,d,β } corresponding to the optimization error.

We show the complete proof in Appendix G.4. In contrast with the existing excess risk studies, our bound does not diverge with increasing t owing to the time-independent generalization bound in Corollary 1.

5 Related studies and discussion

In this section, we compare our generalization bounds with those in related studies. Table 1 shows the order of each bound value along with its assumptions for a loss function.

SGLD analysis with/without changing losses. The existing generalization error bounds in Table 1 are time-dependent; namely, we need to impose restrictive conditions for the step size η in terms of t to achieve a generalization bound that decays to zero with increasing sample size [29, 28, 25, 35] (see the right column in Table 1). Some important applications of SGLD do not satisfy these conditions. For instance, the short-run Markov chain Monte Carlo [26] method used in energy-based models [17] adopts SGLD with a fixed step size. Another example is the cyclic SGLD [42] used in deep learning, where the step size is periodically increased or decreased to facilitate escape from local optima.

Farghly and Rebeschini [10] first analyzed the generalization error of SGLD by using smoothness and dissipative assumptions, which are broadly used in sampling and non-convex optimization studies [29, 41, 4, 44]. Their bound is time-independent; the bound does not diverge with time and achieves the order O ( n -1 / 2 ) . However, the bound depends on the inverse of step size η -1 / 2 owing to the reflection coupling [9], which results in the unnatural behavior of decreasing η with increasing t . Farghly and Rebeschini [10] also derived a bound that does not suffer from this problem by assuming the Lipschitz loss function with weight decay; however, these assumptions excessively restrict the

class of loss functions and algorithms. In contrast to these bounds, our bound is time-independent and does not require scaling η , t , and n to achieve O ( n -1 / 2 ) .

Time evolution analysis of MI via FP equation. The analysis of SGLD using the FP equation has been successfully used in the convergence analysis of SGLD [32, 19]. These studies present analyses of the discretization errors and convergence properties of the unadjusted Langevin algorithm, SGLD, and variance-reduction SGLD (SVRG-LD) [8], comparing them with the continuous Langevin dynamics through the FP equation.

In generalization error analysis, the FP equation is mainly used to analyze the time evolution of KL divergence appearing in a generalization bound on the basis of the stability approach. Li et al. [20] analyzed the time evolution of the KL divergence between the probability densities of the parameters obtained from two training datasets that differ by only one data point under the bounded loss assumption. Mou et al. [23] also studied the KL divergence and Hellinger divergence, and they derived a generalization error bound on the basis of the PAC-Bayes notion [22]. Our idea is similar to these: we analyze the time evolution of the KL divergence between the probability densities of the parameters obtained from two training datasets. The differences between our approach and other approaches are twofold. First, we do not assume weight decay or Lipschitz continuity but instead derive our analysis assuming smoothness and dissipativity. Second, the KL divergence we analyzed is tighter than that of the PAC-Bayes bound with data-independent prior dealt by Mou et al. [23].

6 Limitations and future work

In this paper, we provide a generalization analysis of SGLD, where Gaussian noise is a fundamental assumption for our theoretical results. Thus, it is difficult to extend our analysis to other noisy iterative algorithm variants with a different noise, such as differentially private SGD with Laplace or uniform noise [36]. Another limitation of this study is that we have estimated the sub-exponential parameter roughly with respect to the dimensions of the model parameters. Further investigation of the sub-exponentiality of smooth and dissipative losses, and improvement of the dependence on dimensionality, are crucial for enhancing the practicality of our generalization bounds. The sub-exponential property of a loss function is expected to be helpful in fields other than generalized error analysis. For example, this property opens up room for new theoretical analysis policies that employ useful concentration and transport inequalities [34] in the sampling and optimization context. We hope that the analysis presented in this paper goes beyond generalization analysis and provides valuable insights into understanding the characteristics of machine learning.

Acknowledgments and Disclosure of Funding

We sincerely appreciate the anonymous reviewers for their insightful feedback. FF was supported by JSPS KAKENHI Grant Number JP23K16948. FF was supported by JST, PRESTO Grant Number JPMJPR22C8, Japan. MF was supported by RIKEN Special Postdoctoral Researcher Program. MF was supported by JST, ACT-X Grant Number JPMJAX210K, Japan.

References

R. Amit, B. Epstein, S. Moran, and R. Meir. Integral probability metrics PAC-bayes bounds. In Advances in Neural Information Processing Systems , 2022.

D. Bakry, I. Gentil, and M. Ledoux. Analysis and Geometry of Markov Diffusion Operators , volume 348. Springer Science & Business Media, 2013.

Y. Bu, S. Zou, and V. V. Veeravalli. Tightening mutual information-based bounds on generalization error. IEEE Journal on Selected Areas in Information Theory , 1(1):121-130, 2020.

N. H. Chau, É. Moulines, M. Rásonyi, S. Sabanis, and Y. Zhang. On stochastic gradient Langevin dynamics with dependent data streams: The fully nonconvex case. SIAM Journal on Mathematics of Data Science , 3(3):959-986, 2021.

T. Chen, E. Fox, and C. Guestrin. Stochastic gradient Hamiltonian Monte Carlo. In Proceedings of the 31st International Conference on Machine Learning , volume 32, pages 1683-1691, 2014.

T. M. Cover and J. A. Thomas. Elements of Information Theory . John Wiley & Sons, 2012.

T. Deck and S. Kruse. Parabolic differential equations with unbounded coefficients - A generalization of the parametrix method. Acta Applicandae Mathematica , 74:71-91, 2002.

K. A. Dubey, S. J. Reddi, S. A. Williamson, B. Poczos, A. J. Smola, and E. P. Xing. Variance reduction in stochastic gradient Langevin dynamics. In Advances in Neural Information Processing Systems , volume 29, pages 1154-1162, 2016.

A. Eberle. Reflection couplings and contraction rates for diffusions. Probability Theory and Related Fields , 166:851-886, 2021.

T. Farghly and P. Rebeschini. Time-independent generalization bounds for SGLD in nonconvex settings. In Advances in Neural Information Processing Systems , volume 34, pages 19836-19846, 2021.

V. Feldman, I. Mironov, K. Talwar, and A. Thakurta. Privacy amplification by iteration. In 2018 IEEE 59th Annual Symposium on Foundations of Computer Science (FOCS) , pages 521-532, 2018.

A. Friedman. Partial Differential Equations of Parabolic Type . Courier Dover Publications, 2008.

S. B. Gelfand and S. K. Mitter. Recursive stochastic algorithms for global optimization in R d . SIAM Journal on Control and Optimization , 29(5):999-1018, 1991.

J. K. Hale. Asymptotic Behavior of Dissipative Systems . Mathematical surveys and monographs. American Mathematical Society, 1988.

H. Harutyunyan, M. Raginsky, G. Ver Steeg, and A. Galstyan. Information-theoretic generalization bounds for black-box learning algorithms. In Advances in Neural Information Processing Systems , pages 24670-24682, 2021.

U. G. Haussmann and E. Pardoux. Time reversal of diffusions. The Annals of Probability , 14 (4):1188-1205, 1986.

G. E. Hinton. Training products of experts by minimizing contrastive divergence. Neural Computation , 14(8):1771-1800, 2002.

Y. Jiang, B. Neyshabur, H. Mobahi, D. Krishnan, and S. Bengio. Fantastic generalization measures and where to find them. In The Tenth International Conference on Learning Representations , 2020.

Y. Kinoshita and T. Suzuki. Improved convergence rate of stochastic gradient Langevin dynamics with variance reduction and its application to optimization. In Advances in Neural Information Processing Systems , pages 19022-19034, 2022.

J. Li, X. Luo, and M. Qiao. On generalization error bounds of noisy gradient methods for non-convex learning. In The Eighth International Conference on Learning Representations , 2020.

R. Livni. Information theoretic lower bounds for information theoretic upper bounds. arXiv preprint arXiv:2302.04925 , 2022.

D. A. McAllester. PAC-Bayesian stochastic model selection. Machine Learning , 51(1):5-21, 2003.

W. Mou, L. Wang, X. Zhai, and K. Zheng. Generalization bounds of SGLD for non-convex learning: Two theoretical viewpoints. In Proceedings of the 31st Conference on Learning Theory , volume 75, pages 605-638, 2018.

W. Mou, N. Flammarion, M. J. Wainwright, and P. L. Bartlett. Improved bounds for discretization of Langevin diffusions: Near-optimal rates without convexity. Bernoulli , 28(3):1577 1601, 2022.

J. Negrea, M. Haghifam, G. K. Dziugaite, A. Khisti, and D. M. Roy. Information-theoretic generalization bounds for SGLD via data-dependent estimates. In Advances in Neural Information Processing Systems , volume 32, pages 11015-11025, 2019.

E. Nijkamp, M. Hill, S.-C. Zhu, and Y. N. Wu. Learning non-convergent non-persistent short-run MCMCtoward energy-based model. In Advances in Neural Information Processing Systems , volume 32, pages 5232-5242, 2019.

G. A. Pavliotis. Stochastic processes and applications: Diffusion processes, the Fokker-Planck and Langevin equations , volume 60. Springer, 2014.

A. Pensia, V. Jog, and P.-L. Loh. Generalization error bounds for noisy, iterative algorithms. In 2018 IEEE International Symposium on Information Theory (ISIT) , pages 546-550, 2018.

M. Raginsky, A. Rakhlin, and M. Telgarsky. Non-convex learning via stochastic gradient Langevin dynamics: A nonasymptotic analysis. In Proceedings of the 30th Conference on Learning Theory , volume 65, pages 1674-1703, 2017.

D. Russo and J. Zou. Controlling bias in adaptive data analysis using information theory. In Proceedings of the 19th International Conference on Artificial Intelligence and Statistics , volume 51, pages 1232-1240, 2016.

S. Shalev-Shwartz, O. Shamir, N. Srebro, and K. Sridharan. Stochastic convex optimization. In Proceedings of the 22nd Conference on Learning Theory , 2009.

S. Vempala and A. Wibisono. Rapid convergence of the unadjusted Langevin algorithm: Isoperimetry suffices. In Advances in Neural Information Processing Systems , volume 32, pages 8094-8106, 2019.

R. Vershynin. High-dimensional probability: An introduction with applications in data science , volume 47. Cambridge university press, 2018.

M. J. Wainwright. High-Dimensional Statistics: A Non-Asymptotic Viewpoint . Cambridge Series in Statistical and Probabilistic Mathematics. Cambridge University Press, 2019.

H. Wang, Y. Huang, R. Gao, and F. Calmon. Analyzing the generalization capability of SGLD using properties of Gaussian channels. In Advances in Neural Information Processing Systems , volume 34, pages 24222-24234, 2021.

H. Wang, R. Gao, and F. P. Calmon. Generalization bounds for noisy iterative algorithms using properties of additive noise channels. Journal of Machine Learning Research , 24(26):1-43, 2023.

Z. Wang and Y. Mao. On the generalization of models trained with SGD: Information-theoretic bounds and implications. In The Tenth International Conference on Learning Representations , 2022.

Z. Wang and Y. Mao. Tighter information-theoretic generalization bounds from supersamples. In Proceedings of the 40th International Conference on Machine Learning , volume 202, pages 36111-36137, 2023.

M. Welling and Y. W. Teh. Bayesian learning via stochastic gradient Langevin dynamics. In Proceedings of the 28th International Conference on International Conference on Machine Learning , pages 681-688, 2011.

A. Xu and M. Raginsky. Information-theoretic analysis of generalization capability of learning algorithms. In Advances in Neural Information Processing Systems , volume 30, pages 25242533, 2017.

P. Xu, J. Chen, D. Zou, and Q. Gu. Global convergence of Langevin dynamics based algorithms for nonconvex optimization. In Advances in Neural Information Processing Systems , volume 31, pages 3126-3137, 2018.

R. Zhang, C. Li, J. Zhang, C. Chen, and A. G. Wilson. Cyclical stochastic gradient MCMC for Bayesian deep learning. In The Eighth International Conference on Learning Representations , 2020.

Y. Zhang, P. Liang, and M. Charikar. A hitting time analysis of stochastic gradient Langevin dynamics. In Proceedings of the 30th Conference on Learning Theory , volume 65, pages 1980-2022, 2017.

Y. Zhang, Ö. D. Akyildiz, T. Damoulas, and S. Sabanis. Nonasymptotic estimates for stochastic gradient Langevin dynamics under local conditions in nonconvex optimization. Applied Mathematics & Optimization , 87(25), 2021.

Peer reviews:
comments: SUMMARY:
The paper provides novel information-theoretic generalization bounds for the SGLD. They derive the generalization error bounds by focusing on the time evolution of the Kullback–Leibler divergence.  The bound is   time-independent and removes the step size dependence in previous works.

STRENGTHS:
1. This paper provides novel generalization bounds for SGLD by using the Kullback–Leibler (KL) divergence between the distributions of parameters learned from different training datasets.

2. The authors provide the first information-theoretic generalization error bound and the excess risk bound without surrogate losses.

WEAKNESSES:
1. Although the authors claim they provide the first information-theoretic generalization error bound and the excess risk bound without surrogate losses, they indeed use the dissipativity to fix the non-convexity in the proofs.
 
2.  I feel a little overclaimed on the stepsize because you use the dissipativity, which can be a kind of weak "strongly convexity". In the strongly convex case of   SGLD,  the stepsize can also be improved.

QUESTIONS:
see weakness part.

LIMITATIONS:
Yes
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 3
DATE: 2023-06-22

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper provides novel generalization bounds for SGLD under smooth and dissipative loss conditions via the information-theoretic approach. The proposed bounds are time-independent and decay to zero as the sample size increases, regardless of the number of iterations and whether the step size is fixed. In addition, the proposed bound also leads to an improved excess risk bound by combining the analysis with the existing non-convex optimization error bounds.\n\nSTRENGTHS:\nThe combination of information-theoretic bounds and Fokker–Planck (FP) equation establishes the tighter bounds provided in the paper, which is impressive.\n\nWEAKNESSES:\nI did not note any significant weakness in the paper.\n\nQUESTIONS:\nIt is shown in [1] that the generalization error of the Gibbs algorithm equals the symmetrized KL information between W and S. Under the sub-Gaussian condition, it can be further proved that the generalization error of the Gibbs algorithm has the order of $O(1/n)$ independent of dimension d and time t.  \n\nAs shown in Table 1, it seems that the bound provided by Farghly and Rebeschini is able to recover this 1/n rate for the generalization error Gibbs distribution but with an extra term $\\eta^{1/2}$ for SGLD.\n\nHowever, the bound provided in this paper only scales in the order of $O(1/\\sqrt{n})$. Could you please provide more insights into whether this gap is due to the mismatch between the Gibbs distribution and the SGLD or it is due to the looseness of mutual information-based bound? My guess would be the issue of mutual information-based bound, and maybe the techniques developed in the following recent paper [2,3] might be helpful.\n\n[1] Aminian, Gholamali, Yuheng Bu, Laura Toni, Miguel Rodrigues, and Gregory Wornell. \"An exact characterization of the generalization error for the Gibbs algorithm.\" Advances in Neural Information Processing Systems 34 (2021): 8106-8118.\n\n[2] Zhou, Ruida, Chao Tian, and Tie Liu. \"Exactly Tight Information-Theoretic Generalization Error Bound for the Quadratic Gaussian Problem.\" arXiv preprint arXiv:2305.00876 (2023).\n\n[3] Wu, Xuetong, Jonathan H. Manton, Uwe Aickelin, and Jingge Zhu. \"Fast rate generalization error bounds: Variations on a theme.\" In 2022 IEEE Information Theory Workshop (ITW), pp. 43-48. IEEE, 2022.\n\nLIMITATIONS:\nThe limitations of the proposed method have been addressed well in the last section of the paper.", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-05"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe paper develops \"time-independent\" information-theoretic generalization bounds for SGLD. The key ingredients are: assuming the smoothness and dissipativity assumptions of the loss function , tracking evolution of the weight densities for two independent training sets, and their KL divergence. Such an approach avoids accumulation of the mutual information quantities over iterations.\n\nSTRENGTHS:\nThe overall approach taken appears correct, and the results, if correct, are interesting and novel. It is unfortunate that the paper is not written clear enough for me assess the correctness of the results.\n\nWEAKNESSES:\nI think some important aspects of the paper are not presented clearly or contain errors.\n\n1. The definition of dissipativity appears wrong in Assumption 3: the right-hand side of the inequality is a not a scalar. Although this is easily fixable, it is upsetting to this reviewer that a central concept of the paper is not treated with care. \n2. In section 4 when the authors try to develop generalization bounds of SGLD without a surrogate loss. It is not clear to me how one would perform SGLD on say, the 0-1 loss, instead of on a differentiable surrogate loss: How would you get a gradient signal? -- In the opening paragraph of section 4, the paper states $l=f$. I speculate that the authors may be considering the case where SGLD is run on a differentiable surrogate loss but the generalization error is evaluated using the true loss (say 0-1 loss). If that is the case, I would imagine that somewhere a relationship between the true loss and the surrogate loss be given or assumed. But nowhere I can find it. That confused me and stopped me from further decoding what the authors intend to mean.  \n\n[After Rebuttal] I have read the authors' rebuttal and raised my rating to Weak Accept.\n\nQUESTIONS:\nI am not familiar with dissapative losses. Assuming dissipativity is defined properly, are dissipative losses practically relevant? Can the authors give some loss functions that we use in practice satisfying this condition?\n\nLIMITATIONS:\nNONE", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 2, "SOUNDNESS_CORRECTNESS": 2, "DATE": "2023-07-05"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nIn the literature, there are two approaches to obtain generalization performance bounds. One approach is to study the algorithmic stability and there are well-known results that can bound the generalization error using the algorithmic stability bounds. Another approach is the information-theoretic analysis, which is the approach adopted in this paper. One major contribution of the paper is that it provides generalization bounds that do not grow linearly in time, and moreover, it also avoids an unnatural dependence on the stepsize. In addition, the paper provides additional results that provide information-theoretic generalization bound and the excess risk bound without relying on the surrogate losses. The author(s) are very familiar with the literature and the comparisons with the existing literature is ample and the contributions of the paper are clear.\n\nSTRENGTHS:\nOne major contribution of the paper is that it provides generalization bounds that do not grow linearly in time, and moreover, it also avoids an unnatural dependence on the stepsize. In addition, the paper provides additional results that provide information-theoretic generalization bound and the excess risk bound without relying on the surrogate losses. The author(s) are very familiar with the literature and the comparisons with the existing literature is ample and the contributions of the paper are clear.\n\nWEAKNESSES:\nOne weakness is that even though the paper claims that the bound is independent of time $T$, I think some additional analysis is needed in Theorem 4 to make it more clear. For example, in Theorem 4, there is the term $V_{\\nabla}:=\\max_{t\\in[0,T]}\\beta\\text{Var}[\\nabla F(W_t,B_t)]$. I think the paper should include a lemma or something like that in the Appendix to explain how to provide an upper bound on $V_{\\nabla}$ that is uniform in $T$. My guess is that Lemma 6 in the Appendix or something like that can help. But it is better to provide an upper bound on $V_{\\nabla}$ to make it more clear your bound is uniform in $T$.\n\nQUESTIONS:\n(1)\tIn the last sentence of Theorem 1, ``an training dataset’’, here ``an’’ should be ``a’’.\n(2)\tIn the paragraph before Theorem 1, can you provide a reference to the claim that ``This assumption is necessary to obtain the information-theoretic generalization bound.’’?\n(3)\tAssumption 4 seems to be a bit restrictive. Is it possible for example to relax this assumption to mixture of Gaussians or relax it in some other directions?\n(4)\tIn eqution (10), use \\left( and \\right).\n(5)\tPage 6. ``By integratng Eq. (12) over $t\\in[0,\\eta]$, we have’’ This does not seem to be rigorous. It seems to be that you should consider $\\partial/\\partial t(KL*e^{\\frac{1}{2\\beta c_{LS}}t})$ in (12) instead and then integrate over it.\n(6)\tIn the statements of Corollary 1 and Corollary 2, ``holds’’ should be ``hold’’.\n(7)\tIn remark 4, can you double check it is $s^2>1$? It seems to me it should be $s^2<1$ instead.\n(8)\tIn Remark 3 ``we often assume’’. It is better to write ``It is often assumed’’.\n(9)\tIn equation (19) on page 14, what is the integral $\\int\\left(\\frac{\\rho_{t}}{\\gamma_{t}}\\frac{\\partial\\gamma_{t}}{\\partial t}\\right)$ w.r.t.? Did you miss $dw$ or what? Please double check the other integrals on page 14 and page 15 as well.\n(10)\tAt the end of the first paragraph on page 17, what is $\\mathcal{W}[A]$? What does $[A]$ stand for?\n(11)\tIn the equation before equation (28), what does the notation $w^{d}$ stand for?\n\nLIMITATIONS:\nThe authors have adequately discussed the limitations.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-06"}