Optimizing Information-theoretical Generalization Bound via Anisotropic Noise of SGLD

1 Introduction . Generalization ability is one of the core questions in learning theory [ 4 ] , but remains unclear for deep learning models [ 17 , 38 ] . Existing generalization bounds based on the capacity control become vacuous for practical deep learning models due to the over-parameterization property [ 1 , 26 , 39 ] . From the information theoretical perspective , recent works [ 30 , 37 ] bound the generalization error by the mutual information between the dataset and the learned parameters . This result reveals that good generalization occurs when the learned parameters do not depend on a specific dataset too much , which is intuitively reasonable and closely related with the idea of algorithm stability [ 6 , 15 , 19 , 24 , 20 ] and differential privacy [ 8 , 10 ] . Meanwhile , based on information-theoretic metrics , one can analyze general classes of updates and models , e.g. , stochastic iterative algorithms for non-convex objectives , hence applicable to deep learning . It has been shown that the information theoretical bounds are non-vacuous and closely related with the real generalization error even in deep learning [ 14 , 26 , 9 , 39 ] . Notably , the information theoretical bound is realized in [ 28 ] via decomposing the mutual information across iterations . This framework can perform a step-wise analysis of Stochastic Gradient Langevin Dynamics ( SGLD ) [ 33 , 29 ] , by evaluating the mutual information conditional on the previous learned ∗Corresponding Author 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . parameters at each step . We note that the noise added in SGLD is critical for both the mutual information evaluation and the empirical risk minimization . Most existing work focus on SGLD with constant and isotropic noise covariance [ 22 , 25 , 14 ] due to the technical difficulty . However , it is observed that the test accuracy of SGLD with isotropic noise has a considerable gap compared to that of the widely-used Stochastic Gradient Descent ( SGD ) [ 40 ] . This empirical gap motivates us to consider the following question : Can we find the optimal noise added in SGLD in terms of generalization ? An affirmative answer will lead to an algorithm imitating SGD better , which helps us to better understand the generalization behavior of SGD . Specifically , we propose to optimize the structure of the noise in SGLD such that the generalization bound is minimized while a low empirical risk is guaranteed . To this end , we first show that the trace of the noise covariance in SGLD is a valid constraint that governs the empirical risk behavior both theoretically and empirically . Then we devise a new information theoretical generalization bound that are parallel to those bounds in [ 25 ] , but facilitate the derivation of the optimal noise . With these technical preparations , we prove that when jointly optimizing both the prior and the posterior , the optimal noise covariance is the square root of the expected gradient covariance , i.e. , their eigenvectors are the same and the corresponding eigenvalues of the former are the square root of the latter , and the optimal prior recovers the prior in [ 25 ] . This indicates the optimal noise covariance of SGLD would be quite close to the empirical gradient covariance , i.e. , the noise covariance of SGD , because of the concentration of measure . Our result lends support to the belief that the noise introduced by Stochastic Gradient Descent ( SGD ) is superior to the isotropic noise , which has been widely observed [ 18 , 36 , 40 , 38 ] . As an illustrative example , we plot the generalization errors of SGD , SGLD with the isotropic noise and SGLD with the optimal noise in Figure 1 , where their training curves behave almost the same ( do not show here ) . We can see the optimal noise captures the behavior of SGD much better than the isotropic noise . Specifically , our contribution can be summarized as follows : 1.1 Related Works . Information-theoretical bounds . Recently , researchers [ 37 , 30 , 28 ] propose to bound the generalization error by the mutual information between output hypothesis and input samples . Negrea et al . [ 25 ] further tighten the bound by designing a data-dependent prior . Following [ 25 ] , Haghifam et al . [ 14 ] obtain comparable results through conditional mutual information [ 31 ] . Other related work [ 2 , 16 , 5 , 13 ] tightens the information theoretical generalization bounds from different perspectives . There are also high probability generalization bound [ 24 , 26 , 21 ] obtained by combining information theory [ 7 ] and PAC Bayesian framework [ 27 , 33 ] . Effects of the noise covariance . SGD achieves excellent performance in terms of generalization error in deep learning . In contrast , one recent work [ 40 ] demonstrates empirically that even with well-tuned scaling , isotropic SGLD still achieves much worse generalization error than mini-batch SGD . They obtain an optimal covariance matrix ( low-rank approximation of Hessian ) in terms of escaping efficiency at local minima . Furthermore , [ 34 ] show empirically large-batch SGD with diagonal Fisher Gaussian noise can recover similar validation performance as small-batch SGD . However , both the noise covariance investigated by [ 40 ] and [ 34 ] has gap with that of SGD . It has also been proven [ 23 ] that the stationary distribution of state-dependent SGLD with the empirical loss approximated by quadratic function obeys the power law , and has a better escaping efficiency than SGLD with the isotropic noise . 2 Preliminaries . Notations . Here we briefly introduce the notations which will be used throughout this paper . • ( Set operation ) For a positive integer N , we use [ N ] to denote the index set { 1 , 2 , · · · , N } , and use V [ T ] to denote the set { Vt } Tt=1 . For a finite set S = { z1 , · · · , zN } , |S| is the cardinality of S , and SJ △ = { zi } i∈J is a subset of S with indices in set J ⊂ [ N ] . • ( Probability ) z ∼ S denotes that z is uniformly sampled from the set S. For two random variables X and Y , we denote the conditional distribution of X given Y as PY ( X ) and denote the conditional expectation of X given Y as EY ( X ) . • ( Matrix ) We use Id×d for the d× d identity matrix , abbreviated as I when dimension is clear . For a differentiable function f , we denote the gradient of f at point W as∇f ( W ) . For a positive semi-definite matrix A ∈ Rd×d , we say B = A 12 ∈ Rd×d if B ·B = A . Supervised Learning . In this paper , we focus on the supervised learning . It conducts the empirical risk minimization ( ERM ) over training data : MinimizeW RS ( W ) △ = 1|S| ∑ z∈S ℓ ( z ; W ) , where W ∈ Rd is the parameter of the model , S is the training set with each data point i.i.d . sampled from a distribution D , ℓ is the ( individual ) loss function . A stochastic algorithm solves the ERM by outputting a distribution QS of parameter W . The generalization error measures the gap between the population risk and the training risk as Gerr △ = ES , QS [ RD ( W ) −RS ( W ) ] . ( 1 ) SGD and SGLD . The update rule of Stochastic Gradient Descent ( SGD ) at step t is defined as Wt ←Wt−1 − ηt∇RSVt ( Wt−1 ) , ( 2 ) where Vt is sampled uniformly without replacement from [ N ] with size bt . Given Wt−1 , Wt is random with E [ Wt|Wt−1 ] = Wt−1 − ηt∇RS ( Wt−1 ) , Cov [ Wt|Wt−1 ] = N − bt bt ( N − 1 ) ΣsdS , Wt−1 , where ΣsdS , Wt−1 △ = Covz∼S [ ∇Rz ( Wt−1 ) ] is the covariance of∇Rz with z single drawn from S. The superscript sd means “ single draw ” . Similarly , we define the population covariance of gradient as ΣpopW △ = Covz∼D [ ∇Rz ( W ) ] . We define the Hessian ofRS ( W ) at point W asHS , W . The Stochastic Gradient Langevine Dynamics ( SGLD ) is given by Wt ←Wt−1 − ηt∇RSVt ( Wt−1 ) +N ( 0 , Σt ( S , Wt−1 ) ) , ( 3 ) where Σt ( S , Wt−1 ) ∈ Rd×d is a positive semi-definite matrix with dependence on S and Wt−1 . SGLD ( Eq . 3 ) with Σt ( S , Wt−1 ) = ctI where ct a positive constant , is called isotropic SGLD , and SGLD with Σt ( S , Wt−1 ) dependent on Wt−1 is called state-dependent SGLD . Statistics for iterative algorithms . Both SGD and SGLD are stochastic iterative algorithms [ 11 ] . Specifically , the update rule at step t of a stochastic iterative algorithm A can be generally characterized as Wt =Mt ( Wt−1 , S , Vt ) , where Vt is the auxiliary random variable at step t , e.g. , Vt in Eq . ( 2 ) and Eq. ( 3 ) . We use QS , V [ T ] to denote the joint distribution of ( Wt ) Tt=0 conditional on { S , V [ T ] } , Q S , V [ T ] i : j to denote the joint distribution of ( Wt ) j t=i conditional on { S , V [ T ] } , and Q S , V [ T ] j|j−1 to denote the distribution of Wj conditional on { Wj−1 , S , V [ T ] } . Decomposition of KL divergence . The following Lemma is extensively used throughout this paper . Lemma 1 ( Proposition 2.6 , [ 25 ] ) . Let Q0 : T and P0 : T are two probability measures on Rd× ( T+1 ) with Q0 = P0 . Then , the KL divergence between Q0 : T and P0 : T can be decomposed into KL ( Q0 : T ||P0 : T ) = T∑ t=1 EQ0 : t−1 [ KL ( Qt| [ t−1 ] ∥Pt| [ t−1 ] ) ] . Information theoretical generalization bound . Several existing information-theoretical bounds [ 14 , 25 ] share similar framework . We state one representative proposed in [ 25 ] . Proposition 1 ( Theorem 2.5 in [ 25 ] ) . Let S be the data set i.i.d . sampled from D , and let the loss function ℓ be [ a1 , a2 ] bounded . Let A be an algorithm with update rule Wt ←Mt ( Wt−1 , S , Vt ) , W0 ∼ W0 . Let B be another stochastic iterative optimization algorithm with update rule Wt ← M̃t ( Wt−1 , SJ , Vt , J ) , W0 ∼ W0 . where J is sampled uniformly without replacement from [ N ] with size N − 1 . Given S , J , and V [ T ] , denote the joint distribution of ( Wt ) Tt=1 of A as QS , V [ T ] , and the joint distribution of ( Wt ) Tt=1 of B as PJ , SJ , V [ T ] . Then , for any M̃ [ T ] , the generalization error of A can be bounded as : ES , V [ T ] [ RD ( Q S , V [ T ] T ) − R̂S ( Q S , V [ T ] T ) ] ≤ ES , V [ T ] , J √ ( a2 − a1 ) 2 2 KL ( QS , V [ T ] ∥∥PJ , SJ , V [ T ] ) . This proposition is used to obtain the generalization error bound for SGLD [ 25 , Theorem 3.1 ] by further combing Lemma 1 . As M̃ [ T ] can be arbitrarily picked , P works as an `` Auxiliary Line '' and is called the prior distribution , while Q is the real distribution of parameters called the posterior distribution . In Section 3.2.1 , we will argue the difficulty of applying Proposition 1 to analyze the effect of noise structures .

Summary: This paper studies optimization of the information theoretical generalization bounds of Stochastic Gradient Langevin Dynamics (SGLD) with respect to the noise covariance. Here are the main findings: * When updating rule is fixed and only optimizing noise covariance under the constraint that the trace of the covariance is fixed, the noise covariance of SGLD is similar to the empirical gradient covariance * When both updating rule and the noise covariance are jointly minimized, the noise covariance of SGLD is the square root of the expected gradient covariance.  The two results can potentially be used to improve information theoretical generalization bounds obtained in the literature. The theoretical results are validated by some numerical experiments.