Flat Seeking Bayesian Neural Networks
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.
Van-Anh Nguyen, Tung-Long Vuong, Hoang Phan, Thanh-Toan Do, Dinh Phung, Trung Le
# Flat Seeking Bayesian Neural Networks

###### Abstract

Bayesian Neural Networks (BNNs) offer a probabilistic interpretation for deep learning models by imposing a prior distribution over model parameters and inferencing 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 a lower sharpness have a better generalization ability. Nonetheless, existing posterior inferences are not aware of sharpness/flatness, hence possibly leading to high-sharpness for the models sampled from it. 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 a better flatness, hence possibly possessing a higher generalization ability. We conduct experiments by leveraging the sharpness-aware posterior with the state-of-the-art Bayesian Neural Networks, showing that the flat-seeking counterparts outperform their baselines in all metrics of interest.

Machine Learning, Bayesian Neural Networks, Bayesian Neural Networks

## 1 Introduction

Bayesian Neural Networks (BNNs) (Neal, 2012) offer a probabilistic interpretation for deep learning models by imposing a prior distribution over model parameters and then making inference posterior distribution over model parameters from observed data. The models sampled from the posterior distribution allow us to not only make predictions but also quantify prediction uncertainty which is valuable for many real-world applications.

To sample deep learning models from complex and complicated posterior distributions, we can use Hamiltonian Monte Carlo (HMC) (Neal, 1996) or other advanced particle-sampling approaches, notably Stochastic Gradient HMC (SGHMC) (Chen et al., 2014), Stochastic Gradient Langevin dynamics (SGLD) (Welling and Teh, 2011), and Stein Variational Gradient Descent (SVGD) (Liu and Wang, 2016). Although showing the ability to relieve the computational burden of HMC, the latter particle-sampling approaches are still computationally expensive, especially when needing to sample many models for a better ensemble, because it is computationally prohibitive to maintain and manipulate many deep learning models in the memory at the same time.

To make it more economic to sample multiple deep learning models from posterior distributions and alleviate the computational burden of BNN inferences, variational inference approaches employ approximate posteriors to estimate a true posterior. The essence is to use sufficiently rich approximate families whose elements are economic and convenient to be sampled from. Pioneering works (Graves, 2011; Blundell et al., 2015; Kingma et al., 2015) in variational inference assume approximate posteriors to be fully factorized distributions, also called mean-field variational inference, hence ignoring the strong statistical dependencies among random weights of neural nets and leading to the inability to capture the complicated structure of the true posterior and to estimate the true model uncertainty. To overcome this issue, latter works (Zhang et al., 2018; Ritter et al., 2018; Rossi et al., 2020; Swiatkowski et al., 2020; Ghosh et al., 2018; Ong et al., 2018; Tomczak et al., 2020; Khan et al., 2018) attempt to provide posterior approximations with a richer expressiveness.

For the standard training, flat minimizers have been found to improve the generalization ability of deep learning models because they enable models to find wider local minima, which is more robust against shifts between train and test sets (Jiang et al., 2020; Petzka et al., 2021; Dziugaite and Roy, 2017). Although seeking flat minimizers is a well-known principle to improve the generalization ability of deep learning models, the posteriors used in existing BNNs cannot be aware of the sharpness/flatness of the models drawn from them. Consequently, the models sampled from these posteriors can fall into regions with high sharpness and low flatness, hence possibly possessing an insufficient generalization ability. Moreover, in variational inference approaches, when dedicating approximate posteriors to estimate these non
sharpness-aware posteriors, the models sampled from the corresponding optimal approximate posterior are infeasible to be aware of the sharpness/flatness, hence again possibly possessing an insufficient generalization ability.

In this paper, in the context of learning BNNs, we aim to develop a sharpness-aware posterior, by which the models sampled from this posterior have high flatness for a better generalization ability. Moreover, we also devise the Bayesian setting and the variational inference approach for this sharpness-aware posterior. Accordingly, the optimal approximate posteriors estimating this sharpness-aware posterior can generate flatter models for a better ensemble and further improve the generalization ability.

Specifically, our development pathway is as follows. In Theorem 3.1, we point out that the standard posterior is the optimal solution of an optimization problem, that balances between the empirical loss induced by the models sampled from an approximate posterior for fitting a training set and a Kullback-Leibler (KL) divergence for encouraging a simple approximate posterior. This motivates us to replace the empirical loss induced by the approximate posterior with the corresponding general loss over the entire data-label distribution for boosting the generalization ability. Inspired by sharpness-aware minimization (Foret et al., 2021), we develop an upper-bound of the general loss in Theorem 3.2, hinting us the formulation of the sharpness-aware posterior in Theorem 3.3. Finally, we develop 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 conduct extensive experiments by leveraging our sharpness-aware posterior with the state-of-the-art and well-known BNNs, including SWAG (Maddox et al., 2019), MC-Dropout (Gal and Ghahramani, 2016), and Bayesian deep ensemble (Lakshminarayanan et al., 2017), 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).

Our paper is organized as follows. In Section 2, we introduce the literature review in Bayesian Neural Networks and flat minima. Section 3 is dedicated to our proposed sharpness-aware posterior. Finally, Section 4 presents the experiments and ablation studies, followed by the conclusion.

## 2 Related Work

### 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) (Neal, 1996). 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) (Chen et al., 2014) 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) (Welling and Teh, 2011) employs first-order Langevin dynamics in the stochastic gradient setting. Additionally, Stein Variational Gradient Descent (SVGD) (Liu and Wang, 2016) 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. (Graves, 2011) suggests fitting a Gaussian variational posterior approximation over the weights of neural networks, which was generalized in (Kingma and Welling, 2013; Kingma et al., 2015; Blundell et al., 2015), using the reparameterization trick for training deep latent variable models. To provide posterior approximations with richer expressiveness, many extensive studies have been proposed. Notably, (Louizos and Welling, 2017) treats the weight matrix as a whole via a matrix variate Gaussian (Gupta and Nagar, 2018) and approximates the posterior based on this parametrization.

Several later works have inspected this distribution to examine different structured representations for the variational Gaussian posterior, such as Kronecker-factored (Zhang et al., 2018; Ritter et al., 2018; Rossi et al., 2020), k-tied distribution (Swiatkowski et al., 2020), non-centered or rank-1 parameterization (Ghosh et al., 2018; Dusenberry et al., 2020). Another recipe to represent the true covariance matrix of Gaussian posterior is through the low-rank approximation (Ong et al., 2018; Tomczak et al., 2020; Khan et al., 2018; Maddox et al., 2019).

**Dropout Variational Inference:** This approach utilizes dropout to characterize approximate posteriors. Typically, (Gal and Ghahramani, 2016) and (Kingma et al., 2015) use this principle to propose Bayesian Dropout inference meth
ods such as MC Dropout and Variational Dropout. Concrete dropout (Gal et al., 2017) extends this idea to optimize the dropout probabilities. Variational Structured Dropout (Nguyen et al., 2021) employs Householder transformation to learn a structured representation for multiplicative Gaussian noise in the Variational Dropout method.

### Flat Minima

Flat minimizers have been found to improve the generalization ability of neural networks because they enable models to find wider local minima, by which the models will be more robust against shifts between train and test sets (Jiang et al., 2020; Petzka et al., 2021; Dziugaite & Roy, 2017). The relationship between generalization ability and the width of minima is theoretically and empirically investigated in many studies, notably (Hochreiter & Schmidhuber, 1994; Neyshabur et al., 2017; Dinh et al., 2017; Fort & Ganguli, 2019). Moreover, a variety of methods seeking flat minima have been proposed in (Pereyra et al., 2017; Chaudhari et al., 2017; Keskar et al., 2017; Izmailov et al., 2018; Foret et al., 2021).

Typically, (Keskar et al., 2017; Jastrzebski et al., 2017; Wei et al., 2020) investigate the impacts of different training factors, such as batch-size, learning rate, the covariance of gradient, dropout, on the flatness of found minima. Additionally, several approaches pursue wide local minima by adding regularization terms to the loss function (Pereyra et al., 2017; Zhang et al., 2018; Zhu et al., 2019; Chaudhari et al., 2017), e.g., softmax output's low entropy penalty, (Pereyra et al., 2017), distillation losses (Zhang et al., 2018; Zhu et al., 2019).

Recently, SAM (Foret et al., 2021), a method that seeks flat regions by explicitly minimizing the worst-case loss around the current model, has received significant attention due to its effectiveness and scalability compared to previous methods. Particularly, it has been exploited in a variety of tasks and domains (Cha et al., 2021; Abbas et al., 2022; Qu et al., 2022; Caldarola et al., 2022; Bahri et al., 2022; Chen et al., 2021). A notable example is the improvement that SAM brings to meta-learning bi-level optimization in (Abbas et al., 2022). Another application of SAM is in federated learning (FL) (Qu et al., 2022) in which the authors achieved tighter convergence rates than existing FL works, and proposed a generalization bound for the global model.

In addition, SAM shows its generalization ability in vision models (Chen et al., 2021), language models (Bahri et al., 2022), domain generalization (Cha et al., 2021), and multi-task learning (Phan et al., 2022). Some other works attempt to improve by exploiting geometry SAM (Kwon et al., 2021), additionally minimizing surrogate gap (Zhuang et al., 2022), and speeding up SAM training time (Du et al., 2022; Liu et al., 2022).

## 3 Proposed Framework

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.

### 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 \(\mathcal{S}=\{(x_{1},y_{1}),...,(x_{n},y_{n})\}\) where \((x_{i},y_{i})\sim\mathcal{D}\). We wish to learn a posterior distribution \(\mathbb{Q}_{\mathcal{S}}^{SA}\) with the density function \(q^{SA}(\theta|\mathcal{S})\) such that any model \(\theta\sim\mathbb{Q}^{SA}\) is aware of the sharpness when predicting over the training set \(\mathcal{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}_{\mathcal{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}}\left(\theta\right)=\mathbb{E}_{(x,y)\sim\mathcal{D}} \left[\ell\left(f_{\theta}\left(x\right),y\right)\right].\]

\[\mathcal{L}_{\mathcal{S}}\left(\theta\right)=\mathbb{E}_{(x,y)\sim\mathcal{S}} \left[\ell\left(f_{\theta}\left(x\right),y\right)\right]=\frac{1}{n}\sum_{i=1}^ {n}\ell\left(f_{\theta}\left(x_{i}\right),y_{i}\right).\]

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 \(\mathcal{S}\).

The standard posterior in Eq. (1) can be rewritten as

\[q(\theta\mid\mathcal{S})\propto\exp\left\{-\lambda\mathcal{L}_{\mathcal{S}} \left(\theta\right)\right\}p(\theta). \tag{2}\]
Given a distribution \(\mathbb{Q}\) with the density function \(q\left(\theta\right)\) 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}}\left(\mathbb{Q}\right) =\int_{\Theta}\mathcal{L}_{\mathcal{S}}\left(\theta\right)d \mathbb{Q}\left(\theta\right)=\int_{\Theta}\mathcal{L}_{\mathcal{S}}\left( \theta\right)q\left(\theta\right)d\theta.\] \[\mathcal{L}_{\mathcal{D}}\left(\mathbb{Q}\right) =\int_{\Theta}\mathcal{L}_{\mathcal{D}}\left(\theta\right)d \mathbb{Q}\left(\theta\right)=\int_{\Theta}\mathcal{L}_{\mathcal{D}}\left( \theta\right)q\left(\theta\right)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.

### Our Theory Development

We now present the theory development for the sharpness-aware posterior. Inspired by the Gibbs form of the standard posterior \(\mathbb{Q}_{\mathcal{S}}\) in Eq. (2), we establish the following theorem to connect the standard posterior \(\mathbb{Q}_{\mathcal{S}}\) with the density \(q(\theta\mid\mathcal{S})\) and the empirical loss \(\mathcal{L}_{\mathcal{S}}\left(\mathbb{Q}\right)\)(Catoni, 2007; Alquier et al., 2016).

**Theorem 3.1**.: _Consider the following optimization problem_

\[\min_{\mathbb{Q}<<\mathbb{P}}\left\{\lambda\mathcal{L}_{\mathcal{S}}\left( \mathbb{Q}\right)+KL\left(\mathbb{Q},\mathbb{P}\right)\right\}, \tag{3}\]

_where we search over \(\mathbb{Q}\) absolutely continuous w.r.t. \(\mathbb{P}\) and \(KL\left(\cdot,\cdot\right)\) is the Kullback-Leibler divergence. This optimization has a closed-form optimal solution \(\mathbb{Q}^{*}\) with the density_

\[q^{*}\left(\theta\right)\propto\exp\left\{-\lambda\mathcal{L}_{\mathcal{S}} \left(\theta\right)\right\}p(\theta),\]

_which is exactly the standard posterior \(\mathbb{Q}_{\mathcal{S}}\) with the density \(q(\theta\mid\mathcal{S})\)._

Theorem 3.1 whose proof can be found in Appendix A.1 reveals that we need to find the posterior \(\mathbb{Q}_{\mathcal{S}}\) that balances between optimizing its empirical loss \(\mathcal{L}_{\mathcal{S}}\left(\mathbb{Q}\right)\) and simplicity via \(KL\left(\mathbb{Q},\mathbb{P}\right)\). However, minimizing the empirical loss \(\mathcal{L}_{\mathcal{S}}\left(\mathbb{Q}\right)\) only ensures the correct predictions for the training examples in \(\mathcal{S}\) and might encounter overfitting. Hence, it is natural to replace the empirical loss by the general loss to combat overfitting.

To mitigate overfitting, in (7), we replace the empirical loss by the general loss and desire to solve the following optimization problem (OP):

\[\min_{\mathbb{Q}<<\mathbb{P}}\left\{\lambda\mathcal{L}_{\mathcal{D}}\left( \mathbb{Q}\right)+KL\left(\mathbb{Q},\mathbb{P}\right)\right\}. \tag{4}\]

However, solving the optimization problem (OP) in (4) is generally intractable. To make it tractable, we find its upperbound which is relevant to the sharpness 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\)._

The proof of Theorem 3.2 can be found in Appendix A.1. Here we note that this proof is not a trivial extension of sharpness-aware minimization because we need to tackle the general and empirical losses over a distribution \(\mathbb{Q}\). 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]+KL\left(\mathbb{ Q},\mathbb{P}\right)\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^{SA}(\theta|\mathcal{S}) \propto\exp\left\{-\lambda\max_{\theta^{\prime}:\|\theta^{\prime} -\theta\|\leq\rho}\mathcal{L}_{\mathcal{S}}\left(\theta^{\prime}\right) \right\}p\left(\theta\right)\] \[=\exp\left\{-\lambda\mathcal{L}_{\mathcal{S}}\left(s\left(\theta \right)\right)\right\}p\left(\theta\right),\]

_where we have defined \(s\left(\theta\right)=\underset{\sigma^{\prime}:\|\theta^{\prime}-\theta\|\leq \rho}{\text{argmax}}\mathcal{L}_{\mathcal{S}}\left(\theta^{\prime}\right)\)._

Theorem 3.3 whose proof can be found in Appendix A.1 describes the close form of the sharpness-aware posterior distribution \(\mathbb{Q}_{\mathcal{S}}^{SA}\) with the density function \(q^{SA}(\theta|\mathcal{S})\). Based on this characterization, in what follows, we introduce the sharpness-aware Bayesian setting that sheds limits on its variational approach.

### 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^{SA}\left(y\mid x,\mathcal{S},\theta\right) \propto\exp\left\{-\frac{\lambda}{|\mathcal{S}|}\ell\left(f_{s \left(\theta\right)}(x),y\right)\right\}\] \[=\exp\left\{-\frac{\lambda}{n}\ell\left(f_{s\left(\theta\right)}( x),y\right)\right\},\]

where \(s\left(\theta\right)=\underset{\theta^{\prime}:\|\theta^{\prime}-\theta\|\leq \rho}{\text{argmax}}\mathcal{L}_{\mathcal{S}}\left(\theta^{\prime}\right)\).
With this predefined sharpness-aware likelihood, we can recover the sharpness-aware posterior distribution \(\mathbb{Q}_{\mathcal{S}}^{SA}\) with the density function \(q^{SA}(\theta|\mathcal{S})\):

\[q^{SA}(\theta|\mathcal{S})\propto\prod_{i=1}^{n}p^{SA}\left(y_{i}\mid x_{i}, \mathcal{S},\theta\right)p\left(\theta\right).\]

Variational inference for the sharpness-aware posterior distribution:We now develop the variational inference for the sharpness-aware posterior distribution. Let denote \(X=\left[x_{1},...,x_{n}\right]\) and \(Y=\left[y_{1},...,y_{n}\right]\). Considering an approximate posterior family \(\left\{q_{\phi}\left(\theta\right):\phi\in\Phi\right\}\), we have

\[\log p^{SA}\left(Y\mid X,\mathcal{S}\right)=\int_{\Theta}q_{\phi} \left(\theta\right)\log p^{SA}\left(Y\mid X,\mathcal{S}\right)d\theta\] \[=\int_{\Theta}q_{\phi}\left(\theta\right)\log\frac{p^{SA}\left(Y \mid x,\mathcal{S}\right)}{q_{\phi}\left(\theta\right)}\frac{q_{\phi}\left( \theta\right)}{q^{SA}(\theta|\mathcal{S})}d\theta\] \[=\int_{\Theta}q_{\phi}\left(\theta\right)\log\frac{p^{SA}\left(Y \mid\theta,X,\mathcal{S}\right)p\left(\theta\right)}{q_{\phi}\left(\theta \right)}\frac{q_{\phi}\left(\theta\right)}{q^{SA}(\theta|\mathcal{S})}d\theta\] \[=\mathbb{E}_{q_{\phi}\left(\theta\right)}\left[\sum_{i=1}^{n}\log p ^{SA}\left(y_{i}\mid x_{i},\mathcal{S},\theta\right)\right]-KL\left(q_{\phi},p\right)\] \[+KL\left(q_{\phi},q^{SA}\right).\]

It is obvious that we need to maximize the following lower bound for maximally reducing the gap \(KL\left(q_{\phi},q^{SA}\right)\):

\[\max_{q_{\phi}}\left\{\mathbb{E}_{q_{\phi}\left(\theta\right)} \left[\sum_{i=1}^{n}\log p^{SA}\left(y_{i}\mid x_{i},\mathcal{S},\theta\right) \right]-KL\left(q_{\phi},p\right)\right\},\]

which can be equivalently rewritten as

\[\min_{q_{\phi}}\left\{\lambda\mathbb{E}_{q_{\phi}\left(\theta \right)}\left[\mathcal{L}_{\mathcal{S}}\left(s\left(\theta\right)\right)+KL \left(q_{\phi},p\right)\right\}\text{or}\] \[\min_{q_{\phi}}\left\{\lambda\mathbb{E}_{q_{\phi}\left(\theta \right)}\left[\max_{\theta^{\prime}:\left\|\theta^{\prime}-\theta\right\|\leq \rho}\mathcal{L}_{\mathcal{S}}\left(\theta^{\prime}\right)\right]+KL\left(q_{ \phi},p\right)\right\}. \tag{6}\]

Furthermore, in the OP in (6), the first term implies that we seek an approximate posterior distribution such that any model sampled from it should be aware of the sharpness. Meanwhile, the second term prefers simpler approximate posterior distributions and can be estimated depending to how to equip approximate posterior distributions.

With the Bayesian setting and variational inference formulation, our proposed sharpness-aware posterior is ready to be incorporated into the MCMC-based, variational inference-based Bayesian Neural Networks.

## 4 Experiments

In this section, we conduct various experiments to demonstrate the effectiveness of the sharpness-aware approach on Bayesian Neural networks, including SWAG (Maddox et al., 2019) (with its variations SWA, SWAG-Diagonal, and SWAG), MC-dropout (Gal & Ghahramani, 2016), and Deep Ensemble (Lakshminarayanan et al., 2017).

The experiments are conducted on three benchmark datasets: CIFAR-10, CIFAR-100, and ImageNet ILSVRC-2012, and report accuracy, negative log-likelihood (NLL), and Expected Calibration Error (ECE) to estimate the calibration capability and uncertainty of our method against baselines. Additionally, we visualize the loss-landscape and sharpness score of models to examine the loss geometry awareness of our method and its effectiveness in the Bayesian setting.

We note that the purpose of the experiments is not to seek the state-of-the-art performances. Alternatively, we wish to demonstrate the usefulness of the sharpness-aware posterior when incorporated with specific Bayesian Neural Networks. We first conduct ablation studies to illustrate that the models sampled from our sharpness-aware posterior are flatter than those sampled from the standard posterior (cf. Section 4.4) and the ensemble models produced by our sharpness-aware posterior have lower sharpness compared to the standard posterior (cf. Section 4.4.2). Resultantly, Bayesian Neural Networks incorporated with our sharpness-aware posterior receive better predictive performances and calibration of uncertainty estimates than their counterpart baselines.

### Dataset and implementation detail

**CIFAR:** We run experiments with PreResNet-164 and WideResNet28x10 on both CIFAR-10 and CIFAR-100. These datasets contain 60,000 images in total with 50,000 instances for training and 10,000 for testing. For each network and dataset pair, we apply Sharpness-Aware Bayesian to various settings: F-SWA, F-SWAG-Diag, F-SWAG, F-SWAG-Dropout, F-MC-Dropout, and F-Deep-Ensemble. Note that we train all models for 300 epochs and start to collect models after epoch 161 for the F-SWA and F-SWAG settings, which is the same in (Maddox et al., 2019). We choose \(\rho=0.05\) for CIFAR10 and \(\rho=0.1\) for CIFAR100 in all experiments. For inference, we evaluate a single model for F-SWA and do an ensemble on 30 sample models for the other settings. For the Deep-Ensemble method, we re-produce results by training three independent models and doing ensembles for inference. All hyper-parameters and processes are the same for F-Deep-Ensemble. To get a stable result, we run three times for each setting and report the mean and standard deviation.

**ImageNet:** this is a large and challenging dataset with 1000 labels. We run experiments with Densenet-161 and ResNet-152 on F-SWA, F-SWAG-Diag, and F-SWAG. The results are shown in Table 3. For each setting, we use the pre-trained weight on ImageNet collected from _torchvision_ package and finetune on 10 epochs with \(\rho=0.05\). We start
collecting 4 times per epoch at the beginning of finetuning process for all settings and evaluate the models the same as experiments for the CIFAR dataset.

The result of SWA, SWAG-Diag, SWAG, and MC-Dropout is collected from the original paper (Maddox et al., 2019). There is no available reported result of Deep-Ensemble on examined networks so we report the re-produced result and marked *. 1

Footnote 1: The implementation is provided in [https://anonymous.4open.science/r/flat_bnn-B46B](https://anonymous.4open.science/r/flat_bnn-B46B)

### Metrics

**Ensemble accuracy \(\uparrow\) (ACC)** is computed by sampling several models from the Bayesian Neural Networks and then doing ensemble the prediction probabilities.

**Negative log-likelihood \(\downarrow\) (NLL)** is computed sampling several models from the Bayesian Neural Networks and evaluating the negative log-likelihood of the prediction probabilities.

**Expected calibration error \(\downarrow\) (ECE)** compares the predicted probability (or confidence) of a model to its accuracy (Naeini et al., 2015; Guo et al., 2017). To compute this error, we first bin the the confidence interval \([0,1]\) into \(M\) equal bins, then categorize data samples into these bins according to their confidence scores.

We finally compute the absolute value of the difference between the average confidence and the average accuracy within each bin, and report the average value over all bins as the ECE. Specifically, let \(B_{m}\) denote the set of indices of samples having their confidence scores belonging to the \(m^{th}\) bin. The average accuracy and the average confidence within this bin are:

\[acc(B_{m}) =\frac{1}{|B_{m}|}\sum_{i\in B_{m}}\mathbf{1}[\hat{y_{i}}=y_{i}],\] \[conf(B_{m}) =\frac{1}{|B_{m}|}\sum_{i\in B_{m}}\hat{p}(x_{i}).\]

Then the ECE of the model is defined as:

\[ECE=\sum_{m=1}^{M}\frac{|B_{m}|}{N}|acc(B_{m})-conf(B_{m})|.\]

### Experimental results

#### 4.3.1 Predictive performance

The experimental results are shown in Table 1 for CIFAR-100, Table 2 for CIFAR-10, and Table 3 for ImageNet dataset. There is an appropriate improvement in all experiments. Especially on ImageNet and CIFAR-10, our method consistently surpasses the baselines on both three calibrations, which demonstrate the stability and reliability of the model in uncertainty matter. On CIFAR-100, we gain from 0.7% to 1.03% on PreResNet-164 and at least 1.2% on the WideResNet28x10 model. We observe that there is a trade-off between accuracy, negative log-likelihood, and expected calibration error. However, the trade-off is miner compared to overall improvement.

The Deep-Ensemble method trains several models of the

Figure 1: Comparing loss landscape of PreResNet-164 on CIFAR-100 dataset training with SWAG and F-SWAG method. For visualization purposes, we sample two models for each SWAG and F-SWAG and then plot the loss landscapes. It can be observed that the loss landscapes of our F-SWAG are flatter, supporting our argument for the flatter sampled models.
same network independently, then does ensemble for inference. We do ensemble three models and get a remarkable result for both baseline and sharpness-aware settings. Additionally, this simple yet effective method takes much longer time to train but is much faster for inference. Again, our approach surpasses the baseline in all metrics of interest.

#### 4.3.2 Calibration of uncertainty estimates

We evaluate the ECE of each setting and compare it to baselines in Tables 1, 2, and 3. This score measures the maximum discrepancy between the accuracy and confidence of the model. To further clarify it, we display the Reliability Diagrams of PreResNet-164 on CIFAR-100 to understand how well the model predicts according to the confidence threshold in Figure 2. To produce this plot, the predicted confidence score is split into 20 bins uniformly. For each bin, we calculate the mean confidence and the accuracy of test samples within the bin. The difference between these two scores is displayed. The plot closer to the zero lines (black dashed line in Figure 2) is better. Also, the plot above this line represents over-confident prediction, while the one below this line represents under-confident prediction. As can be seen, our methods F-SWA, F-SWAG-Diag, and F-SWAG are closer to the zeros line in most bins, while the corresponding baselines tend to be over-confident. The flatter model could help to generate smoother outcomes, therefore mitigating the over or under-confidence problems. This demonstrates the sharpness-aware effect on reliability

\begin{table}
\begin{tabular}{l c c c|c c c} \hline \hline  & \multicolumn{3}{c}{**PreResNet-164**} & \multicolumn{3}{c}{**WideResNet28x10**} \\ Method & ACC \(\uparrow\) & NLL \(\downarrow\) & ECE \(\downarrow\) & ACC \(\uparrow\) & NLL \(\downarrow\) & ECE \(\downarrow\) \\ \hline \hline SWA & 80.19 \(\pm\) 0.52 & 0.7370 \(\pm\) 0.0265 & 0.0700 \(\pm\) 0.0056 & 82.40 \(\pm\) 0.16 & 0.6684 \(\pm\) 0.0034 & 0.0684 \(\pm\) 0.0022 \\ F-SWA & **80.94 \(\pm\) 0.23** & **0.6839 \(\pm\) 0.0047** & **0.0482 \(\pm\) 0.0025** & **83.61 \(\pm\) 0.29** & **0.5801 \(\pm\) 0.0174** & **0.0283 \(\pm\) 0.0042** \\ \(\hat{\text{SWAG-Diag}}\) & 80.18 \(\pm\) 0.50 & 0.6837 \(\pm\) 0.0186 & **0.0239 \(\pm\) 0.0047** & 82.40 \(\pm\) 0.09 & 0.6150 \(\pm\) 0.0029 & 0.0322 \(\pm\) 0.0018 \\ F-SWAG-Diag & **81.01 \(\pm\) 0.29** & **0.6645 \(\pm\) 0.0050** & 0.0242 \(\pm\) 0.0039 & **83.50 \(\pm\) 0.29** & **0.5763 \(\pm\) 0.0120** & **0.0151 \(\pm\) 0.0020** \\ \(\hat{\text{SWAG}}\) & 79.90 \(\pm\) 0.50 & **0.6595 \(\pm\) 0.019** & 0.0587 \(\pm\) 0.0048 & 82.23 \(\pm\) 0.19 & 0.6078 \(\pm\) 0.0066 & **0.0113 \(\pm\) 0.0026** \\ F-SWAG & **80.93 \(\pm\) 0.27** & 0.6704 \(\pm\) 0.0049 & **0.0350 \(\pm\) 0.0025** & **83.57 \(\pm\) 0.26** & **0.5757 \(\pm\) 0.0136** & 0.0196 \(\pm\) 0.0015 \\ \hline MC-Dropout & - & - & 82.30 \(\pm\) 0.19 & 0.6500 \(\pm\) 0.0049 & 0.0574 \(\pm\) 0.0028 \\ F-MC-Dropout & 81.06 \(\pm\) 0.44 & 0.7027 \(\pm\) 0.0049 & 0.0514 \(\pm\) 0.0047 & **83.24 \(\pm\) 0.11** & **0.6144 \(\pm\) 0.0068** & **0.0250 \(\pm\) 0.0027** \\ \hline Deep-ens* & 82.08 \(\pm\) 0.42 & 0.7189 \(\pm\) 0.0108 & 0.0334 \(\pm\) 0.0064 & 83.04 \(\pm\) 0.15 & 0.6958 \(\pm\) 0.0335 & 0.0483 \(\pm\) 0.0017 \\ F-Deep-ens & **82.54 \(\pm\) 0.10** & **0.6286 \(\pm\) 0.0022** & **0.0143 \(\pm\) 0.0041** & **84.52 \(\pm\) 0.03** & **0.5644 \(\pm\) 0.0106** & **0.0191 \(\pm\) 0.0039** \\ \hline \hline \end{tabular}
\end{table}
Table 1: Classification score on CIFAR100 dataset. * Experiments are reproduced with the same hyper-parameters with our corresponding method for a fair comparison

\begin{table}
\begin{tabular}{l c c c|c c c} \hline \hline  & \multicolumn{3}{c}{**PreResNet-161**} & \multicolumn{3}{c}{**ResNet-152**} \\ Model & ACC \(\uparrow\) & NLL \(\downarrow\) & ECE \(\downarrow\) & ACC \(\uparrow\) & NLL \(\downarrow\) & ECE \(\downarrow\) \\ \hline \hline SWA & 78.60 & 0.8655 & 0.0509 & 78.92 & 0.8682 & 0.0605 \\ F-SWA & **78.72** & **0.8269** & **0.0201** & **79.15** & **0.8080** & **0.0211** \\ \(\hat{\text{SWAG-Diag}}\) & 78.59 & 0.8559 & 0.0459 & 78.96 & 0.8584 & 0.0566 \\ F-SWAG-Diag & **78.71** & **0.8267** & **0.0194** & **79.20** & **0.8065** & **0.0199** \\ \(\hat{\text{SWAG}}\) & 78.59 & 0.8303 & 0.0204 & 79.08 & 0.8205 & 0.0279 \\ F-SWAG & **78.70** & **0.8262** & **0.0185** & **79.17** & **0.8078** & **0.0208** \\ \hline \hline \end{tabular}
\end{table}
Table 3: Classification score on ImageNet dataset

\begin{table}
\begin{tabular}{l c c c|c c c} \hline \hline  & \multicolumn{3}{c}{**Densenet-161**} & \multicolumn{3}{c}{**ResNet-152**} \\ Model & ACC \(\uparrow\) & NLL \(\downarrow\) & ECE \(\downarrow\) & ACC \(\uparrow\) & NLL \(\downarrow\) & ECE \(\downarrow\) \\ \hline \hline SWA & 78.60 & 0.8655 & 0.0509 & 78.92 & 0.8682 & 0.0605 \\ F-SWA & **78.72** & **0.8269** & **0.0201** & **79.15** & **0.8080** & **0.0211** \\ \(\hat{\text{SWAG-Diag}}\) & 78.59 & 0.8559 & 0.0459 & 78.96 & 0.8584 & 0.0566 \\ F-SWAG-Diag & **78.71** & **0.8267** & **0.0194** & **79.20** & **0.8065** & **0.0199** \\ \(\hat{\text{SWAG}}\) & 78.59 & 0.8303 & 0.0204 & 79.08 & 0.8205 & 0.0279 \\ F-SWAG & **78.70** & **0.8262** & **0.0185** & **79.17** & **0.8078** & **0.0208** \\ \hline \hline \end{tabular}
\end{table}
Table 2: Classification score on CIFAR10 dataset. * Experiments is re-produced with the same hyper-parameters with our corresponding method for a fair comparison
for BNNs, reflected by the better ECE score overall.

### Ablation studies

#### 4.4.1 Loss landscape

In Figure 1, we plot the loss-landscape of the models sampled from our proposal of sharpness-aware posterior against the non-sharpness-aware one. Particularly, we compare two methods F-SWAG and SWAG by selecting four random models sampled from the posterior distribution of each method under the same hyper-parameter settings. As observed, our method not only improves the generalization of ensemble inference, demonstrated by classification results in Section 4 and sharpness in Section 4.4.2, but also the individual sampled model is flatter itself.

#### 4.4.2 Sharpness evaluation

We measure and visualize the sharpness of the models. To this end, we sample five models from the approximate posteriors and then take the average of the sharpness of these models. For a model \(\theta\), the sharpness is evaluated as \(\max\limits_{||\epsilon||_{2}\leq\rho}\mathcal{L}_{\mathcal{S}}(\theta+ \epsilon)-\mathcal{L}_{\mathcal{S}}(\theta)\) to measure the change of loss value around \(\theta\).

We calculate the sharpness score of PreResNet-164 network for SWAG-Diag, SWAG, and flat settings F-SWAG-Diag, F-SWAG training on CIFAR100 dataset and visualize them in Figure 3. The sharpness-aware settings generate smaller _sharpness_ scores compared to the corresponding baseline, indicating that our models get into flatter regions.

## 5 Conclusion

In this paper, we devise theories, the Bayesian setting, and the variational inference for the sharpness-aware posterior in the context of Bayesian Neural Networks (BNNs). The sharpness-aware posterior enables the models sampled from it and the optimal approximate posterior estimating it to have a higher flatness, hence possibly possessing a better generalization ability. We conduct extensive experiments by leveraging the sharpness-aware posterior with the state-of-the-art Bayesian Neural Networks, showing that the flat-seeking counterparts outperform their baselines in all metrics of interest, including the ensemble accuracy, ECE, and NLL.

Figure 3: Evaluate the sharpness over training epoch. The average sharpness of the models sampled from our sharpness-aware posterior is smaller than the corresponding baseline.

Figure 2: Reliability diagrams for PreResNet164 on CIFAR-100. The confidence is split into 20 bins and plot the gap between confidence and accuracy is in each bin. The best case is the black dashed line when this gap is zeros. The plots of F-SWAG get closer to the zero lines, implying our F-SWAG can calibrate the uncertainty better.