summaries: This paper study the interplay between generalization bound and the covariance of perturbed gradient descent. The authors invoke a recent (information-theoretic) bound for stochastic optimization methods. Then, they optimize the upper bound with respect to trace constraint on the covariance of noise in each step and experimentally show that this optimization may be effective for training neural networks.  This paper considers the problem of understanding the generalization of the SGLD algorithms using information-theoretic techniques. This line of work initiates by [27],[24], and [13]. The goal in this paper is to "modify" the SGLD algorithms such that it has a good generalization properties and good performance on the training set. Basically, the modification that the authors considered is to make the covariance of the Gaussian noise in the SGLD update to depend on the trajectory. Note that in the usual setting the noise is isotropic. Then, authors claim that controlling the trace of the SGLD ensures we have good performance on training set and good generalization. Then the authors provide a generalization bound which uses reverse KL between posterior and the prior distribution. I think this bound is not true as stated in my review. Since this bound has been used in the paper, I think the authors should comment on the validity of the theorem. This interesting paper studies the connection between generalization abilities of of SGLD and the covariance structure of its noise term. It first proves that with constraint to guarantee low empirical risk and commonly used data-dependent priors, the optimal noise covariance of SGLD in terms of the bound is similar to the empirical gradient covariance. The paper further proves that if the generalization bound is jointly optimized with respect to both prior and posterior, the optimal noise covariance is the square root of the expected gradient covariance. These facts can be used to further support the belief of the superiority of SGD noise over isotropic noise. 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. 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.  
abstractText: Recently, the information-theoretical framework has been proven to be able to obtain non-vacuous generalization bounds for large models trained by Stochastic Gradient Langevin Dynamics (SGLD) with isotropic noise. In this paper, we optimize the information-theoretical generalization bound by manipulating the noise structure in SGLD. We prove that with constraint to guarantee low empirical risk, the optimal noise covariance is the square root of the expected gradient covariance if both the prior and the posterior are jointly optimized. This validates that the optimal noise is quite close to the empirical gradient covariance. Technically, we develop a new information-theoretical bound that enables such an optimization analysis. We then apply matrix analysis to derive the form of optimal noise covariance. Presented constraint and results are validated by the empirical observations.
authors.name: Bohan Wang
authors.name: Huishuai Zhang
authors.name: Jieyu Zhang
authors.name: Wei Chen
references.authors: S. Arora R. Ge B. Neyshabur Y. Zhang
references.title: Stronger generalization bounds for deep nets via a compression approach
references.venue: International Conference on Machine Learning, pages 254–263. PMLR
references.authors: A.R. Asadi E. Abbe S. Verdú
references.title: Chaining mutual information and tightening generalization bounds
references.venue: Advances in Neural Information Processing Systems (NeurIPS)
references.authors: S. Boucheron G. Lugosi P. Massart
references.title: Concentration inequalities: A nonasymptotic theory of independence
references.venue: Oxford university press
references.authors: O. Bousquet S. Boucheron G. Lugosi
references.title: Introduction to statistical learning theory
references.venue: Summer School on Machine Learning, pages 169–207. Springer
references.authors: Y. Bu S. Zou V.V. Veeravalli
references.title: Tightening mutual information-based bounds on generalization error
references.venue: IEEE Journal on Selected Areas in Information Theory, 1(1):121–130
references.authors: Z. Charles D. Papailiopoulos
references.title: Stability and generalization of learning algorithms that converge to global optima
references.venue: International Conference on Machine Learning, pages 745–754. PMLR
references.authors: T.M. Cover J.A. Thomas
references.title: Elements of information theory
references.venue: Wiley
references.authors: C. Dwork V. Feldman M. Hardt T. Pitassi O. Reingold A. Roth
references.title: Generalization in adaptive data analysis and holdout reuse
references.venue: Advances in Neural Information Processing Systems
references.authors: G.K. Dziugaite D.M. Roy
references.title: Computing nonvacuous generalization bounds for deep (stochastic) neural networks with many more parameters than training data
references.venue: arXiv preprint arXiv:1703.11008
references.authors: V. Feldman T. Steinke
references.title: Calibrating noise to variance in adaptive data analysis
references.venue: Conference On Learning Theory, pages 535–544. PMLR
references.authors: S.B. Gelfand S.K. Mitter
references.title: Recursive stochastic algorithms for global optimization in rˆd
references.venue: SIAM Journal on Control and Optimization, 29(5):999–1018
references.authors: N. Gergely G.K. Dziugaite M. Haghifam D.M. Roy
references.title: Information-theoretic generalization bounds for stochastic gradient descent
references.venue: COLT
references.authors: H. Hafez-Kolahi Z. Golgooni S. Kasaei M. Soleymani
references.title: Conditioning and processing: Techniques to improve information-theoretic generalization bounds
references.venue: Advances in Neural Information Processing Systems (NeurIPS), volume 33
references.authors: M. Haghifam J. Negrea A. Khisti D.M. Roy G.K. Dziugaite
references.title: Sharpened generalization bounds based on conditional mutual information and an application to noisy
references.venue: iterative algorithms. arXiv preprint arXiv:2004.12983
references.authors: M. Hardt B. Recht Y. Singer
references.title: Train faster
references.venue: generalize better: Stability of stochastic gradient descent. In International Conference on Machine Learning, pages 1225–1234. PMLR
references.authors: I. Issa A.R. Esposito M. Gastpar
references.title: Strengthened information-theoretic bounds on the generalization error
references.venue: 2019 IEEE International Symposium on Information Theory (ISIT), pages 582–586. IEEE
references.authors: Y. Jiang B. Neyshabur H. Mobahi D. Krishnan S. Bengio
references.title: Fantastic generalization measures and where to find them
references.venue: International Conference on Learning Representations
references.authors: N.S. Keskar D. Mudigere J. Nocedal M. Smelyanskiy P.T.P. Tang
references.title: On large-batch training for deep learning: Generalization gap and sharp minima
references.venue: arXiv preprint arXiv:1609.04836
references.authors: I. Kuzborskij C. Lampert
references.title: Data-dependent stability of stochastic gradient descent
references.venue: International Conference on Machine Learning, pages 2815–2824. PMLR
references.authors: Y. Lei Y. Ying
references.title: Sharper generalization bounds for learning with gradient-dominated objective functions
references.venue: International Conference on Learning Representations
references.authors: J. Li X. Luo M. Qiao
references.title: On generalization error bounds of noisy gradient methods for non-convex learning
references.venue: International Conference on Learning Representations
references.authors: S. Mandt M.D. Hoffman D.M. Blei
references.title: Stochastic gradient descent as approximate Bayesian inference
references.venue: The Journal of Machine Learning Research, 18(1):4873–4907
references.authors: Q. Meng S. Gong W. Chen Z.-M. Ma T.-Y. Liu
references.title: Dynamic of stochastic gradient descent with state-dependent noise
references.venue: arXiv preprint arXiv:2006.13719
references.authors: W. Mou L. Wang X. Zhai K. Zheng
references.title: Generalization bounds of SGLD for non-convex learning: Two theoretical viewpoints
references.venue: Conference on Learning Theory, pages 605–638. PMLR
references.authors: J. Negrea M. Haghifam G.K. Dziugaite A. Khisti D.M. Roy
references.title: Information-theoretic generalization bounds for SGLD via data-dependent estimates
references.venue: arXiv preprint arXiv:1911.02151
references.authors: B. Neyshabur S. Bhojanapalli N. Srebro
references.title: A PAC-Bayesian approach to spectrallynormalized margin bounds for neural networks
references.venue: International Conference on Learning Representations
references.authors: E. Parrado-Hernández A. Ambroladze J. Shawe-Taylor S. Sun
references.title: PAC-Bayes bounds with data dependent priors
references.venue: The Journal of Machine Learning Research, 13(1):3507–3531
references.authors: A. Pensia V. Jog P.-L. Loh
references.title: Generalization error bounds for noisy
references.venue: iterative algorithms. In 2018 IEEE International Symposium on Information Theory (ISIT), pages 546–550. IEEE
references.authors: M. Raginsky A. Rakhlin M. Telgarsky
references.title: Non-convex learning via stochastic gradient langevin dynamics: a nonasymptotic analysis
references.venue: Conference on Learning Theory, pages 1674–1703. PMLR
references.authors: D. Russo J. Zou
references.title: Controlling bias in adaptive data analysis using information theory
references.venue: Artificial Intelligence and Statistics, pages 1232–1240. PMLR
references.authors: T. Steinke L. Zakynthinou
references.title: Reasoning about generalization via conditional mutual information
references.venue: Conference on Learning Theory, pages 3437–3452. PMLR
references.authors: G. Strang G. Strang G. Strang G. Strang
references.title: Introduction to linear algebra
references.venue: volume 3. Wellesley-Cambridge Press Wellesley, MA
references.authors: M. Welling Y.W. Teh
references.title: Bayesian learning via stochastic gradient langevin dynamics
references.venue: Proceedings of the 28th international conference on machine learning (ICML-11), pages 681–688. Citeseer
references.authors: Y. Wen K. Luk M. Gazeau G. Zhang H. Chan J. Ba
references.title: An empirical study of large-batch stochastic gradient descent with structured covariance noise
references.venue: arXiv e-prints, pages arXiv–1902
references.authors: H. Xiao K. Rasul R. Vollgraf
references.title: Fashion-MNIST: a novel image dataset for benchmarking machine learning algorithms
references.venue: arXiv preprint arXiv:1708.07747
references.authors: Z. Xie I. Sato M. Sugiyama
references.title: A diffusion theory for deep learning dynamics: Stochastic gradient descent escapes from sharp minima exponentially fast
references.venue: arXiv preprint arXiv:2002.03495
references.authors: A. Xu M. Raginsky
references.title: Information-theoretic analysis of generalization capability of learning algorithms
references.venue: arXiv preprint arXiv:1705.07809
references.authors: C. Zhang S. Bengio M. Hardt B. Recht O. Vinyals
references.title: Understanding deep learning requires rethinking generalization
references.venue: International Conference on Learning Representations
references.authors: W. Zhou V. Veitch M. Austern R.P. Adams P. Orbanz
references.title: Non-vacuous generalization bounds at the imagenet scale: a PAC-Bayesian compression approach
references.venue: International Conference on Learning Representations
references.authors: Z. Zhu J. Wu B. Yu L. Wu J. Ma
references.title: The anisotropic noise in stochastic gradient descent: Its behavior of escaping from sharp minima and regularization effects
references.venue: arXiv preprint arXiv:1803.00195
sections.heading: 1 Introduction
sections.text: 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:
sections.heading: 1.1 Related Works
sections.text: 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.
sections.heading: 2 Preliminaries
sections.text: 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.
sections.heading: 3 Formulate the Problem: Proper Constraints and Optimization Target
sections.text: In this part, we formulate the optimization problem, i.e., finding the optimal noise covariance of SGLD in terms of the information-theoretical generalization bound, by selecting the proper optimization constraint and optimization target. Specifically, in Section 3.1, we argue that the trace of the noise covariance is a proper constraint to ensure the same optimization error; in Section 3.2, we argue that existing generalization bounds are not proper candidates as the optimization target, and propose a new information-theoretical bound parallel to existing ones but easier to analyze.
sections.heading: 3.1 Constraint on the Covariance to Control the Empirical Risk
sections.text: We first derive the constraint of noise covariance Σt in Eq. (3) from the perspective of training performance, under which we optimize the generalization bounds in the rest of this paper.
  Without any constraint on the noise covariance, optimizing generalization error is trivial but meaningless: a direct combination of Theorem 1 of [37] and Theorem 1 of [28] shows that for isotropic SGLD with Σt = σtI, the generalization error after T step satisfies Gerr ≤ O(( ∑T t=1 log(1 + 1/σt))
  1/2). Therefore, as σt →∞ for t ∈ [T ], we have Gerr→ 0, but then the update of SGLD is dominated by the noise, leading to arbitrary bad empirical risk. Hence, we need constraints on the covariance in order to control the empirical risk when minimizing the generalization error. Specifically, the expected decrease of the empirical risk for one iteration can be bounded as follows. Lemma 2. Let empirical riskRS(W ) be β-smooth w.r.t. W . Let W[T ] be given by state-dependent SGLD (Eq. (3)). Then, Et+1|tRS(Wt+1)−RS(Wt) ≤ − ( 1 − βηt+1
  2
  ) ηt+1 ∥∇RS (Wt)∥2+ β
  2 tr ( η2t+1(N − bt) (N − 1)bt Σ sd S,Wt + Σt+1(S,Wt) ) .
  The proof can be obtained by a standard analysis in optimization, and we defer it to Appendix B. By Lemma 2, the noise covariance Σt affects the upper bound of the empirical risk by its trace. Therefore, it is reasonable to keep tr(Σt(S,Wt−1)) unchanged while seeking the optimal Σt(S,Wt−1) to minimize the generalization error. The constraint is given formally as follows: Constraint 1. The trace of Σt(S,Wt−1) is fixed when optimizing the generalization error. That is, there exist positive constants ct(S,Wt−1) depending on S and Wt−1, such that,
  tr(Σt(S,Wt−1)) = ct(S,Wt−1).
  We do not put any constraint on the value of ct in our latter analyses. Therefore, it is also possible to manipulate ct in order to jointly optimize the empirical risk and the generalization bound, which
  however, is beyond the scope of this paper and we defer it to future works. Similar constraint is also proposed by [40] from the standpoint of kinetic energy when analyzing the effect of noise structure on the escaping efficiency from saddle points.
  We next verify this constraint empirically. We run SGLD with different covariances on a four-layer neural network for the Fashion-MNIST classification problem. Concretely, the noise covariances are chosen respectively as Σ(1)t = Σ sd S,Wt−1 , i.e., “EC-SGLD” ( Empirical Covariance SGLD) curve in Fig.2, and Σ(2)t = 1 d tr(Σ sd S,Wt−1
  )I, i.e., the “Iso-SGLD (C)” curve in Fig.2. It is easy to verify that tr(Σ
  (1) t ) = tr(Σ (2) t ), which is exactly the Constraint 1. We can see from Fig.2 that the convergence curves corresponding to Σ(1)t and Σ (2) t almost coincide with each other, validating the Constraint 1.
  In comparison, another quantity tr(HS,Wt−1Σt), the trace of the empirical risk’s Hessian times the empirical covariance, has also been proposed to govern the convergence behavior [34, Theorem 4.1]. We plot the convergence curve of SGLD with noise covariance
  Σ (3) t =
  tr(HS,Wt−1Σ sd S,Wt−1 )
  tr(HS,Wt−1 ) I ("Iso-SGLD (H)" in Fig.2). While
  tr(HS,Wt−1Σ (1) t ) = tr(HS,Wt−1Σ (3) t ), there is a significant gap between curves of “EC-SGLD ” and “Iso-SGLD (H)” in Fig.2. This implies that tr(HS,Wt−1Σt) is not a good constraint for the empirical risk of SGLD, validating the Constraint 1 from another side.
  In the rest of this paper, we will optimize the information-theoretical bound under Constraint 1 to ensure low empirical risk.
sections.heading: 3.2 New Information-theoretical Bounds as the Optimization Target
sections.text: We first demonstrate that existing information-theoretical bounds are not suitable for the optimization target in Section 3.2.1. Then we propose a new information-theoretical bound as the optimization target in Section 3.2.2.
sections.heading: 3.2.1 Difficulties When Applying Traditional Information-theoretical Bounds
sections.text: By Lemma 1, the generalization error bound in Proposition 1 can be rewritten as
  ES,V[T ],J √√√√ (a2 − a1)2 2 T∑ s=1 E Q S,V[s−1] s−1 KL ( QS,Vss|(s−1) ∥∥∥PJ,SJ ,Vss|(s−1) ). (4) For Problem 1, there are two difficulties. For a fixed QS,V[T ] , in order to obtain the optimal PJ,SJ ,Vss|(s−1) for any s ∈ [T ], one would first calculate the outside expectation of Eq. (4) over SJc , as PSJ ,V[T ] is independent of SJc . However, for each term E
  Q S,V[s−1] s−1
  KL (QS,Vss|(s−1)∥P J,SJ ,Vs s|(s−1) ) where s ∈ [T ],
  both the probability measure Q S,V[s−1] s−1 and the function KL(Q S,Vs s|(s−1)∥P J,SJ ,Vs s|(s−1) ) has dependency on SJc , which makes evaluating the generalization bound with respect to PSJ ,V[T ] extremely hard. Furthermore, when we come to optimize the bound w.r.t. Σ[T ], for each Σi where i ∈ [T − 1], the Eq. (4) depends on Σi via two terms A and B, where
  A = E Q
  S,V[T ] i−1
  KL ( Q
  S,V[T ] i|(i−1)∥P J,SJ ,V[T ] i|(i−1)
  ) , B = T∑ s=i+1 E Q S,V[s−1] s−1 KL ( QS,Vss|(s−1)∥P J,SJ ,Vs s|(s−1) ) .
  While the first term is easy to deal with as Q S,V[T ] i−1 is irrelevant of Σi, the second term depends on Σi via the distribution Q S,V[s−1] s−1 for s ≥ i+ 1, and can be very complex (please refer to Appendix C.1).
sections.heading: 3.2.2 Information-theoretical Generalization Bound for Identifying the Noise Effect
sections.text: We then establish a new information-theoretical generalization bound as the optimization target that are parallel to the results in [25] but more suitable for analyzing the effect of noise structure. The
  basic idea is to reverse the order of prior and posterior in the KL divergence to let Σt only affect one term in the generalization bound for each t when finding the optimal point. The formal theorem is stated as follows: Theorem 1. Let sample set S, mini-batch V[T ], random subset J , the posterior distribution Q output by Algorithm A with update ruleM and the prior distribution P output by Algorithm B with update rule M̃ be defined as Proposition 1. Let [a1, a2] be the range of loss. Then, the generalization error of A can be bounded as
  Gerr ≤ ES,V[T ],J
  √ (a2 − a1)2
  2 KL
  ( PJ,SJ ,V[T ] ∥∥QS,V[T ] ). (5) Compared to Proposition 1, Theorem 1 can be viewed as a parallel version with the positions of the prior distribution P and the posterior distribution Q reversed. This reverse benefits the optimization of the generalization error bound with respect to the noise covariance under Constraint 1 in two ways. First, by Lemma 1, the KL term in Eq. (5) can be further decomposed into
  KL ( PSJ ,V[T ] ∥∥QS,V[T ] ) = T∑ s=1 E P J,SJ ,V[s−1] 1:s−1 KL ( PJ,SJ ,Vss|(s−1) ||Q S,Vs s|(s−1) ) ,
  in which only each KL term depends on each Σs, respectively. Secondly, for each summand in the above decomposition, only QS,Vss|(s−1) depends on SJc , making it easy to compute the optimal prior in a "greedy" sense (see details in Section 4).
  Hence in the rest of this paper, we study the optimal structure of Σt in terms of the generalization bound in Theorem 1 under the Constraint 1. Specifically, our ultimate goal can be stated as follows: Problem 1. What is the optimal structure of the noise covariance Σ[T ] for SGLD in terms of the generalization bound in Theorem 1 under the Constraint 1?
  In the rest of this paper, we focus on solving Problem 1 and its variants.
sections.heading: 4 Obtain the Optimal Noise Covariance with the Greedy Prior
sections.text: The generalization bound in Theorem 1 depends on both the prior distribution P and the posterior distribution Q (consequently on Σ[T ]). Therefore it requires searching P and Σ[T ] jointly to optimize the generalization error.
  In this section, we solve Problem 1 with greedily selected priors. Due to that the square root in the bound Eq. (5) make the dependency on Σ[T ] even more complex, we optimize a slightly different version of the bound by taking the expectation with respect to SJc into the square root, i.e.,
  GenT △ = ESJ ,V[T ],J
  √ (a2 − a1)2
  2 ESJc KL
  ( PJ,SJ ,V[T ] ∥∥QS,V[T ] ). (6) By Jensen’s Inequality, Eq. (6) is still a generalization bound, but allows the expectation with respect to SJc to interact with the KL divergence directly. One can easily observe that GenT is a mapping from P and Q (and thus from P and Σ[T ]) to a positive real. Therefore, we can reformulate Problem 1 mathematically as the following optimization problem (P1):
  (P1). min P,Σ[T ] GenT (P,Σ[T ]), subject to: Constraint 1.
  For simplicity, we restrict that the considered noise covariance Σ[T ] only depends on the parameter W , and is independent of the sample S. By Lemma 1, the KL term in the optimization target GenT can be decomposed into
  KL ( PJ,SJ ,V[T ] ∥∥∥QS,V[T ] ) = T∑ s=1 E P J,SJ ,V[s−1] s−1 KL ( PJ,SJ ,Vss|(s−1) ∥∥∥QS,Vss|(s−1) ) . (7) Therefore, when optimizing GenT with respect to prior P , both KL(P J,SJ ,Vs s|(s−1) ||Q S,Vs s|(s−1)) and P J,SJ ,V[i−1] i−1 (i > s) have dependence on P J,SJ ,Vs s|(s−1) . Similar to the discussion in Section 3.2.1, the dependence of P J,SJ ,V[i−1] i−1 with i > s on P J,SJ ,Vs s|(s−1) can be very complex. Therefore, we approximate the optimal PJ,SJ ,Vss|(s−1) by the greedy prior which is defined as follows:
  Definition 1 (Greedy Prior). We say P ∗ is the optimal prior in the greedy sense, or the greedy prior for brevity, if for any 1 ≤ s ≤ T and any SJ and V[T ], P ∗J,SJ ,Vss|(s−1) = P sJ,SJ ,Vs s|(s−1) , where P
  s is defined as follows:
  P s △ = argmin
  P ( min Σ[s] Gens(P,Σ[s]) ) , subject to: Constraint 1.
  Intuitively, the conditional probability of P ∗ of the step s is the optimal one if we only consider the generalization bound for steps up to s, and a special case is that the step T conditional probability of P ∗ agrees with the step T conditional probability of PT , which is the desired optimal prior. This is why we call P ∗ "greedy" and use it to approximate the optimal prior.
  With the greedy prior, we characterize the optimal noise covariance by the following theorem.
  Theorem 2. Let the iteration of SGLD with state-dependent noise QS,V[T ] be given as Eq.(3). Under Constraint 1, the greedy optimal prior of step t is given by
  Wt = Wt−1 − ηt ( |Vt ∩ J | |Vt| ∇RSVt∩J (Wt−1) + |Vt ∩ Jc| |Vt| ∇RD (Wt−1) ) +N (0,Σ∗t (Wt−1)) ,
  while the optimal covariance of noise Σ∗[T ] for GenT with the greedy prior is given by Σ ∗ t (W ) =
  λt(W ) (Σ pop W ) 1 2 (∀t ∈ [T ]), where λt(W ) = ct(W )/ tr((ΣpopW ) 1 2 ).
  As the sample size is large enough, we have ΣsdS,W → Σ pop W almost surely, which demonstrates the similarity between the solution of Problem 1 and the noise covariance of SGD. Also, by the Law of Large Numbers, ∇RSJ ≈ ∇RD, and the mean of the greedy optimal prior recovers the mean of the prior used in [25] (one can also refer to Eq. (10) in this paper for the form).
  We briefly state the proof skeleton of Theorem 2, with the proof details deferred to Appendix D. To obtain the final optimal noise covariance in Theorem 2, we need to first derive the greedy prior, i.e., the optimal conditional distribution Ps|(s−1) of the prior of step s in terms of the generalization bound Gens, which has the form
  ESJ ,V[s],J √√√√ (a2 − a1)2 2 ESJc s∑
  t=1
  E P
  J,SJ ,V[t−1] t−1
  KL ( PJ,SJ ,Vtt|(t−1) ∥∥∥QS,Vtt|(t−1)). (8) Typically, solving the optimal Ps|(s−1) requires optimizing Gens with respect to all Pt|(t−1) t ∈ [s] and Σ[s], which is still very complex. However, we can tackle this problem in a rather elegant way. We first investigate the optimal noise covariance in terms of a single KL divergence term in Gens.
  Lemma 3. Under Constraint 1, the optimal noise covariance of the following problem
  min Σs  min P
  J,SJ ,Vs s|(s−1)
  ESJc∼D KL ( PJ,SJ ,Vss|(s−1) ∥∥∥QS,Vss|(s−1))  . (9)
  is attained at Σ∗t (W ), where Σ ∗ t (W ) is defined as Theorem 2.
  By Lemma 3, the optimal solution of Eq. (9) doesn’t rely on J ,SJ , or V[s]. On the other hand, by Eq. (8), ESJc∼D KL ( PJ,SJ ,Vss|(s−1) ∥∥∥QS,Vss|(s−1)) (∀J ,SJ ,V[s]) are the only terms depending on PJ,SJ ,Vss|(s−1) and Σs. Therefore, the optimal Σs is also λt(W ) (Σ pop W ) 1 2 , which is formally stated as the following lemma.
  Lemma 4. The optimal Σs and Ps|(s−1) in terms of Gens are the same as Σ∗s and P ∗s|(s−1) given by Theorem 2, respectively.
  With the greedy prior derived, we apply it back to the generalization bound GenT . As GenT depends on Σs also through ESJc∼D KL ( PJ,SJ ,Vss|(s−1) ∥∥∥QS,Vss|(s−1)), by applying Lemma 3 again, we derive Theorem 2.
sections.heading: 5 Extension: Optimal Noise Covariance with Fixed Priors
sections.text: In existing works [25, 12], the prior distribution is set to be the SGLD with isotropic noise, which (with the notations in Theorem 1) is given by
  M̃t(Wt−1,SJ ,Vt,J) = Wt−1 − ηt ( |Vt ∩ J | |Vt| ∇RSVt∩J (Wt−1) + |Vt ∩ Jc| |Vt| ∇RSJ (Wt−1) ) +N (0, σtId) , (10)
  where σt > 0 is the noise scale of prior noise covariance. In our latter analysis, we generalize the prior by allowing σt depend on Wt−1. The formal description of the iteration of the prior is deferred to Appendix E.1 for completeness.
  Therefore, it is also interesting to see what the optimal noise covariance looks like if the prior is fixed as the one commonly adopted in the existing analyses, e.g., Eq. (10). We still set the optimization constraint the same as Section 4, but change the optimization target a little to
  G̃enT △ = ES
  √ (a2 − a1)2
  2 EV[T ],J KL
  ( PJ,SJ ,V[T ] ∥∥QS,V[T ] ). By Jensen’s Inequality and Theorem 1, G̃enT is still a generalization bound, but allows us to treat the expectation of the KL divergence with respect to V[T ] and J for given S as a whole in optimization. Similar trick is also adopted in [25] to obtain the final generalization error of SGLD. The problem can be mathematically formulated as the following optimization problem (P2):
  (P2) min Σ[T ] G̃enT (P,Σ[T ]), subject to: Constraint 1,
  where P is given by the update rule Eq. (10). As SJ is obtained by removing only one sample from the dataset S and the size of S is usually large in practice, it is reasonable to make the assumption that for any fixed V[T ] and S, the prior distribution is the same regardless of J .
  Assumption 1. For any fixed dataset S and mini-batches V[T ], the distribution PJ,SJ ,V[T ] is invariant of J .
  We also restrict that ct(S,W ) ≤ dσt in order to guarantee the noise scale of the prior comparable to that of the posterior. The optimal noise covariance with prior fixed can then be characterized by the following theorem. Theorem 3. Let prior and posterior be defined as Eq.(10) and Eq.(3), respectively. Let Assumption 1 hold. Then, the solution of (P2) is given by
  Σ∗t (S,W ) = O sd S,W Diag(ω̃ S,W 1 , · · · , ω̃ S,W d ) ( OsdS,W )⊤ ,
  where ω̃S,Wi ≥ 0 (i ∈ [d]) (the exact form is omitted here) and ( OsdS,W ) is the orthogonal matrix
  that diagonalizes ΣsdS,W as Σ sd S,W = O sd S,W Diag(ω S,W 1 , · · · , ω S,W d ) ( OsdS,W )⊤ . Moreover, for any i ̸= j, ω̃S,Wi ≥ ω̃ S,W j if and only if ω S,W i ≥ ω S,W j .
  The proof together with the the exact formula of ω̃S,Wi is deferred to Appendix 5. From Theorem 3, the optimal point (Σ∗t ) T t=1 is similar to the empirical gradient covariance matrix Σ
  sd in two ways: 1) {Σ∗t }Tt=1 share the same eigenvectors with Σsd; 2) the corresponding eigenvalues of Σ∗t has the same order as Σsd. Though the value of ω̃S,Wi is not comparable to ω S,W i because Σ ∗ t is affected by the prior noise and the posterior noise covariance, which are freely chosen, it can be shown that the condition number of Σ∗t is smaller than Σ sd S,W (please refer to Appendix E.4 for details).
  Theorem 3 also reveals an interesting correlation between the noise covariance matrices of the prior and posterior. While the optimal noise covariance is affected by the prior noise covariance, it also depends on the distance between the means of the prior and the posterior (see Lemma 13 in Appendix E.2), which brings the information of empirical gradient covariance into the optimal noise structure. That being said, the optimal posterior noise covariance is biased to the empirical gradient covariance from prior covariance. We note that such analysis can be easily extended to arbitrary priors.
sections.heading: 6 Empirical Verification
sections.text: In this section, we support our theoretical findings with some experiments. We adopt the setting in [40] where a four-layer neural network with 11330 parameters is used to conduct the classification task on the Fashion-MNIST except that we use 10000 training samples instead of 1200 used in [40]. We defer detailed settings of the experiments to Appendix F.
  We first verify if the empirical gradient covariance is far from isotropic Gaussian distribution. We plot the ratios of the 1st eigenvalue to the 500th largest eigenvalue of the empirical gradient covariance along the training trajectory of both SGD and Iso-SGLD in Figure 4. We can see that throughout the training procedure, the ratios of empirical gradient covariance stay around 107 for both SGD and Iso-SGLD. In constrast, for the isotropic Gaussian, the distribution of eigenvalues follows semi-circle law, and the ratio would be around 1.1 for dimension 11330. This demonstrates that the energy of empirical gradient covariance concentrates in a very small subspace, less than 5% of the total 11330 dimensions. Hence the empirical gradient covariance is highly anisotropic.
  We next verify our main claim Theorem 2 that under Constraint 1, the generalization error for SGLD with the optimal noise covariance is much smaller than that for Iso-SGLD. Specifically, we consider the SGLD with noise covariance equal to the (scaled) square root of the empirical gradient covariance, coined “SREC-SGLD”. We compare the training loss and the generalization error of SREC-SGLD and Iso-SGLD under Constraint 1 on the noise covariance trace. From Figure 3, we can see that the generalization error of SREC-SGLD is smaller than Iso-SGLD for different noise trace scales, which supports Theorem 2. Moreover, if we look at the training loss curves, the SREC-SGLD and Iso-SGLD behave almost the same for the same noise trace scale, which supports Lemma 2.
  Finally though our analysis does not cover the generalization bound for SGD, we further empirically show how SREC-SGLD performs in comparison with SGD and the Iso-SGLD in Figure 1. We can see that with the same noise trace scale, SREC indeed provides a more accurate characterization of SGD compared to Iso-SGLD.
sections.heading: 7 Conclusion and Future Direction
sections.text: In this paper, we study the optimal noise covariance of SGLD in terms of its generalization ability. Specifically, we first formulate the optimization problem both by deriving constraints from the optimization performance and proposing the optimization target by constructing a new information theoretical bound. We then solve the problem with both greedy optimal prior and fixed prior. Interestingly, we observe that the optimal noise covariance aligns with the empirical gradient covariance, which indicates the superiority of the noise covariance of SGD in terms of generalization.
sections.heading: Acknowledgement
sections.text: The authors would like to thank Mr. Ziming Liu for helpful theoretical discussions.
title: Optimizing Information-theoretical Generalization Bounds via Anisotropic Noise in SGLD