SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Van-Anh Nguyen
Tung-Long Vuong1,2
Hoang Phan2,3
Dinh Phung 1,2
Trung Le
Thanh-Toan Do
$^{1}$ Department of Data Science and AI, Monash University, Australia
$^{2}$ VinAI, Vietnam
$^{3}$ New York University, United States
{van-anh.nguyen, tung-long.vuong, toan.do, dinh.phung, trunglm}@monash.edu
hvp2011@nyu.edu
Bayesian Neural Networks (BNNs) provide a probabilistic interpretation for deep learning models by imposing a prior distribution over model parameters and inferring a posterior distribution based on observed data. The model sampled from the posterior distribution can be used for providing ensemble predictions and quantifying prediction uncertainty. It is well-known that deep learning models with lower sharpness have better generalization ability. However, existing posterior inferences are not aware of sharpness/flatness in terms of formulation, possibly leading to high sharpness for the models sampled from them. In this paper, we develop theories, the Bayesian setting, and the variational inference approach for the sharpness-aware posterior. Specifically, the models sampled from our sharpness-aware posterior, and the optimal approximate posterior estimating this sharpness-aware posterior, have better flatness, hence possibly possessing higher generalization ability. We conduct experiments by leveraging the sharpness-aware posterior with state-of-the-art Bayesian Neural Networks, showing that the flat-seeking counterparts outperform their baselines in all metrics of interest.
Bayesian Neural Networks (BNNs) provide a way to interpret deep learning models probabilistically. This is done by setting a prior distribution over model parameters and then inferring a posterior distribution over model parameters based on observed data. This allows us to not only make predictions, but also quantify prediction uncertainty, which is useful for many real-world applications. To sample deep learning models from complex and complicated posterior distributions, advanced particle-sampling approaches such as Hamiltonian Monte Carlo (HMC) [41], Stochastic Gradient HMC (SGHMC) [10], Stochastic Gradient Langevin dynamics (SGLD) [58], and Stein Variational Gradient Descent (SVGD) [36] are often used. However, these methods can be computationally expensive, particularly when many models need to be sampled for better ensembles.
To alleviate this computational burden and enable the sampling of multiple deep learning models from posterior distributions, variational inference approaches employ approximate posteriors to estimate the true posterior. These methods utilize approximate posteriors that belong to sufficiently rich families, which are both economical and convenient to sample from. However, the pioneering works in variational inference, such as [21, 5, 33], assume approximate posteriors to be fully factorized distributions, also known as mean-field variational inference. This approach fails to account for the strong statistical dependencies among random weights of neural networks, limiting its ability to capture the complex structure of the true posterior and estimate the true model uncertainty. To overcome this issue, latter works have attempted to provide posterior approximations with richer
expressiveness [61, 52, 53, 54, 20, 45, 55, 30, 48]. These approaches aim to improve the accuracy of the posterior approximation and enable more effective uncertainty quantification.
In the context of standard deep network training, it has been observed that flat minimizers can enhance the generalization capability of models. This is achieved by enabling them to locate wider local minima that are more robust to shifts between train and test sets. Several studies, including [27, 47, 15], have shown evidence to support this principle. However, the posteriors used in existing Bayesian neural networks (BNNs) do not account for the sharpness/flatness of the models derived from them in terms of model formulation. As a result, the sampled models can be located in regions of high sharpness and low flatness, leading to poor generalization ability. Moreover, in variational inference methods, using approximate posteriors to estimate these non-sharpness-aware posteriors can result in sampled models from the corresponding optimal approximate posterior lacking awareness of sharpness/flatness, hence causing them to suffer from poor generalization ability.
In this paper, our objective is to propose a sharpness-aware posterior for learning BNNs, which samples models with high flatness for better generalization ability. To achieve this, we devise both a Bayesian setting and a variational inference approach for the proposed posterior. By estimating the optimal approximate posteriors, we can generate flatter models that improve the generalization ability. Our approach is as follows: In Theorem 3.1, we show that the standard posterior is the optimal solution to an optimization problem that balances the empirical loss induced by models sampled from an approximate posterior for fitting a training set with a Kullback-Leibler (KL) divergence, which encourages a simple approximate posterior. Based on this insight, we replace the empirical loss induced by the approximate posterior with the general loss over the entire data-label distribution in Theorem 3.2 to improve the generalization ability. Inspired by sharpness-aware minimization [16], we develop an upper-bound of the general loss in Theorem 3.2, leading us to formulate the sharpness-aware posterior in Theorem 3.3. Finally, we devise the Bayesian setting and variational approach for the sharpness-aware posterior. Overall, our contributions in this paper can be summarized as follows:
- We propose and develop theories, the Bayesian setting, and the variational inference approach for the sharpness-aware posterior. This posterior enables us to sample a set of flat models that improve the model generalization ability. We note that SAM [16] only considers the sharpness for a single model, while ours is the first work studying the concept and theory of the sharpness for a distribution  $\mathbb{Q}$  over models. Additionally, the proof of Theorem 3.2 is very challenging, elegant, and complicated because of the infinite number of models in the support of  $\mathbb{Q}$ .
- We conduct extensive experiments by leveraging our sharpness-aware posterior with the state-of-the-art and well-known BNNs, including BNNs with an approximate Gaussian distribution [33], BNNs with stochastic gradient Langevin dynamics (SGLD) [58], MC-Dropout [18], Bayesian deep ensemble [35], and SWAG [39] to demonstrate that the flat-seeking counterparts consistently outperform the corresponding approaches in all metrics of interest, including the ensemble accuracy, expected calibration error (ECE), and negative log-likelihood (NLL).
2.1 Bayesian Neural Networks
Markov chain Monte Carlo (MCMC): This approach allows us to sample multiple models from the posterior distribution and was well-known for inference with neural networks through the Hamiltonian Monte Carlo (HMC) [41]. However, HMC requires the estimation of full gradients, which is computationally expensive for neural networks. To make the HMC framework practical, Stochastic Gradient HMC (SGHMC) [10] enables stochastic gradients to be used in Bayesian inference, crucial for both scalability and exploring a space of solutions. Alternatively, stochastic gradient Langevin dynamics (SGLD) [58] employs first-order Langevin dynamics in the stochastic gradient setting. Additionally, Stein Variational Gradient Descent (SVGD) [36] maintains a set of particles to gradually approach a posterior distribution. Theoretically, all SGHMC, SGLD, and SVGD asymptotically sample from the posterior in the limit of infinitely small step sizes.
Variational Inference: This approach uses an approximate posterior distribution in a family to estimate the true posterior distribution by maximizing a variational lower bound. [21] suggests fitting
a Gaussian variational posterior approximation over the weights of neural networks, which was generalized in [32, 33, 5], using the reparameterization trick for training deep latent variable models. To provide posterior approximations with richer expressiveness, many extensive studies have been proposed. Notably, [38] treats the weight matrix as a whole via a matrix variate Gaussian [22] and approximates the posterior based on this parameterization. Several later works have inspected this distribution to examine different structured representations for the variational Gaussian posterior, such as Kronecker-factored [59, 52, 53], k-tied distribution [54], non-centered or rank-1 parameterization [20, 14]. Another recipe to represent the true covariance matrix of Gaussian posterior is through the low-rank approximation [45, 55, 30, 39].
Dropout Variational Inference: This approach utilizes dropout to characterize approximate posteriors. Typically, [18] and [33] use this principle to propose Bayesian Dropout inference methods such as MC Dropout and Variational Dropout. Concrete dropout [19] extends this idea to optimize the dropout probabilities. Variational Structured Dropout [43] employs Householder transformation to learn a structured representation for multiplicative Gaussian noise in the Variational Dropout method.
Flat minimizers have been found to improve the generalization ability of neural networks. This is because they enable models to find wider local minima, which makes them more robust against shifts between train and test sets [27, 47, 15, 44]. The relationship between generalization ability and the width of minima has been investigated theoretically and empirically in many studies, notably [23, 42, 12, 17]. Moreover, various methods seeking flat minima have been proposed in [46, 9, 29, 25, 16, 44]. Typically, [29, 26, 57] investigate the impacts of different training factors such as batch size, learning rate, covariance of gradient, and dropout on the flatness of found minima. Additionally, several approaches pursue wide local minima by adding regularization terms to the loss function [46, 61, 60, 9]. Examples of such regularization terms include softmax output's low entropy penalty [46] and distillation losses [61, 60].
SAM, a method that aims to minimize the worst-case loss around the current model by seeking flat regions, has recently gained attention due to its scalability and effectiveness compared to previous methods [16, 56]. SAM has been widely applied in various domains and tasks, such as meta-learning bi-level optimization [1], federated learning [51], multi-task learning [50], where it achieved tighter convergence rates and proposed generalization bounds. SAM has also demonstrated its generalization ability in vision models [11], language models [3], domain generalization [8], and multi-task learning [50]. Some researchers have attempted to improve SAM by exploiting its geometry [34, 31], additionally minimizing the surrogate gap [62], and speeding up its training time [13, 37]. Regarding the behavior of SAM, [28] empirically studied the difference in sharpness obtained by SAM [16] and SWA [24], [40] showed that SAM is an optimal Bayes relaxation of the standard Bayesian inference with a normal posterior, while [44] proved that distribution robustness [4, 49] is a probabilistic extension of SAM.
In what follows, we present the technicality of our proposed sharpness-aware posterior. Particularly, Section 3.1 introduces the problem setting and motivation for our sharpness-aware posterior. Section 3.2 is dedicated to our theory development, while Section 3.3 is used to describe the Bayesian setting and variational inference approach for our sharpness-aware posterior.
3.1 Problem Setting and Motivation
We aim to develop Sharpness-Aware Bayesian Neural Networks (SA-BNN). Consider a family of neural networks  $f_{\theta}(x)$  with  $\theta \in \Theta$  and a training set  $S = \{(x_1, y_1), \ldots, (x_n, y_n)\}$  where  $(x_i, y_i) \sim \mathcal{D}$ . We wish to learn a posterior distribution  $\mathbb{Q}_S^{SA}$  with the density function  $q^{SA}(\theta | S)$  such that any model  $\theta \sim \mathbb{Q}_S^{SA}$  is aware of the sharpness when predicting over the training set  $S$ .
We depart with the standard posterior
$$
q (\theta \mid \mathcal {S}) \propto \prod_ {i = 1} ^ {n} p (y _ {i} \mid x _ {i}, \mathcal {S}, \theta) p (\theta),
$$
where the prior distribution  $\mathbb{P}$  has the density function  $p(\theta)$  and the likelihood has the form
$$
p \left(y \mid x, \mathcal {S}, \theta\right) \propto \exp \left\{- \frac {\lambda}{| \mathcal {S} |} \ell \left(f _ {\theta} (x), y\right) \right\} = \exp \left\{- \frac {\lambda}{n} \ell \left(f _ {\theta} (x), y\right) \right\}
$$
with the loss function  $\ell$ . The standard posterior  $\mathbb{Q}_S$  has the density function defined as
$$
q (\theta \mid \mathcal {S}) \propto \exp \left\{- \frac {\lambda}{n} \sum_ {i = 1} ^ {n} \ell \left(f _ {\theta} \left(x _ {i}\right), y _ {i}\right) \right\} p (\theta), \tag {1}
$$
where  $\lambda \geq 0$  is a regularization parameter.
We define the general and empirical losses as follows:
$$
\mathcal {L} _ {\mathcal {D}} (\theta) = \mathbb {E} _ {(x, y) \sim \mathcal {D}} [ \ell (f _ {\theta} (x), y) ].
$$
$$
\mathcal {L} _ {\mathcal {S}} (\theta) = \mathbb {E} _ {(x, y) \sim \mathcal {S}} [ \ell (f _ {\theta} (x), y) ] = \frac {1}{n} \sum_ {i = 1} ^ {n} \ell (f _ {\theta} (x _ {i}), y _ {i}).
$$
Basically, the general loss is defined as the expected loss over the entire data-label distribution  $\mathcal{D}$ , while the empirical loss is defined as the empirical loss over a specific training set  $S$ .
The standard posterior in Eq. (1) can be rewritten as
$$
q (\theta \mid \mathcal {S}) \propto \exp \left\{- \lambda \mathcal {L} _ {\mathcal {S}} (\theta) \right\} p (\theta). \tag {2}
$$
Given a distribution  $\mathbb{Q}$  with the density function  $q(\theta)$  over the model parameters  $\theta \in \Theta$ , we define the empirical and general losses over this model distribution  $\mathbb{Q}$  as
$$
\mathcal {L} _ {\mathcal {S}} (\mathbb {Q}) = \int_ {\Theta} \mathcal {L} _ {\mathcal {S}} (\theta) d \mathbb {Q} (\theta) = \int_ {\Theta} \mathcal {L} _ {\mathcal {S}} (\theta) q (\theta) d \theta .
$$
$$
\mathcal {L} _ {\mathcal {D}} (\mathbb {Q}) = \int_ {\Theta} \mathcal {L} _ {\mathcal {D}} (\theta) d \mathbb {Q} (\theta) = \int_ {\Theta} \mathcal {L} _ {\mathcal {D}} (\theta) q (\theta) d \theta .
$$
Specifically, the general loss over the model distribution  $\mathbb{Q}$  is defined as the expectation of the general losses incurred by the models sampled from this distribution, while the empirical loss over the model distribution  $\mathbb{Q}$  is defined as the expectation of the empirical losses incurred by the models sampled from this distribution.
3.2 Our Theory Development
We now present the theory development for the sharpness-aware posterior whose proofs can be found in the supplementary material. Inspired by the Gibbs form of the standard posterior  $\mathbb{Q}_S$  in Eq. (2), we establish the following theorem to connect the standard posterior  $\mathbb{Q}_S$  with the density  $q(\theta \mid S)$  and the empirical loss  $\mathcal{L}_S(\mathbb{Q})$  [7, 2].
Theorem 3.1. Consider the following optimization problem
$$
\min  _ {\mathbb {Q} <   <   \mathbb {P}} \left\{\lambda \mathcal {L} _ {S} (\mathbb {Q}) + K L (\mathbb {Q}, \mathbb {P}) \right\}, \tag {3}
$$
where we search over  $\mathbb{Q}$  absolutely continuous w.r.t.  $\mathbb{P}$  and  $KL(\cdot ,\cdot)$  is the Kullback-Leibler divergence. This optimization has a closed-form optimal solution  $\mathbb{Q}^*$  with the density
$$
q ^ {*} (\theta) \propto \exp \left\{- \lambda \mathcal {L} _ {\mathcal {S}} (\theta) \right\} p (\theta),
$$
which is exactly the standard posterior  $\mathbb{Q}_S$  with the density  $q(\theta \mid S)$ .
Theorem 3.1 reveals that we need to find the posterior  $\mathbb{Q}_S$  balancing between optimizing its empirical loss  $\mathcal{L}_S(\mathbb{Q})$  and simplicity via  $KL(\mathbb{Q},\mathbb{P})$ . However, minimizing the empirical loss  $\mathcal{L}_S(\mathbb{Q})$  only ensures the correct predictions for the training examples in  $S$ , hence possibly encountering overfitting. Therefore, it is desirable to replace the empirical loss by the general loss to combat overfitting.
To mitigate overfitting, in (3), we replace the empirical loss by the general loss and solve the following optimization problem (OP):
$$
\min  _ {\mathbb {Q} <   <   \mathbb {P}} \left\{\lambda \mathcal {L} _ {\mathcal {D}} (\mathbb {Q}) + K L (\mathbb {Q}, \mathbb {P}) \right\}. \tag {4}
$$
Notably, solving the optimization problem (OP) in (4) is generally intractable. To make it tractable, we find its upper-bound which is relevant to the sharpness of a distribution  $\mathbb{Q}$  over models as shown in the following theorem.
Theorem 3.2. Assume that  $\Theta$  is a compact set. Under some mild conditions, given any  $\delta \in [0;1]$  with the probability at least  $1 - \delta$  over the choice of  $\mathcal{S} \sim \mathcal{D}^n$ , for any distribution  $\mathbb{Q}$ , we have
$$
\mathcal {L} _ {\mathcal {D}} \left(\mathbb {Q}\right) \leq \mathbb {E} _ {\theta \sim \mathbb {Q}} \left[ \max  _ {\theta^ {\prime}: \| \theta^ {\prime} - \theta \| \leq \rho} \mathcal {L} _ {\mathcal {S}} \left(\theta^ {\prime}\right) \right] + f \left(\max  _ {\theta \in \Theta} \| \theta \| ^ {2}, n\right),
$$
where  $f$  is a non-decreasing function w.r.t. the first variable and approaches 0 when the training size  $n$  approaches  $\infty$ .
We note that the proof of Theorem 3.2 is not a trivial extension of sharpness-aware minimization because we need to tackle the general and empirical losses over a distribution  $\mathbb{Q}$ . To make explicit our sharpness over a distribution  $\mathbb{Q}$  on models, we rewrite the upper-bound of the inequality as
$$
\mathbb {E} _ {\theta \sim \mathbb {Q}} \left[ \max _ {\theta^ {\prime}: \| \theta^ {\prime} - \theta \| \leq \rho} \mathcal {L} _ {\mathcal {S}} (\theta^ {\prime}) - \mathcal {L} _ {\mathcal {S}} (\theta) \right] + \mathcal {L} _ {\mathcal {S}} (\mathbb {Q}) + f \left(\max _ {\theta \in \Theta} \| \theta \| ^ {2}, n\right),
$$
where the first term  $\mathbb{E}_{\theta \sim \mathbb{Q}}\left[\max_{\theta^{\prime}:\| \theta^{\prime} - \theta \| \leq \rho}\mathcal{L}_S(\theta^{\prime}) - \mathcal{L}_S(\theta)\right]$  can be regarded as the sharpness over the distribution  $\mathbb{Q}$  on the model space and the last term  $f\left(\max_{\theta \in \Theta}\|\theta\|^2,n\right)$  is a constant.
Moreover, inspired by Theorem 3.2, we propose solving the following OP which forms an upper-bound of the desirable OP in (4)
$$
\min  _ {\mathbb {Q} <   <   \mathbb {P}} \left\{\lambda \mathbb {E} _ {\theta \sim \mathbb {Q}} \left[ \max  _ {\theta^ {\prime}: \| \theta^ {\prime} - \theta \| \leq \rho} \mathcal {L} _ {\mathcal {S}} \left(\theta^ {\prime}\right) \right] + K L (\mathbb {Q}, \mathbb {P}) \right\}. \tag {5}
$$
The following theorem characterizes the optimal solution of the OP in (5).
Theorem 3.3. The optimal solution the OP in (5) is the sharpness-aware posterior distribution  $\mathbb{Q}_S^{SA}$  with the density function  $q^{SA}(\theta | \mathcal{S})$ :
$$
q ^ {S A} (\theta | \mathcal {S}) \propto \exp \left\{- \lambda \max _ {\theta^ {\prime}: \| \theta^ {\prime} - \theta \| \leq \rho} \mathcal {L} _ {\mathcal {S}} (\theta^ {\prime}) \right\} p (\theta) = \exp \left\{- \lambda \mathcal {L} _ {\mathcal {S}} (s (\theta)) \right\} p (\theta),
$$
where we have defined  $s(\theta) = \arg\max \mathcal{L}_S(\theta')$ .
$$
\theta^ {\prime}: \| \theta^ {\prime} - \theta \| \leq \rho
$$
Theorem 3.3 describes the close form of the sharpness-aware posterior distribution  $\mathbb{Q}_S^{SA}$  with the density function  $q^{SA}(\theta | \mathcal{S})$ . Based on this characterization, in what follows, we introduce the SA Bayesian setting that sheds lights on its variational approach.
3.3 Sharpness-Aware Bayesian Setting and Its Variational Approach
Bayesian Setting: To promote the Bayesian setting for sharpness-aware posterior distribution  $\mathbb{Q}_S^{SA}$ , we examine the sharpness-aware likelihood
$$
p ^ {S A} \left(y \mid x, \mathcal {S}, \theta\right) \propto \exp \left\{- \frac {\lambda}{| \mathcal {S} |} \ell \left(f _ {s (\theta)} (x), y\right) \right\} = \exp \left\{- \frac {\lambda}{n} \ell \left(f _ {s (\theta)} (x), y\right) \right\},
$$
where  $s(\theta) = \mathrm{argmax}\mathcal{L}_S(\theta ')$
With this predefined sharpness-aware likelihood, we can recover the sharpness-aware posterior distribution  $\mathbb{Q}_S^{SA}$  with the density function  $q^{SA}(\theta|\mathcal{S})$ :
$$
q ^ {S A} (\theta | \mathcal {S}) \propto \prod_ {i = 1} ^ {n} p ^ {S A} \left(y _ {i} \mid x _ {i}, \mathcal {S}, \theta\right) p (\theta).
$$
Variational inference for the sharpness-aware posterior distribution: We now develop the variational inference for the sharpness-aware posterior distribution. Let denote  $X = [x_{1},\dots,x_{n}]$  and
$Y = [y_{1},\dots,y_{n}]$ . Considering an approximate posterior family  $\{q_{\phi}(\theta):\phi \in \Phi \}$ , we have
$$
\begin{array}{l} \log p ^ {S A} (Y \mid X, \mathcal {S}) = \int_ {\Theta} q _ {\phi} (\theta) \log p ^ {S A} (Y \mid X, \mathcal {S}) d \theta \\ = \int_ {\Theta} q _ {\phi} (\theta) \log \frac {p ^ {S A} (Y | \theta , X , \mathcal {S}) p (\theta)}{q _ {\phi} (\theta)} \frac {q _ {\phi} (\theta)}{q ^ {S A} (\theta | \mathcal {S})} d \theta \\ = \mathbb {E} _ {q _ {\phi} (\theta)} \left[ \sum_ {i = 1} ^ {n} \log p ^ {S A} \left(y _ {i} \mid x _ {i}, \mathcal {S}, \theta\right) \right] - K L \left(q _ {\phi}, p\right) + K L \left(q _ {\phi}, q ^ {S A}\right). \\ \end{array}
$$
It is obvious that we need to maximize the following lower bound for maximally reducing the gap  $KL(q_{\phi},q^{SA})$ :
$$
\max _ {q _ {\phi}} \left\{\mathbb {E} _ {q _ {\phi} (\theta)} \left[ \sum_ {i = 1} ^ {n} \log p ^ {S A} (y _ {i} | x _ {i}, \mathcal {S}, \theta) \right] - K L (q _ {\phi}, p) \right\},
$$
which can be equivalently rewritten as
$$
\min  _ {q _ {\phi}} \left\{\lambda \mathbb {E} _ {q _ {\phi} (\theta)} \left[ \mathcal {L} _ {\mathcal {S}} (s (\theta)) \right] + K L (q _ {\phi}, p) \right\} \text {o r}
$$
$$
\min  _ {q _ {\phi}} \left\{\lambda \mathbb {E} _ {q _ {\phi} (\theta)} \left[ \max  _ {\theta^ {\prime}: \| \theta^ {\prime} - \theta \| \leq \rho} \mathcal {L} _ {\mathcal {S}} \left(\theta^ {\prime}\right) \right] + K L \left(q _ {\phi}, p\right) \right\}. \tag {6}
$$
Derivation for Variational Approach with A Gaussian Approximate Posterior: Inspired by the geometry-based SAM approaches [34, 31], we incorporate the geometry to the SA variational approach via the distance to define the ball for the sharpness as  $\| \theta^{\prime} - \theta \|_{\mathrm{diag}(T_{\theta})} = \sqrt{(\theta^{\prime} - \theta)^{T}\mathrm{diag}(T_{\theta})^{-1}(\theta^{\prime} - \theta)}$  as
$$
\min  _ {q _ {\phi}} \left\{\lambda \mathbb {E} _ {q _ {\phi} (\theta)} \left[ \max  _ {\theta^ {\prime}: \| \theta^ {\prime} - \theta \| _ {\operatorname {d i a g} (T _ {\theta})} \leq \rho} \mathcal {L} _ {\mathcal {S}} \left(\theta^ {\prime}\right) \right] + K L \left(q _ {\phi}, p\right) \right\}.
$$
To further clarify, we consider our SA posterior distribution to Bayesian NNs, wherein we impose the Gaussian distributions to its weight matrices  $W_{i} \sim \mathcal{N}\left(\mu_{i},\sigma_{i}^{2}\mathbb{I}\right)$ ,  $i = 1,\dots ,L^{1}$ . The parameter  $\phi$  consists of  $\mu_{i},\sigma_{i},i = 1,\ldots ,L$ . For  $\theta = W_{1:L} \sim q_{\phi}$ , using the reparameterization trick  $W_{i} = \mu_{i} + \mathrm{diag}(\sigma_{i})\epsilon_{i},\epsilon_{i} \sim \mathcal{N}(0,\mathbb{I})$  and by searching  $\theta^{\prime} = W_{1:L}^{\prime}$  with  $W_{i}^{\prime} = \mu_{i}^{\prime} + \mathrm{diag}(\sigma_{i})\epsilon_{i},\epsilon_{i} \sim \mathcal{N}(0,\mathbb{I})$ , the constraint  $\| \theta -\theta^{\prime}\|_{\mathrm{diag}(T_{\theta})} = \| \mu -\mu^{\prime}\|_{\mathrm{diag}(T_{\theta})}$  with  $\mu = \mu_{1:L}$  and  $\mu^{\prime} = \mu_{1:L}^{\prime}$ . Thus, the OP in (6) reads