# Optimizing Information-theoretical Generalization Bounds via Anisotropic
Noise in SGLD

Bohan Wang
University of Science & Technology of China
Microsoft Research Asia &Huishuai Zhang
Microsoft Research Asia &Jieyu Zhang
University of Washington
Microsoft Research Asia &Qi Meng
Microsoft Research Asia &Wei Chen
Microsoft Research Asia &Tie-Yan Liu
Microsoft Research Asia Corresponding Author

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

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

Figure 1: Generalization errors of SGD, SGLD with the isotropic noise (Iso-
SGLD) and SGLD with the optimal noise (SREC-SGLD).

Specifically, our contribution can be summarized as follows: 1\. We formulate
a problem of finding the optimal noise covariance by optimizing an
information-theoretical bound;
2\. We develop a new information-theoretical bound to facilitate the analysis
of the above optimization problem;
3\. We obtain the optimal structure of the noise covariances, and demonstrate
the similarity to empirical/expected gradient covariance;

###  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,\cdots,N\\}$, and use $\bm{V}_{[T]}$ to denote the set
$\\{\bm{V}_{t}\\}_{t=1}^{T}$. For a finite set
$\bm{S}=\\{\bm{z}_{1},\cdots,\bm{z}_{N}\\}$, $|\bm{S}|$ is the cardinality of
$\bm{S}$, and ${\small\bm{S}_{J}\overset{\triangle}{=}\\{\bm{z}_{i}\\}_{i\in
J}}$ is a subset of $\bm{S}$ with indices in set $J\subset[N]$.

  * •

(Probability) $\bm{z}\sim\bm{S}$ denotes that $\bm{z}$ is uniformly sampled
from the set $\bm{S}$. For two random variables $\bm{X}$ and $\bm{Y}$, we
denote the conditional distribution of $\bm{X}$ given $\bm{Y}$ as
$\mathbb{P}^{\bm{Y}}(\bm{X})$ and denote the conditional expectation of
$\bm{X}$ given $\bm{Y}$ as $\mathbb{E}^{\bm{Y}}(\bm{X})$.

  * •

(Matrix) We use $\mathbb{I}_{d\times d}$ for the $d\times d$ identity matrix,
abbreviated as $\mathbb{I}$ when dimension is clear. For a differentiable
function $f$, we denote the gradient of $f$ at point $\bm{W}$ as $\nabla
f(\bm{W})$. For a positive semi-definite matrix $\bm{A}\in\mathbb{R}^{d\times
d}$, we say $\bm{B}=\bm{A}^{\frac{1}{2}}\in\mathbb{R}^{d\times d}$ if
$\bm{B}\cdot\bm{B}=\bm{A}$.

Supervised Learning. In this paper, we focus on the supervised learning. It
conducts the empirical risk minimization (ERM) over training data:
$\operatorname{Minimize}_{\bm{W}}\mathcal{R}_{\bm{S}}(\bm{W})\overset{\triangle}{=}\frac{1}{|\bm{S}|}\sum_{\bm{z}\in\bm{S}}\ell(\bm{z};\bm{W}),$
where $\bm{W}\in\mathbb{R}^{d}$ is the parameter of the model, $\bm{S}$ is the
training set with each data point i.i.d. sampled from a distribution
$\mathcal{D}$, $\ell$ is the (individual) loss function. A _stochastic
algorithm_ solves the ERM by outputting a distribution $Q^{\bm{S}}$ of
parameter $\bm{W}$. The _generalization error_ measures the gap between the
population risk and the training risk as

$\operatorname{Gerr}\overset{\triangle}{=}\mathbb{E}_{\bm{S},Q^{\bm{S}}}\left[\mathcal{R}_{\mathcal{D}}(\bm{W})-\mathcal{R}_{\bm{S}}(\bm{W})\right].$
(1)

SGD and SGLD.  The update rule of Stochastic Gradient Descent (SGD) at step
$t$ is defined as

$\bm{W}_{t}\leftarrow\bm{W}_{t-1}-\eta_{t}\nabla\mathcal{R}_{\bm{S}_{\bm{V}_{t}}}(\bm{W}_{t-1}),$
(2)

where $\bm{V}_{t}$ is sampled uniformly without replacement from $[N]$ with
size $b_{t}$. Given $\bm{W}_{t-1}$, $\bm{W}_{t}$ is random with

$\mathbb{E}\left[\bm{W}_{t}|\bm{W}_{t-1}\right]=\bm{W}_{t-1}-\eta_{t}\nabla\mathcal{R}_{\bm{S}}\left(\bm{W}_{t-1}\right),\text{
}\mathrm{Cov}\left[\bm{W}_{t}|\bm{W}_{t-1}\right]=\frac{N-b_{t}}{b_{t}(N-1)}\bm{\Sigma}^{sd}_{\bm{S},\bm{W}_{t-1}},$

where
${\small\bm{\Sigma}^{sd}_{\bm{S},\bm{W}_{t-1}}\overset{\triangle}{=}\mathrm{Cov}_{\bm{z}\sim\bm{S}}\left[\nabla\mathcal{R}_{\bm{z}}(\bm{W}_{t-1})\right]}$
is the covariance of $\nabla\mathcal{R}_{\bm{z}}$ with $\bm{z}$ single drawn
from $\bm{S}$. The superscript sd means “single draw”. Similarly, we define
the population covariance of gradient as
$\bm{\Sigma}^{\operatorname{pop}}_{\bm{W}}\overset{\triangle}{=}\mathrm{Cov}_{\bm{z}\sim\mathcal{D}}\left[\nabla\mathcal{R}_{\bm{z}}(\bm{W})\right]$.
We define the Hessian of $\mathcal{R}_{\bm{S}}(\bm{W})$ at point $\bm{W}$ as
$\mathcal{H}_{\bm{S},\bm{W}}$.

The Stochastic Gradient Langevine Dynamics (SGLD) is given by

$\bm{W}_{t}\leftarrow\bm{W}_{t-1}-\eta_{t}\nabla\mathcal{R}_{\bm{S}_{\bm{V}_{t}}}(\bm{W}_{t-1})+\mathcal{N}(\bm{0},\bm{\Sigma}_{t}(\bm{S},\bm{W}_{t-1})),$
(3)

where $\bm{\Sigma}_{t}(\bm{S},\bm{W}_{t-1})\in\mathbb{R}^{d\times d}$ is a
positive semi-definite matrix with dependence on $\bm{S}$ and $\bm{W}_{t-1}$.
SGLD (Eq. 3) with $\bm{\Sigma}_{t}(\bm{S},\bm{W}_{t-1})=c_{t}\mathbb{I}$ where
$c_{t}$ a positive constant, is called _isotropic_ SGLD, and SGLD with
$\bm{\Sigma}_{t}(\bm{S},\bm{W}_{t-1})$ dependent on $\bm{W}_{t-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 $\mathcal{A}$ can be generally characterized as
$\bm{W}_{t}=\mathcal{M}_{t}(\bm{W}_{t-1},\bm{S},\bm{V}_{t})$, where
$\bm{V}_{t}$ is the _auxiliary random variable_ at step $t$, e.g.,
$\bm{V}_{t}$ in Eq.(2) and Eq.(3). We use $Q^{\bm{S},\bm{V}_{[T]}}$ to denote
the joint distribution of $(\bm{W}_{t})_{t=0}^{T}$ conditional on
$\\{\bm{S},\bm{V}_{[T]}\\}$, $Q_{i:j}^{\bm{S},\bm{V}_{[T]}}$ to denote the
joint distribution of $(\bm{W}_{t})_{t=i}^{j}$ conditional on
$\\{\bm{S},\bm{V}_{[T]}\\}$, and $Q_{j|j-1}^{\bm{S},\bm{V}_{[T]}}$ to denote
the distribution of $\bm{W}_{j}$ conditional on
$\\{\bm{W}_{j-1},\bm{S},\bm{V}_{[T]}\\}$.

Decomposition of KL divergence. The following Lemma is extensively used
throughout this paper.

######  Lemma 1 (Proposition 2.6, [25]).

Let $Q_{0:T}$ and $P_{0:T}$ are two probability measures on
$\mathbb{R}^{d\times(T+1)}$ with $Q_{0}=P_{0}$. Then, the KL divergence
between $Q_{0:T}$ and $P_{0:T}$ can be decomposed into

$\operatorname{KL}(Q_{0:T}||P_{0:T})=\sum_{t=1}^{T}\mathbb{E}_{Q_{0:t-1}}\left[\operatorname{KL}\left(Q_{t|[t-1]}\|P_{t|[t-1]}\right)\right].$

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 $\bm{S}$ be the data set i.i.d. sampled from $\mathcal{D}$, and let the
loss function $\ell$ be $[a_{1},a_{2}]$ bounded. Let $\mathcal{A}$ be an
algorithm with update rule

$\bm{W}_{t}\leftarrow\mathcal{M}_{t}(\bm{W}_{t-1},\bm{S},\bm{V}_{t}),\bm{W}_{0}\sim\mathcal{W}_{0}.$

Let $\mathcal{B}$ be another stochastic iterative optimization algorithm with
update rule

$\bm{W}_{t}\leftarrow\tilde{\mathcal{M}}_{t}(\bm{W}_{t-1},\bm{S}_{\bm{J}},\bm{V}_{t},\bm{J}),\bm{W}_{0}\sim\mathcal{W}_{0}.$

where $\bm{J}$ is sampled uniformly without replacement from $[N]$ with size
$N-1$. Given $\bm{S}$, $\bm{J}$, and $\bm{V}_{[T]}$, denote the joint
distribution of $(\bm{W}_{t})_{t=1}^{T}$ of $\mathcal{A}$ as
$Q^{\bm{S},\bm{V}_{[T]}}$, and the joint distribution of
$(\bm{W}_{t})_{t=1}^{T}$ of $\mathcal{B}$ as
$P^{\bm{J},\bm{S}_{\bm{J}},\bm{V}_{[T]}}$. Then, for any
$\tilde{\mathcal{M}}_{[T]}$, the generalization error of $\mathcal{A}$ can be
bounded as:

$\displaystyle\mathbb{E}_{\bm{S},\bm{V}_{[T]}}\left[\mathcal{R}_{\mathcal{D}}\left(Q_{T}^{\bm{S},\bm{V}_{[T]}}\right)-\hat{\mathcal{R}}_{S}\left(Q_{T}^{\bm{S},\bm{V}_{[T]}}\right)\right]\leq\mathbb{E}_{\bm{S},\bm{V}_{[T]},\bm{J}}\sqrt{\frac{(a_{2}-a_{1})^{2}}{2}\operatorname{KL}\left(Q^{\bm{S},\bm{V}_{[T]}}\left\|P^{\bm{J},\bm{S}_{\bm{J}},\bm{V}_{[T]}}\right.\right)}.$

This proposition is used to obtain the generalization error bound for SGLD
[25, Theorem 3.1] by further combing Lemma 1. As $\tilde{\mathcal{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.

##  3 Formulate the Problem: Proper Constraints and Optimization Target

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.

###  3.1 Constraint on the Covariance to Control the Empirical Risk

We first derive the constraint of noise covariance $\bm{\Sigma}_{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
$\bm{\Sigma}_{t}=\sigma_{t}\mathbb{I}$, the generalization error after $T$
step satisfies
$\operatorname{Gerr}\leq\mathcal{O}((\sum_{t=1}^{T}\log(1+1/\sigma_{t}))^{1/2})$.
Therefore, as $\sigma_{t}\rightarrow\infty$ for $t\in[T]$, we have
$\operatorname{Gerr}\rightarrow 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 risk $\mathcal{R}_{\bm{S}}(\bm{W})$ be $\beta$-smooth w.r.t.
$\bm{W}$. Let $\bm{W}_{[T]}$ be given by state-dependent SGLD (Eq. (3)). Then,

$\mathbb{E}_{t+1|t}\mathcal{R}_{\bm{S}}(\bm{W}_{t+1})-\mathcal{R}_{\bm{S}}(\bm{W}_{t})\leq-\left(1-\frac{\beta\eta_{t+1}}{2}\right)\eta_{t+1}\left\|\nabla\mathcal{R}_{\bm{S}}\left(\bm{W}_{t}\right)\right\|^{2}+\frac{\beta}{2}\operatorname{tr}\left(\frac{\eta_{t+1}^{2}(N-b_{t})}{(N-1)b_{t}}\bm{\Sigma}^{sd}_{\bm{S},\bm{W}_{t}}+\bm{\Sigma}_{t+1}(\bm{S},\bm{W}_{t})\right).$

The proof can be obtained by a standard analysis in optimization, and we defer
it to Appendix B. By Lemma 2, the noise covariance $\bm{\Sigma}_{t}$ affects
the upper bound of the empirical risk by its _trace_. Therefore, it is
reasonable to keep $\operatorname{tr}(\bm{\Sigma}_{t}(\bm{S},\bm{W}_{t-1}))$
unchanged while seeking the optimal $\bm{\Sigma}_{t}(\bm{S},\bm{W}_{t-1})$ to
minimize the generalization error. The constraint is given formally as
follows:

######  Constraint 1.

The trace of $\bm{\Sigma}_{t}(\bm{S},\bm{W}_{t-1})$ is fixed when optimizing
the generalization error. That is, there exist positive constants
$c_{t}(\bm{S},\bm{W}_{t-1})$ depending on $\bm{S}$ and $\bm{W}_{t-1}$, such
that,

$\operatorname{tr}(\bm{\Sigma}_{t}(\bm{S},\bm{W}_{t-1}))=c_{t}(\bm{S},\bm{W}_{t-1}).$

We do not put any constraint on the value of $c_{t}$ in our latter analyses.
Therefore, it is also possible to manipulate $c_{t}$ 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
$\bm{\Sigma}^{(1)}_{t}=\bm{\Sigma}^{sd}_{\bm{S},\bm{W}_{t-1}}$, i.e., “EC-
SGLD” ( Empirical Covariance SGLD) curve in Fig.2, and
$\bm{\Sigma}^{(2)}_{t}=$
$\frac{1}{d}\operatorname{tr}(\bm{\Sigma}^{sd}_{\bm{S},\bm{W}_{t-1}})\mathbb{I}$,
i.e., the “Iso-SGLD (C)” curve in Fig.2. It is easy to verify that
$\operatorname{tr}(\bm{\Sigma}^{(1)}_{t})=\operatorname{tr}(\bm{\Sigma}^{(2)}_{t})$,
which is exactly the Constraint 1. We can see from Fig.2 that the convergence
curves corresponding to $\bm{\Sigma}^{(1)}_{t}$ and $\bm{\Sigma}^{(2)}_{t}$
almost coincide with each other, validating the Constraint 1.

Figure 2: Training errors of SGLD with different noise covariances. The
experiment is run on the Fashion-MNIST dataset with a four-layer neural
network (see Appendix F).

In comparison, another quantity
$\operatorname{tr}(\mathcal{H}_{\bm{S},\bm{W}_{t-1}}\bm{\Sigma}_{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
$\bm{\Sigma}^{(3)}_{t}=\frac{\operatorname{tr}(\mathcal{H}_{\bm{S},\bm{W}_{t-1}}\bm{\Sigma}^{sd}_{\bm{S},\bm{W}_{t-1}})}{\operatorname{tr}(\mathcal{H}_{\bm{S},\bm{W}_{t-1}})}\mathbb{I}$
("Iso-SGLD (H)" in Fig.2). While
$\operatorname{tr}(\mathcal{H}_{\bm{S},\bm{W}_{t-1}}\bm{\Sigma}^{(1)}_{t})=\operatorname{tr}(\mathcal{H}_{\bm{S},\bm{W}_{t-1}}\bm{\Sigma}^{(3)}_{t})$,
there is a significant gap between curves of “EC-SGLD ” and “Iso-SGLD (H)” in
Fig.2. This implies that
$\operatorname{tr}(\mathcal{H}_{\bm{S},\bm{W}_{t-1}}\bm{\Sigma}_{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.

###  3.2 New Information-theoretical Bounds as the Optimization Target

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.

####  3.2.1 Difficulties When Applying Traditional Information-theoretical
Bounds

By Lemma 1, the generalization error bound in Proposition 1 can be rewritten
as

$\mathbb{E}_{\bm{S},\bm{V}_{[T]},\bm{J}}\sqrt{\frac{(a_{2}-a_{1})^{2}}{2}\sum_{s=1}^{T}\mathbb{E}_{Q_{s-1}^{\bm{S},\bm{V}_{[s-1]}}}\operatorname{KL}\left(Q_{s|(s-1)}^{\bm{S},\bm{V}_{s}}\left\|P_{s|(s-1)}^{\bm{J},\bm{S}_{\bm{J}},\bm{V}_{s}}\right.\right)}.$
(4)

For Problem 1, there are two difficulties. For a fixed
$Q^{\bm{S},\bm{V}_{[T]}}$, in order to obtain the optimal
$P_{s|(s-1)}^{\bm{J},\bm{S}_{\bm{J}},\bm{V}_{s}}$ for any $s\in[T]$, one would
first calculate the outside expectation of Eq. (4) over $\bm{S}_{\bm{J}^{c}}$
, as $P^{\bm{S}_{\bm{J}},\bm{V}_{[T]}}$ is independent of
$\bm{S}_{\bm{J}^{c}}$. However, for each term
$\mathbb{E}_{Q_{s-1}^{\bm{S},\bm{V}_{[s-1]}}}\operatorname{KL}$
$(Q_{s|(s-1)}^{\bm{S},\bm{V}_{s}}\|P_{s|(s-1)}^{\bm{J},\bm{S}_{\bm{J}},\bm{V}_{s}})$
where $s\in[T]$, both the probability measure
$Q_{s-1}^{\bm{S},\bm{V}_{[s-1]}}$ and the function
$\operatorname{KL}(Q_{s|(s-1)}^{\bm{S},\bm{V}_{s}}\|P_{s|(s-1)}^{\bm{J},\bm{S}_{\bm{J}},\bm{V}_{s}})$
has dependency on $\bm{S}_{\bm{J}^{c}}$ , which makes evaluating the
generalization bound with respect to $P^{\bm{S}_{\bm{J}},\bm{V}_{[T]}}$
extremely hard. Furthermore, when we come to optimize the bound w.r.t.
$\bm{\Sigma}_{[T]}$, for each $\bm{\Sigma}_{i}$ where $i\in[T-1]$, the Eq. (4)
depends on $\bm{\Sigma}_{i}$ via two terms $A$ and $B$, where

$A=\mathbb{E}_{Q_{i-1}^{\bm{S},\bm{V}_{[T]}}}\operatorname{KL}\left(Q_{i|(i-1)}^{\bm{S},\bm{V}_{[T]}}\|P_{i|(i-1)}^{\bm{J},\bm{S}_{\bm{J}},\bm{V}_{[T]}}\right),~{}B=\sum_{s=i+1}^{T}\mathbb{E}_{Q_{s-1}^{\bm{S},\bm{V}_{[s-1]}}}\operatorname{KL}\left(Q_{s|(s-1)}^{\bm{S},\bm{V}_{s}}\|P_{s|(s-1)}^{\bm{J},\bm{S}_{\bm{J}},\bm{V}_{s}}\right).$

While the first term is easy to deal with as $Q_{i-1}^{\bm{S},\bm{V}_{[T]}}$
is irrelevant of $\bm{\Sigma}_{i}$, the second term depends on
$\bm{\Sigma}_{i}$ via the distribution $Q_{s-1}^{\bm{S},\bm{V}_{[s-1]}}$ for
$s\geq i+1$, and can be very complex (please refer to Appendix C.1).

####  3.2.2 Information-theoretical Generalization Bound for Identifying the
Noise Effect

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 $\bm{\Sigma}_{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 $\bm{S}$, mini-batch $\bm{V}_{[T]}$, random subset $\bm{J}$,
the posterior distribution $Q$ output by Algorithm $\mathcal{A}$ with update
rule $\mathcal{M}$ and the prior distribution $P$ output by Algorithm
$\mathcal{B}$ with update rule $\tilde{\mathcal{M}}$ be defined as Proposition
1. Let $[a_{1},a_{2}]$ be the range of loss. Then, the generalization error of
$\mathcal{A}$ can be bounded as

$\displaystyle\operatorname{Gerr}\leq\mathbb{E}_{\bm{S},\bm{V}_{[T]},\bm{J}}\sqrt{\frac{(a_{2}-a_{1})^{2}}{2}\operatorname{KL}\left(P^{\bm{J},\bm{S}_{\bm{J}},\bm{V}_{[T]}}\left\|Q^{\bm{S},\bm{V}_{[T]}}\right.\right)}.$
(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

$\operatorname{KL}\left(P^{\bm{S}_{\bm{J}},\bm{V}_{[T]}}\left\|Q^{\bm{S},\bm{V}_{[T]}}\right.\right)=\sum_{s=1}^{T}\mathbb{E}_{P_{1:s-1}^{\bm{J},\bm{S}_{\bm{J}},\bm{V}_{[s-1]}}}\operatorname{KL}\left(P_{s|(s-1)}^{\bm{J},\bm{S}_{\bm{J}},\bm{V}_{s}}||Q_{s|(s-1)}^{\bm{S},\bm{V}_{s}}\right),$

in which only each KL term depends on each $\bm{\Sigma}_{s}$, respectively.
Secondly, for each summand in the above decomposition, only
$Q_{s|(s-1)}^{\bm{S},\bm{V}_{s}}$ depends on $\bm{S}_{\bm{J}^{c}}$, 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
$\bm{\Sigma}_{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 $\bm{\Sigma}_{[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.

##  4 Obtain the Optimal Noise Covariance with the Greedy Prior

The generalization bound in Theorem 1 depends on both the prior distribution
$P$ and the posterior distribution $Q$ (consequently on $\bm{\Sigma}_{[T]}$).
Therefore it requires searching $P$ and $\bm{\Sigma}_{[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
$\bm{\Sigma}_{[T]}$ even more complex, we optimize a slightly different
version of the bound by taking the expectation with respect to
$\bm{S}_{\bm{J}^{c}}$ into the square root, i.e.,

${\operatorname{Gen}}_{T}\overset{\triangle}{=}\mathbb{E}_{\bm{S}_{\bm{J}},\bm{V}_{[T]},\bm{J}}\sqrt{\frac{(a_{2}-a_{1})^{2}}{2}\mathbb{E}_{\bm{S}_{\bm{J}^{c}}}\operatorname{KL}\left(P^{\bm{J},\bm{S}_{\bm{J}},\bm{V}_{[T]}}\left\|Q^{\bm{S},\bm{V}_{[T]}}\right.\right)}.$
(6)

By Jensen’s Inequality, Eq. (6) is still a generalization bound, but allows
the expectation with respect to $\bm{S}_{\bm{J}^{c}}$ to interact with the KL
divergence directly. One can easily observe that $\operatorname{Gen}_{T}$ is a
mapping from $P$ and $Q$ (and thus from $P$ and $\bm{\Sigma}_{[T]}$) to a
positive real. Therefore, we can reformulate Problem 1 mathematically as the
following optimization problem (P1):

$\textbf{(P1)}.\min_{P,\bm{\Sigma}_{[T]}}\operatorname{Gen}_{T}(P,\bm{\Sigma}_{[T]}),\text{
subject to: Constraint \ref{constraint: trace}}.$

For simplicity, we restrict that the considered noise covariance
$\bm{\Sigma}_{[T]}$ only depends on the parameter $\bm{W}$, and is independent
of the sample $\bm{S}$. By Lemma 1, the KL term in the optimization target
$\operatorname{Gen}_{T}$ can be decomposed into

$\operatorname{KL}\left(P^{\bm{J},\bm{S}_{\bm{J}},\bm{V}_{[T]}}\left\|Q^{\bm{S},\bm{V}_{[T]}}\right.\right)=\sum_{s=1}^{T}\mathbb{E}_{P_{s-1}^{\bm{J},\bm{S}_{\bm{J}},\bm{V}_{[s-1]}}}\operatorname{KL}\left(P_{s|(s-1)}^{\bm{J},\bm{S}_{\bm{J}},\bm{V}_{s}}\left\|Q_{s|(s-1)}^{\bm{S},\bm{V}_{s}}\right.\right).$
(7)

Therefore, when optimizing $\operatorname{Gen}_{T}$ with respect to prior $P$,
both
$\operatorname{KL}(P_{s|(s-1)}^{\bm{J},\bm{S}_{\bm{J}},\bm{V}_{s}}||Q_{s|(s-1)}^{\bm{S},\bm{V}_{s}})$
and $P_{i-1}^{\bm{J},\bm{S}_{\bm{J}},\bm{V}_{[i-1]}}$ $(i>s)$ have dependence
on $P_{s|(s-1)}^{\bm{J},\bm{S}_{\bm{J}},\bm{V}_{s}}$. Similar to the
discussion in Section 3.2.1, the dependence of
$P_{i-1}^{\bm{J},\bm{S}_{\bm{J}},\bm{V}_{[i-1]}}$ with $i>s$ on
$P_{s|(s-1)}^{\bm{J},\bm{S}_{\bm{J}},\bm{V}_{s}}$ can be very complex.
Therefore, we approximate the optimal
$P_{s|(s-1)}^{\bm{J},\bm{S}_{\bm{J}},\bm{V}_{s}}$ 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\leq s\leq T$ and any $\bm{S}_{\bm{J}}$ and
$\bm{V}_{[T]}$,
${P^{*}}^{\bm{J},\bm{S}_{\bm{J}},\bm{V}_{s}}_{s|(s-1)}={P^{s}}^{\bm{J},\bm{S}_{\bm{J}},\bm{V}_{s}}_{s|(s-1)}$,
where $P^{s}$ is defined as follows:

$P^{s}\overset{\triangle}{=}\arg\min_{P}\left(\min_{\bm{\Sigma}_{[s]}}\operatorname{Gen}_{s}(P,\bm{\Sigma}_{[s]})\right),\text{
subject to: Constraint \ref{constraint: trace}}.$

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 $P^{T}$, 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 $Q^{\bm{S},\bm{V}_{[T]}}$
be given as Eq.(3). Under Constraint 1, the greedy optimal prior of step $t$
is given by

$\bm{W}_{t}=\bm{W}_{t-1}-\eta_{t}\left(\frac{|\bm{V}_{t}\cap\bm{J}|}{|\bm{V}_{t}|}\nabla\mathcal{R}_{\bm{S}_{\bm{V}_{t}\cap\bm{J}}}\left(\bm{W}_{t-1}\right)+\frac{|\bm{V}_{t}\cap\bm{J}^{c}|}{|\bm{V}_{t}|}\nabla\mathcal{R}_{\mathcal{D}}\left(\bm{W}_{t-1}\right)\right)+\mathcal{N}\left(\bm{0},\bm{\Sigma}^{*}_{t}(\bm{W}_{t-1})\right),$

while the optimal covariance of noise $\bm{\Sigma}^{*}_{[T]}$ for
$\operatorname{Gen}_{T}$ with the greedy prior is given by
$\bm{\Sigma}^{*}_{t}(\bm{W})=\lambda_{t}(\bm{W})\left(\bm{\Sigma}^{pop}_{\bm{W}}\right)^{\frac{1}{2}}$
$(\forall t\in[T])$, where
$\lambda_{t}(\bm{W})=c_{t}(\bm{W})/\operatorname{tr}(\left(\bm{\Sigma}^{pop}_{\bm{W}}\right)^{\frac{1}{2}})$.

As the sample size is large enough, we have
$\bm{\Sigma}^{sd}_{\bm{S},\bm{W}}\rightarrow\bm{\Sigma}^{pop}_{\bm{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,
$\nabla\mathcal{R}_{\bm{S}_{\bm{J}}}\approx\nabla\mathcal{R}_{\mathcal{D}}$,
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 $P_{s|(s-1)}$ of the prior of step $s$ in terms of
the generalization bound $\operatorname{Gen}_{s}$, which has the form

$\mathbb{E}_{\bm{S}_{\bm{J}},\bm{V}_{[s]},\bm{J}}\sqrt{\frac{(a_{2}-a_{1})^{2}}{2}\mathbb{E}_{\bm{S}_{\bm{J}^{c}}}\sum_{t=1}^{s}\mathbb{E}_{P_{t-1}^{\bm{J},\bm{S}_{\bm{J}},\bm{V}_{[t-1]}}}\operatorname{KL}\left(P_{t|(t-1)}^{\bm{J},\bm{S}_{\bm{J}},\bm{V}_{t}}\left\|Q_{t|(t-1)}^{\bm{S},\bm{V}_{t}}\right.\right)}.$
(8)

Typically, solving the optimal $P_{s|(s-1)}$ requires optimizing
$\operatorname{Gen}_{s}$ with respect to all $P_{t|(t-1)}$ $t\in[s]$ and
$\bm{\Sigma}_{[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
$\operatorname{Gen}_{s}$.

######  Lemma 3.

Under Constraint 1, the optimal noise covariance of the following problem

$\min_{\bm{\Sigma}_{s}}\left(\min_{{P}_{s|(s-1)}^{\bm{J},\bm{S}_{\bm{J}},\bm{V}_{s}}}\mathbb{E}_{\bm{S}_{\bm{J}^{c}}\sim\mathcal{D}}\operatorname{KL}\left({P}_{s|(s-1)}^{\bm{J},\bm{S}_{\bm{J}},\bm{V}_{s}}\left\|Q_{s|(s-1)}^{\bm{S},\bm{V}_{s}}\right.\right)\right).$
(9)

is attained at $\bm{\Sigma}^{*}_{t}(\bm{W})$, where
$\bm{\Sigma}^{*}_{t}(\bm{W})$ is defined as Theorem 2.

By Lemma 3, the optimal solution of Eq. (9) doesn’t rely on
$\bm{J},\bm{S}_{\bm{J}}$, or $\bm{V}_{[s]}$. On the other hand, by Eq. (8),
$\mathbb{E}_{\bm{S}_{\bm{J}^{c}}\sim\mathcal{D}}\operatorname{KL}\left({P}_{s|(s-1)}^{\bm{J},\bm{S}_{\bm{J}},\bm{V}_{s}}\left\|Q_{s|(s-1)}^{\bm{S},\bm{V}_{s}}\right.\right)$
$(\forall\bm{J},\bm{S}_{\bm{J}},\bm{V}_{[s]})$ are the only terms depending on
${P}_{s|(s-1)}^{\bm{J},\bm{S}_{\bm{J}},\bm{V}_{s}}$ and $\bm{\Sigma}_{s}$.
Therefore, the optimal $\bm{\Sigma}_{s}$ is also
$\lambda_{t}(\bm{W})\left(\bm{\Sigma}^{pop}_{\bm{W}}\right)^{\frac{1}{2}}$,
which is formally stated as the following lemma.

######  Lemma 4.

The optimal $\bm{\Sigma}_{s}$ and $P_{s|(s-1)}$ in terms of
$\operatorname{Gen}_{s}$ are the same as $\bm{\Sigma}^{*}_{s}$ and
$P^{*}_{s|(s-1)}$ given by Theorem 2, respectively.

With the greedy prior derived, we apply it back to the generalization bound
$\operatorname{Gen}_{T}$. As $\operatorname{Gen}_{T}$ depends on
$\bm{\Sigma}_{s}$ also through
$\mathbb{E}_{\bm{S}_{\bm{J}^{c}}\sim\mathcal{D}}\operatorname{KL}\left({P}_{s|(s-1)}^{\bm{J},\bm{S}_{\bm{J}},\bm{V}_{s}}\left\|Q_{s|(s-1)}^{\bm{S},\bm{V}_{s}}\right.\right)$,
by applying Lemma 3 again, we derive Theorem 2.

##  5 Extension: Optimal Noise Covariance with Fixed Priors

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

$\displaystyle\tilde{\mathcal{M}}_{t}(\bm{W}_{t-1},\bm{S}_{\bm{J}},\bm{V}_{t},\bm{J})$
$\displaystyle~{}~{}~{}~{}~{}~{}~{}~{}~{}=\bm{W}_{t-1}-\eta_{t}\left(\frac{|\bm{V}_{t}\cap\bm{J}|}{|\bm{V}_{t}|}\nabla\mathcal{R}_{\bm{S}_{\bm{V}_{t}\cap\bm{J}}}\left(\bm{W}_{t-1}\right)+\frac{|\bm{V}_{t}\cap\bm{J}^{c}|}{|\bm{V}_{t}|}\nabla\mathcal{R}_{\bm{S}_{\bm{J}}}\left(\bm{W}_{t-1}\right)\right)+\mathcal{N}\left(\bm{0},\sigma_{t}\mathbb{I}_{d}\right),$
(10)

where $\sigma_{t}>0$ is the noise scale of prior noise covariance. In our
latter analysis, we generalize the prior by allowing $\sigma_{t}$ depend on
$\bm{W}_{t-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

$\widetilde{\operatorname{Gen}}_{T}\overset{\triangle}{=}\mathbb{E}_{\bm{S}}\sqrt{\frac{(a_{2}-a_{1})^{2}}{2}\mathbb{E}_{\bm{V}_{[T]},\bm{J}}\operatorname{KL}\left(P^{\bm{J},\bm{S}_{\bm{J}},\bm{V}_{[T]}}\left\|Q^{\bm{S},\bm{V}_{[T]}}\right.\right)}.$

By Jensen’s Inequality and Theorem 1, $\widetilde{\operatorname{Gen}}_{T}$ is
still a generalization bound, but allows us to treat the expectation of the KL
divergence with respect to $\bm{V}_{[T]}$ and $\bm{J}$ for given $\bm{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):

$\textbf{(P2)}\;\;\;\;\;\min_{\bm{\Sigma}_{[T]}}\widetilde{\operatorname{Gen}}_{T}(P,\bm{\Sigma}_{[T]}),\text{
subject to: Constraint \ref{constraint: trace}},$

where $P$ is given by the update rule Eq. (10). As $\bm{S}_{\bm{J}}$ is
obtained by removing only one sample from the dataset $\bm{S}$ and the size of
$\bm{S}$ is usually large in practice, it is reasonable to make the assumption
that for any fixed $\bm{V}_{[T]}$ and $\bm{S}$, the prior distribution is the
same regardless of $\bm{J}$.

######  Assumption 1.

For any fixed dataset $\bm{S}$ and mini-batches $\bm{V}_{[T]}$, the
distribution $P^{\bm{J},\bm{S}_{\bm{J}},\bm{V}_{[T]}}$ is invariant of
$\bm{J}$.

We also restrict that $c_{t}(\bm{S},\bm{W})\leq d\sigma_{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

$\bm{\Sigma}^{*}_{t}(\bm{S},\bm{W})=\bm{O}_{\bm{S},\bm{W}}^{sd}\operatorname{Diag}(\tilde{\omega}^{\bm{S},\bm{W}}_{1},\cdots,\tilde{\omega}^{\bm{S},\bm{W}}_{d})\left(\bm{O}_{\bm{S},\bm{W}}^{sd}\right)^{\top},$

where $\tilde{\omega}^{\bm{S},\bm{W}}_{i}\geq 0$ $(i\in[d])$ (the exact form
is omitted here) and $\left(\bm{O}_{\bm{S},\bm{W}}^{sd}\right)$ is the
orthogonal matrix that diagonalizes $\bm{\Sigma}^{sd}_{\bm{S},\bm{W}}$ as
$\bm{\Sigma}^{sd}_{\bm{S},\bm{W}}=\bm{O}_{\bm{S},\bm{W}}^{sd}\operatorname{Diag}(\omega^{\bm{S},\bm{W}}_{1},\cdots,\omega^{\bm{S},\bm{W}}_{d})\left(\bm{O}_{\bm{S},\bm{W}}^{sd}\right)^{\top}.$
Moreover, for any $i\neq j$,
$\tilde{\omega}^{\bm{S},\bm{W}}_{i}\geq\tilde{\omega}^{\bm{S},\bm{W}}_{j}$ if
and only if $\omega^{\bm{S},\bm{W}}_{i}\geq\omega^{\bm{S},\bm{W}}_{j}$.

The proof together with the the exact formula of
$\tilde{\omega}_{i}^{\bm{S},\bm{W}}$ is deferred to Appendix 5. From Theorem
3, the optimal point $(\bm{\Sigma}^{*}_{t})_{t=1}^{T}$ is similar to the
empirical gradient covariance matrix $\bm{\Sigma}^{sd}$ in two ways: 1)
$\\{\bm{\Sigma}^{*}_{t}\\}_{t=1}^{T}$ share the same eigenvectors with
$\bm{\Sigma}^{sd}$; 2) the corresponding eigenvalues of $\bm{\Sigma}^{*}_{t}$
has the same order as $\bm{\Sigma}^{sd}$. Though the value of
$\tilde{\omega}_{i}^{S,W}$ is not comparable to $\omega_{i}^{S,W}$ because
$\Sigma_{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 $\Sigma_{t}^{*}$ is smaller than $\Sigma^{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.

##  6 Empirical Verification

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 $500$th
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 $10^{7}$ 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 4, 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.

(a) $\text{trace scale}=1$

(b) $\text{trace scale}=5$

Figure 3: The training loss and and the generalization error (test loss $-$
training loss) of the Iso-SGLD and SREC-SGLD. Traces of the noise covariance
in (a) and (b) are 1 and 5 times of
$\operatorname{tr}((\bm{\Sigma}^{sd}_{\bm{S},\bm{W}})^{1/2})$, respectively.

Figure 4: Ratio of the 1st to the 500th largest eigenvalue of the empirical
gradient covariance for two SGLDs.

##  7 Conclusion and Future Direction

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.

## Acknowledgement

The authors would like to thank Mr. Ziming Liu for helpful theoretical
discussions.

## References

  * [1] S. Arora, R. Ge, B. Neyshabur, and Y. Zhang.  Stronger generalization bounds for deep nets via a compression approach.  In International Conference on Machine Learning, pages 254–263. PMLR, 2018.
  * [2] A. R. Asadi, E. Abbe, and S. Verdú.  Chaining mutual information and tightening generalization bounds.  In Advances in Neural Information Processing Systems (NeurIPS), 2018\.
  * [3] S. Boucheron, G. Lugosi, and P. Massart.  Concentration inequalities: A nonasymptotic theory of independence.  Oxford university press, 2013.
  * [4] O. Bousquet, S. Boucheron, and G. Lugosi.  Introduction to statistical learning theory.  In Summer School on Machine Learning, pages 169–207. Springer, 2003\.
  * [5] Y. Bu, S. Zou, and V. V. Veeravalli.  Tightening mutual information-based bounds on generalization error.  IEEE Journal on Selected Areas in Information Theory, 1(1):121–130, 2020.
  * [6] Z. Charles and D. Papailiopoulos.  Stability and generalization of learning algorithms that converge to global optima.  In International Conference on Machine Learning, pages 745–754. PMLR, 2018.
  * [7] T. M. Cover and J. A. Thomas.  Elements of information theory.  Wiley, 2001.
  * [8] C. Dwork, V. Feldman, M. Hardt, T. Pitassi, O. Reingold, and A. Roth.  Generalization in adaptive data analysis and holdout reuse.  In Advances in Neural Information Processing Systems, 2015.
  * [9] G. K. Dziugaite and D. M. Roy.  Computing nonvacuous generalization bounds for deep (stochastic) neural networks with many more parameters than training data.  arXiv preprint arXiv:1703.11008, 2017.
  * [10] V. Feldman and T. Steinke.  Calibrating noise to variance in adaptive data analysis.  In Conference On Learning Theory, pages 535–544. PMLR, 2018.
  * [11] S. B. Gelfand and S. K. Mitter.  Recursive stochastic algorithms for global optimization in r^d.  SIAM Journal on Control and Optimization, 29(5):999–1018, 1991\.
  * [12] N. Gergely, G. K. Dziugaite, M. Haghifam, and D. M. Roy.  Information-theoretic generalization bounds for stochastic gradient descent.  In COLT, 2021.
  * [13] H. Hafez-Kolahi, Z. Golgooni, S. Kasaei, and M. Soleymani.  Conditioning and processing: Techniques to improve information-theoretic generalization bounds.  In Advances in Neural Information Processing Systems (NeurIPS), volume 33, 2020.
  * [14] M. Haghifam, J. Negrea, A. Khisti, D. M. Roy, and G. K. Dziugaite.  Sharpened generalization bounds based on conditional mutual information and an application to noisy, iterative algorithms.  arXiv preprint arXiv:2004.12983, 2020.
  * [15] M. Hardt, B. Recht, and Y. Singer.  Train faster, generalize better: Stability of stochastic gradient descent.  In International Conference on Machine Learning, pages 1225–1234. PMLR, 2016.
  * [16] I. Issa, A. R. Esposito, and M. Gastpar.  Strengthened information-theoretic bounds on the generalization error.  In 2019 IEEE International Symposium on Information Theory (ISIT), pages 582–586. IEEE, 2019.
  * [17] Y. Jiang, B. Neyshabur, H. Mobahi, D. Krishnan, and S. Bengio.  Fantastic generalization measures and where to find them.  In International Conference on Learning Representations, 2019.
  * [18] N. S. Keskar, D. Mudigere, J. Nocedal, M. Smelyanskiy, and P. T. P. Tang.  On large-batch training for deep learning: Generalization gap and sharp minima.  arXiv preprint arXiv:1609.04836, 2016.
  * [19] I. Kuzborskij and C. Lampert.  Data-dependent stability of stochastic gradient descent.  In International Conference on Machine Learning, pages 2815–2824. PMLR, 2018.
  * [20] Y. Lei and Y. Ying.  Sharper generalization bounds for learning with gradient-dominated objective functions.  In International Conference on Learning Representations, 2021.
  * [21] J. Li, X. Luo, and M. Qiao.  On generalization error bounds of noisy gradient methods for non-convex learning.  In International Conference on Learning Representations, 2019.
  * [22] S. Mandt, M. D. Hoffman, and D. M. Blei.  Stochastic gradient descent as approximate Bayesian inference.  The Journal of Machine Learning Research, 18(1):4873–4907, 2017\.
  * [23] Q. Meng, S. Gong, W. Chen, Z.-M. Ma, and T.-Y. Liu.  Dynamic of stochastic gradient descent with state-dependent noise.  arXiv preprint arXiv:2006.13719, 2020.
  * [24] W. Mou, L. Wang, X. Zhai, and K. Zheng.  Generalization bounds of SGLD for non-convex learning: Two theoretical viewpoints.  In Conference on Learning Theory, pages 605–638. PMLR, 2018.
  * [25] J. Negrea, M. Haghifam, G. K. Dziugaite, A. Khisti, and D. M. Roy.  Information-theoretic generalization bounds for SGLD via data-dependent estimates.  arXiv preprint arXiv:1911.02151, 2019.
  * [26] B. Neyshabur, S. Bhojanapalli, and N. Srebro.  A PAC-Bayesian approach to spectrally-normalized margin bounds for neural networks.  In International Conference on Learning Representations, 2018.
  * [27] E. Parrado-Hernández, A. Ambroladze, J. Shawe-Taylor, and S. Sun.  PAC-Bayes bounds with data dependent priors.  The Journal of Machine Learning Research, 13(1):3507–3531, 2012\.
  * [28] A. Pensia, V. Jog, and P.-L. Loh.  Generalization error bounds for noisy, iterative algorithms.  In 2018 IEEE International Symposium on Information Theory (ISIT), pages 546–550. IEEE, 2018.
  * [29] M. Raginsky, A. Rakhlin, and M. Telgarsky.  Non-convex learning via stochastic gradient langevin dynamics: a nonasymptotic analysis.  In Conference on Learning Theory, pages 1674–1703. PMLR, 2017.
  * [30] D. Russo and J. Zou.  Controlling bias in adaptive data analysis using information theory.  In Artificial Intelligence and Statistics, pages 1232–1240. PMLR, 2016.
  * [31] T. Steinke and L. Zakynthinou.  Reasoning about generalization via conditional mutual information.  In Conference on Learning Theory, pages 3437–3452. PMLR, 2020.
  * [32] G. Strang, G. Strang, G. Strang, and G. Strang.  Introduction to linear algebra, volume 3.  Wellesley-Cambridge Press Wellesley, MA, 1993.
  * [33] M. Welling and Y. W. Teh.  Bayesian learning via stochastic gradient langevin dynamics.  In Proceedings of the 28th international conference on machine learning (ICML-11), pages 681–688. Citeseer, 2011.
  * [34] Y. Wen, K. Luk, M. Gazeau, G. Zhang, H. Chan, and J. Ba.  An empirical study of large-batch stochastic gradient descent with structured covariance noise.  arXiv e-prints, pages arXiv–1902, 2019.
  * [35] H. Xiao, K. Rasul, and R. Vollgraf.  Fashion-MNIST: a novel image dataset for benchmarking machine learning algorithms.  arXiv preprint arXiv:1708.07747, 2017.
  * [36] Z. Xie, I. Sato, and M. Sugiyama.  A diffusion theory for deep learning dynamics: Stochastic gradient descent escapes from sharp minima exponentially fast.  arXiv preprint arXiv:2002.03495, 2020.
  * [37] A. Xu and M. Raginsky.  Information-theoretic analysis of generalization capability of learning algorithms.  arXiv preprint arXiv:1705.07809, 2017.
  * [38] C. Zhang, S. Bengio, M. Hardt, B. Recht, and O. Vinyals.  Understanding deep learning requires rethinking generalization.  In International Conference on Learning Representations, 2017.
  * [39] W. Zhou, V. Veitch, M. Austern, R. P. Adams, and P. Orbanz.  Non-vacuous generalization bounds at the imagenet scale: a PAC-Bayesian compression approach.  In International Conference on Learning Representations, 2018.
  * [40] Z. Zhu, J. Wu, B. Yu, L. Wu, and J. Ma.  The anisotropic noise in stochastic gradient descent: Its behavior of escaping from sharp minima and regularization effects.  arXiv preprint arXiv:1803.00195, 2018.

Supplementary materials for
“Optimizing Information-theoretical Generalization Bound via Anisotropic Noise
in SGLD”

The supplementary materials are organized as follows. In Appendix A, we
provide some basic lemmas which are used throughout the proofs in the rest of
the materials. In Appendix B, we provide the proof of Lemma 2. In Appendix C,
D, and E, we provide the detailed proofs of Lemmas and Theorems respectively
in Section 3.2, Section 4 , and Section 5. In Appendix F, we provide the
detailed settings of the experiments in the main text together with an
additional experiments to justify the result of Theorem 2.

##  Appendix A Preliminaries

In this section, we provide some basic lemmas that will be used throughout the
proof both from probability theory and from matrix analysis.

###  A.1 Preparations in Probability Theory

The first lemma is a standard result characterizing the KL divergence between
two Gaussian distributions.

######  Lemma 5 (KL divergence between Gaussian distributions).

Let $\mathrm{P}_{1}$ and $\mathrm{P}_{2}$ are multivariate Gaussian
distributions on $\mathbb{R}^{d}$ with mean and covariance respectively
$\mu_{1}$, $\bm{\Sigma}_{1}$ and $\mu_{2}$, $\bm{\Sigma}_{2}$. Then the KL
divergence between $\mathrm{P}_{1}$ and $\mathrm{P}_{2}$ are given as follows:

$\mathrm{KL}(\mathrm{P}_{1}||\mathrm{P}_{2})=\frac{1}{2}\left(\operatorname{tr}\left(\bm{\Sigma}_{2}^{-1}\bm{\Sigma}_{1}\right)+\left(\mu_{2}-\mu_{1}\right)^{\top}\bm{\Sigma}_{2}^{-1}\left(\mu_{2}-\mu_{1}\right)-d+\ln\left(\frac{\operatorname{det}\bm{\Sigma}_{2}}{\operatorname{det}\bm{\Sigma}_{1}}\right)\right).$

We then provide a lemma which gives the expected difference between two
uniform sampling variables.

######  Lemma 6 (Two step sampling).

Suppose $\bm{z}$ is a discrete random variable with
$\mathbb{P}(\bm{z}=\bm{z}_{i})=\frac{1}{N}$, $\forall i=1,2,\cdots,N$, where
the support set is
$\mathcal{Z}=\\{\bm{z}_{1},\cdots,\bm{z}_{N}\\}\subset\mathbb{R}^{d}$. Suppose
further $\bm{U}$ is a random index set with size $b$ and sampled uniformly
without replacement from $[N]$. Suppose $\bm{V}$ is another random index set
independent of $\bm{U}$ with size $N-1$ and sampled uniformly without
replacement from $[N]$. Denote subset of $\mathcal{Z}$ with index in
$\bm{U}\cap\bm{V}^{c}$ and $\bm{V}$ respectively as
$\mathcal{Z}_{\bm{U}\cap\bm{V}^{c}}=\\{\bm{z}_{i},i\in\bm{U}\cap\bm{V}^{c}\\}$,
$\mathcal{Z}_{\bm{V}}=\\{\bm{z}_{i},i\in\bm{V}\\}$, and the average of
$\mathcal{Z}_{\bm{U}\cap\bm{V}^{c}}$ and $\mathcal{Z}_{\bm{V}}$ respectively
as $\bar{\mathcal{Z}}_{\bm{U}\cap\bm{V}^{c}}$ and
$\bar{\mathcal{Z}}_{\bm{V}}$. Then the following equation holds:

$\mathbb{E}_{\bm{U},\bm{V}}\left(\frac{(b-|\bm{U}\cap\bm{V}|)^{2}}{b^{2}}(\bar{\mathcal{Z}}_{\bm{V}}-\bar{\mathcal{Z}}_{\bm{U}\cap\bm{V}^{c}})(\bar{\mathcal{Z}}_{\bm{V}}-\bar{\mathcal{Z}}_{\bm{U}\cap\bm{V}^{c}})^{\top}\right)=\frac{1}{Nb}\left(\frac{N}{N-1}\right)^{2}\mathrm{Cov}(\bm{z}).$

###### Proof.

We rewrite
$\mathbb{E}_{\bm{U},\bm{V}}\left(\frac{(b-|\bm{U}\cap\bm{V}|)^{2}}{b^{2}}(\bar{\mathcal{Z}}_{\bm{V}}-\bar{\mathcal{Z}}_{\bm{U}\cap\bm{V}^{c}})(\bar{\mathcal{Z}}_{\bm{V}}-\bar{\mathcal{Z}}_{\bm{U}\cap\bm{V}^{c}})^{\top}\right)$
by taking conditional expectation with respect to $|\bm{U}\cap\bm{V}|$ as
follows:

$\displaystyle\mathbb{E}_{\bm{U},\bm{V}}\left(\frac{(b-|\bm{U}\cap\bm{V}|)^{2}}{b^{2}}(\bar{\mathcal{Z}}_{\bm{V}}-\bar{\mathcal{Z}}_{\bm{U}\cap\bm{V}^{c}})(\bar{\mathcal{Z}}_{\bm{V}}-\bar{\mathcal{Z}}_{\bm{U}\cap\bm{V}^{c}})^{\top}\right)$
$\displaystyle=\mathbb{E}_{|\bm{U}\cap\bm{V}^{c}|}\mathbb{E}^{|\bm{U}\cap\bm{V}^{c}|}_{\bm{U},\bm{V}}\left(\frac{|\bm{U}\cap\bm{V}^{c}|^{2}}{b^{2}}(\bar{\mathcal{Z}}_{\bm{V}}-\bar{\mathcal{Z}}_{\bm{U}\cap\bm{V}^{c}})(\bar{\mathcal{Z}}_{\bm{V}}-\bar{\mathcal{Z}}_{\bm{U}\cap\bm{V}^{c}})^{\top}\right)$
$\displaystyle=\mathbb{P}(|\bm{U}\cap\bm{V}^{c}|=1)\mathbb{E}^{|\bm{U}\cap\bm{V}^{c}|=1}_{\bm{U},\bm{V}}\left(\frac{1}{b^{2}}(\bar{\mathcal{Z}}_{\bm{V}}-\bar{\mathcal{Z}}_{\bm{V}^{c}})(\bar{\mathcal{Z}}_{\bm{V}}-\bar{\mathcal{Z}}_{\bm{V}^{c}})^{\top}\right)$
$\displaystyle=\mathbb{P}(|\bm{U}\cap\bm{V}^{c}|=1)\mathbb{E}_{\bm{V}}\left(\frac{1}{b^{2}}(\bar{\mathcal{Z}}_{\bm{V}}-\bar{\mathcal{Z}}_{\bm{V}^{c}})(\bar{\mathcal{Z}}_{\bm{V}}-\bar{\mathcal{Z}}_{\bm{V}^{c}})^{\top}\right)$
$\displaystyle=\frac{1}{Nb}\left(\frac{N}{N-1}\right)^{2}\mathrm{Cov}(\bm{z}).$

The proof is completed. ∎

We provide the following lemma for computing the KL divergence between two
joint distributions.

######  Lemma 7.

Let $\bm{X}$, $\bm{Y}$, and $\bm{Z}$ be three random variables with $\bm{X}$
and $\bm{Y}$ having the same support set. Then the KL divergence between the
joint distribution of $(\bm{X},\bm{Z})$ and $(\bm{Y},\bm{Z})$ can be
decomposed into

$\displaystyle\operatorname{KL}\left((\bm{X},\bm{Z})\|(\bm{Y},\bm{Z})\right)=\mathbb{E}_{\bm{Z}}\operatorname{KL}\left((\bm{X}|\bm{Z})\|(\bm{Y}|\bm{Z})\right).$

###### Proof.

By the definition of KL divergence,

$\displaystyle\operatorname{KL}\left((\bm{X},\bm{Z})\|(\bm{Y},\bm{Z})\right)=$
$\displaystyle\int\mathbb{P}(\bm{X},\bm{Z})\log\frac{\mathbb{P}(\bm{X},\bm{Z})}{\mathbb{P}(\bm{Y},\bm{Z})}$
$\displaystyle=$
$\displaystyle\int\mathbb{P}(\bm{Z})\mathbb{P}^{\bm{Z}}(\bm{X})\log\frac{\mathbb{P}^{\bm{Z}}(\bm{X})\mathbb{P}(\bm{Z})}{\mathbb{P}^{\bm{Z}}(\bm{Y})\mathbb{P}(\bm{Z})}$
$\displaystyle=$
$\displaystyle\int\mathbb{P}(\bm{Z})\int\mathbb{P}^{\bm{Z}}(\bm{X})\log\frac{\mathbb{P}^{\bm{Z}}(\bm{X})}{\mathbb{P}^{\bm{Z}}(\bm{Y})}$
$\displaystyle=$
$\displaystyle\mathbb{E}_{\bm{Z}}\operatorname{KL}\left((\bm{X}|\bm{Z})\|(\bm{Y}|\bm{Z})\right).$

The proof is completed. ∎

In the end of this section, we provide a proof of Lemma 1 using Lemma 7 for
the completeness of this paper.

###### Proof of Lemma 1.

By Lemma 7, we have

$\displaystyle\operatorname{KL}(Q_{0:T}\|P_{0:T})$
$\displaystyle=\operatorname{KL}(Q_{0:T}\|P_{0:T})-\operatorname{KL}(Q_{0:T}\|(Q_{0:T-1},P_{T|[T-1]}))+\operatorname{KL}(Q_{0:T}\|(Q_{0:T-1},P_{T|[T-1]}))$
$\displaystyle=\int Q_{0:T}\log\frac{Q_{0:T}}{P_{0:T}}-\int
Q_{0:T}\log\frac{Q_{0:T}}{Q_{0:T-1}P_{T|[T-1]}}+\mathbb{E}_{Q_{T-1}}\operatorname{KL}\left(Q_{T|[T-1]}\|P_{T|[T-1]}\right)$
$\displaystyle=\int
Q_{0:T}\log\frac{Q_{0:T-1}}{P_{0:T-1}}+\mathbb{E}_{Q_{T-1}}\operatorname{KL}\left(Q_{T|[T-1]}\|P_{T|[T-1]}\right)$
$\displaystyle=\int
Q_{0:T-1}\log\frac{Q_{0:T-1}}{P_{0:T-1}}+\mathbb{E}_{Q_{T-1}}\operatorname{KL}\left(Q_{T|[T-1]}\|P_{T|[T-1]}\right)$
$\displaystyle=\operatorname{KL}\left(Q_{0:T-1}\left\|P_{0:T-1}\right.\right)+\mathbb{E}_{Q_{T-1}}\operatorname{KL}\left(Q_{T|[T-1]}\|P_{T|[T-1]}\right).$

The proof is then completed by induction. ∎

######  Remark 1.

In this paper, we focus on the case where $Q_{0:T}$ and $P_{0:T}$ obeys the
_Markov Property_ , i.e., for any $t$,

$Q_{t|[t-1]}=Q_{t|(t-1)},\text{ }P_{t|[t-1]}=P_{t|(t-1)}.$

Therefore, the result in Lemma 1 becomes

$\operatorname{KL}(Q_{0:T}||P_{0:T})=\sum_{t=1}^{T}\mathbb{E}_{Q_{0:t-1}}\left[\operatorname{KL}\left(Q_{t|(t-1)}\|P_{t|(t-1)}\right)\right].$

###  A.2 Technical Lemmas in Matrix Analysis

We first provide a sufficient and necessary condition of that two symmetric
matrices commute, and the proof can be found from any Linear Algebra Textbook
(e.g. [32]).

######  Lemma 8.

Let $\bm{A}$ and $\bm{B}$ be two $d\times d$ real symmetric matrices. Then,
$\bm{A}$ and $\bm{B}$ commute (i.e., $\bm{A}\bm{B}=\bm{B}\bm{A}$), if and only
if there exists an orthogonal matrix $\bm{O}$ which can diagonalize $\bm{A}$
and $\bm{B}$ simultaneously, i.e., both $\bm{O}^{\top}\bm{A}\bm{O}$ and
$\bm{O}^{\top}\bm{B}\bm{O}$ are diagonal.

The next lemma is a key technique to obtain the optimal noise covariance of
Theorem 2 and Theorem 3.

######  Lemma 9.

Let $\bm{B}\in\mathbb{R}^{d\times d}$ be a (fixed) positive definite matrix
with eigenvalues $(\beta_{1},\cdots,\beta_{d})$, where $\beta_{i}\geq 0$. Let
$\bm{G}\in\mathbb{R}^{d\times d}$ be a positive definite matrix variable with
fixed trace $\operatorname{tr}{\bm{G}}=c$, where $c$ is a positive constant
and $c\leq\operatorname{tr}(\bm{B})$. Then the minimum of
$\operatorname{tr}(\bm{G}^{-1}\bm{B})+\ln\operatorname{det}(\bm{G})$ is
achieved at
$\bm{G}=\bm{O}^{\top}\operatorname{Diag}(\alpha_{1},\cdots,\alpha_{d})\bm{O}$,
where

$\alpha_{i}^{*}=\frac{\sqrt{1-4\lambda^{*}\beta_{i}}-1}{-2\lambda^{*}},$

$\bm{O}$ is any orthogonal matrix which diagonalize $\bm{B}$ as

$\bm{B}=\bm{O}^{\top}\operatorname{Diag}(\beta_{1},\cdots,\beta_{d})\bm{O},$

and $\lambda^{*}\leq 0$ is the unique solution of

$\sum_{i=1}^{d}\frac{2\beta_{i}}{1+\sqrt{1-4\lambda^{*}\beta_{i}}}=c.$ (11)

######  Remark 2.

$f(\lambda)=\sum_{i=1}^{d}\frac{2\beta_{i}}{1+\sqrt{1-4\lambda\beta_{i}}}$ is
a monotonously increasing function with respect to $\lambda$, which guarantees
the uniqueness of the solution of $f(\lambda)=c$.

Lemma 9 is proved via two steps: 1) we first prove for $\bm{G}$ with fixed
eigenvalues,
$\operatorname{tr}{\bm{G}^{-1}\bm{B}}+\ln\operatorname{det}(\bm{G})$ is
optimized if and only if $\bm{G}$ and $\bm{B}$ share the same eigenvectors; 2)
we then calculate the eigenvalues of the optimal $\bm{G}$ using the method of
Lagrange multipliers. Theorem 3 can then be obtained by applying Lemma 9 and
setting $\bm{G}=\bm{\Sigma}_{t}(\bm{S},\bm{W})$ and
$\bm{B}=\sigma_{t}\mathbb{I}+\frac{\eta_{t}^{2}}{Nb_{t}}\left(\frac{N}{N-1}\right)^{2}\bm{\Sigma}^{sd}_{\bm{S},\bm{W}}$.
We first prove the eigenvectors of $\bm{G}$ agree with those of $\bm{B}$.

######  Lemma 10.

Let $\bm{G}\in\mathbb{R}^{d\times d}$ be a positive definite matrix variable
with fixed eigenvalues $(\alpha_{i})_{i=1}^{d}$. Specifically, let
$\alpha_{1}\geq\alpha_{2}\geq\cdots\geq\alpha_{d}>0$ be all the eigenvalues of
$\bm{G}$, and $\bm{G}$ can be any element from the following set

$\\{\bm{Q}^{\top}\operatorname{Diag}(\alpha_{1},\cdots,\alpha_{d})\bm{Q}:\bm{Q}\text{
is orthogonal}\\}.$

Let $\bm{B}$ be a fixed positive semi-definite matrix, with eigenvalues
$(\beta_{i})_{i=1}^{d}$ satisfies
$\beta_{1}\geq\beta_{2}\cdots\geq\beta_{d}\geq 0$. Then, the optimal (minimal)
value of $g(\bm{G})=\operatorname{tr}(\bm{G}^{-1}\bm{B})$ is achieved when

$\bm{G}^{*}=\bm{O}^{\top}\operatorname{Diag}(\alpha_{1},\cdots,\alpha_{d})\bm{O},$

where $\bm{O}$ is any orthogonal matrix which diagonalizes $\bm{B}$ as

$\bm{B}=\bm{O}^{\top}\operatorname{Diag}(\beta_{1},\cdots,\beta_{d})\bm{O}$

and the optimal value of $g(\bm{G})$ is
$\sum_{i=1}^{d}\frac{\beta_{i}}{\alpha_{i}}$.

###### Proof.

Let $\bm{G}^{*}$ be a optimal point of $\operatorname{tr}(\bm{G}^{-1}\bm{B})$.
We will then obtain the condition of $\bm{G}^{*}$ by adding a disturbance.
Specifically, let $\bm{A}$ be an anti-symmetric matrix. Then,

$(\mathbb{I}-\varepsilon\bm{A})(\mathbb{I}-\varepsilon\bm{A})^{T}=\mathbb{I}+\varepsilon^{2}\bm{A}\bm{A}^{\top}.$

As $\varepsilon$ is small enough,
$\mathbb{I}+\varepsilon^{2}\bm{A}\bm{A}^{\top}$ is inevitable, and positive
definite. Therefore,
$(\mathbb{I}-\varepsilon\bm{A})(\mathbb{I}+\varepsilon^{2}\bm{A}\bm{A}^{\top})^{-\frac{1}{2}}$
is orthogonal. As $\varepsilon\rightarrow 0$,

$\lim_{\varepsilon\rightarrow
0}(\mathbb{I}-\varepsilon\bm{A})(\mathbb{I}+\varepsilon^{2}\bm{A}\bm{A}^{\top})^{-\frac{1}{2}}=\mathbb{I},$

and

$(\mathbb{I}-\varepsilon\bm{A})(\mathbb{I}+\varepsilon^{2}\bm{A}\bm{A}^{\top})^{-\frac{1}{2}}-\mathbb{I}=-\varepsilon\bm{A}+\bm{o}(\varepsilon).$

Since $\bm{G}^{*}$ is an optimal point of
$\operatorname{tr}(\bm{G}^{-1}\bm{B})$, we have

$\displaystyle\operatorname{tr}((\bm{G}^{*})^{-1}\bm{B})\leq$
$\displaystyle\operatorname{tr}\left((\mathbb{I}-\varepsilon\bm{A})(\mathbb{I}+\varepsilon^{2}\bm{A}\bm{A}^{\top})^{-\frac{1}{2}}(\bm{G}^{*})^{-1}\left((\mathbb{I}-\varepsilon\bm{A})(\mathbb{I}+\varepsilon^{2}\bm{A}\bm{A}^{\top})^{-\frac{1}{2}}\right)^{\top}\bm{B}\right)$
$\displaystyle=$
$\displaystyle\operatorname{tr}\left((\mathbb{I}-\varepsilon\bm{A})(\mathbb{I}+\varepsilon^{2}\bm{A}\bm{A}^{\top})^{-\frac{1}{2}}(\bm{G}^{*})^{-1}(\mathbb{I}+\varepsilon^{2}\bm{A}\bm{A}^{\top})^{-\frac{1}{2}}(\mathbb{I}+\varepsilon\bm{A})\bm{B}\right),$

which further leads to

$\displaystyle-\varepsilon\operatorname{tr}\left(\bm{A}(\bm{G}^{*})^{-1}\bm{B}\right)+\varepsilon\operatorname{tr}\left((\bm{G}^{*})^{-1}\bm{A}\bm{B}\right)+\bm{o}(\varepsilon)\geq
0.$

By letting $\varepsilon\rightarrow 0$, we further have

$\displaystyle-\operatorname{tr}\left(\bm{A}(\bm{G}^{*})^{-1}\bm{B}\right)+\operatorname{tr}\left((\bm{G}^{*})^{-1}\bm{A}\bm{B}\right)=0,$

which further leads to

$\displaystyle 0=$
$\displaystyle-\operatorname{tr}\left(\bm{A}(\bm{G}^{*})^{-1}\bm{B}\right)+\operatorname{tr}\left((\bm{G}^{*})^{-1}\bm{A}\bm{B}\right)$
$\displaystyle=$
$\displaystyle\operatorname{tr}\left(\bm{A}^{\top}(\bm{G}^{*})^{-1}\bm{B}\right)+\operatorname{tr}\left(\bm{B}(\bm{G}^{*})^{-1}\bm{A}\right)$
$\displaystyle=$ $\displaystyle
2\operatorname{tr}\left(\bm{B}(\bm{G}^{*})^{-1}\bm{A}\right).$ (12)

Since Eq.(12) holds for any anti-symmetry matrix $\bm{A}$, let
$\bm{A}=\bm{E}_{i,j}-\bm{E}_{j,i}$, where $i,j\in[d]$ and $i\neq j$. By
Eq.(12), we have

$\displaystyle\left(\bm{B}(\bm{G}^{*})^{-1}\right)_{i,j}=\left(\bm{B}(\bm{G}^{*})^{-1}\right)_{j,i},$

which further leads to

$\displaystyle\bm{B}(\bm{G}^{*})^{-1}=\left(\bm{B}(\bm{G}^{*})^{-1}\right)^{\top}=(\bm{G}^{*})^{-\top}\bm{B}^{\top}=(\bm{G}^{*})^{-1}\bm{B}.$

By simple rearranging, we have

$\displaystyle\bm{G}^{*}\bm{B}=\bm{B}\bm{G}^{*}.$

Therefore, by Lemma 8, we have that there exists an orthogonal matrix
$\bm{O}_{0}$, such that both $\bm{O}_{0}\bm{G}^{*}\bm{O}_{0}^{\top}$ and
$\bm{O}_{0}\bm{B}\bm{O}_{0}^{\top}$ are diagonal. By multiplying a permutation
matrix, we further have there exists an orthogonal matrix $\tilde{\bm{O}}$
such that $\tilde{\bm{O}}\bm{G}^{*}\tilde{\bm{O}}^{\top}$ is diagonal, and

$\tilde{\bm{O}}\bm{B}\tilde{\bm{O}}^{\top}=\operatorname{Diag}(\beta_{1},\cdots,\beta_{d}).$
(13)

Since $\tilde{\bm{O}}\bm{G}^{*}\tilde{\bm{O}}^{\top}$ is diagonal, there
exists a permutation mapping $\mathcal{T}:[d]\rightarrow[d]$, such that

$\tilde{\bm{O}}\bm{G}^{*}\tilde{\bm{O}}^{\top}=\operatorname{Diag}\left(\alpha_{\mathcal{T}(1)},\cdots,\alpha_{\mathcal{T}(d)}\right).$
(14)

Denote the order of $\beta_{i}$ $(i=1,2,\cdots,d)$ as

$\displaystyle\beta_{1}=\cdots=\beta_{s_{1}}>\beta_{s_{1}+1}=\cdots=\beta_{s_{1}+s_{2}}>\cdots>\beta_{\sum_{i=1}^{k-1}s_{i}+1}=\cdots=\beta_{\sum_{i=1}^{k}s_{i}}>0,$
(15)

where $\sum_{i=1}^{k}s_{i}=d$, and we denote $s_{0}=0$. Since $\bm{G}^{*}$ is
the optimal point of $\operatorname{tr}((\bm{G}^{*})^{-1}\bm{B})$, for any
$1\leq i<j\leq d$ and $\beta_{i}>\beta_{j}$, we have
$\alpha_{\mathcal{T}(i)}>\alpha_{\mathcal{T}(j)}$: otherwise, let

$\bm{G}^{\prime}=\tilde{\bm{O}}^{\top}\operatorname{Diag}\left(\alpha_{\mathcal{T}(1)},\cdots,\alpha_{\mathcal{T}(i-1)},\alpha_{\mathcal{T}(j)},\alpha_{\mathcal{T}(i+1)},\cdots,\alpha_{\mathcal{T}(j-1)},\alpha_{\mathcal{T}(i)},\alpha_{\mathcal{T}(j+1)},\cdots,\alpha_{\mathcal{T}(d)}\right)\tilde{\bm{O}},$

we have

$\operatorname{tr}((\bm{G}^{*})^{-1}\bm{B})>\operatorname{tr}((\bm{G}^{\prime})^{-1}\bm{B}),$

which contradicts that $\bm{G}^{*}$ is optimal.

Therefore,
$\mathcal{T}(\sum_{i=1}^{j}s_{i}+1),\cdots,\mathcal{T}(\sum_{i=1}^{j+1}s_{i})$
is then a permutation of $\sum_{i=1}^{j}s_{i}+1,\cdots,\sum_{i=1}^{j+1}s_{i}$,
and there exists permutation matrix $\bm{Q}$ such that

$\bm{Q}=\operatorname{Diag}\left(\bm{Q}_{1},\cdots,\bm{Q}_{k}\right),$ (16)

where $\bm{Q}_{i}$ is a $s_{i}\times s_{i}$ permutation sub-matrix, such that,

$\bm{Q}\operatorname{Diag}\left(\alpha_{\mathcal{T}(1)},\cdots,\alpha_{\mathcal{T}(d)}\right)\bm{Q}^{\top}=\operatorname{Diag}\left(\alpha_{1},\cdots,\alpha_{d}\right).$
(17)

Furthermore, by Eq.(16) and Eq.(15), we have

$\bm{Q}\operatorname{Diag}\left(\beta_{1},\cdots,\beta_{d}\right)\bm{Q}^{\top}=\operatorname{Diag}\left(\beta_{1},\cdots,\beta_{d}\right).$
(18)

Therefore, by Eqs.(13), (14), (17), and (18), we have

$\displaystyle\bm{Q}\tilde{\bm{O}}\bm{B}\left(\bm{Q}\tilde{\bm{O}}\right)^{\top}=\operatorname{Diag}\left(\beta_{1},\cdots,\beta_{d}\right),$
$\displaystyle\bm{Q}\tilde{\bm{O}}\bm{G}^{*}\left(\bm{Q}\tilde{\bm{O}}\right)^{\top}=\operatorname{Diag}\left(\alpha_{1},\cdots,\alpha_{d}\right).$

Furthermore,

$\displaystyle\operatorname{tr}\left(\bm{G}^{-1}\bm{B}\right)$
$\displaystyle=\operatorname{tr}\left(\left(\bm{Q}\tilde{\bm{O}}\right)^{\top}\operatorname{Diag}\left(\alpha^{-1}_{1},\cdots,\alpha^{-1}_{d}\right)\left(\bm{Q}\tilde{\bm{O}}\right)\left(\bm{Q}\tilde{\bm{O}}\right)^{\top}\operatorname{Diag}\left(\beta_{1},\cdots,\beta_{d}\right)\left(\bm{Q}\tilde{\bm{O}}\right)\right)$
$\displaystyle=\sum_{i=1}^{d}\frac{\beta_{i}}{\alpha_{i}}.$

Therefore, the optimal value of $\operatorname{tr}(\bm{G}^{-1}\bm{B})$ is
$\sum_{i=1}^{d}\frac{\beta_{i}}{\alpha_{i}}$, and the corresponding optimal
point $\bm{G}^{*}$ belongs to the following set

$\mathcal{G}=\\{\bm{O}^{\top}\operatorname{Diag}(\alpha_{1},\cdots,\alpha_{d})\bm{O}:\bm{B}=\bm{O}^{\top}\operatorname{Diag}(\beta_{1},\cdots,\beta_{d})\bm{O}\\}.$

On the other hand, it is easy to verify that for any element
$\bm{G}\in\mathcal{G}$,

$\operatorname{tr}(\bm{G}^{-1}\bm{B})=\sum_{i=1}^{d}\frac{\beta_{i}}{\alpha_{i}}.$

The proof is completed. ∎

Lemma 10 indicates that with eigenvalues fixed, the eigenvectors of $\bm{G}$
should agree with those of $\bm{B}$ by the order of eigenvalues. We then
provide the following lemma to determine the optimal eigenvalues.

######  Lemma 11.

Let $\beta_{1},\beta_{2},\cdots,\beta_{d}$ be a series of fixed positive
reals. Let $\alpha_{1},\alpha_{2},\cdots,\alpha_{d}\in\mathbb{R}^{+}$ be a
series of real variables with constraint $\sum_{i=1}^{d}\alpha_{i}=c$, where
$c$ is a positive real constant which satisfies
$c\leq\sum_{i=1}^{d}\beta_{i}$. Then the minimum of function

$f(\alpha_{1},\cdots,\alpha_{d})=\sum_{i=1}^{d}\frac{\beta_{i}}{\alpha_{i}}+\sum_{i=1}^{d}\ln\alpha_{i}$

is achieved at

$\alpha_{i}^{*}=\frac{\sqrt{1-4\lambda^{*}\beta_{i}}-1}{-2\lambda^{*}},$

where $\lambda^{*}\leq 0$ is the unique solution of

$\sum_{i=1}^{d}\frac{2\beta_{i}}{1+\sqrt{1-4\lambda^{*}\beta_{i}}}=c.$

###### Proof.

We find the minimum of $f$ under the constraint that
$\alpha_{1}+\cdots+\alpha_{d}=c$ by the method of Lagrange Multiplier.
Specifically, as for any $i\in[d]$, $\alpha_{i}\rightarrow 0^{+}$ or
$\alpha_{i}\rightarrow c^{-}$ will lead to
$f(\alpha_{1},\cdots,\alpha_{d})\rightarrow\infty$, we have that for any
global optimal (minimal) point $(\alpha_{1}^{*},\cdots,\alpha_{d}^{*})$ of $f$
under the constraint $\alpha_{1}+\cdots+\alpha_{d}=c$, we have that there
exist a real $\lambda^{*}$, such that
$((\alpha_{1}^{*},\cdots,\alpha_{d}^{*}),\lambda^{*})$ is a saddle point of
$\mathcal{L}((\alpha_{1},\cdots,\alpha_{d}),\lambda)$, which is defined as

$\mathcal{L}((\alpha_{1},\cdots,\alpha_{d}),\lambda)=f(\alpha_{1},\cdots,\alpha_{d})+\lambda(c-\alpha_{1}-\cdots-\alpha_{d}).$

By taking partial derivative of $\mathcal{L}$ with respect to $\alpha_{i}$, we
have

$-\lambda^{*}=-\frac{1}{\alpha_{i}}+\frac{\beta_{i}}{\alpha_{i}^{2}}=\frac{\beta_{i}-\alpha_{i}}{\alpha_{i}^{2}},$
(19)

which further leads to

$\sum_{i=1}^{d}\beta_{i}-c=\sum_{i=1}^{d}\left(\beta_{i}-\alpha_{i}\right)=-\lambda^{*}\left(\sum_{i=1}^{d}\alpha_{i}^{2}\right).$

Since $\sum_{i=1}^{d}\beta_{i}\geq c$, we have $\lambda^{*}\leq 0$. Therefore,
for any $i\in[d]$, the quadratic equation $\beta_{i}x^{2}-x+\lambda^{*}=0$ has
only one positive solution
$\frac{1+\sqrt{1-4\lambda^{*}\beta_{i}}}{2\beta_{i}}$, and

$\alpha_{i}^{*}=\frac{2\beta_{i}}{1+\sqrt{1-4\lambda^{*}\beta_{i}}}=\frac{\sqrt{1-4\lambda^{*}\beta_{i}}-1}{-2\lambda^{*}}.$

On the other hand, by taking derivative of $\mathcal{L}$ with respect to
$\lambda^{*}$, we have

$\sum_{i=1}^{d}\alpha_{i}^{*}=\sum_{i=1}^{d}\frac{2\beta_{i}}{1+\sqrt{1-4\lambda^{*}\beta_{i}}}=c.$
(20)

Since
$\sum_{i=1}^{d}\alpha_{i}^{*}=\sum_{i=1}^{d}\frac{2\beta_{i}}{1+\sqrt{1-4\lambda^{*}\beta_{i}}}$
is a monotonously increasing function of $\lambda^{*}$, there is only one
solution of $\lambda^{*}$ of Eq.(20).

The proof is completed. ∎

The proof of Lemma 9 can then be obtained by combining Lemma 10 and Lemma 11
together.

###### Proof of Lemma 9.

The original optimization problem can be written as

$\min_{\operatorname{tr}(\bm{G})=c}\operatorname{tr}\left(\bm{G}^{-1}\bm{B}\right)+\ln\left(\det\bm{G}\right),$

which can be further decomposed into

$\displaystyle\min_{\operatorname{tr}(\bm{G})=c}\operatorname{tr}\left(\bm{G}^{-1}\bm{B}\right)+\ln\left(\det\bm{G}\right)$
$\displaystyle=\min_{\begin{subarray}{c}\sum_{i=1}^{d}\alpha_{i}=c\\\
\alpha_{1}\geq\cdots\geq\alpha_{d}>0\end{subarray}}\min_{\bm{O}\in\mathcal{O}(d)}\left(\operatorname{tr}\left(\bm{O}^{\top}\operatorname{Diag}\left(\alpha_{1}^{-1},\cdots,\alpha_{d}^{-1}\right)\bm{O}\bm{B}\right)+\sum_{i=1}^{d}\ln\alpha_{i}\right)$
$\displaystyle\overset{(*)}{=}\min_{\begin{subarray}{c}\sum_{i=1}^{d}\alpha_{i}=c\end{subarray}}\left(\sum_{i=1}^{d}\frac{\beta_{i}}{\alpha_{i}}+\sum_{i=1}^{d}\ln\alpha_{i}\right)$
$\displaystyle\overset{(**)}{=}\sum_{i=1}^{d}\frac{1+\sqrt{1-4\lambda^{*}\beta_{i}}}{2}+\sum_{i=1}^{d}\ln\frac{2\beta_{i}}{1+\sqrt{1-4\lambda^{*}\beta_{i}}},$

where Eq. $(*)$ is due to Lemma 10, Eq. $(**)$ is due to Lemma 11, and
$\lambda^{*}\leq 0$ is the unique solution of

$\sum_{i=1}^{d}\frac{2\beta_{i}}{1+\sqrt{1-4\lambda^{*}\beta_{i}}}=c.$

Furthermore, the optimal point of
$\operatorname{tr}(\bm{G}^{-1}\bm{B})+\ln(\operatorname{det}\bm{G})$ can be
calculated as

$\displaystyle\arg\min_{\operatorname{tr}(\bm{G})=c}\operatorname{tr}\left(\bm{G}^{-1}\bm{B}\right)+\ln\left(\det\bm{G}\right)$
$\displaystyle=\left\\{\bm{O}^{\top}\operatorname{Diag}\left(\frac{\sqrt{1-4\lambda^{*}\beta_{1}}-1}{-2\lambda^{*}},\cdots,\frac{\sqrt{1-4\lambda^{*}\beta_{d}}-1}{-2\lambda^{*}}\right)\bm{O}:\bm{B}=\bm{O}^{\top}\operatorname{Diag}(\beta_{1},\cdots,\beta_{d})\bm{O},\right.$
$\displaystyle\;\;\left.\lambda^{*}=\arg_{\lambda}\left(\sum_{i=1}^{d}\frac{2\beta_{i}}{1+\sqrt{1-4\lambda\beta_{i}}}=c\right)\right\\}.$

The proof is completed. ∎

##  Appendix B Supplementary Materials of Section 3.1

###### Proof of Lemma 2.

The $\beta$-smooth condition gives

$\mathcal{R}_{\bm{S}}(\bm{W}_{t+1})\leq\mathcal{R}_{\bm{S}}(\bm{W}_{t})+\langle\nabla\mathcal{R}_{\bm{S}}(\bm{W}_{t}),\bm{W}_{t+1}-\bm{W}_{t}\rangle+\frac{\beta}{2}\|\bm{W}_{t+1}-\bm{W}_{t}\|^{2}.$
(21)

Based on the update rule Eq.(3), we have

$\displaystyle\bm{W}_{t+1}-\bm{W}_{t}=-\eta_{t+1}\nabla\mathcal{R}_{\bm{S}_{\bm{V}_{t+1}}}(\bm{W}_{t})+\varepsilon_{t+1},$
(22)

where
$\varepsilon_{t+1}\sim\mathcal{N}(0,\bm{\Sigma}_{t+1}(\bm{S},\bm{W}_{t}))$.

Take expectation on Eq.(21) with respect to $\bm{W}_{t+1}|\bm{W}_{t}$, by
$\mathbb{E}^{\bm{W}_{t}}(\nabla\mathcal{R}_{\bm{S}_{\bm{V}_{t+1}}}(\bm{W}_{t}))=\nabla\mathcal{R}_{\bm{S}}(\bm{W}_{t})$,

$\displaystyle\mathbb{E}^{\bm{W}_{t}}[\mathcal{R}_{\bm{S}}(\bm{W}_{t+1})]\leq\mathcal{R}_{\bm{S}}(\bm{W}_{t})-\eta_{t+1}\|\nabla\mathcal{R}_{\bm{S}}(\bm{W}_{t})\|^{2}+\frac{\beta}{2}\mathbb{E}^{\bm{W}_{t}}\|\bm{W}_{t+1}-\bm{W}_{t}\|^{2}.$
(23)

Furthermore,

$\displaystyle\mathbb{E}^{\bm{W}_{t}}\|\bm{W}_{t+1}-\bm{W}_{t}\|^{2}$
$\displaystyle=\mathbb{E}^{\bm{W}_{t}}\|-\eta_{t+1}\nabla\mathcal{R}_{\bm{S}_{\bm{V}_{t+1}}}(\bm{W}_{t})+\varepsilon_{t+1}\|^{2}$
$\displaystyle\overset{(*)}{=}\mathbb{E}^{\bm{W}_{t}}\|-\eta_{t+1}\nabla\mathcal{R}_{\bm{S}_{\bm{V}_{t}}}(\bm{W}_{t})\|+\mathbb{E}^{\bm{W}_{t}}\|\varepsilon_{t+1}\|^{2}$
$\displaystyle=\eta_{t+1}^{2}\mathbb{E}^{\bm{W}_{t}}\|\nabla\mathcal{R}_{\bm{S}_{\bm{V}_{t+1}}}(\bm{W}_{t})-\nabla\mathcal{R}_{\bm{S}}(\bm{W}_{t})+\nabla\mathcal{R}_{\bm{S}}(\bm{W}_{t})\|^{2}+\bm{\Sigma}_{t+1}(\bm{S},\bm{W}_{t})$
$\displaystyle=\eta_{t+1}^{2}\mathbb{E}^{\bm{W}_{t}}\|\nabla\mathcal{R}_{\bm{S}_{\bm{V}_{t+1}}}(\bm{W}_{t})-\nabla\mathcal{R}_{\bm{S}}(\bm{W}_{t})\|^{2}+\eta_{t+1}^{2}\|\nabla\mathcal{R}_{\bm{S}}(\bm{W}_{t})\|^{2}+\bm{\Sigma}_{t+1}(\bm{S},\bm{W}_{t})$
$\displaystyle=\frac{\eta_{t+1}^{2}}{N-1}\frac{N-b_{t+1}}{b_{t+1}}\bm{\Sigma}^{sd}_{\bm{S},\bm{W}_{t}}+\eta_{t+1}^{2}\|\nabla\mathcal{R}_{\bm{S}}(\bm{W}_{t})\|^{2}+\bm{\Sigma}_{t+1}(\bm{S},\bm{W}_{t}),$
(24)

where Eq.$(*)$ is due to $\varepsilon_{t+1}$ is independent of $\bm{V}_{t+1}$,
and $\mathbb{E}^{\bm{W}_{t}}\varepsilon_{t+1}=0$.

Applying Eq.(24) back to Eq.(23) completes the proof.

∎

##  Appendix C Supplementary Materials of Section 3.2

###  C.1 Example to illustrate the difficulty to apply Proposition 1 to solve
Problem 1

In this section, we show an example to demonstrate the difficulty for tackling
Problem 1 through Proposition 1. To start with, by the definition of state-
dependent SGLD (Eq.(3)), covariance $\bm{\Sigma}_{[T]}$ is independent of
$\bm{J}$ and $\bm{V}_{[T]}$. Therefore, the square root separates the
expectation with respect to $\bm{V}_{[T]}$ and $\bm{J}$ from the KL divergence
term in the generalization bound

$\mathbb{E}_{\bm{S},\bm{V}_{[T]},\bm{J}}\sqrt{\frac{(a_{2}-a_{1})^{2}}{2}\sum_{s=1}^{T}\mathbb{E}_{Q_{s-1}^{\bm{S},\bm{V}_{[T]}}}\operatorname{KL}\left(Q_{s|(s-1)}^{\bm{S},\bm{V}_{[T]}}\left\|P_{s|(s-1)}^{\bm{J},\bm{S}_{\bm{J}},\bm{V}_{[T]}}\right.\right)},$

which makes the dependency of the bound on $\bm{\Sigma}_{[T]}$ even more
complex. However, even though we change the optimization target into

$\mathbb{E}_{\bm{S}}\sqrt{\frac{(a_{2}-a_{1})^{2}}{2}\mathbb{E}_{\bm{V}_{[T]},\bm{J}}\sum_{s=1}^{T}\mathbb{E}_{Q_{s-1}^{\bm{S},\bm{V}_{[T]}}}\operatorname{KL}\left(Q_{s|(s-1)}^{\bm{S},\bm{V}_{[T]}}\left\|P_{s|(s-1)}^{\bm{J},\bm{S}_{\bm{J}},\bm{V}_{[T]}}\right.\right)},$
(25)

which is still a generalization bound by Jensen’s Inequality, we demonstrate
that the dependency on $\bm{\Sigma}_{[T]}$ is still too complex to tackle as
follows.

To optimize Eq.(25) with respect to $\bm{\Sigma}_{[T]}(\bm{S},\cdot)$ for
fixed $\bm{S}$, we are actually seeking the optimal point of the following
optimization problem:

$\bm{\Sigma}^{*}_{[T]}(\bm{S},\cdot)=\arg\min_{\bm{\Sigma}_{[T]}(\bm{S},\cdot)}\sqrt{\mathbb{E}_{\bm{V}_{[T]},\bm{J}}\sum_{s=1}^{T}\mathbb{E}_{Q_{s-1}^{\bm{S},\bm{V}_{[T]}}}\operatorname{KL}\left(Q_{s|(s-1)}^{\bm{S},\bm{V}_{[T]}}\left\|P_{s|(s-1)}^{\bm{J},\bm{S}_{\bm{J}},\bm{V}_{[T]}}\right.\right)}.$
(26)

However, we will show it is technically hard to solve Eq. (26). As discussed
in Section 3.2.1, for any fixed index $i\in[T]$, Eq.(26) depends on
$\bm{\Sigma}_{s}(\bm{S},\cdot)$ through both
$\mathbb{E}_{\bm{V}_{[T]},\bm{J}}\mathbb{E}_{Q_{s-1}^{\bm{S},\bm{V}_{[T]}}}$
$\operatorname{KL}(Q_{s|(s-1)}^{\bm{S},\bm{V}_{[T]}}\|P_{s|(s-1)}^{\bm{J},\bm{S}_{\bm{J}},\bm{V}_{[T]}})$
and
$\mathbb{E}_{\bm{V}_{[T]},\bm{J}}\mathbb{E}_{Q_{i-1}^{\bm{S},\bm{V}_{[T]}}}\operatorname{KL}(Q_{i|(i-1)}^{\bm{S},\bm{V}_{[T]}}\|P_{i|(i-1)}^{\bm{J},\bm{S}_{\bm{J}},\bm{V}_{[T]}})$
for $\forall i>s$. Specifically, we adopt the update rule for prior for all
the steps and posterior for all steps $t\neq s$ to be the isotropic SGLD in
[25], i.e.,

$\displaystyle\text{Posterior:
}\bm{W}_{t}=\bm{W}_{t-1}-\eta_{t}\nabla\mathcal{R}_{\bm{S}_{\bm{V}_{t}}}(\bm{W}_{t-1})+\mathcal{N}(\bm{0},\sigma_{t}\mathbb{I})$
$\displaystyle\text{Prior:
}\bm{W}_{t}=\bm{W}_{t-1}-\eta_{t}\left(\frac{|\bm{V}_{t}\cap\bm{J}|}{|\bm{V}_{t}|}\nabla\mathcal{R}_{\bm{S}_{\bm{V}_{t}\cap\bm{J}}}\left(\bm{W}_{t-1}\right)+\frac{|\bm{V}_{t}\cap\bm{J}^{c}|}{|\bm{V}_{t}|}\nabla\mathcal{R}_{\bm{S}_{\bm{J}}}\left(\bm{W}_{t-1}\right)\right)+\mathcal{N}\left(\bm{0},\sigma_{t}\mathbb{I}\right),$

while we only optimize the noise covariance $\bm{\Sigma}_{s}(\bm{S},\cdot)$ of
step $s$:

$\bm{W}_{s}=\bm{W}_{s-1}-\eta_{s}\nabla\mathcal{R}_{\bm{S}_{\bm{V}_{s}}}(\bm{W}_{s-1})+\mathcal{N}(\bm{0},\bm{\Sigma}_{s}(\bm{S},\bm{W}_{s-1})).$

By simple calculation, for any step $t\in[T]$, given the same $\bm{W}_{t-1}$,
$\bm{V}_{t}$, $\bm{J}$, and $\bm{S}$, the mean between the prior and posterior
can be calculated as

$\displaystyle\mu^{\bm{S},\bm{V}_{t},\bm{J},\bm{W}_{t-1}}$
$\displaystyle=-\eta_{t}\left(\frac{|\bm{V}_{t}\cap\bm{J}|}{|\bm{V}_{t}|}\nabla\mathcal{R}_{\bm{S}_{\bm{V}_{t}\cap\bm{J}}}\left(\bm{W}_{t-1}\right)+\frac{|\bm{V}_{t}\cap\bm{J}^{c}|}{|\bm{V}_{t}|}\nabla\mathcal{R}_{\bm{S}_{\bm{J}}}\left(\bm{W}_{t-1}\right)\right)+\eta_{t}\nabla\mathcal{R}_{\bm{S}_{\bm{V}_{t}}}(\bm{W}_{t-1})$
$\displaystyle=\eta_{t}\frac{|\bm{V}_{t}\cap\bm{J}^{c}|}{|\bm{V}_{t}|}\left(\nabla\mathcal{R}_{\bm{S}_{\bm{V}_{t}\cap\bm{J}^{c}}}(\bm{W}_{t-1})-\nabla\mathcal{R}_{\bm{S}_{\bm{J}}}(\bm{W}_{t-1})\right).$
(27)

Therefore, by Lemma 5 and Lemma 6, the expected KL divergence
$\mathbb{E}_{\bm{V}_{[T]},\bm{J}}\mathbb{E}_{Q_{i-1}^{\bm{S},\bm{V}_{[T]}}}\operatorname{KL}$
$(Q_{i|(i-1)}^{\bm{S},\bm{V}_{[T]}}\|P_{i|(i-1)}^{\bm{J},\bm{S}_{\bm{J}},\bm{V}_{[T]}})$
can be calculated as

$\displaystyle\mathbb{E}_{\bm{V}_{[T]},\bm{J}}\mathbb{E}_{Q_{i-1}^{\bm{S},\bm{V}_{[T]}}}\operatorname{KL}\left(Q_{i|(i-1)}^{\bm{S},\bm{V}_{[T]}}\left\|P_{i|(i-1)}^{\bm{J},\bm{S}_{\bm{J}},\bm{V}_{[T]}}\right.\right)$
$\displaystyle=\frac{1}{2}\mathbb{E}_{\bm{V}_{[T]},\bm{J}}\mathbb{E}_{Q_{i-1}^{\bm{S},\bm{V}_{[T]}}}\left(\sigma_{i}^{-1}\mu^{\bm{S},\bm{V}_{t},\bm{J},\bm{W}_{i-1}}\left(\mu^{\bm{S},\bm{V}_{t},\bm{J},\bm{W}_{i-1}}\right)^{\top}\right)$
$\displaystyle=\frac{1}{2\sigma_{i}}\frac{1}{Nb_{i}}\left(\frac{N}{N-1}\right)^{2}\mathbb{E}_{\bm{V}_{[i-1]}}\mathbb{E}_{Q_{i-1}^{\bm{S},\bm{V}_{[i-1]}}}\bm{\Sigma}^{sd}_{\bm{S},\bm{W}_{i-1}}.$

Therefore, the exact form of
$\mathbb{E}_{\bm{V}_{[T]},\bm{J}}\mathbb{E}_{Q_{i-1}^{\bm{S},\bm{V}_{[T]}}}\operatorname{KL}(Q_{i|(i-1)}^{\bm{S},\bm{V}_{[T]}}\|P_{i|(i-1)}^{\bm{J},\bm{S}_{\bm{J}},\bm{V}_{[T]}})$
requires taking expectation to $\bm{\Sigma}^{sd}_{\bm{S},\bm{W}_{i-1}}$ with
respect to Gaussian distribution with covariance $\bm{\Sigma}_{s}$, and can be
complex due to the complex structure of the model. Specifically, if $i=s+1$,
then
$\mathbb{E}_{\bm{V}_{[T]},\bm{J}}\mathbb{E}_{Q_{i-1}^{\bm{S},\bm{V}_{[T]}}}$
$\operatorname{KL}(Q_{i|(i-1)}^{\bm{S},\bm{V}_{[T]}}\|P_{i|(i-1)}^{\bm{J},\bm{S}_{\bm{J}},\bm{V}_{[T]}})$
can be further written as

$\displaystyle\mathbb{E}_{\bm{V}_{[T]},\bm{J}}\mathbb{E}_{Q_{s}^{\bm{S},\bm{V}_{[T]}}}\operatorname{KL}\left(Q_{s+1|s}^{\bm{S},\bm{V}_{[T]}}\left\|P_{s+1|s}^{\bm{J},\bm{S}_{\bm{J}},\bm{V}_{[T]}}\right.\right)$
$\displaystyle=$
$\displaystyle\frac{1}{2}\frac{1}{Nb_{s+1}}\left(\frac{N}{N-1}\right)^{2}\mathbb{E}_{\bm{V}_{[s]}}\mathbb{E}_{Q_{s-1}^{\bm{S},\bm{V}_{[s-1]}}}\mathbb{E}_{Q_{s|(s-1)}^{\bm{S},\bm{V}_{s}}}\bm{\Sigma}^{sd}_{\bm{S},\bm{W}_{s}}.$

Therefore, we need to optimize
$\mathbb{E}_{\bm{V}_{s}}\mathbb{E}_{Q_{s|(s-1)}^{\bm{S},\bm{V}_{s}}}\bm{\Sigma}^{sd}_{\bm{S},\bm{W}_{s}}$,
which can be further written as

$\displaystyle\mathbb{E}_{\bm{V}_{s}}\mathbb{E}_{Q_{s|(s-1)}^{\bm{S},\bm{V}_{s}}}\bm{\Sigma}^{sd}_{\bm{S},\bm{W}_{s}}=\mathbb{E}_{\bm{V}_{s}}\mathbb{E}_{\mathcal{N}(-\eta_{s}\nabla\mathcal{R}_{\bm{S}_{\bm{V}_{s}}}(\bm{W}_{s-1}),\bm{\Sigma}_{s}(\bm{S},\bm{W}_{s-1}))}\bm{\Sigma}^{sd}_{\bm{S},\bm{W}_{s}}.$

The explicit form of
$\mathbb{E}_{\bm{V}_{s}}\mathbb{E}_{\mathcal{N}(-\eta_{s}\nabla\mathcal{R}_{\bm{S}_{\bm{V}_{s}}}(\bm{W}_{s-1}),\bm{\Sigma}_{s}(\bm{S},\bm{W}_{s-1}))}\bm{\Sigma}^{sd}_{\bm{S},\bm{W}_{s}}$
can be obtained only when $\bm{\Sigma}^{sd}_{\bm{S},\bm{W}_{s}}$ is some
simple functions with respect to $\bm{W}_{s}$ (e.g. quadratic functions),
which makes the optimal of
$\mathbb{E}_{\bm{V}_{s}}\mathbb{E}_{\mathcal{N}(-\eta_{s}\nabla\mathcal{R}_{\bm{S}_{\bm{V}_{s}}}(\bm{W}_{s-1}),\bm{\Sigma}_{s}(\bm{S},\bm{W}_{s-1}))}\bm{\Sigma}^{sd}_{\bm{S},\bm{W}_{s}}$
complicated due to the complex structure of $\mathcal{R}_{\bm{S}}$ and
$\bm{\Sigma}^{sd}_{\bm{S},\bm{W}}$ in pratical learning problems.

###  C.2 Proof of Theorem 1

###### Proof of Theorem 1.

For any two random measures
$P^{\bm{J},\bm{S}_{\bm{J}},\bm{V}_{[T]}},Q^{\bm{S},\bm{V}_{[T]}}$, by the
Donsker-Varadhan variational formula [3], for any function $g$ satisfying
$Q^{\bm{S},\bm{V}_{[T]}}(\exp g)<\infty$, we have

$\displaystyle\operatorname{KL}(P^{\bm{J},\bm{S}_{\bm{J}},\bm{V}_{[T]}}||Q^{\bm{S},\bm{V}_{[T]}})\geq
P^{\bm{J},\bm{S}_{\bm{J}},\bm{V}_{[T]}}(g)-Q^{\bm{S},\bm{V}_{[T]}}(g)-\log
Q^{\bm{S},\bm{V}_{[T]}}\left(\exp(g-Q^{\bm{S},\bm{V}_{[T]}}(g))\right).$

Letting
$g(\bm{W})=\lambda\left(\hat{\mathcal{R}}_{\bm{S}_{\bm{J}^{c}}}(\bm{W})-\mathcal{R}_{\mathcal{D}}(\bm{W})\right)$,
we further have

$\displaystyle\operatorname{KL}(P^{\bm{J},\bm{S}_{\bm{J}},\bm{V}_{[T]}}||Q^{\bm{S},\bm{V}_{[T]}})$
$\displaystyle\geq\lambda\left(\mathcal{R}_{\mathcal{D}}(Q^{\bm{S},\bm{V}_{[T]}})-\hat{\mathcal{R}}_{\bm{S}_{\bm{J}^{c}}}(Q^{\bm{S},\bm{V}_{[T]}})-\left(\mathcal{R}_{\mathcal{D}}(P^{\bm{J},\bm{S}_{\bm{J}},\bm{V}_{[T]}})-\hat{\mathcal{R}}_{\bm{S}_{\bm{J}^{c}}}(P^{\bm{J},\bm{S}_{\bm{J}},\bm{V}_{[T]}})\right)\right)$
$\displaystyle\;\;-\log
Q^{\bm{S},\bm{V}_{[T]}}\left(\exp\left(\lambda\left(\hat{\mathcal{R}}_{\bm{S}_{\bm{J}^{c}}}-\mathcal{R}_{\mathcal{D}}-\left(\hat{\mathcal{R}}_{\bm{S}_{\bm{J}^{c}}}(Q^{\bm{S},\bm{V}_{[T]}})-\mathcal{R}_{\mathcal{D}}(Q^{\bm{S},\bm{V}_{[T]}})\right)\right)\right)\right).$

On the other hand, since $\ell\in[a_{1},a_{2}]$,
$\lambda\left(\hat{\mathcal{R}}_{\bm{S}_{\bm{J}^{c}}}(\bm{W})-\mathcal{R}_{\mathcal{D}}(\bm{W})\right)$
is $\frac{\lambda(a_{2}-a_{1})}{2}$ subgaussian. Therefore,

$\displaystyle\left(\mathcal{R}_{\mathcal{D}}(Q^{\bm{S},\bm{V}_{[T]}})-\hat{\mathcal{R}}_{\bm{S}_{\bm{J}^{c}}}(Q^{\bm{S},\bm{V}_{[T]}})\right)-\left(\mathcal{R}_{\mathcal{D}}(P^{\bm{J},\bm{S}_{\bm{J}},\bm{V}_{[T]}})-\hat{\mathcal{R}}_{\bm{S}_{\bm{J}^{c}}}(P^{\bm{J},\bm{S}_{\bm{J}},\bm{V}_{[T]}})\right)$
$\displaystyle\leq\inf_{\lambda>0}\frac{\operatorname{KL}(P^{\bm{J},\bm{S}_{\bm{J}},\bm{V}_{[T]}}||Q^{\bm{S},\bm{V}_{[T]}})+\frac{1}{8}\lambda^{2}(a_{2}-a_{1})^{2}}{\lambda}.$

Since $P^{\bm{J},\bm{S}_{\bm{J}},\bm{V}_{[T]}}$ is independent of
$\bm{S}_{\bm{J}^{c}}$ then we have
$\mathbb{E}^{\bm{S}_{\bm{J}},\bm{J},\bm{V}_{[T]}}\left[\mathcal{R}_{\mathcal{D}}(P^{\bm{J},\bm{S}_{\bm{J}},\bm{V}_{[T]}})-\hat{\mathcal{R}}_{\bm{S}_{\bm{J}^{c}}}(P^{\bm{J},\bm{S}_{\bm{J}},\bm{V}_{[T]}})\right]=0$.
Hence, by averaging over $\bm{S}_{\bm{J}^{c}}$ (equivalently, taking the
conditional expectation conditional on
$(\bm{S}_{\bm{J}},\bm{J},\bm{V}_{[T]})$) we have, with probability one

$\displaystyle\mathbb{E}^{\bm{S}_{\bm{J}},\bm{J},\bm{V}_{[T]}}\left[\mathcal{R}_{\mathcal{D}}(Q^{\bm{S},\bm{V}_{[T]}})-\hat{\mathcal{R}}_{\bm{S}_{\bm{J}^{c}}}(Q^{\bm{S},\bm{V}_{[T]}})\right]$
$\displaystyle=\mathbb{E}^{\bm{S}_{\bm{J}},\bm{J},\bm{V}_{[T]}}\left[\mathcal{R}_{\mathcal{D}}(Q^{\bm{S},\bm{V}_{[T]}})-\hat{\mathcal{R}}_{\bm{S}_{\bm{J}^{c}}}(Q^{\bm{S},\bm{V}_{[T]}})-\left(\mathcal{R}_{\mathcal{D}}(P^{\bm{J},\bm{S}_{\bm{J}},\bm{V}_{[T]}})-\hat{\mathcal{R}}_{\bm{S}_{\bm{J}^{c}}}(P^{\bm{J},\bm{S}_{\bm{J}},\bm{V}_{[T]}})\right)\right]$
$\displaystyle\leq\mathbb{E}^{\bm{S}_{\bm{J}},\bm{J},\bm{V}_{[T]}}\left(\inf_{\lambda>0}\frac{\operatorname{KL}(P^{\bm{J},\bm{S}_{\bm{J}},\bm{V}_{[T]}}||Q^{\bm{S},\bm{V}_{[T]}})+\frac{1}{8}\lambda^{2}(a_{2}-a_{1})^{2}}{\lambda}\right)$

Finally, by taking the full expectation, since $J\perp\\!\\!\\!\perp
Q^{\bm{S},\bm{V}_{[T]}}$ we get:

$\mathbb{E}_{\bm{S},\bm{V}_{[T]}}\left[\mathcal{R}_{\mathcal{D}}(Q^{\bm{S},\bm{V}_{[T]}})-\hat{\mathcal{R}}_{\bm{S}}(Q^{\bm{S},\bm{V}_{[T]}})\right]\leq\mathbb{E}_{\bm{S},\bm{V}_{[T]},\bm{J}}\left[\inf_{\lambda>0}\frac{\operatorname{KL}(P^{\bm{J},\bm{S}_{\bm{J}},\bm{V}_{[T]}}||Q^{\bm{S},\bm{V}_{[T]}})+\frac{1}{8}\lambda^{2}(a_{2}-a_{1})^{2}}{\lambda}\right]$

where the final
$\operatorname{KL}(P^{\bm{J},\bm{S}_{\bm{J}},\bm{V}_{[T]}}||Q^{\bm{S},\bm{V}_{[T]}})$
on the right hand side is between two random measures, and hence is a random
variable depending on $(\bm{S},\bm{J},\bm{V}_{[T]})$; and the expectation on
the right hand side integrates over $(\bm{S},\bm{J},\bm{V}_{[T]})$.

Since

$\frac{\operatorname{KL}(P^{\bm{J},\bm{S}_{\bm{J}},\bm{V}_{[T]}}||Q^{\bm{S},\bm{V}_{[T]}})+\frac{1}{8}\lambda^{2}(a_{2}-a_{1})^{2}}{\lambda}\geq\sqrt{\frac{1}{2}(a_{2}-a_{1})^{2}\operatorname{KL}(P^{\bm{J},\bm{S}_{\bm{J}},\bm{V}_{[T]}}||Q^{\bm{S},\bm{V}_{[T]}})},$

the proof is completed. ∎

##  Appendix D Supplementary of Section 4

In this section, we provide the proof of Theorem 2. Specifically, as mentioned
in the main body, optimizing $\operatorname{Gen}_{T}$ with greedily selected
prior involves three steps. (1). we first prove Lemma 3, which provides the
optimal solution of noise covariance and prior for one single KL divergence
term in the generalization bound $\operatorname{Gen}_{T}$; (2). as the optimal
solution of noise covariance in Lemma 3 is independent of $\bm{S}_{\bm{J}}$,
$\bm{V}_{[T]}$, and $\bm{V}_{[T]}$, we are then able to obtain the greedy
prior by Lemma 4; (3). applying the greedy prior back to
$\operatorname{Gen}_{T}$, we are finally able to derive Theorem 2.

We start by restating Lemma 12 and providing its proof.

######  Lemma 12 (Lemma 3, restated).

For any $s\in[T]$, $\bm{J}$, $\bm{S}_{\bm{J}}$, and $\bm{V}_{[T]}$, under
Constraint 1,

$\min_{{P}_{s|(s-1)}^{\bm{J},\bm{S}_{\bm{J}},\bm{V}_{s}}}\mathbb{E}_{\bm{S}_{\bm{J}^{c}}\sim\mathcal{D}}\operatorname{KL}\left({P}_{s|(s-1)}^{\bm{J},\bm{S}_{\bm{J}},\bm{V}_{s}}\left\|Q_{s|(s-1)}^{\bm{S},\bm{V}_{s}}\right.\right)$
(28)

(1). is independent of $\bm{\Sigma}_{s}$ when
$\bm{V}_{s}\cap\bm{J}^{c}=\emptyset$, and (2). is minimized at
$\bm{\Sigma}_{s}(\bm{W})=\lambda_{s}(\bm{W})\left(\bm{\Sigma}^{pop}_{\bm{W}}\right)^{\frac{1}{2}}$,
$\forall\bm{W}$, when $\bm{V}_{s}\cap\bm{J}^{c}\neq\emptyset$, where
$\lambda_{s}(\bm{W})=c_{s}(\bm{W})/\operatorname{tr}(\left(\bm{\Sigma}^{pop}_{\bm{W}}\right)^{\frac{1}{2}})$.

###### Proof.

We first calculate
$\min_{{P}_{s|(s-1)}^{\bm{J},\bm{S}_{\bm{J}},\bm{V}_{s}}}\mathbb{E}_{\bm{S}_{\bm{J}^{c}}\sim\mathcal{D}}\operatorname{KL}\left({P}_{s|(s-1)}^{\bm{J},\bm{S}_{\bm{J}},\bm{V}_{s}}\left\|Q_{s|(s-1)}^{\bm{S},\bm{V}_{s}}\right.\right)$
for any $\bm{\Sigma}_{s}$. By applying the definition of the KL divergence, we
have

$\displaystyle\arg\min_{{P}_{s|(s-1)}^{\bm{J},\bm{S}_{\bm{J}},\bm{V}_{s}}}\mathbb{E}_{\bm{S}_{\bm{J}^{c}}\sim\mathcal{D}}\operatorname{KL}\left({P}_{s|(s-1)}^{\bm{J},\bm{S}_{\bm{J}},\bm{V}_{s}}\left\|Q_{s|(s-1)}^{\bm{S},\bm{V}_{[T]}}\right.\right)$
$\displaystyle=\arg\min_{{P}_{s|(s-1)}^{\bm{J},\bm{S}_{\bm{J}},\bm{V}_{s}}}\mathbb{E}_{\bm{S}_{\bm{J}^{c}}\sim\mathcal{D}}\int{P}_{s|(s-1)}^{\bm{J},\bm{S}_{\bm{J}},\bm{V}_{s}}(\bm{W}_{s})\log\frac{{P}_{s|(s-1)}^{\bm{J},\bm{S}_{\bm{J}},\bm{V}_{s}}(\bm{W}_{s})}{Q_{s|(s-1)}^{\bm{S},\bm{V}_{[T]}}(\bm{W}_{s})}\mathrm{d}\bm{W}_{s}$
$\displaystyle\overset{(*)}{=}\arg\min_{{P}_{s|(s-1)}^{\bm{J},\bm{S}_{\bm{J}},\bm{V}_{s}}}\int{P}_{s|(s-1)}^{\bm{J},\bm{S}_{\bm{J}},\bm{V}_{s}}(\bm{W}_{s})\log\frac{{P}_{s|(s-1)}^{\bm{J},\bm{S}_{\bm{J}},\bm{V}_{s}}(\bm{W}_{s})}{e^{\mathbb{E}_{\bm{S}_{\bm{J}^{c}}\sim\mathcal{D}}\log
Q_{s|(s-1)}^{\bm{S},\bm{V}_{[T]}}(\bm{W}_{s})}}\mathrm{d}\bm{W}_{s},$ (29)

where Eq. $(*)$ is due to the independence of $P$ on $\bm{S}_{\bm{J}^{c}}$.

Let

$\displaystyle\tilde{Q}_{s|(s-1)}^{\bm{J},\bm{S}_{\bm{J}},\bm{V}_{[T]}}(\bm{W})=\frac{e^{\mathbb{E}_{\bm{S}_{\bm{J}^{c}}\sim\mathcal{D}}\log
Q_{s|(s-1)}^{\bm{S},\bm{V}_{[T]}}(\bm{W})}}{\int
e^{\mathbb{E}_{\bm{S}_{\bm{J}^{c}}\sim\mathcal{D}}\log
Q_{s|(s-1)}^{\bm{S},\bm{V}_{[T]}}(\tilde{\bm{W}})}\mathrm{d}\tilde{\bm{W}}},$
(30)

and $\tilde{Q}_{s|(s-1)}^{\bm{J},\bm{S}_{\bm{J}},\bm{V}_{[T]}}(\bm{W})$ is
then a probability measure on $\mathbb{R}^{d}$. Applying Eq. (30) back to Eq.
(29), we obtain

$\displaystyle\arg\min_{{P}_{s|(s-1)}^{\bm{J},\bm{S}_{\bm{J}},\bm{V}_{s}}}\left(\int{P}_{s|(s-1)}^{\bm{J},\bm{S}_{\bm{J}},\bm{V}_{s}}(\bm{W}_{s})\log\frac{{P}_{s|(s-1)}^{\bm{J},\bm{S}_{\bm{J}},\bm{V}_{s}}(\bm{W}_{s})}{\tilde{Q}_{s|(s-1)}^{\bm{J},\bm{S}_{\bm{J}},\bm{V}_{[T]}}(\bm{W}_{s})}\mathrm{d}\bm{W}_{s}\right.$
$\displaystyle~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}-\left.\int{P}_{s|(s-1)}^{\bm{J},\bm{S}_{\bm{J}},\bm{V}_{s}}(\bm{W}_{s})\log\left(\int
e^{\mathbb{E}_{\bm{S}_{\bm{J}^{c}}\sim\mathcal{D}}\log
Q_{s|(s-1)}^{\bm{S},\bm{V}_{[T]}}(\tilde{\bm{W}})}\mathrm{d}\tilde{\bm{W}}\right)\mathrm{d}\bm{W}_{s}\right)$
$\displaystyle=$
$\displaystyle\arg\min_{{P}_{s|(s-1)}^{\bm{J},\bm{S}_{\bm{J}},\bm{V}_{s}}}\left(\int{P}_{s|(s-1)}^{\bm{J},\bm{S}_{\bm{J}},\bm{V}_{s}}(\bm{W}_{s})\log\frac{{P}_{s|(s-1)}^{\bm{J},\bm{S}_{\bm{J}},\bm{V}_{s}}(\bm{W}_{s})}{\tilde{Q}_{s|(s-1)}^{\bm{J},\bm{S}_{\bm{J}},\bm{V}_{[T]}}(\bm{W}_{s})}\mathrm{d}\bm{W}_{s}\right.-\left.\log\left(\int
e^{\mathbb{E}_{\bm{S}_{\bm{J}^{c}}\sim\mathcal{D}}\log
Q_{s|(s-1)}^{\bm{S},\bm{V}_{[T]}}(\tilde{\bm{W}})}\mathrm{d}\tilde{\bm{W}}\right)\right)$
$\displaystyle=$
$\displaystyle\arg\min_{{P}_{s|(s-1)}^{\bm{J},\bm{S}_{\bm{J}},\bm{V}_{s}}}\left(\int{P}_{s|(s-1)}^{\bm{J},\bm{S}_{\bm{J}},\bm{V}_{s}}(\bm{W}_{s})\log\frac{{P}_{s|(s-1)}^{\bm{J},\bm{S}_{\bm{J}},\bm{V}_{s}}(\bm{W}_{s})}{\tilde{Q}_{s|(s-1)}^{\bm{J},\bm{S}_{\bm{J}},\bm{V}_{[T]}}(\bm{W}_{s})}\mathrm{d}\bm{W}_{s}\right)$
$\displaystyle=$
$\displaystyle\arg\min_{{P}_{s|(s-1)}^{\bm{J},\bm{S}_{\bm{J}},\bm{V}_{s}}}\operatorname{KL}\left(P\|\tilde{Q}_{s|(s-1)}^{\bm{J},\bm{S}_{\bm{J}},\bm{V}_{[T]}}\right).$
(31)

The minimum of Eq.(31) is achieved if and only if
${P}_{s|(s-1)}^{\bm{J},\bm{S}_{\bm{J}},\bm{V}_{s}}=\tilde{Q}_{s|(s-1)}^{\bm{J},\bm{S}_{\bm{J}},\bm{V}_{[T]}}$,
and we only need to calculate the exact form of
$\tilde{Q}_{s|(s-1)}^{\bm{J},\bm{S}_{\bm{J}},\bm{V}_{[T]}}$. Since
$\bm{W}_{s}|(\bm{W}_{s-1},\bm{S},\bm{V}_{s})\sim\mathcal{N}(\bm{W}_{s-1}-\eta_{s}\nabla_{\bm{W}_{s-1}}\mathcal{R}_{\bm{S}_{\bm{V}_{s}}}(\bm{W}_{s-1}),$
$\bm{\Sigma}_{s}(\bm{W}_{s-1}))$, we have

$\displaystyle\exp{\mathbb{E}_{\bm{S}_{\bm{J}^{c}}\sim\mathcal{D}}\log
Q_{s|(s-1)}^{\bm{S},\bm{V}_{[T]}}(\bm{W})}$
$\displaystyle=\exp\left(\mathbb{E}_{\bm{S}_{\bm{J}^{c}}\sim\mathcal{D}}\left(-\frac{1}{2}(\bm{W}-\bm{W}_{s-1}+\eta_{s}\nabla_{\bm{W}_{s-1}}\mathcal{R}_{\bm{S}_{\bm{V}_{s}}}(\bm{W}_{s-1}))^{\top}\bm{\Sigma}_{s}(\bm{W}_{s-1})^{-1}(\bm{W}-\bm{W}_{s-1}\right.\right.$
$\displaystyle\;\;\left.\left.+\eta_{s}\nabla_{\bm{W}_{s-1}}\mathcal{R}_{\bm{S}_{\bm{V}_{s}}}(\bm{W}_{s-1}))-\frac{d}{2}\log
2\pi-\frac{1}{2}\log\det(\bm{\Sigma}_{s}(\bm{W}_{s-1}))\right)\right)$
$\displaystyle=\exp\left(\mathbb{E}_{\bm{S}_{\bm{J}^{c}}\sim\mathcal{D}}\bigg{(}-\frac{1}{2}(\bm{W}-\bm{W}_{s-1}+\eta_{s}\nabla_{\bm{W}_{s-1}}\mathcal{R}_{\bm{S}_{\bm{V}_{s}}}(\bm{W}_{s-1}))^{\top}\bm{\Sigma}_{s}(\bm{W}_{s-1})^{-1}(\bm{W}-\bm{W}_{s-1}\right.$
$\displaystyle\;\;\left.+\eta_{s}\nabla_{\bm{W}_{s-1}}\mathcal{R}_{\bm{S}_{\bm{V}_{s}}}(\bm{W}_{s-1})\bigg{)}-\frac{d}{2}\log
2\pi-\frac{1}{2}\log\det(\bm{\Sigma}_{s}(\bm{W}_{s-1}))\right).$ (32)

On the other hand,

$\displaystyle\mathbb{E}_{\bm{S}_{\bm{J}^{c}}\sim\mathcal{D}}\left(-\frac{1}{2}(\bm{W}-\bm{W}_{s-1}+\eta_{s}\nabla_{\bm{W}_{s-1}}\mathcal{R}_{\bm{S}_{\bm{V}_{s}}}(\bm{W}_{s-1}))^{\top}\bm{\Sigma}_{s}(\bm{W}_{s-1})^{-1}(\bm{W}-\bm{W}_{s-1}\right.$
$\displaystyle\;\;\left.+\eta_{s}\nabla_{\bm{W}_{s-1}}\mathcal{R}_{\bm{S}_{\bm{V}_{s}}}(\bm{W}_{s-1}))\right)$
$\displaystyle=-\frac{1}{2}\mathbb{E}_{\bm{S}_{\bm{J}^{c}}\sim\mathcal{D}}\left(\bm{W}-\bm{W}_{s-1}+\eta_{s}\left(\frac{|\bm{V}_{s}\cap\bm{J}|}{|\bm{V}_{s}|}\nabla\mathcal{R}_{\bm{S}_{\bm{V}_{s}\cap\bm{J}}}\left(\bm{W}_{s-1}\right)+\frac{|\bm{V}_{s}\cap\bm{J}^{c}|}{|\bm{V}_{s}|}\nabla\mathcal{R}_{\bm{S}_{\bm{V}_{s}\cap\bm{J}^{c}}}\left(\bm{W}_{s-1}\right)\right)\right)^{\top}$
$\displaystyle\;\;\cdot\bm{\Sigma}_{s}(\bm{W}_{s-1})^{-1}\left(\bm{W}-\bm{W}_{s-1}+\eta_{s}\left(\frac{|\bm{V}_{s}\cap\bm{J}|}{|\bm{V}_{s}|}\nabla\mathcal{R}_{\bm{S}_{\bm{V}_{s}\cap\bm{J}}}\left(\bm{W}_{s-1}\right)+\frac{|\bm{V}_{s}\cap\bm{J}^{c}|}{|\bm{V}_{s}|}\nabla\mathcal{R}_{\bm{S}_{\bm{V}_{s}\cap\bm{J}^{c}}}\left(\bm{W}_{s-1}\right)\right)\right)$
$\displaystyle=-\frac{1}{2}\left(\bm{W}-\bm{W}_{s-1}+\eta_{s}\left(\frac{|\bm{V}_{s}\cap\bm{J}|}{|\bm{V}_{s}|}\nabla\mathcal{R}_{\bm{S}_{\bm{V}_{s}\cap\bm{J}}}\left(\bm{W}_{s-1}\right)+\frac{|\bm{V}_{s}\cap\bm{J}^{c}|}{|\bm{V}_{s}|}\nabla\mathcal{R}_{\mathcal{D}}\left(\bm{W}_{s-1}\right)\right)\right)^{\top}$
$\displaystyle\;\;\cdot\bm{\Sigma}_{s}(\bm{W}_{s-1})^{-1}\left(\bm{W}-\bm{W}_{s-1}+\eta_{s}\left(\frac{|\bm{V}_{s}\cap\bm{J}|}{|\bm{V}_{s}|}\nabla\mathcal{R}_{\bm{S}_{\bm{V}_{s}\cap\bm{J}}}\left(\bm{W}_{s-1}\right)+\frac{|\bm{V}_{s}\cap\bm{J}^{c}|}{|\bm{V}_{s}|}\nabla\mathcal{R}_{\mathcal{D}}\left(\bm{W}_{s-1}\right)\right)\right)$
$\displaystyle\;\;-\frac{1}{2}\mathbb{E}_{\bm{S}_{\bm{J}^{c}}\sim\mathcal{D}}\eta_{s}^{2}\frac{|\bm{V}_{s}\cap\bm{J}^{c}|^{2}}{|\bm{V}_{s}|^{2}}\left(\nabla\mathcal{R}_{\mathcal{D}}\left(\bm{W}_{s-1}\right)-\nabla\mathcal{R}_{\bm{S}_{\bm{V}_{s}\cap\bm{J}^{c}}}\left(\bm{W}_{s-1}\right)\right)^{\top}\bm{\Sigma}_{s}(\bm{W}_{s-1})^{-1}$
$\displaystyle\;\;\cdot\left(\nabla\mathcal{R}_{\mathcal{D}}\left(\bm{W}_{s-1}\right)-\nabla\mathcal{R}_{\bm{S}_{\bm{V}_{s}\cap\bm{J}^{c}}}\left(\bm{W}_{s-1}\right)\right).$
(33)

By combining Eq.(32) and Eq.(33), we further have

$\displaystyle\exp{\mathbb{E}_{\bm{S}_{\bm{J}^{c}}\sim\mathcal{D}}\log
Q_{s|(s-1)}^{\bm{S},\bm{V}_{[T]}}(\bm{W})}$
$\displaystyle=\frac{1}{(2\pi)^{-\frac{d}{2}}\det(\bm{\Sigma}_{s}(\bm{W}_{s-1}))^{\frac{1}{2}}}\exp\left(-\frac{1}{2}\left(\bm{W}-\bm{W}_{s-1}+\eta_{s}\left(\frac{|\bm{V}_{s}\cap\bm{J}|}{|\bm{V}_{s}|}\nabla\mathcal{R}_{\bm{S}_{\bm{V}_{s}\cap\bm{J}}}\left(\bm{W}_{s-1}\right)\right.\right.\right.$
$\displaystyle\;\;+\left.\left.\left.\frac{|\bm{V}_{s}\cap\bm{J}^{c}|}{|\bm{V}_{s}|}\nabla\mathcal{R}_{\mathcal{D}}\left(\bm{W}_{s-1}\right)\right)\right)^{\top}\bm{\Sigma}_{s}(\bm{W}_{s-1})^{-1}\left(\bm{W}-\bm{W}_{s-1}+\eta_{s}\left(\frac{|\bm{V}_{s}\cap\bm{J}|}{|\bm{V}_{s}|}\nabla\mathcal{R}_{\bm{S}_{\bm{V}_{s}\cap\bm{J}}}\left(\bm{W}_{s-1}\right)\right.\right.\right.$
$\displaystyle\;\;\left.\left.+\frac{|\bm{V}_{s}\cap\bm{J}^{c}|}{|\bm{V}_{s}|}\nabla\mathcal{R}_{\mathcal{D}}\left(\bm{W}_{s-1}\right)\right)\right)\exp\mathbb{E}_{\bm{S}_{\bm{J}^{c}}}\bigg{(}-\frac{1}{2}\eta_{s}^{2}\frac{|\bm{V}_{s}\cap\bm{J}^{c}|^{2}}{|\bm{V}_{s}|^{2}}\left(\nabla\mathcal{R}_{\mathcal{D}}\left(\bm{W}_{s-1}\right)-\nabla\mathcal{R}_{\bm{S}_{\bm{V}_{s}\cap\bm{J}^{c}}}\left(\bm{W}_{s-1}\right)\right)^{\top}$
$\displaystyle\;\;\cdot\bm{\Sigma}_{s}(\bm{W}_{s-1})^{-1}\left(\nabla\mathcal{R}_{\mathcal{D}}\left(\bm{W}_{s-1}\right)-\nabla\mathcal{R}_{\bm{S}_{\bm{V}_{s}\cap\bm{J}^{c}}}\left(\bm{W}_{s-1}\right)\right)\bigg{)}.$
(34)

Therefore, by taking integration with respect to $\tilde{\bm{W}}$, we have,

$\displaystyle\int e^{\mathbb{E}_{\bm{S}_{\bm{J}^{c}}\sim\mathcal{D}}\log
Q_{s|(s-1)}^{\bm{S},\bm{V}_{[T]}}(\tilde{\bm{W}})}\mathrm{d}\tilde{\bm{W}}$
$\displaystyle=\exp\mathbb{E}_{\bm{S}_{\bm{J}^{c}}}\left(-\frac{1}{2}\eta_{s}^{2}\frac{|\bm{V}_{s}\cap\bm{J}^{c}|^{2}}{|\bm{V}_{s}|^{2}}\left(\nabla\mathcal{R}_{\mathcal{D}}\left(\bm{W}_{s-1}\right)-\nabla\mathcal{R}_{\bm{S}_{\bm{V}_{s}\cap\bm{J}^{c}}}\left(\bm{W}_{s-1}\right)\right)^{\top}\bm{\Sigma}_{s}(\bm{W}_{s-1})^{-1}\right.$
$\displaystyle\;\;\left.\cdot\left(\nabla\mathcal{R}_{\mathcal{D}}\left(\bm{W}_{s-1}\right)-\nabla\mathcal{R}_{\bm{S}_{\bm{V}_{s}\cap\bm{J}^{c}}}\left(\bm{W}_{s-1}\right)\right)\right).$
(35)

Therefore, by Eq.(30), Eq.(34), and Eq.(35), we have

$\displaystyle\arg\min_{{P}_{s|(s-1)}^{\bm{J},\bm{S}_{\bm{J}},\bm{V}_{s}}}\mathbb{E}_{\bm{S}_{\bm{J}^{c}}\sim\mathcal{D}}\operatorname{KL}\left({P}_{s|(s-1)}^{\bm{J},\bm{S}_{\bm{J}},\bm{V}_{s}}\left\|Q_{s|(s-1)}^{\bm{S},\bm{V}_{[T]}}\right.\right)=\tilde{Q}_{s|(s-1)}^{\bm{J},\bm{S}_{\bm{J}},\bm{V}_{[T]}}$
(36)
$\displaystyle\sim\mathcal{N}\left(\bm{W}_{s-1}-\eta_{s}\left(\frac{|\bm{V}_{s}\cap\bm{J}|}{|\bm{V}_{s}|}\nabla\mathcal{R}_{\bm{S}_{\bm{V}_{s}\cap\bm{J}}}\left(\bm{W}_{s-1}\right)+\frac{|\bm{V}_{s}\cap\bm{J}^{c}|}{|\bm{V}_{s}|}\nabla\mathcal{R}_{\mathcal{D}}\left(\bm{W}_{s-1}\right)\right),\bm{\Sigma}_{s}(\bm{W}_{s-1})\right).$

Applying Eq. (36) back to
$\mathbb{E}_{\bm{S}_{\bm{J}^{c}}\sim\mathcal{D}}\operatorname{KL}\left({P}_{s|(s-1)}^{\bm{J},\bm{S}_{\bm{J}},\bm{V}_{s}}\left\|Q_{s|(s-1)}^{\bm{S},\bm{V}_{[T]}}\right.\right)$,
we obtain

$\displaystyle\min_{{P}_{s|(s-1)}^{\bm{J},\bm{S}_{\bm{J}},\bm{V}_{s}}}\mathbb{E}_{\bm{S}_{\bm{J}^{c}}\sim\mathcal{D}}\operatorname{KL}\left({P}_{s|(s-1)}^{\bm{J},\bm{S}_{\bm{J}},\bm{V}_{s}}\left\|Q_{s|(s-1)}^{\bm{S},\bm{V}_{[T]}}\right.\right)=\mathbb{E}_{\bm{S}_{\bm{J}^{c}}\sim\mathcal{D}}\operatorname{KL}\left(\tilde{Q}_{s|(s-1)}^{\bm{J},\bm{S}_{\bm{J}},\bm{V}_{[T]}}\left\|Q_{s|(s-1)}^{\bm{S},\bm{V}_{[T]}}\right.\right)$
$\displaystyle=\int\tilde{Q}_{t|(t-1)}^{\bm{S}_{\bm{J}},\bm{V}_{[s]}}(\bm{W}_{s})\log\frac{\tilde{Q}_{t|(t-1)}^{\bm{S}_{\bm{J}},\bm{V}_{[s]}}(\bm{W}_{s})}{e^{\mathbb{E}_{\bm{S}_{\bm{J}^{c}}\sim\mathcal{D}}\log
Q_{t|(t-1)}^{\bm{S},\bm{V}_{[s]}}(\bm{W}_{s})}}\mathrm{d}\bm{W}_{s}$
$\displaystyle\overset{(\circ)}{=}-\int\tilde{Q}_{t|(t-1)}^{\bm{S}_{\bm{J}},\bm{V}_{[s]}}(\bm{W}_{s})\log\int
e^{\mathbb{E}_{\bm{S}_{\bm{J}^{c}}\sim\mathcal{D}}\log
Q_{t|(t-1)}^{\bm{S},\bm{V}_{[s]}}(\tilde{\bm{W}})}\mathrm{d}\tilde{\bm{W}}\mathrm{d}\bm{W}_{s}$
$\displaystyle=-\log\int
e^{\mathbb{E}_{\bm{S}_{\bm{J}^{c}}\sim\mathcal{D}}\log
Q_{t|(t-1)}^{\bm{S},\bm{V}_{[s]}}(\tilde{\bm{W}})}\mathrm{d}\tilde{\bm{W}}$
$\displaystyle\overset{(\bullet)}{=}\frac{1}{2}\eta_{t}^{2}\mathbb{E}_{\bm{S}_{\bm{J}^{c}}\sim\mathcal{D}}\frac{|\bm{V}_{t}\cap\bm{J}^{c}|^{2}}{|\bm{V}_{t}|^{2}}\left(\nabla\mathcal{R}_{\mathcal{D}}\left(\bm{W}_{t-1}\right)-\nabla\mathcal{R}_{\bm{S}_{\bm{V}_{t}\cap\bm{J}^{c}}}\left(\bm{W}_{t-1}\right)\right)^{\top}\bm{\Sigma}_{t}(\bm{W}_{t-1})^{-1}$
$\displaystyle\;\;\cdot\left(\nabla\mathcal{R}_{\mathcal{D}}\left(\bm{W}_{t-1}\right)-\nabla\mathcal{R}_{\bm{S}_{\bm{V}_{t}\cap\bm{J}^{c}}}\left(\bm{W}_{t-1}\right)\right)$
$\displaystyle=\frac{1}{2}\eta_{t}^{2}\mathbb{E}_{\bm{S}_{\bm{J}^{c}}\sim\mathcal{D}}\operatorname{tr}\bigg{(}\bm{\Sigma}_{t}(\bm{W}_{t-1})^{-1}\frac{|\bm{V}_{t}\cap\bm{J}^{c}|^{2}}{|\bm{V}_{t}|^{2}}\left(\nabla\mathcal{R}_{\mathcal{D}}\left(\bm{W}_{t-1}\right)-\nabla\mathcal{R}_{\bm{S}_{\bm{V}_{t}\cap\bm{J}^{c}}}\left(\bm{W}_{t-1}\right)\right)^{\top}$
$\displaystyle\;\;\cdot\left(\nabla\mathcal{R}_{\mathcal{D}}\left(\bm{W}_{t-1}\right)-\nabla\mathcal{R}_{\bm{S}_{\bm{V}_{t}\cap\bm{J}^{c}}}\left(\bm{W}_{t-1}\right)\right)\bigg{)}$
$\displaystyle\overset{(\Diamond)}{=}\left\\{\begin{aligned}
&0\qquad\qquad\qquad\qquad\qquad\qquad\qquad\qquad,\bm{V}_{t}\cap\bm{J}^{c}=\emptyset;\\\
&\frac{1}{2}\frac{\eta_{t}^{2}N}{b_{t}(N-1)^{2}}\operatorname{tr}\bigg{(}\bm{\Sigma}_{t}(\bm{W}_{t-1})^{-1}\bm{\Sigma}^{pop}_{\bm{W}_{t-1}}\bigg{)},\bm{V}_{t}\cap\bm{J}^{c}\neq\emptyset.\end{aligned}\right.$
(37)

where Eq. ($\circ$) is due to the definition of
$\tilde{Q}_{t|(t-1)}^{\bm{S}_{\bm{J}},\bm{V}_{[s]}}$ (Eq.(30)), Eq.
($\bullet$) is due to Eq.(35) and Eq. ($\Diamond$) is due to Lemma 6.

Therefore, when $\bm{V}_{t}\cap\bm{J}^{c}=\emptyset$, Eq.(28) is independent
of $\bm{\Sigma}_{s}$. On the other hand, if
$\bm{V}_{t}\cap\bm{J}^{c}\neq\emptyset$, we only need to solve

$\bm{\Sigma}_{s}(\bm{W})^{*}=\arg\min_{\operatorname{tr}(\bm{\Sigma}_{s}(\bm{W}))=c_{s}(\bm{W})}\operatorname{tr}\bigg{(}\bm{\Sigma}_{s}(\bm{W})^{-1}\bm{\Sigma}^{pop}_{\bm{W}}\bigg{)},\text{
subject to Constraint \ref{constraint: trace}}.$ (38)

We complete the proof by solving Problem (38). Specifically, let the
eigenvalues of $\bm{\Sigma}_{\bm{W}}^{pop}$ be $(\omega_{i}^{pop})_{i=1}^{d}$
(the value is by non-increasing order with respect to index) we first fix the
eigenvalues of $\bm{\Sigma}_{s}(\bm{W})$ to be $\alpha_{[d]}$ with
$\alpha_{i}\geq 0$ (the value is by non-increasing order with respect to
index), $i\in[d]$. Then, by Lemma 10, the minimum of
$\operatorname{tr}\bigg{(}\bm{\Sigma}_{s}(\bm{W})^{-1}\bm{\Sigma}^{pop}_{\bm{W}}\bigg{)}$
is achieved when

$\bm{\Sigma}_{s}(\bm{W})\in\left\\{P^{\top}\left(\alpha_{[d]}\right)P:P\text{
is orthogonal and
}\bm{\Sigma}^{pop}_{\bm{W}}=P^{\top}\left(\omega_{[d]}^{pop}\right)P\right\\},$
(39)

and

$\operatorname{tr}\bigg{(}\bm{\Sigma}_{s}(\bm{W})^{-1}\bm{\Sigma}^{pop}_{\bm{W}}\bigg{)}=\sum_{i=1}^{d}\frac{\omega_{i}^{pop}}{\alpha_{i}}.$

We then optimize $\sum_{i=1}^{d}\frac{\omega_{i}^{pop}}{\alpha_{i}}$ under the
constraint $\sum_{i=1}^{d}\alpha_{i}=c_{s}(\bm{W}_{s-1})$. By the Cauchy-
Schwarz inequality,

$c_{s}(\bm{W}_{s-1})\left(\sum_{i=1}^{d}\frac{\omega_{i}^{pop}}{\alpha_{i}}\right)=\left(\sum_{i=1}^{d}\frac{\omega_{i}^{pop}}{\alpha_{i}}\right)\left(\sum_{i=1}^{d}\alpha_{i}\right)\overset{(*)}{\geq}\left(\sum_{i=1}^{d}\sqrt{\omega_{i}^{pop}}\right)^{2},$
(40)

where equality in inequality $(*)$ holds when
$\alpha_{i}^{2}/\omega^{pop}_{i}$ is invariant of $i$. By combining Eq.(39)
and Eq.(40), the proof is completed. ∎

By Lemma 12, the optimal noise covariances $\bm{\Sigma}_{s}$ of all KL
divergence terms
$\mathbb{E}_{\bm{S}_{\bm{J}^{c}}\sim\mathcal{D}}\operatorname{KL}\left({P}_{s|(s-1)}^{\bm{J},\bm{S}_{\bm{J}},\bm{V}_{s}}\left\|Q_{s|(s-1)}^{\bm{S},\bm{V}_{s}}\right.\right)$
are the same regardless of $\bm{V}_{s}$, $\bm{J}$, and $\bm{S}_{\bm{J}}$,
which helps us to obtain Lemma 4.

###### Proof of Lemma 4.

To begin with, denote the optimal noise covariance of first $s$-step in terms
of the generalization bound $\operatorname{Gen}_{s}$ as $\bm{\Sigma}^{s}$
under Constraint 1, i.e.,

$\bm{\Sigma}^{s}_{[s]}\overset{\triangle}{=}\arg\min_{\bm{\Sigma}_{[s]}}\left(\min_{P}\operatorname{Gen}_{s}(P,\bm{\Sigma}_{[s]})\right),\text{
subject to: Constraint \ref{constraint: trace}},$

we also define $Q^{s}$ accordingly as the posterior distribution with noise
covariance $\bm{\Sigma}^{s}$. Also, recall that $P^{s}$ is the optimal prior
in terms of the generalization bound $\operatorname{Gen}_{s}$ under Constraint
1, i.e.,

$P^{s}=\arg\min_{P}\left(\min_{\bm{\Sigma}_{[s]}}\operatorname{Gen}_{s}(P,\bm{\Sigma}_{[s]})\right),\text{
subject to: Constraint \ref{constraint: trace}}.$

We would like to derive the form of $\bm{\Sigma}_{s}^{s}$ and
$P_{s|(s-1)}^{s}$.

Specifically, we have

$\displaystyle
P^{s}_{s|(s-1)}=\operatornamewithlimits{\arg\min}_{P_{s|(s-1)}}\left(\operatorname{Gen}_{s}(P,\bm{\Sigma}^{s}_{[s]})\right),\text{
subject to: $P_{t|(t-1)}=P^{s}_{t|(t-1)}(t<s)$},$

and

$\displaystyle\bm{\Sigma}^{s}_{s}=\operatornamewithlimits{\arg\min}_{\bm{\Sigma}_{s}}\left(\operatorname{Gen}_{s}(P^{s},\bm{\Sigma}_{[s]})\right),\text{
subject to: Constraint \ref{constraint: trace} and
$\bm{\Sigma}_{t}=\bm{\Sigma}^{s}_{t|(t-1)}(t<s)$}.$

That is, to obtain the desired $\bm{\Sigma}^{s}_{s}$ and $P^{s}_{s|(s-1)}$, we
only need to solve

$\min_{\bm{\Sigma}_{s},P^{s}_{s|(s-1)}}\operatorname{Gen}_{s}(P,\bm{\Sigma}_{[s]}),\text{
subject to: $P_{t|(t-1)}=P^{s}_{t|(t-1)}(t<s)$ and
$\bm{\Sigma}_{t}=\bm{\Sigma}^{s}_{t|(t-1)}(t<s)$}.$

On the other hand, with $P_{t|(t-1)}=P^{s}_{t|(t-1)}(t<s)$ and
$\bm{\Sigma}_{t}=\bm{\Sigma}^{s}_{t|(t-1)}(t<s)$ and under Constraint 1, we
have

$\displaystyle\min_{\bm{\Sigma}_{s},P_{s|(s-1)}}\operatorname{Gen}_{s}(P,\bm{\Sigma}_{[s]})$
$\displaystyle=\min_{\bm{\Sigma}_{s},P_{s|(s-1)}}\mathbb{E}_{\bm{S}_{\bm{J}},\bm{V}_{[s]},\bm{J}}\sqrt{\frac{(a_{2}-a_{1})^{2}}{2}\mathbb{E}_{\bm{S}_{\bm{J}^{c}}}\operatorname{KL}\left(P^{\bm{J},\bm{S}_{\bm{J}},\bm{V}_{[s]}}\left\|Q^{\bm{S},\bm{V}_{[s]}}\right.\right)}$
$\displaystyle=\min_{\bm{\Sigma}_{s},P_{s|(s-1)}}\mathbb{E}_{\bm{S}_{\bm{J}},\bm{V}_{[s]},\bm{J}}\sqrt{\frac{(a_{2}-a_{1})^{2}}{2}\mathbb{E}_{\bm{S}_{\bm{J}^{c}}}\sum_{t=1}^{s}\mathbb{E}_{P_{t-1}^{\bm{J},\bm{S}_{\bm{J}},\bm{V}_{[s]}}}\operatorname{KL}\left(P_{t|(t-1)}^{\bm{J},\bm{S}_{\bm{J}},\bm{V}_{s}}\left\|Q_{t|(t-1)}^{\bm{S},\bm{V}_{s}}\right.\right)}$
$\displaystyle=\min_{\bm{\Sigma}_{s},P_{s|(s-1)}}\mathbb{E}_{\bm{S}_{\bm{J}},\bm{V}_{[s]},\bm{J}}\left[\sqrt{\frac{(a_{2}-a_{1})^{2}}{2}\mathbb{E}_{\bm{S}_{\bm{J}^{c}}}\sum_{t=1}^{s}\mathbb{E}_{{P^{s}}_{t-1}^{\bm{J},\bm{S}_{\bm{J}},\bm{V}_{[s]}}}\operatorname{KL}\left({P^{s}}_{t|(t-1)}^{\bm{J},\bm{S}_{\bm{J}},\bm{V}_{s}}\left\|{Q^{s}}_{t|(t-1)}^{\bm{S},\bm{V}_{s}}\right.\right)}\right.$
$\displaystyle\;\;\left.\overline{+\frac{(a_{2}-a_{1})^{2}}{2}\mathbb{E}_{\bm{S}_{\bm{J}^{c}}}\mathbb{E}_{P_{s-1}^{\bm{J},\bm{S}_{\bm{J}},\bm{V}_{[s]}}}\operatorname{KL}\left(P_{s|(s-1)}^{\bm{J},\bm{S}_{\bm{J}},\bm{V}_{s}}\left\|Q_{s|(s-1)}^{\bm{S},\bm{V}_{s}}\right.\right)}\right]$
$\displaystyle\overset{(*)}{\geq}\mathbb{E}_{\bm{S}_{\bm{J}},\bm{V}_{[s]},\bm{J}}\left[\sqrt{\frac{(a_{2}-a_{1})^{2}}{2}\mathbb{E}_{\bm{S}_{\bm{J}^{c}}}\sum_{t=1}^{s}\mathbb{E}_{{P^{s}}_{t-1}^{\bm{J},\bm{S}_{\bm{J}},\bm{V}_{[s]}}}\operatorname{KL}\left({P^{s}}_{t|(t-1)}^{\bm{J},\bm{S}_{\bm{J}},\bm{V}_{s}}\left\|{Q^{s}}_{t|(t-1)}^{\bm{S},\bm{V}_{s}}\right.\right)}\right.$
$\displaystyle\;\;\left.\overline{+\frac{(a_{2}-a_{1})^{2}}{2}\mathbb{E}_{\bm{S}_{\bm{J}^{c}}}\mathbb{E}_{P_{s-1}^{\bm{J},\bm{S}_{\bm{J}},\bm{V}_{[s]}}}\min_{\bm{\Sigma}_{s},P^{\bm{J},\bm{S}_{\bm{J}},\bm{V}_{s}}_{s|(s-1)}}\operatorname{KL}\left(P_{s|(s-1)}^{\bm{J},\bm{S}_{\bm{J}},\bm{V}_{s}}\left\|Q_{s|(s-1)}^{\bm{S},\bm{V}_{s}}\right.\right)}\right].$

By Lemma 12,
$\min_{\bm{\Sigma}_{s},P^{\bm{J},\bm{S}_{\bm{J}},\bm{V}_{s}}_{s|(s-1)}}\operatorname{KL}\left(P_{s|(s-1)}^{\bm{J},\bm{S}_{\bm{J}},\bm{V}_{s}}\left\|Q_{s|(s-1)}^{\bm{S},\bm{V}_{s}}\right.\right)$
is attained at
$\bm{\Sigma}_{s}(\bm{W})=\lambda_{s}(\bm{W})\left(\bm{\Sigma}^{pop}_{\bm{W}}\right)^{\frac{1}{2}}$,
which is not dependent on $\bm{J},\bm{S}_{\bm{J}},\bm{V}_{s}$, and

$P^{\bm{J},\bm{S}_{\bm{J}},\bm{V}_{s}}_{s|(s-1)}\sim\mathcal{N}\left(\bm{W}_{s-1}-\eta_{s}\left(\frac{|\bm{V}_{s}\cap\bm{J}|}{|\bm{V}_{s}|}\nabla\mathcal{R}_{\bm{S}_{\bm{V}_{s}\cap\bm{J}}}\left(\bm{W}_{s-1}\right)+\frac{|\bm{V}_{s}\cap\bm{J}^{c}|}{|\bm{V}_{s}|}\nabla\mathcal{R}_{\mathcal{D}}\left(\bm{W}_{s-1}\right)\right),\lambda_{s}(\bm{W})\left(\bm{\Sigma}^{pop}_{\bm{W}}\right)^{\frac{1}{2}}\right).$

Therefore, Inequality $(*)$ holds, and the proof is completed.

∎

By Lemma 4, we obtain the form of $P^{*}$, i.e.,

${P^{*}}^{\bm{J},\bm{S}_{\bm{J}},\bm{V}_{s}}_{s|(s-1)}\sim\mathcal{N}\left(\bm{W}_{s-1}-\eta_{s}\left(\frac{|\bm{V}_{s}\cap\bm{J}|}{|\bm{V}_{s}|}\nabla\mathcal{R}_{\bm{S}_{\bm{V}_{s}\cap\bm{J}}}\left(\bm{W}_{s-1}\right)+\frac{|\bm{V}_{s}\cap\bm{J}^{c}|}{|\bm{V}_{s}|}\nabla\mathcal{R}_{\mathcal{D}}\left(\bm{W}_{s-1}\right)\right),\lambda_{s}(\bm{W})\left(\bm{\Sigma}^{pop}_{\bm{W}}\right)^{\frac{1}{2}}\right),$

which allows us to further derive Theorem 2.

###### Proof of Theorem 2.

By the definition of $\operatorname{Gen}_{T}$, with prior the greedy prior and
under Constraint 1, we have

$\displaystyle\min_{\bm{\Sigma}_{[T]}}\operatorname{Gen}_{T}(P^{*},\bm{\Sigma}_{[T]})$
$\displaystyle=\min_{\bm{\Sigma}_{[T]}}\mathbb{E}_{\bm{S}_{\bm{J}},\bm{V}_{[T]},\bm{J}}\sqrt{\frac{(a_{2}-a_{1})^{2}}{2}\mathbb{E}_{\bm{S}_{\bm{J}^{c}}}\operatorname{KL}\left({P^{*}}^{\bm{J},\bm{S}_{\bm{J}},\bm{V}_{[T]}}\left\|Q^{\bm{S},\bm{V}_{[T]}}\right.\right)}$
$\displaystyle=\min_{\bm{\Sigma}_{[T]}}\mathbb{E}_{\bm{S}_{\bm{J}},\bm{V}_{[T]},\bm{J}}\sqrt{\frac{(a_{2}-a_{1})^{2}}{2}\mathbb{E}_{\bm{S}_{\bm{J}^{c}}}\sum_{t=1}^{T}\mathbb{E}_{{P^{*}}_{t-1}^{\bm{J},\bm{S}_{\bm{J}},\bm{V}_{[t-1]}}}\operatorname{KL}\left({P^{*}}_{t|(t-1)}^{\bm{J},\bm{S}_{\bm{J}},\bm{V}_{t}}\left\|Q_{t|(t-1)}^{\bm{S},\bm{V}_{t}}\right.\right)}$
$\displaystyle\overset{(\bullet)}{\geq}\mathbb{E}_{\bm{S}_{\bm{J}},\bm{V}_{[T]},\bm{J}}\sqrt{\frac{(a_{2}-a_{1})^{2}}{2}\sum_{t=1}^{T}\mathbb{E}_{{P^{*}}_{t-1}^{\bm{J},\bm{S}_{\bm{J}},\bm{V}_{[t-1]}}}\min_{\bm{\Sigma}_{t}}\mathbb{E}_{\bm{S}_{\bm{J}^{c}}}\operatorname{KL}\left({P^{*}}_{t|(t-1)}^{\bm{J},\bm{S}_{\bm{J}},\bm{V}_{t}}\left\|Q_{t|(t-1)}^{\bm{S},\bm{V}_{t}}\right.\right)}$
$\displaystyle\overset{(*)}{=}\mathbb{E}_{\bm{S}_{\bm{J}},\bm{V}_{[T]},\bm{J}}\sqrt{\frac{(a_{2}-a_{1})^{2}}{2}\sum_{t=1}^{T}\mathbb{E}_{{P^{*}}_{t-1}^{\bm{J},\bm{S}_{\bm{J}},\bm{V}_{[t-1]}}}\min_{\bm{\Sigma}_{t},P_{t|(t-1)}^{\bm{J},\bm{S}_{\bm{J}},\bm{V}_{t}}}\mathbb{E}_{\bm{S}_{\bm{J}^{c}}}\operatorname{KL}\left(P_{t|(t-1)}^{\bm{J},\bm{S}_{\bm{J}},\bm{V}_{t}}\left\|Q_{t|(t-1)}^{\bm{S},\bm{V}_{t}}\right.\right)},$

where Eq. $(*)$ is due to that by the proof of Lemma 4,
${P^{*}}^{\bm{J},\bm{S}_{\bm{J}},\bm{V}_{s}}_{s|(s-1)}$ is the same as the
prior minimizing
$\mathbb{E}_{\bm{S}_{\bm{J}^{c}}\sim\mathcal{D}}\operatorname{KL}\left({P}_{s|(s-1)}^{\bm{J},\bm{S}_{\bm{J}},\bm{V}_{s}}\left\|Q_{s|(s-1)}^{\bm{S},\bm{V}_{s}}\right.\right)$
for any given ${\bm{J},\bm{S}_{\bm{J}},\bm{V}_{s}}$. Therefore,
$\min_{\bm{\Sigma}_{t}}\mathbb{E}_{\bm{S}_{\bm{J}^{c}}}\operatorname{KL}\left({P^{*}}_{t|(t-1)}^{\bm{J},\bm{S}_{\bm{J}},\bm{V}_{t}}\left\|Q_{t|(t-1)}^{\bm{S},\bm{V}_{t}}\right.\right)$
is attained when
$\bm{\Sigma}_{t}(\bm{W})=\lambda_{s}(\bm{W})(\bm{\Sigma}^{pop}_{\bm{W}})^{\frac{1}{2}}$,
which is independent of ${\bm{J},\bm{S}_{\bm{J}},\bm{V}_{s}}$, and Inequality
($\bullet$) holds. Therefore,
$\min_{\bm{\Sigma}_{[T]}}\operatorname{Gen}_{T}(P^{*},\bm{\Sigma}_{[T]})$ is
also attained at
$\bm{\Sigma}_{t}(\bm{W})=\lambda_{s}(\bm{W})(\bm{\Sigma}^{pop}_{\bm{W}})^{\frac{1}{2}}$.

The proof is completed.

∎

##  Appendix E Supplementary materials of Section 5

###  E.1 Formal Description of the Prior in Section 5

In this section, we provide a detailed description of the update rule of the
prior defined by Eq.(10).

Input: Sample set $\bm{S}$ with size $N$, initialization distribution
$\mathcal{W}_{0}$, total step $T$, learning rate $(\eta_{t})_{t=1}^{T}$

Output: $\bm{W}_{[T]}$, $\bm{J}$

1 Initialize $\bm{W}_{0}$ according to $\mathcal{W}_{0}$; initialize $\bm{J}$
by uniformly sampling $N-1$ elements from $[N]$ without replacement; set $t=0$

2 while  _$t <T$_ do

3 Uniformly sample index set $\bm{V}_{t}\subset[N]$ such that
$|\bm{V}_{t}|=b_{t}$ without replacement and independent of $\bm{J}$

4 if  _$\bm{V}_{t}\subset\bm{J}$_ then

5
$\bm{W}_{t}=\bm{W}_{t-1}-\eta_{t}\nabla\mathcal{R}_{\bm{S}_{\bm{V}_{t}}}(\bm{W}_{t-1})+\mathcal{N}\left(\bm{0},\sigma_{t}\mathbb{I}_{d}\right)$

6

7 else

8
$\bm{W}_{t}=\bm{W}_{t-1}-\eta_{t}\frac{b_{t}-1}{b_{t}}\nabla\mathcal{R}_{\bm{S}_{\bm{V}_{t}\cap\bm{J}}}(\bm{W}_{t-1})-\eta_{t}\frac{1}{b_{t}}\nabla\mathcal{R}_{\bm{S}_{\bm{J}}}(\bm{W}_{t-1})+\mathcal{N}\left(\bm{0},\sigma_{t}\mathbb{I}_{d}\right)$

9

10 $t=t+1$

11

Algorithm 1 Iteration of Prior

###  E.2 Calculation of the Generalization Bound

To obtain the optimal noise covariance of (P2), we first derive the explicit
form of the generalization bound $\widetilde{\operatorname{Gen}}_{T}$ with the
prior given by Eq. (10) as the following lemma:

######  Lemma 13 (Calculate $\widetilde{\operatorname{Gen}}_{T}$).

Let Assumption 1 hold. Let the prior $P$ is given by the update rule Eq. (10).
Then, the generalization bound $\widetilde{\operatorname{Gen}}_{T}$ can be
represented as

$\widetilde{\operatorname{Gen}}_{T}=\mathbb{E}_{\bm{S}_{\bm{J}}}\sqrt{\frac{(a_{2}-a_{1})^{2}}{2}\sum_{t=1}^{T}\mathbb{E}_{\bm{S}_{\bm{J}^{c}},\bm{V}_{[t-1]}}\mathbb{E}_{P_{s-1}^{\bm{J},\bm{S}_{\bm{J}},\bm{V}_{[s-1]}}}A_{t}(\bm{S},\bm{W}_{t-1})},$

where $A_{t}(\bm{S},\bm{W})$ is given as

$\displaystyle A_{t}(\bm{S},\bm{W})\overset{\triangle}{=}$
$\displaystyle\frac{1}{2}\left(\sigma_{t}(\bm{W})\operatorname{tr}\left(\bm{\Sigma}_{t}(\bm{S},\bm{W})^{-1}\right)+\ln\left(\operatorname{det}\bm{\Sigma}_{t}(\bm{S},\bm{W})\right)-d\right)$
$\displaystyle+\frac{\eta_{t}^{2}}{2Nb_{t}}\left(\frac{N}{N-1}\right)^{2}\operatorname{tr}\left(\bm{\Sigma}_{t}(\bm{S})^{-1}\bm{\Sigma}^{sd}_{\bm{S},\bm{W}_{t-1}}\right)-\frac{1}{2}d\ln\sigma_{t}(\bm{W}).$

###### Proof.

By the definition of $\widetilde{\operatorname{Gen}}_{T}$, we have

$\widetilde{\operatorname{Gen}}_{T}=\mathbb{E}_{\bm{S}}\sqrt{\frac{(a_{2}-a_{1})^{2}}{2}\mathbb{E}_{\bm{V}_{[T]},\bm{J}}\operatorname{KL}\left(P^{\bm{J},\bm{S}_{\bm{J}},\bm{V}_{[T]}}\left\|Q^{\bm{S},\bm{V}_{[T]}}\right.\right)},$

which by the decomposition of KL divergence (Lemma 1) further leads to

$\displaystyle\widetilde{\operatorname{Gen}}_{T}=$
$\displaystyle\mathbb{E}_{\bm{S}}\sqrt{\frac{(a_{2}-a_{1})^{2}}{2}\mathbb{E}_{\bm{V}_{[T]},\bm{J}}\sum_{s=1}^{T}\mathbb{E}_{P_{s-1}^{\bm{J},\bm{S}_{\bm{J}},\bm{V}_{[T]}}}\operatorname{KL}\left(P_{s|(s-1)}^{\bm{J},\bm{S}_{\bm{J}},\bm{V}_{[T]}}\left\|Q_{s|(s-1)}^{\bm{S},\bm{V}_{[T]}}\right.\right)}$
$\displaystyle=$
$\displaystyle\mathbb{E}_{\bm{S}}\sqrt{\frac{(a_{2}-a_{1})^{2}}{2}\sum_{s=1}^{T}\mathbb{E}_{\bm{V}_{[s]},\bm{J}}\mathbb{E}_{P_{s-1}^{\bm{J},\bm{S}_{\bm{J}},\bm{V}_{[s-1]}}}\operatorname{KL}\left(P_{s|(s-1)}^{\bm{J},\bm{S}_{\bm{J}},\bm{V}_{s}}\left\|Q_{s|(s-1)}^{\bm{S},\bm{V}_{s}}\right.\right)}$
$\displaystyle\overset{(*)}{=}$
$\displaystyle\mathbb{E}_{\bm{S}}\sqrt{\frac{(a_{2}-a_{1})^{2}}{2}\sum_{s=1}^{T}\mathbb{E}_{\bm{V}_{[s-1]}}\mathbb{E}_{P_{s-1}^{\bm{J},\bm{S}_{\bm{J}},\bm{V}_{[s-1]}}}\mathbb{E}_{\bm{V}_{s},\bm{J}}\operatorname{KL}\left(P_{s|(s-1)}^{\bm{J},\bm{S}_{\bm{J}},\bm{V}_{s}}\left\|Q_{s|(s-1)}^{\bm{S},\bm{V}_{s}}\right.\right)}$
$\displaystyle\overset{(**)}{=}$
$\displaystyle\mathbb{E}_{\bm{S}}\sqrt{\frac{(a_{2}-a_{1})^{2}}{2}\sum_{s=1}^{T}\mathbb{E}_{\bm{V}_{[s-1]},\bm{J}}\mathbb{E}_{P_{s-1}^{\bm{J},\bm{S}_{\bm{J}},\bm{V}_{[s-1]}}}\mathbb{E}_{\bm{V}_{s},\bm{J}}\operatorname{KL}\left(P_{s|(s-1)}^{\bm{J},\bm{S}_{\bm{J}},\bm{V}_{s}}\left\|Q_{s|(s-1)}^{\bm{S},\bm{V}_{s}}\right.\right)},$

where in Eq. $(*)$ we exchange the order between
$\mathbb{E}_{P_{s-1}^{\bm{J},\bm{S}_{\bm{J}},\bm{V}_{[s-1]}}}$ and
$\mathbb{E}_{\bm{V}_{s},\bm{J}}$ due to Assumption 1, and Eq. $(**)$ is due to
that
$\mathbb{E}_{P_{s-1}^{\bm{J},\bm{S}_{\bm{J}},\bm{V}_{[s-1]}}}\mathbb{E}_{\bm{V}_{s},\bm{J}}\operatorname{KL}\left(P_{s|(s-1)}^{\bm{J},\bm{S}_{\bm{J}},\bm{V}_{s}}\left\|Q_{s|(s-1)}^{\bm{S},\bm{V}_{s}}\right.\right)$
is independent of $\bm{J}$ by Assumption 1.

Therefore, we only need to prove
$\mathbb{E}_{\bm{V}_{t},\bm{J}}\operatorname{KL}\left(P_{t|(t-1)}^{\bm{J},\bm{S}_{\bm{J}},\bm{V}_{t}}\left\|Q_{t|(t-1)}^{\bm{S},\bm{V}_{t}}\right.\right)=A(t)$,
which can be obtained by

$\displaystyle\mathbb{E}_{\bm{V}_{t},\bm{J}}\operatorname{KL}\left(P_{t|(t-1)}^{\bm{J},\bm{S}_{\bm{J}},\bm{V}_{t}}\left\|Q_{t|(t-1)}^{\bm{S},\bm{V}_{t}}\right.\right)$
$\displaystyle\overset{(\bullet)}{=}\frac{1}{2}\mathbb{E}_{\bm{V}_{t},\bm{J}}\bigg{(}\left(\mu^{\bm{S},\bm{V}_{t},\bm{J},\bm{W}_{t-1}}\right)^{\top}\bm{\Sigma}_{t}(\bm{S},\bm{W}_{t-1})^{-1}\mu^{\bm{S},\bm{V}_{t},\bm{J},\bm{W}_{t-1}}+\ln\frac{\det\bm{\Sigma}_{t}(\bm{S},\bm{W}_{t-1})}{\sigma_{t}(\bm{W}_{t-1})^{d}}$
$\displaystyle\;\;+\operatorname{tr}\left(\sigma_{t}(\bm{W}_{t-1})\bm{\Sigma}_{t}(\bm{S},\bm{W}_{t-1})^{-1}\right)\bigg{)}-\frac{d}{2}$
$\displaystyle=\frac{1}{2}\mathbb{E}_{\bm{V}_{t},\bm{J}}\bigg{(}\operatorname{tr}\left(\bm{\Sigma}_{t}(\bm{S},\bm{W}_{t-1})^{-1}\mu^{\bm{S},\bm{V}_{t},\bm{J},\bm{W}_{t-1}}\left(\mu^{\bm{S},\bm{V}_{t},\bm{J},\bm{W}_{t-1}}\right)^{\top}\right)+\ln\frac{\det\bm{\Sigma}_{t}(\bm{S},\bm{W}_{t-1})}{\sigma_{t}(\bm{W}_{t-1})^{d}}$
$\displaystyle\;\;+\operatorname{tr}\left(\sigma_{t}(\bm{W}_{t-1})\bm{\Sigma}_{t}(\bm{S},\bm{W}_{t-1})^{-1}\right)\bigg{)}-\frac{d}{2}$
$\displaystyle=\frac{1}{2}\operatorname{tr}\left(\bm{\Sigma}_{t}(\bm{S},\bm{W}_{t-1})^{-1}\mathbb{E}_{\bm{J},\bm{V}_{t}}\mu^{\bm{S},\bm{V}_{t},\bm{J},\bm{W}_{t-1}}\left(\mu^{\bm{S},\bm{V}_{t},\bm{J},\bm{W}_{t-1}}\right)^{\top}\right)+\frac{1}{2}\ln\frac{\det\bm{\Sigma}_{t}(\bm{S},\bm{W}_{t-1})}{\sigma_{t}(\bm{W}_{t-1})^{d}}$
$\displaystyle\;\;+\frac{1}{2}\operatorname{tr}\left(\sigma_{t}(\bm{W}_{t-1})\bm{\Sigma}_{t}(\bm{S},\bm{W}_{t-1})^{-1}\right)-\frac{d}{2}$
$\displaystyle\overset{(\circ)}{=}\frac{1}{2}\left(\sigma_{t}(\bm{W}_{t-1})\operatorname{tr}\left(\bm{\Sigma}_{t}(\bm{S},\bm{W}_{t-1})^{-1}\right)+\ln\left(\operatorname{det}\bm{\Sigma}_{t}(\bm{S},\bm{W}_{t-1})\right)-d\right)-\frac{1}{2}d\ln\sigma_{t}(\bm{W}_{t-1})$
$\displaystyle\;\;+\frac{\eta_{t}^{2}}{2Nb_{t}}\left(\frac{N}{N-1}\right)^{2}\operatorname{tr}\left(\bm{\Sigma}_{t}(\bm{S},\bm{W}_{t-1})^{-1}\bm{\Sigma}^{sd}_{\bm{S},\bm{W}_{t-1}}\right),$

where Eq. $(\bullet)$ is due to Lemma 5, where
$\mu^{\bm{S},\bm{V}_{t},\bm{J},\bm{W}_{t-1}}$ is defined by Eq.(27), and Eq.
$(\circ)$ is obtained by Lemma 6.

The proof is completed. ∎

By Lemma 13, for any $t\in[T]$, $\bm{S}$, and $\bm{W}_{t-1}$,
$\operatorname{Gen}_{[T]}$ depend on
$\bm{\Sigma}_{t}(\bm{W},\operatorname{Gen}_{[T]})$ only through
$A_{t}(\bm{S},\bm{W})$, and the solution of optimizing $A_{t}$ with respect to
$\bm{\Sigma}_{t}$ under Constraint 1 has already been given by Lemma 9. We
then complete the proof of Theorem 3 in the next section by combining Lemma 13
and Lemma 9 together.

###  E.3 Proof of Theorem 3

In this section, we first restate Theorem 3 with explicit form of
$\tilde{\omega}_{i}^{\bm{S},\bm{W}}$ (omitted in the main text). We then
provide the proof of the theorem by Lemma 13 and Lemma 9.

######  Theorem 4.

Let prior and posterior be defined as Eq.(10) and Eq.(3), respectively. Then,
with Assumption 1, the solution of (P2) is given by

$\bm{\Sigma}^{*}_{t}(\bm{S},\bm{W})=\bm{Q}_{\bm{S},\bm{W}}^{sd}\operatorname{Diag}(\tilde{\omega}^{\bm{S},\bm{W}}_{t,1},\cdots,\tilde{\omega}^{\bm{S},\bm{W}}_{t,d})\left(\bm{Q}_{\bm{S},\bm{W}}^{sd}\right)^{\top},$

where

$\tilde{\omega}^{\bm{S},\bm{W}}_{t,i}=\frac{\sqrt{1-4\lambda^{*}\left(\mathbb{E}_{\bm{J}}\sigma_{t}(\bm{S}_{\bm{J}},\bm{W})+\frac{\eta_{t}^{2}}{Nb_{t}}\left(\frac{N}{N-1}\right)^{2}\omega^{\bm{S},\bm{W}}_{i}\right)}-1}{-2\lambda^{*}},$

$\lambda^{*}$ is determined by
$\sum_{i=1}^{d}\tilde{\omega}^{\bm{S},\bm{W}}_{t,i}=c_{t}\left(\bm{S},\bm{W}\right)$,
and $\bm{Q}_{\bm{S},\bm{W}}^{sd}$ is the orthogonal matrix that diagonalizes
$\bm{\Sigma}^{sd}_{\bm{S},\bm{W}}$ as

$\bm{\Sigma}^{sd}_{\bm{S},\bm{W}}=\bm{Q}_{\bm{S},\bm{W}}^{sd}\operatorname{Diag}(\omega^{\bm{S},\bm{W}}_{1},\cdots,\omega^{\bm{S},\bm{W}}_{d})\left(\bm{Q}_{\bm{S},\bm{W}}^{sd}\right)^{\top}.$

###### Proof of Theorem 3.

By Lemma 13, $\operatorname{Gen}_{T}$ depends on
$\bm{\Sigma}_{t}(\bm{S},\bm{W})$ only through $A_{t}(\bm{S},\bm{W})$, and we
have

$\displaystyle\bm{\Sigma}^{*}_{s}(\bm{S},\bm{W})$
$\displaystyle=\arg\min_{\text{Constraint }\ref{constraint:
trace}}A_{s}(\bm{S},\bm{W})$
$\displaystyle=\arg\min_{\operatorname{tr}(\bm{\Sigma})=c_{s}(\bm{S},\bm{W})}\frac{1}{2}\bigg{(}\operatorname{tr}\left(\bm{\Sigma}^{-1}\left(\sigma_{s}(\bm{S}_{\bm{J}},\bm{W}_{s-1})\mathbb{I}+\frac{\eta_{s}^{2}}{Nb_{s}}\left(\frac{N}{N-1}\right)^{2}\bm{\Sigma}^{sd}_{\bm{S},\bm{W}_{s-1}}\right)\right)$
$\displaystyle\;\;-d\ln\sigma_{s}(\bm{S}_{\bm{J}},\bm{W}_{s-1})-d+\ln\left(\operatorname{det}\bm{\Sigma}\right)\bigg{)}$
$\displaystyle=\arg\min_{\operatorname{tr}(\bm{\Sigma})=c_{s}(\bm{S},\bm{W})}\frac{1}{2}\bigg{(}\operatorname{tr}\left(\bm{\Sigma}^{-1}\left(\sigma_{s}(\bm{S}_{\bm{J}},\bm{W}_{s-1})\mathbb{I}+\frac{\eta_{s}^{2}}{Nb_{s}}\left(\frac{N}{N-1}\right)^{2}\bm{\Sigma}^{sd}_{\bm{S},\bm{W}_{s-1}}\right)\right)$
$\displaystyle\;\;+\ln\left(\operatorname{det}\bm{\Sigma}\right)\bigg{)}.$

Applying Lemma 9 completes the proof.

∎

###  E.4 Smaller Condition Number

In this section, we demonstrate why the optimal noise of Theorem 3 has smaller
condition number than $\bm{\Sigma}^{sd}$ as the following corollary.

######  Corollary 1.

The optimal noise covariance $\bm{\Sigma}^{*}$ given by Theorem 3 has smaller
condition number than $\bm{\Sigma}^{sd}$.

###### Proof.

We prove this claim following two steps.

Firstly, the noise covariance of the prior is isotropic, has condition number
$1$, and push the condition number of
$\sigma_{t}\mathbb{I}+\frac{\eta_{t}^{2}}{Nb_{t}}\left(\frac{N}{N-1}\right)^{2}\Sigma^{sd}_{\bm{S},\bm{W}}$
smaller than $\Sigma^{sd}_{\bm{S},\bm{W}}$.

Secondly, the optimal solution $G$ of Lemma 9 always has a smaller condition
number than $B$, which implies that $\Sigma_{t}^{*}(\bm{S},\bm{W})$ has
smaller condition number than
$B=\sigma_{t}\mathbb{I}+\frac{\eta_{t}^{2}}{Nb_{t}}\left(\frac{N}{N-1}\right)^{2}\Sigma^{sd}_{\bm{S},\bm{W}}$.
Hence the condition number of $\Sigma_{t}^{*}(\bm{S},\bm{W})$ is smaller than
$\Sigma^{sd}_{\bm{S},\bm{W}}$. ∎

##  Appendix F Experiments

In this section, we introduce the settings of the experiments in Fig. (1)
Fig.(2), Fig.(4), and Fig.(4). We further include an additional experiment
comparing the generalization error between SGLD with square rooted empirical
gradient covariance (SREC-SGLD) (the optimal noise covariance in Theorem 2)
and SGLD with empirical gradient covariance (EC-SGLD) subject to Constraint 1.

###  F.1 Experiment settings

For both Fig. (1), Fig.(2), Fig. (4), and Fig. (4), we adopt the same setting
as the Fashion-MNIST experiment of [40, Section D.3] despite enlarging the
training set. Specifically, we use the $4$-layer convolutional neural network
as our model to conduct multi-class classification on Fashion-MNIST [35].
Concretely, this convolutional neural network can be expressed in order as:
convolutional layer with $10$ channel and filter size $5\times 5$, max-pool
layer with kernel size $2$ and stride $2$, convolutional layer with $10$
channel and filter size $5\times 5$, max-pool with kernel size $2$, two fully
connected layer with width $50$. Our training set consists of 10,000 examples
uniformly sampled without replacement from the Fashion-MNIST dataset. Our
training set is larger than that in [40] (which only contains 1200 samples),
but is still one sixth of the whole Fashion-MNIST dataset due to the
computational burden of gradient descent (without mini-batch) in the SGLD. The
learning rates of all SGLD are set to $0.07$, which is exactly the same as
[40]. We also set the learning rate of SGD in Fig. (1) to $0.07$ for fair
comparison with SGLD.

Empirical gradient covariance: We use top $100$ eigenvalues to approximate the
empirical gradient covariance matrix. Specifically, we decompose the matrix
$\bm{\Sigma}^{sd}_{\bm{S},\bm{W}}$ into
$(Q_{\bm{S},\bm{W}})^{\top}(\omega^{\bm{S},\bm{W}}_{[d]})Q_{\bm{S},\bm{W}}$,
and use
$(Q_{\bm{S},\bm{W}})^{\top}(\omega^{\bm{S},\bm{W}}_{[100]},\bm{0}_{d-100})Q_{\bm{S},\bm{W}}$
to approximate $\bm{\Sigma}^{sd}_{\bm{S},\bm{W}}$.

Noise Scale: In Fig. (1) and Fig.(2), the traces of all SGLDs are set to be
$\operatorname{tr}(\bm{\Sigma}^{sd}_{\bm{S},\bm{W}})$; in Fig. (4), the traces
are set to be $\operatorname{tr}((\bm{\Sigma}_{\bm{S},\bm{W}}^{sd})^{1/2})$
and $5\operatorname{tr}((\bm{\Sigma}_{\bm{S},\bm{W}}^{sd})^{1/2})$,
respectively in (a) and (b); in Fig. (4), the traces are set to be
$\operatorname{tr}((\bm{\Sigma}_{\bm{S},\bm{W}}^{sd})^{1/2})$.

Noise frequency: Similar to [40], we re-estimate the noise structure of all
SGLDs every 10 epochs to ease the computational burden.

###  F.2 Comparison between EC-SGLD and SREC-SGLD

We further conduct an experiment to compare the generalization performance
between Iso-SGLD, EC-SGLD and SREC-SGLD, with the traces of the covariance are
all set to be $5\operatorname{tr}((\bm{\Sigma}_{\bm{S},\bm{W}}^{sd})^{1/2})$,
and all other settings consistent with Appendix F.1. The generalization error
along the iteration of SREC-SGLD, Iso-SGLD, and EC-SGLD is plotted as Fig. 5,
where one can easily observe the generalization error of SREC-SGLD is the
smallest, which supports Theorem 2.

Figure 5: Comparison of generalization error for SGLDs with different noise
structures. Traces of the covariances are all set to be
$5\operatorname{tr}((\bm{\Sigma}_{\bm{S},\bm{W}}^{sd})^{1/2})$.