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: From usual isentropic noise to data-dependent priors, the authors extend the existing information-theoretic analysis of SGLD. Additionally, they identify the optimal choice of prior distribution covariance. Last but not least, the paper includes empirical observations to validate its technical findings.