Optimizing Information-theoretical Generalization Bounds via Anisotropic Noise in SGLD
Year: 2022

Abstract:
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.

heading: 1 Introduction
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:

heading: 1.1 Related Works
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.

heading: 2 Preliminaries
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.

heading: 3 Formulate the Problem: Proper Constraints and Optimization Target
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.

heading: 3.1 Constraint on the Covariance to Control the Empirical Risk
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.

heading: 3.2 New Information-theoretical Bounds as the Optimization Target
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.

heading: 3.2.1 Difficulties When Applying Traditional Information-theoretical Bounds
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).

heading: 3.2.2 Information-theoretical Generalization Bound for Identifying the Noise Effect
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.

heading: 4 Obtain the Optimal Noise Covariance with the Greedy Prior
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.

heading: 5 Extension: Optimal Noise Covariance with Fixed Priors
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.

heading: 6 Empirical Verification
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.

heading: 7 Conclusion and Future Direction
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.

heading: Acknowledgement
text: The authors would like to thank Mr. Ziming Liu for helpful theoretical discussions.

Review 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.