Analyzing the Generalization Capability of SGLD Using Properties of Gaussian Channels

1 Introduction . Modern deep neural networks ( DNNs ) are highly expressive : they can memorize an entire training dataset and still generalize well to unseen data ( Zhang et al. , 2016 ) . This empirical observation is not captured by traditional generalization bounds found in statistical learning theory , which attribute the generalization ability to the use of a hypothesis class with constrained complexity ( Vapnik and Chervonenkis , 1971 ; Valiant , 1984 ) . Recent studies demonstrate that different algorithmic choices and data distributions may yield DNNs with contrasting generalization behaviors ( Hardt et al. , 2016 ; Neyshabur et al. , 2017 ; Bartlett et al. , 2017 ) . In this paper , we study how one optimization method used for training DNNs , namely the stochastic gradient Langevin dynamics ( SGLD ) algorithm ( Gelfand and Mitter , 1991 ; Welling and Teh , 2011 ) , may influence their generalization . The SGLD algorithm is used in different practical settings . For example , it has been implemented in open-source libraries ( Facebook AI , 2020 ; Radebaugh and Erlingsson , 2019 ) for training models with differential privacy guarantees ( Dwork et al. , 2006 ; Song et al. , 2013 ; Abadi et al. , 2016 ) . The additive noise in the SGLD algorithm can also mitigate overfitting for DNNs ( Neelakantan et al. , 2015 ) . Recently , there is an increasing number of efforts ( see e.g. , Raginsky et al. , 2017 ; Mou et al. , 2018 ; Li et al. , 2019 ; Pensia et al. , 2018 ; Negrea et al. , 2019 ) that investigate the generalization properties of the SGLD algorithm . It is within this body of work that the present paper is inscribed . 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . We derive a new generalization bound ( Theorem 1 ) for the SGLD algorithm in Section 3 . Our bound ( see Figure 1 for an illustration ) incorporates the variance of gradients , which can be estimated from the training data and captures a particular kind of “ sharpness ” of the loss landscape ( Keskar et al. , 2016 ) . This variance term can also be predictive of the generalization gap as shown in a recent empirical study ( Jiang et al. , 2019 ) . We consider a general setting in which the output from the SGLD algorithm can be any function of the iterates . This is crucial since many theoretical analyses ( see e.g. , Zhang et al. , 2017 ; Jin et al. , 2019 ) require the final output to be the iterate that achieves the smallest value of the loss function or an average of the iterates , but not necessarily the last iterate . Finally , the numerical experiments in Section 5 suggest that our generalization bound is highly correlated with the true generalization gap . We also investigate the DP-SGD algorithm in Section 4 . In particular , we prove that if it is known a prior , that the algorithm outputs the last iterate rather than an arbitrary function of all iterates , then our bound can be further tightened by incorporating a time-decaying factor . Our analysis is motivated by a line of recent works ( Feldman et al. , 2018 ; Balle et al. , 2019 ; Asoodeh et al. , 2020 ) on privacy amplification by iteration . Specifically , the original work by Feldman et al . ( 2018 ) provided two intertwined observations of the DP-SGD algorithm : ( i ) not releasing the intermediate steps can amplify the privacy guarantees and ( ii ) data points used in the early iterations get stronger privacy protection than those occurring late . In this paper , we establish two analogous results : ( i ) our generalization bound can be sharpened by incorporating a time-decaying factor if DP-SGD only outputs the last iterate ( Theorem 2 ) and ( ii ) this decay factor enables the impact of early iterations on our bound to reduce with time ( Lemma 4 ) . The proof techniques of this paper are based on fundamental tools from information theory . We first use an information-theoretic framework , proposed by Russo and Zou ( 2016 ) and Xu and Raginsky ( 2017 ) and further tightened by Bu et al . ( 2020 ) , for deriving an algorithmic generalization bound . This framework relates the generalization gap with the mutual information I ( W ; Zi ) between the output parameter W from the SGLD algorithm and each individual data point Zi . However , estimating the mutual information from data is often intractable . Given this major challenge , our key contribution is to connect the SGLD algorithm with a well-understood notion in data transmission , namely additive white Gaussian noise ( AWGN ) channels . This connection allows us to use properties of Gaussian channels for analyzing the mutual information . First , we upper bound I ( W ; Zi ) using the variance of gradients by exploring the input-output mutual information of a Gaussian channel . This variance term can be estimated from the training data and is highly correlated with the true generalization gap . Second , we incorporate a time-decaying factor into our bound . This factor is established by strong data processing inequalities ( Dobrushin , 1956 ; Cohen et al. , 1998 ) and has an intuitive interpretation : if a data point is used at an early iteration , its impact on the generalization gap reduces with time due to external Gaussian noise . The above two aspects correspond to Lemma 4 and Lemma 5 which , in turn , are the basis of our main results in Theorem 1 and Theorem 2 . The supplementary material of this paper includes : ( i ) omitted proofs of all theoretical results and ( ii ) supporting experimental results . Related Works We contextualize our contributions in regard to existing literature ( Mou et al. , 2018 ; Li et al. , 2019 ; Pensia et al. , 2018 ; Bu et al. , 2020 ; Negrea et al. , 2019 ; Haghifam et al. , 2020 ; Rodríguez-Gálvez et al. , 2020 ; Neu et al. , 2021 ) . Among them , Mou et al . ( 2018 ) introduced two generalization bounds . The first one ( Theorem 1 of Mou et al. , 2018 ) , a stability-based bound , achieves O ( 1/n ) rate in terms of the sample size n but relies on the Lipschitz constant of the loss function which makes it distribution-independent . A distribution-independent bound can be potentially loose and may not capture empirical observations ( e.g. , a network trained using true labels generalizes better than a network trained using corrupted labels as shown in Figure 1 ) . The second one ( Theorem 2 of Mou et al. , 2018 ) , a PAC-Bayes bound , replaces the Lipschitz constant by an expected-squared gradient norm but suffers from a slower rate O ( 1/ √ n ) . In contrast , our bound has order O ( 1/n ) and tightens the expected-squared gradient norm by the variance of gradients . The PAC-Bayes bound in Mou et al . ( 2018 ) is the only SGLD bound , besides ours , which incorporates an explicit time-decaying factor . However , their decay factor is incomparable with ours since they are established through distinct proof techniques and rely on different assumptions : their factor relies on an L2 regularization1 whereas ours is focused on the DP-SGD algorithm . A follow-up work by Li et al . ( 2019 ) combined the algorithmic stability approach with PAC-Bayesian theory and presented a bound which achieves order O ( 1/n ) . However , their bound requires the scale of the learning rate to be upper bounded by the inverse Lipschitz constant of the loss function . In contrast , we do not need any assumptions on the learning rate . There are significant recent works that adopt the information-theoretic framework ( Xu and Raginsky , 2017 ) for bounding the SGLD generalization gap . Among them , Pensia et al . ( 2018 ) initially proposed a bound in Corollary 1 for analyzing a class of noisy iterative algorithms , including SGLD , and their bound was extended in Proposition 3 of Bu et al . ( 2020 ) . However , the bounds in Pensia et al . ( 2018 ) ; Bu et al . ( 2020 ) are distribution-independent . Recently , Negrea et al . ( 2019 ) improved these bounds by replacing the Lipschitz constant with a gradient prediction residual , which is distribution-dependent . To compare with the generalization bound in Negrea et al . ( 2019 ) , we incorporate a time-decaying factor into our bound under additional assumptions . Furthermore , the numerical experiments ( Section 5 ) suggest that our bound is more favourably correlated with the true generalization gap . Negrea et al . ( 2019 ) provided another bound in their Theorem 3.1 which involves a variance term but it suffers from a slower rate O ( 1/ √ n ) . More broadly , there are several recent works considering algorithms closely related to SGLD . For example , Haghifam et al . ( 2020 ) investigated the Langevin dynamics algorithm ( i.e. , full batch SGLD ) , which was later extended by Rodríguez-Gálvez et al . ( 2020 ) to SGLD , and observed a time-decaying phenomenon in their experiments . Specifically , Haghifam et al . ( 2020 ) incorporated a quantity , namely the squared error probability of the hypothesis test , into their bound in Theorem 4.2 and this quantity decays with the number of iterations . This seems to suggest that earlier iterations have a larger impact on their generalization bound . In contrast , our decay factor and the counterpart in Mou et al . ( 2018 ) indicate that the impact of earlier iterations is reducing with the total number of iterations . A recent work by Neu et al . ( 2021 ) investigated the generalization properties of SGD and provided a bound that also involves the variance of gradients . Note that the generalization bound in their Proposition 3 suffers from a weaker order O ( 1/ √ n ) when the analysis is applied to SGLD . To summarize , our main contribution is to provide a generalization bound for SGLD that incorporates the variance of gradients for quantifying the “ sharpness ” of the loss landscape and hasO ( 1/n ) sample size dependence . Moreover , our bound holds under mild conditions and is applicable to a general setting in which the output from SGLD can be any function of the iterations . Under additional assumptions that are satisfied by DP-SGD , our bound includes an explicit time-decaying factor . Although each of the above contributions has been discussed in existing works , to the best of our knowledge , our bound is the first one that simultaneously combines all these aspects . Moreover , our proof techniques based on information-theoretic tools ( e.g. , strong data processing inequalities and properties of Gaussian channels ) may be of broader interest to the community studying generalization theory . 1We note that Mou et al . ( 2018 ) also provided a bound in Theorem 23 without requiring regularization but the rate of the decay factor is slower than our factor in this case . 2 Preliminaries . Consider the following ( possibly non-convex ) optimization problem : min w∈W Lµ ( w ) , E [ ` ( w , Z ) ] = ∫ Z ` ( w , z ) dµ ( z ) , where w ∈ W ⊆ Rd is the parameter ( e.g. , weights of a neural network ) to optimize ; µ is the underlying data distribution that generates Z ; and ` : W ×Z → R+ is the loss function ( e.g. , 0-1 loss ) . Since µ is unknown , Lµ ( w ) can not be computed directly . Hence , one can instead minimize the empirical risk using a dataset S , ( Z1 , · · · , Zn ) which contains n i.i.d . points Zi ∼ µ : min w∈W LS ( w ) , 1 n n∑ i=1 ` ( w , Zi ) . Stochastic gradient Langevin dynamics . We consider the stochastic gradient Langevin dynamics ( SGLD ) algorithm ( Gelfand and Mitter , 1991 ; Welling and Teh , 2011 ) for solving this empirical risk optimization . The dataset S is first divided into m disjoint mini-batches : S = m⋃ j=1 Sj , where |Sj | = b and Sj ∩ Sk = ∅ for j 6= k. We initialize the parameter with a random point W0 ∈ W and update using the following rule : Wt = Wt−1 − ηt∇w ˆ̀ ( Wt−1 , SBt ) + √ 2ηt βt N , ( 1 ) where ηt is the learning rate ; βt is the inverse temperature ; N ∼ N ( 0 , Id ) is a random variable drawn independently from a standard Gaussian distribution ; Bt ∈ [ m ] is the mini-batch index2 ; ˆ̀ is a surrogate loss ( e.g. , hinge loss ) ; and ∇w ˆ̀ ( Wt−1 , SBt ) , 1 b ∑ Z∈SBt ∇w ˆ̀ ( Wt−1 , Z ) . ( 2 ) The recursion in ( 1 ) runs for T iterations and the final output is W = f ( W1 , · · · , WT ) which is a function of the parameters across all iterations . For example , the output can be the parameter at the last iteration W = WT , the one which achieves the smallest value of the loss function W = argminWt Lµ ( Wt ) , or an average of the parameters ( i.e. , Polyak averaging ) W = 1 T ∑ t Wt . Information-theoretic generalization bounds . The goal of this paper is to derive an upper bound for the expected generalization gap : E [ Lµ ( W ) − LS ( W ) ] . ( 3 ) A recent work by Xu and Raginsky ( 2017 ) provided a new method for bounding the expected generalization gap in terms of the mutual information between the input dataset S and the output parameter W. The bound in Xu and Raginsky ( 2017 ) was later tightened by Bu et al . ( 2020 ) . Lemma 1 ( Bu et al . ( 2020 ) Proposition 1 ) . Let the loss function ` ( w , Z ) be σ-sub-Gaussian under Z ∼ µ for all w ∈ W . For any learning algorithm which takes a dataset S = ( Z1 , · · · , Zn ) as input and outputs W , |E [ Lµ ( W ) − LS ( W ) ] | ≤ 1 n n∑ i=1 √ 2σ2I ( W ; Zi ) , ( 4 ) where I ( W ; Zi ) is the mutual information between the learning algorithm ’ s output W and an individual data point Zi . 2For the sake of illustration , we assume that the mini-batch indices are specified before the SGLD is run . Strong data processing inequalities . In order to characterize the time-decaying phenomenon , we use an information-theoretic tool : strong data processing inequalities ( Dobrushin , 1956 ; Cohen et al. , 1998 ) . The data processing inequality ( Cover and Thomas , 2012 ) states that if a Markov chain U→ X→ Y holds , then I ( U ; Y ) ≤ I ( U ; X ) . In other words , no post-processing of X can increase the information about U . Under certain conditions , the data processing inequality can be sharpened , which leads to a strong data processing inequality , often cast in terms of a contraction coefficient . Next , we recall the contraction coefficients of f -divergences and show their connection with strong data processing inequalities . Let f : ( 0 , ∞ ) → R be a convex function with f ( 1 ) = 0 and P , Q be two probability distributions over a set X ⊆ Rd . The f -divergence ( Csiszár , 1967 ) between P and Q is defined as Df ( P‖Q ) , ∫ X f ( dP dQ ) dQ . ( 5 ) Examples of f -divergence include KL-divergence ( f ( t ) = t log t ) and total variation distance ( f ( t ) = |t − 1|/2 ) . For a given transition probability kernel PY|X : X → Y , let PY|X ◦ P be the distribution on Y induced by the push-forward of the distribution P ( i.e. , the distribution of Y when the distribution of X is P ) . The contraction coefficient of PY|X for Df is defined as ηf ( PY|X ) , sup P , Q : P 6=Q Df ( PY|X ◦ P‖PY|X ◦Q ) Df ( P‖Q ) ∈ [ 0 , 1 ] . In particular , when the total variation distance is used , the corresponding contraction coefficient ηTV ( PY|X ) is known as the Dobrushin ’ s coefficient ( Dobrushin , 1956 ) , which upper bounds all other contraction coefficients ( Cohen et al. , 1998 ) : ηf ( PY|X ) ≤ ηTV ( PY|X ) . Furthermore , for any Markov chain U → X → Y , the contraction coefficient of KL-divergence satisfies ( Ahlswede and Gács , 1976 ) I ( U ; Y ) ≤ ηKL ( PY|X ) · I ( U ; X ) . ( 6 ) When ηKL ( PY|X ) < 1 , the strict inequality I ( U ; Y ) < I ( U ; X ) improves the data processing inequality and , hence , is referred to as a strong data processing inequality . We refer the reader to Polyanskiy and Wu ( 2016 ) and Raginsky ( 2016 ) for a more comprehensive review on strong data processing inequalities and Calmon et al . ( 2017 ) for non-linear strong data processing inequalities in Gaussian channels . Gaussian channels . We describe next a few fundamental properties of Gaussian channels . They will be used to derive a closed-form expression of the decay factor and to upper bound the mutual information by a quantity that can be estimated from data . Consider a pair of random variables ( X , Y ) related by Y = X +mN where X is lying on X ; m > 0 is a constant ; and N ∼ N ( 0 , Id ) follows a standard Gaussian distribution . This model can be regarded as a single use of a Gaussian channel , which has a long history in information theory and possesses many interesting properties . For example , if X is a compact set , the contraction coefficients have a non-trivial upper bound ηKL ( PY|X ) ≤ ηTV ( PY|X ) = 1− 2Φ̄ ( diam ( X ) 2m ) , ( 7 ) where diam ( X ) , supx , x′∈X ‖x− x′‖2 is the diameter of X and Φ̄ ( t ) , ∫∞ t 1√ 2π exp ( −v2/2 ) dv is the Gaussian complementary cumulative distribution function ( CCDF ) . Another useful property is the following inequality ( see Lemma 3.4.2 in Raginsky and Sason , 2012 , for a proof ) which upper bounds the KL-divergence of the output distributions from the Gaussian channel by the Wasserstein distance of their input distributions . It also serves as a fundamental lemma for proving Otto-Villani ’ s HWI inequality ( Otto and Villani , 2000 ) in the Gaussian case . Lemma 2 . Let X and X′ be a pair of random variables which are independent of N ∼ N ( 0 , Id ) . Then for any m > 0 DKL ( PX+mN‖PX′+mN ) ≤ 1 2m2 W22 ( PX , PX′ ) . ( 8 ) Here W2 ( PX , PX′ ) is the 2-Wasserstein distance equipped with the L2 cost function : W22 ( PX , PX′ ) , inf E [ ‖X− X′‖22 ] , where the infimum is taken over all couplings ( i.e. , joint distributions ) of the random variables X and X′ with marginals PX and PX′ , respectively . We recall an analogous result ( Guo et al. , 2005 ) which is also used in our proof . It gives an upper bound for the input-output mutual information of a Gaussian channel . Lemma 3 . Let X be a random variable which is independent of N ∼ N ( 0 , Id ) . Then for any m > 0 I ( X +mN ; X ) ≤ 1 2m2 Var ( X ) . ( 9 )

Summary:  The authors study the generalization properties of stochastic gradient Langevin dynamics (SGLD),  using results about Gaussian channels from info theory.   The paper derives a generalization bound that involves the variance of the mini batch gradients in various steps. The variance terms are multiplied by a decaying factor (Thm 1) when one is interested in generalization bound for the last iterate, and the factor is (not surprisingly) absent when the function under consideration depends on all the iterates (Thm 2). The results are established by building on the recent results relating generalization bounds for sub-Gaussian loss functions in terms of the mutual information between the parameter iterates and samples used at various iterations (Bu et al. 2020 result), and upper bounding the mutual information in terms of the variance of gradient of the loss function, and the SGLD step size eta, and temperature beta (Lemma 5). The results scale as 1/n with sample size n---which is better than some of the previous results that try to remove the Lip constant dependence.   The authors present some numerical experiments that demonstrate that their generalization bounds correlate well with generalization gaps observed with small versions of MNIST and CIFAR 10 datasets (5k training examples): (a) with the amount of corrupted labels---increasing corruption increases the generalization error, (b) the width of the network---increasing width decreases the generalization. Their bounds (Table 1) correlate with generalization bounds better compared with Negrea et al. 2019 results, when one considers various settings across learning rate, width, and depth hyper-parameters.  The paper is very well written and fairly easy to follow---a feature missing in several submissions in recent times. It summarizes some useful results in a reader-friendly way.