Flat Seeking Bayesian Neural Networks
Authors: Van-Anh Nguyen
Long Tung Vuong
Hoang Phan
Thanh-Toan Do
Dinh Phung
Trung Le
Venue: NeurIPS 2023 poster (2023)
Keywords: Bayesian
sharpness-aware
posterior
TL;DR: Sharpness-aware posterior for BNN models

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

Flat Seeking Bayesian Neural Networks
Van-Anh Nguyen'
Tung-Long Vuong 1.2
Hoang Phan
Thanh-Tone Do
Dinh Phung 1,2
Trung Le
DDeparment of Data Science and AI, Monash University, Australia
VinAl, Vietnam
'New York University, United States
{van-anh. nguyen, tung-long vuong, toan.do, dinh. phung, trungim) @monash edu 1vp20110nyu.edu
Abstract
Bayesian Neural Networks (BNNs) provide a probabilistic interpretation for deep learning models by imposing a prior distribution over model parameters and infer- ring posterior distribution based on observed data. The model sampled from the posterior distribution can be used for providing ensemble predictions and quantify- ing 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/fltnesss 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-awe posterior. Specifically, the models sampled from our sharpness-aware posterior, and the optimal approximate posterior estimating this sharpness-awart posterior, have better flatness, hence possibly possessing higher generalization ability. con- duct experiments by leveraging the sharpness-awee posterior with state-of-ther -ht Bayesian Neural Networks, showing that the flat-sceking counterparts outperform their baselines in all metrics of interest.
Introduction
Bayesian Neural Networks (BNNs) provide 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 observed data. This allows us to not only make predictions, but also quantify prediction uncertainty, which useful for many real-world applications. To sample deep learning models from complex and complicated posterior distributions, advanced particle-san saailiin approaches such as Hamiltonian Monte Carlo (HMC) [41]. Stochastic Gradient HMC (SGHMC) [10], Stochastic Gradient Langevin dynamics (SGLD) [58], and Stein Variational Gradient Descent (S (SVGD) [36] are often used. However, these methods can be computationally expensive, particularly when many models need 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 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 mean-field variational inference. This approach fails to account for the strong statistical dependencies among random weights of neural networks, limiting its ability 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
37th Conference on Neural Information Processing Systems (NcurIPS 2023).

---

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 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/Rlatness 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 estimate these non-sharpness-awee aware posteriors can result in sampled models from the corresponding optimal approximate posterior lacking awareness of sharpness/Ratnnsss 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 which encourages 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-awwn aware posterior. Overall, our contributions in this paper can be summarized as follows:
(KL) divergence,
We propose and develop theories, the Bayesian setting, and the variational inference approach for the sharpness-awant posterior. This posterior enables us sample set of flat models that improve the model generalization ability. We note that SAM [16] only considers the sharpness for single model, while ours is the first work studying the concept and theory of the sharpness for a distribution 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 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 Related Work
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 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 family to estimate the true posterior distribution by maximizing a variational lower bound. [21] suggests fitting

---

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 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 posteri- ors. 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.
2.2 Flat Minima
Flat minimizers have been found to improve the generalization ability of neural networks. This is because they enable models 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, 17J. Moreover, various methods seeking flat minima have been proposed in [46, 9, 2,,9, 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,
9J. Examples of such regularization terms include softmax output's 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, 37J. 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 probabilistic extension of SAM.
3 Proposed Framework
In what follows, we present the technicality of our proposed sharpness-awar posterior. Particularly, Section 3.1 introduces the problem setting and motivation for our sharpness-aware posterior. Section 3.2 is dedicated to theory development, while Section 3.3 is used to describe the Bayesian setting and variational inference approach for our sharpness-awe 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 fo( with E and training set = ((),,))),.... (Ins W/r) where (xi, yr) D. We wish to learn a posterior distribution Q& with the density function ( such that any model f ~ Q is aware of the sharpness when predicting over the training set S.
We depart with the standard posterior
n1
q(
S) x
p(u x. S)0(())
3

---

where the prior distribution P has the density function p(0) and the likelihood has the form
[ ()o(#)) - exp { { (e {
(y x, S,0) x exp
with the loss function l. The standard posterior Qs has the density function defined as
() exp
(Jo (2x)) (1a)) jo(0),
(I)
where À 2 0 is regularization parameter.
We define the general and empirical losses as follows:
LD (0) = B(aw)>D (e v)]-
Ls (0) = B(.,))= /e (Jo (x), 4)1
Σ (So (r.)
Basically, the general loss is defined the expected loss over the entire data-label distribution 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(0 S) x exp {-ALs (0)) p0))
Given a distribution Q with the density function (0) over the model parameters 0 E O we define the empirical and general losses over this model distribution Q as
Ls (Q) / Ls (0) dQ (0) / Cs (0)
@
/
CD (Q)
CD (0)
CD (0) d0) d0.
Specifically, the general loss over the model distribution Q is defined as the expectation the general losses incurred by the models sampled from this distribution, while the empirical loss the model distribution 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-award posterior whose proofs can be found in the supplementary material. Inspired by the Gibbs form of the standard posterior Qs in Eq. (2), we establish the following theorem connect the standard posterior Qs with the density g(0 S)
and the empirical loss Ls (Q) [7,2]. 2].
Theorem 3.1. Consider the following optimization problem
min (ASs (Q) KL (Q,
Q<<T
where we search over Q absolutely continuous w.r.t. and KL )) is the Kallback-Leible diver- gence. This optimization has a closed-form optimal solution Q" with the density
q' (0) x exp (- (0)) p(0),
which is exactly the standard posterior Qs with the density q(0 S).
Theorem 3.1 reveals that we need find the posterior Qs balancing between optimizing its empirical loss Ls (Q) and simplicity via KL P). However, minimizing the empirical loss Ls (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):
(ACD (Q) <<<P
min
KL (Q.)
Notably, solving the optimization problem (OP) in (4) is generally intractable. To make it tractable, we find its bpper-bound which relevant the sharpness of a distribution Q over models as shown in the following theorem.

---

Theorem 3.2. Assume that @ is compact set. Under some mild conditions, given any d E [0; 1), with the probability at least ó over the choice of DD,, foor any distribution Q, we have
CD (Q) S Ee~Q
Es (0)) -
max
+
max |00 .n)
where f is a non-decreasing function w.r.. the first variable and approaches 0 when the training size approaches
We note that the proof of Theorem 3.2 is not trivial extension of sharpness because we need to tackle the general and empirical losses over a distribution Q. To make explicit sharpness over a distribution Q on models, we rewrite the upper-bound of the inequality as
minimization
E
Ls (s) Lss (0) + Ls (Q)) -
max
+
max |001 ,m)
BE0
where the first term E [maxp: - Ls (8') - Ls (0)] can be regarded as the sharpness over the distribution Q on the model space and the last term f (maxxee |101]2, is constant.
Moreover, inspired by Theorem 3.2, we propose solving the following OP which forms an upper- bound of the desirable OP in (4)
Ls (07) + (L
min QC<P
XEEa-
max 00:100'00
(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 Q with the density function q (0)S))
-x max Ls (s) (((0) 3 0':10'00SSp
(A1S)) & exp
- exp (xp(
where we have defined 8 (0) argmax Ls (0').
Theorem 3.3 describes the close form of the sharpness-aware posterior distribution Q with the density function (0 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 Q we examine the sharpness-aware likelihood
SA (y x. 0, x exp ( (lac) exp (-
where (0) argmax Ls (8') *' lot
SP
With this predefined sharpness-awan likelihood, we can recover the sharpness-aware posterior distribution Q with the density function q (B)S):
q (A(())
II SA (yi I.. S,0)
Variational inference for the sharpness-aware posterior distribution: We now develop the variational inference for the sharpness-aware posterior distribution. Let denote X [21, Tn] and
5

---

Y = Mi,... Br]. Considering an approximate posterior family {96 (0) @ E D), we have
/ 2 log 2 (Y) X, de
log SA (Y X, 5)
:/8 (0) log
S4 (Y 0 S)
qs (0) d0
q (0)
(S)
Eg (0)
logp
SA (ys 24 2 S.
KL (qo. p) + KL (AA
is obvious that we need to maximize the following lower bound for maximally reducing the gap
KL (90.9 SAA):
1 log Log (ys ry S,0)
max 9+
E (o)
KL (90.
which can be equivalently rewritten as
min
{ (o)
[Cs (s (0)) KL (go,
min 2
(0) max
Cs (8") 01:10r-01Spp
+ KL (gp. p)
(6)
Derivation for Variational Approach with A Gaussian Approximate Posterior: Inspired by geometry-based SAM approaches [34, 31], we incorporate the geometry to the SA vari- ational approach via the distance to define the ball for the sharpness 110' -Ilaiug(To))
V(Ør 8) diag(To) (0' 0) as
min
AB
max amer Sp
Ls (0')
+ KK (2
0':10''
To further clarify, we consider our SA posterior distribution Bayesian NNs, wherein we impose the Gaussian distributions to its weight matrices W; N (14, af , 1.... L'. The parameter consists of μ4, o,,
1.... L. For
WiL qa. using the reparameterization trick W,
J4s + diag(o; (aag(()))ss) (3, ~ N (0, I) and by searching = WiL with W. = + diag((a)) ~ N (0, IT), the constraint ||0 6' l' dldg((())) = || p²l p'll (la)g((T)) with μ = fi. and Mi:L- Thus, the OP in (6)
reads
LS (1r diag(o;) (o:)) )c: i=1 }}
min ΛE
max
||
Baac Sp
where σ = 01:L, C1:L. and we define diag(To) diag(Tp.o) in the distance of the geometry.
To solve the OP in (7), sample c1.t. from the standard Gaussian distributions, employ one-step gradient ascent to find μr, and use the gradient μ update μ.. Specifically, find j [6] (Chapter 9) as
diag(T (
[uu diag(o,)
1' = +
diag(T V (14 diag (o,) (o.)ci)
The diagnose of diag(Tp.r) specifies the importance level of the model weights, i.e., the weight with a higher importance level is encouraged have higher sharpness via a smaller absolute partial derivative of the loss w.r.t. this weight. We consider diag(T),.>) I (i.c., the standard SA
BNN) and diag (aag(T,..) (T,..) diag element- wise division.
[Md] (i.e., the geometry SA BNN). Here we note that : represents the
We absorb the biases the weight matrices.

---

Table 1: Classification score on CIFAR- 100 dataset. Each experiment is repeated three times with different random seeds and reports the mean and standard deviation.
Pre N
WideResNet28x10 N
Method
ACC 1
BCE 4
ACC
ECE
Variational inference
MC-Dropout
79.50 = 0.37
0.9162 1 0.0103
0.0993
0.0033
82.30 +: 0.19
0.6500
0.0049
0.0574 : 0.0028
F-MC- Pro-orop
81,06 0.44
0.7027
0.0049
0.0514 0.00447
83.24 1: 0.11
0.6144
0.0068
0.0250 0.0027
Deep-ons
82.08 = 0.42
0.7189 + 0.0108
0.0334 + 0.0064
83.04 + 0.15
0.6958 + 0.0335
0.0483 = 0.0017
F-Deep-ens
82.54 + 0.10
0.6286 0.0022
0.0143 0.0041
84.52 + 0.03
0.5644 0.0106
0.0191 0.0039
Markov chain Monte Carlo
SGLD
80.13 0.01
0.7604 0.0010
0.1161 0.0031
81.38 +- 0.10
0.7123 0.0204
0.0958 + 0.0004
F-SGLD
80.82 + 0.02
0.7276 +: 0.0012
0.1085 + 0.0008
82.12 + 0.16
0.6722 + 0.0112
0.0820 + 0.0021
Sample
SWAG-Diag
80.18 = 0.50
0.6837 1 0.0186
0.0239 0.0047
82,40 1: 0.09
0,6150 0.0029
0.0322 : 0.0018
F-SWAG Diag
81,01
0.29
0.6645 0.0050
0.0242
0.0039
83,50 0.29
0.5763
0.0120
0.0151
0,0020
SWAG
79.90 1: 0.50
0,6595 1 0.0019
0.0587 1 0.0048
82.23
1: 0.19
0.6078 1: 0.0006
0.0113 1: 0.0020
F-SWAG
80.93 0.27
0.6704 0.0049
0.0350 0.0025
83.57
0.5757 0.0136
0.0196
Table 2: Classification score on CIFAR-10 Lataset.Fach experiment is repeated three times with different random seeds and reports the mean and standard deviation.
PreResNet-164 NLL
Ner2Res NLL 4
Method
1CC
ECE
ACC
ECE
Variational inference
MC-Dropout
96.18 + 0.02
0.1270 + 0.0030
0.0162 + 0.0007
96.39 + 0.09
0,1094
0.0021
0.0094 + 0.0014
F-MC-Dropout 96.39 = 0,18
0.1137 0.0024
0,0118 + 0.0006
97.10 0.12
0.0966 0.0047
0.0095 + 0.0008
Deop-ens
96.39 + 0.09 96.70
0.1277
= 0.0030
0.0108 + 0.0015
96.96 + 0.10
0.0076 0.0851
0,1031
0,0087 + 0.0018
F-Deep-ens
0,1031 0.0016
0.0057 0.0031
97.11 0.10
0.0011
0.0059
Markov chain Monte Carlo
SGLD
94.79 + 0.10
0.2089
0.0021
0.0711 + 0.0061
95.87 + 0.08
0.1573
0.0190
0.0463 0.0050
F-SGLD
95.04 :: 0.06
0,1912 0.0080
0.0601 1 0.0002
96.43 1 0.05
0.1336 1 0.004
0.0385 1 0.0003
Sample
SWAG-Ding
96.03 + 0.10
0,1251 + 0.0029
0.0082 + 0.0008
96.41 + 0.05
0.1077
0.0009
0.0047 + 0.0013
F.SWAG Diag
96.23 0.01
0.1108
0.0013
0.0M44
0.0005
97.05
0.08
0.0888
0.0052
0.0043
0.0004
SWAG
96.03 1. 0.02
0,1232
1. 0.0022
0.0053 1. 0.0004
96.32
1. 0.08
0.1122
0.0009
0.0088 1 0.0006
F-SWAG
96.25
0.11062
0.0014
0.0056 0.0002
97.09
0.0883
0.0IMI4
0.0M008
Finally, the objective function in (6) indicates that we aim find an approximate posterior distribution that ensures any model sampled from it is aware of the sharpness, while also preferring simpler approximate posterior distributions. This preference can be estimated based on how we equip these distributions. With the Bayesian setting and variational inference formulation, our proposed sharpness- aware posterior can be integrated into MCMC-based and variational infrrence-based Bayesian Neural Networks. The supplementary material contains the details on how to derive variational approaches and incorporate the saarpness-awarenes into the BNNs used in our experiments including BNNs with approximate Gaussian distribution [33], BNNs with stochastic gradient Langevin dynamics (SGLD) [58], MC-Dropout [18], Bayesian deep ensemble [35], and SWAG [39].
4 Experiments
In this section, we conduct various experiments to demonstrate the effectiveness of the sharpness- aware approach on Bayesian Neural networks, including BNNs with an approximate Gaussian distribution [33] (i.e.. SGVB for model's reparameterization trick and SGVB-LRT for representation's 's reparameterization trick), BNNs with stochastic gradient Langevin dynamics (SGLD) [58], MC- Dropout [18], Bayesian deep ensemble [35], and SWAG [39]. The experiments are conducted on three benchmark datasets: CIFAR-10, CIFAR- and ImageNet ILSVRC-2012, and report accuracy, negative log- likeeiidddood (NLL), and Expected Calibration Error (ECE) to estimate the calibration capability and uncertainty of our method against baselines. The details of the dataset and implementation are described in the supplementary material?.
"The implementation is provided in httpp::/ggttt
bit.git
7

---

Table 3: Classification scores of approximate the Gaussian posterior the CIFAR datasets. Each experiment is repeated three times with different random seeds and reports the mean and standard deviation.
Resnet10 NLL 4
Resnet18 NLL
Method
ACC
ECE
ACC 1
ECE І
Experiments
Cifar-100 dataset
SGVB-LRT
61.75 + 0.75
1.534 + 0.03
0.0676 + 0.01
68.95 + 1.20
1.140 + 0.21
0.063 + 0.04
F-SGVB-LRT
62.25 = 0.57
1.4001
+ 0.04
0.0642 = 0.01
70.00 = 1.42
1.127 + 0.25
0.022 = 0.05
+ Geometry
62.54 i 0.67
1.3704 0.01
0.0301 i 0.03
70.12 1 1.02
1.121
1 0.23
0.036 0.06
SGVB
54,40 + 0.98
1.968 + 0.05
0.214 +: 0.00
0.00
60.91 +: 2.31
1.746 +: 0.15
0.246 + 0.03
F-SGVB
54.53 + 0.33
1.967 + 0.00
0.212 + 0.00
61.54
+: 2.23
1.695 +: 0.15
0.242 + 0.03
+ Geometry
55.53 1: 0.65
1.906 1: 0.02
0.207 1: 0.00
62.58 1: 0.53
1.612 1: 0.03
0.224 0.00
Experiments on Cifar-10 dataset
SGVB-LRT
84,98 +: 1.87
0.422 1: 0.10
0.043 0.04
89.10 1 1.32
0.344
0.02
0.033 1: 0.02
F-SGVB-IRRT
86.32 1.34
0.409
0.03
0.017 0.06
90.00 1.10
0.291
0.02
0.019 + 0.01
+ Geometry
86.44 1.12
0.403
0.06
0.025 0.03
90.31 1.11
0.262
0.01
0.02 0.014
SGVB
80.52 + 2.10
0.781 + 0.23
0.237 + 0.06
86.74 = 1.25
0.541 + 0.01
0.181 = 0.02
F-SGVB
80.60 + 1.88
0.776 0.13
0.223 + 0.05
87.01 + 0.91
0.534 + 0.01
0.183 + 0.01
+ Geometry
82.05 + 0.47
0.704 + 0.01
0.206 + 0.00
86.80 = 1.30
0.531 + 0.01
0.175 = 0.01
Table 4: Classification
ImageNet dataset
Densenet-
ResNet-152
Model
ACC
NLL +
ECE 4
ACC
NLL +
ECE 1
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
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
SGLD
78. 50
0.8317
0.0157
79.00
0.8165
0.0220
F-SGLD
78.64
0.8236
0.0166
79.16
0,8050
4.1 Experimental results
4.1.1 Predictive performance
Our experimental results, presented in Tables 1, 2, 3 for CIFAR-100 and CIFAR-10 dataset, and Table 4 for the ImageNet dataset, indicate notable improvement across all experiments. It is worth noting that there is trade-off between accuracy, negative log-likelihood and expected calibration eroo. Nonetheless, approach obtains fine balance between these factors compared to the overall improvement.
4.2 Effectiveness of sharpness-aware posterior
Calibration of uncertainty estimates: We evaluate the ECE of each setting and compare it to baselines in Tables 1, 2, and 4. This score measures the maximum discrepancy between the accuracy
Figure 1: Comparing loss landscape PreResNet-164 on CIFAR- 100 dataset training with SWAG and F-SWAG method. For visualization purposes, sample two models for each SWAG and F. SWAG and then plot the loss landscapes. can be observed that the loss landscapes our F-SWAG are flatter, supporting our argument for the flatter sampled models.

---

0.9
Contitetee -erarary
vccurbry
Cartilerre -vcuureyy
0.4
011
000
-0.35
251A
3-5005
0.10
1-50RG Dvp3
Widds 64Ng
PONAAO
0.2
0.2
0u4
0.6
60
0.3
0ut
044
08
6.0
0.22
Contiderce EITIK probi
Contidence tmax prob)
Confidennce OTTe SPDD)
Figure 2: Reliability diagrams for PreResNet164 on CIFAR-ION The confidence is split into 20 bins and plots the gap between confidence and accuracy 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.
Table 5: Classification
CIFAR-10-C, PreResNet-164 model when training with CIFAR-10.
The full result on each type of corruption is displayed the supplemetary material.
ECE L
Accuracy
Coruptton
SWAG-D
F-SWAG
SWAG
F-SWAG
SWAG-D
F-SWAG-D
SWAG
F-SWAG
Noise
0.0729
0.0701
0.0958
0.0078
74,26
75.59
74.02
75.08
Blur
0.0121
0.0090
0.0202
0.0273
91.13
90.55
91.03
90.93
Weather
0.018
0.0142
0.0272
0.0240
89.47
89.18
89.42
89.11
Digital and others
0.0277
0.0229
0.0384
0.0209
87.03
86.94
86.93
87.19
Average
0.0328
0.0290
0.0454
0.0200
85.47
85.56
85.58
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. The experiments is detailed in the supplementary material.
Out-of-distribution prediction: The effectiveness of the sharpness-awars Bayesian neural network (BNN) is demonstrated in the above experiments, particularly in comparison to non-flat methods. In this section, we extend the evaluation to an out-of-distribution setting. Specifically, we utilize the BNN models trained on the CIFAR-10 dataset to assess their performance on the CIFAR-10-4 dataset. This an extension of the CIFAR-10 designed evaluate the robustness of machine learning models against common corruptions and perturbations in the input data. The corruptions include various forms of noise, blur, weather conditions, and digital distortions. We conduct an ensemble of 30 models sampled from the flat-posterior distribution and compared them with non-flat ones. We present the average result of each corruption group and the average result on the whole dataset in Table 5, the detailed result of each corruption form is displayed in the supplementary material. Remarkably, the flat BNN models consistently surpass their non-flat counterparts with respect to average ECE and accuracy metrics. This finding is additional evidence of the generalization ability of the tharpness-awan posterior.
4.3 Ablation studies
In Figure 1, we plot the loss-landscape of the models sampled from our proposal of sharpness-aware posterior against the non-sharpness-art 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 paremmrer settings. As observed, our method not only improves the generalization of ensemble inference, demonstrated by classification results in Section 4.1 and sharpness in Section 4.2, but also the individual sampled model is flatter itself.
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 model 0, the sharpness is evaluated as max Ls(0 + E) (s( to measure the change of loss value |<llaSp
around calculate the sharpness score of PreResNet-166 network for SWAG, and F-SWAG training CIFAR- 100 dataset and visualize them in the supplementary material. As shown there, the sharpness-aware versions produce smaller sharpness scores compared to the corresponding baselines, indicating that our models get into flatter regions.

---

Conclusion
In this paper, we introduce theories in the Bayesian setting and discuss variational inference for the sharpness- aappne posterior in the context of Bayesian Neural Networks (BNNs). The sharpness-aware posterior results in models that are less sensitive to noise and have a better generalization ability, as it enables the models sampled from it and the optimal approximate posterior estimates to have higher flatness. We conducted extensive experiments that leveraged the sharpness-awand posterior with state-of-the-art Bayesian Neural Networks. Our main results show that the models sampled from proposed posterior outperform their baselines terms of ensemble accuracy, expected calibration error (ECE), and negative log- ligeeiilioood (NLL). This indicates that the flat-sceking counterparts are better at capturing the true distribution of weights in neural networks and providing accurate probabilistic predictions. Furthermore, we performed ablation studies to showcase the effectiveness of the flat posterior distribution on various factors such as uncertainty estimation, loss landscape, and outoo--distribution prediction. Overall, the sharpness aware posterior presents promising approach for improving the generalization performance Bayesian neural networks.
This work was partly supported by ARC DP23 grant DP230101 176 and by the Air Force Office Scientific Research under award number FA2386-23-1-4044,
Acknowledgements.
References
[1] Momin Abbas, Quan Xiao, Lisha Chen, Pin-Yu Chen, and Tianyi Chen. Sharp-maml: Sharpness- aware model-agnostic meta learning. arXiv preprint arXiv:2206.0399 2022. 3
[2] Pierre Alquier, James Ridgway, and Nicolas Chopin. On the properties of variational approx- imations of gibbs posteriors. The Journal of Machine Learning Research, 17(1):837484414 2016.
[3] Dara Bahri, Hossein Mobahi, and Yi Tay. Sharpness-aware minimization improves language model generalization. In Proceedings of the 60th Annual Meeting of the Association for Computational Linguistics (Volume Long Papers), pages 7360-7371, Dublin, Ireland, May 2022. Association for Computational Linguistics. 3
[4] Jose Blanchet and Karthyek Murthy. Quantifying distributional model risk via optimal transport. Mathematics of Operations Research, 44(2):565-600, 2019. 3
[5] Charles Blundell, Julien Cornebise, Koray Kavukcuoglu, and Daan Wierstra. Weight uncertainty in neural network. In International conference on machine learning, pages 1613-1622, PMLR, 2015. 3.3
[6] Stephen Boyd and Lieven Vandenberghe. Convex Optimization. Cambridge University Press, March.
[7] Olivier Catoni. Pac-bayesian supervised classification: the thermodynamics statistical learning, arXiv preprint arXiv:0712.02488. 2007. 4
[8] Junbum Cha. Sanghyuk Chun, Kyungjae Lee, HlnnChholl Cho, Scunghyun Park, Yunsung Lee, and Sungrae Park. Swad: Domain generalization by seeking flat minima. Advances in Neural Information Processing Systems, 34:22405-2222 2021.
[9] Pratik Chaudhari, Anna Choromańska, Stefano Soatto, Yann LeCun, Carlo Baldassi, Christian Borgs, Jennifer T. Chayes, Levent Sagun, and Riccardo Zecchina. Entropy-sgd. biasing gradient descent into wide valleys. Journal of Statistical Mechanics: Theory and Experiment, 2019, 2017.
[10] Tianqi Chen, Emily Fox, and Carlos Guestrin. Stochastic gradient hamiltonian monte carlo. In International conference on machine learning, pages 1683-1691, PMLR, 2014. 1,2
[11] Xiangning Chen, Cho-Jui Hsieh, and Boging Gong. When vision transformers outperform resnets without pre-training or strong data augmentations, arXiv preprint arX.v:2106.0154. 2021. 3
10

---

[12] Laurent Dinh, Razvan Pascanu, Samy Bengio, and Yoshua Bengio. Sharp minima can generalize for deep nets. In International Conference on Machine Learning, pages 1019-1028. PMLR, 2017. 3
[13] Jiawei Du, Daquan Zhou, Jiashi Feng, Vincent YF Tan, and Joey Tianyi Zhou. Sharpness-aware training for free. arXiv preprint arXiv:2205. 1400,,, 2022. 3
[14] Michael Dusenberry, Ghassen Jerfel, Yeming Wen, Yian Ma, Jasper Snoek, Katherine Heller, Balaji Lakshminarayanan, and Dustin Tran. Efficient and scalable bayesian neural nets with rank-1 factors. In International conference machine learning, pages 2782-2792. PMLR, 2020.
[15] Gintare Karolina Dziugaite and Daniel M. Roy. Computing nonvacuous generalization bounds for deep (stochastic) neural networks with many more parameters than training data. In UAL. AUAI Press, 2017. 2,
[16] Pierre Foret, Ariel Kleiner, Hossein Mobahi, and Behnam Neyshabur. Sharpness- aware min- imization for efficiently improving generalization. In International Conference on Learning Representations, 2021. 3,3
[17] Stanislav Fort and Surya Ganguli. Emergent properties of the local geometry of neural loss landscapes. arXiv preprint arXiv: 1910.09929, 2019.
[18] Yarin Gal and Zoubin Ghahramani. Dropout as bayesian approximation: Representing model uncertainty in deep learning. In Maria Florina Balcan and Kilian Q. Weinberger, editors, Proceedings of The 33rd International Conference on Machine Learning, volume 48 of Proceedings of Machine Learning Research, pages 1050-1059, New York, New York, USA, 20-22 Jun 2016. PMLR. 2,
[19] Yarin Gal, Jiri Hron, and Alex Kendall. Concrete dropout. Advances in neural information processing systems, 30, 2
[20] Soumya Ghosh, Jiayu Yao, and Finale Doshi-Velez, Structured variational learning of bayesian neural networks with horseshoe priors. In International Conference on Machine Learning, pages 1744-1753. PMLR, 2018. 2.
[21] Alex Graves. Practical variational inference for neural networks. Advances in neural information processing systems, 24, 2011. 1,2
[22] Arjun K Gupta and Daya K Nagar. Matrix variate distributions. Chapman and Hal//CRC, 2018.
3
[23] Sepp Hochreiter and Jürgen Schmidhuher: Simplifying neural nets by discovering flat minima. In NIPS, pages 329-536. MIT Press, 1994. 3
[24] Pavel Izmailov, Dmitrii Podoprikhin, Timur Garipov, Dmitry Vetrov, and Andrew Gordon Wilson. Averaging weights leads to wider optima and better generalization. arXiv preprint arXiv:1803.05407, 2018.
[25] Pavel Izmailov, Dmitrii Podoprikhin, Timur Garipov, Dmitry P. Vetrov, and Andrew Gordon Wilson. Averaging weights leads to wider optima and better generalization. In UAI, pages 876-885, AUAI Press, 2018.
[26] Stanislaw Jastrzebski, Zachary Kenton, Devansh Arpit, Nicolas Ballas, Asja Fischer, Yoshua Bengio, and Amos J. Storkey. Three factors influencing minima in sgd. ArXiv, abs/1711.04623, 2017.
[27] Yiding Jiang, Behnam Neyshabur, Hossein Mobahi, Dilip Krishnan, and Samy Bengio. Fantastic generalization measures and where to find them. In ICLR. OpenReview.ned. 2020. 2,
[28] Jean Kaddour, Linqing Liu, Ricardo Silva, and Matt Kusner. When do flat minima optimizers work? In Alice H. Oh, Alekh Agarwal, Danielle Belgrave, and Kyunghyun Cho, editors, Advances in Neural Information Processing Systems, 2022.
11

---

[29] Nitish Shirish Keskar, Dheevatsa Mudigere, Jorge Nocedal, Mikhail Smelyanskiy, and Ping Tak Peter Tang. On turge-batch training for deep learning: Generalization gap and sharp minima. In ICLR. OvenReview.net, 2017.
[30]
Mohammad Khan, Didrik Nielsen, Voot Tangkaratt, Wu Lin, Yarin Gal, and Akash Srivastava. Fast and scalable bayesian deep learning by weight-pertuubation in adam. In International Conference on Machine Learning, pages 2611-2622 PMLR, 2018. 2,
[31] Minyoung Kim, Da Li, Shell X Hu, and Timothy Hospedales. Fisher SAM: Information geometry and sharpness aware minimisation. In Kamalika Chaudhuri, Stefanie Jegelka, Le Song, Csaba Szepesvari, Gang Niu, and Sivan Sabato, editors, Proceedings of the 39th International Conference on Machine Learning, volume 162 of Proceedings of Machine Learning Research, pages 11 11148-11166 161. PMLR, 17-23 Jul 2022. 3,6 6
[32] Diederik P Kingma and Max Welling. Auto-encoding variational bayes. arXiv preprint arXiv:1312.6114 2013. 3
[33] Durk P Kingma, Tim Salimans, and Max Welling. Variational dropout and the local repa- rameterization trick. Advances neural information processing systems, 28, 2015. 1, 2, 3,
7
[34] Jungmin Kwon, Jeongseop Kim, Hyunseo Park, and In Kwon Choi. Asam: Adaptive sharpness- aware minimization for scale-invariant learning of deep neural networks. In International Conference on Machine Learning, pages 5905-5914. PMLR, 2021. 3,
[35] Balaji Lakshminarayanan, Alexander Pritzel, and Charles Blundell. Simple and scalable predictive uncertainty estimation using deep ensembles. Advances in neural information processing systems, 30, 2017. 7.7
[36] Qiang Liu and Dilin Wang. Stein variational gradient descent: A general purpose bayesian inference algorithm. Advances in neural information processing systems, 29, 2016. 1,2
[37] Yong Liu, Siqi Mai, Xiangning Chen, Cho-Jui Hsieh, and Yang You. Towards efficient and scalable sharpness-aware minimization. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pages 12360-12370, 2022. 3
[38] Christos Louizos and Max Welling. Multiplicative normalizing flows for variational bayesian neural networks. In International Conference on Machine Learning, pages 2218-2227. PMLR,
2017. 3
[39] Wesley J. Maddox, Timur Garipov, Pavel Izmailov, Dmitry Vetrov, and Andrew Gordon Wilson. A Simple Baseline for Bayesian Uncertainty in Deep Leamning. Curran Associates Inc., Red Hook, NY, USA, 2019. 2, 3,7
[40] Thomas Möllenhoff and Mohammad Emtiyaz Khan. Sam as an optimal relaxation of bayes. arXiv preprint arXhv:2210.01620 2022. 3
[41] Radford M. Neal. Bayesian Learning for Neural Networks. Springer-Verlag, Berlin, Heidelberg, 1996.
1,2
[42] Behnam Neyshabur, Srinadh Bhojanapalli, David McAllester, and Nati Srebro. Exploring generalization in deep learning. Advances in neural information processing systems, 30, 2017.
[43] Son Nguyen, Duong Nguyen, Khai Nguyen, Khoat Than, Hung Bui, and Nhat Ho. Structured dropout variational inference for bayesian neural networks. Advances in Neural Information Processing Systems, 34:15188-15201 2021.
[44] Van-Anh Nguyen, Trung Le, Anh Bui, Thanh-Tour Do, and Dinh Phung. Optimal transport model distributional robustness. In Advances in Neural Information Processing Systems, 2023.
3
[45] Victor M-H Ong, David J Nott, and Michael S Smith. Gaussian variational approximation with factor covariance structure. Journal of Computational and Graphical Statistics, 2763)-465-478 2018. 2, 3
12

---

[46] Gabriel Pereyra, George Tucker, Jan Chorowski, Lukasz Kaiser, and Geoffrey E. Hinton. Regularizing neural networks by penalizing confident output distributions. In ICLR (Workshop). OpenReview.net, 2017. 3
[47] Henning Petzka, Michael Kamp, Linara Adilova, Cristian Sminchisescu, and Mario Boley. Relative flatness and generalization. In NeurIPS, pages 18420-18432, 2021. 2,3
[48] Cuong Pham, C. Cuong Nguyen, Trung Le, Phung Dinh, Gustavo Carneio, and Thanh-Toan Do. Model and feature diversity for bayesian neural networks in mutual learning, In Advances in Neural Information Processing Systems, 2023.
[49] Hoang Phan, Trung Le, Trung Phung, Anh Tuan Bui, Nhat Ho, and Dinh Phung. Global- local regularization via distributional robustness. In Francisco Ruiz, Jennifer Dy, and Jan- Willem van de Meent, editors. Proceedings of The 26th International Conference on Artificial Intelligence and Statistics, volume 206 of Proceedings of Machine Learning Research, pages 7644-7664. PMLR, 25-27 Apr 2023. 3
[50] Hoang Phan, Lam Tran, Ngoc N Tran, Nhat Ho, Dinh Phung, and Trung Le. Improving multi-task learning via seeking task-based flat regions. arXiv preprint arXiv:2211.13723, 2022. 3
[51] Zhe Qu, Xingyu Li, Rui Duan, Yao Liu, Bo Tang, and Zhuo Lu, Generalized federated learning via sharpness aware minimization. arXiv preprint arXiv:2206. 02618, 2022.
[52] Hippolyt Ritter, Aleksandar Botev, and David Barber. A scalable laplace approximation for neural networks. In 6th International Conference on Learning Representations, ICLR 2018- Conference Track Proceedings, volume 6. International Conference on Representation Learning, 2018. 2,3
[53] Simone Rossi, Sebastien Marmin, and Maurizio Filippone. Walsh-hadamard variational in- ference for bayesian deep learning. Advances in Neural Information Processing Systems, 9686, 2020. 2,
[54] Jakub Swiatkowski, Kevin Roth, Bastiaan Veeling, Linh Tran, Joshua Dillon, Jasper Snoek, Stephan Mandt, Tim Salimans, Rodolphe Jenatton, and Sebastian Nowozin. The k-tied normal distribution: A compact parameterization of gaussian field posteriors in bayesian neural networks. In International Conference on Machine Learning, pages 9289-9299. PMLR, 2020. 2, 3
[55] Marcin Tomczak, Siddharth Swaroop, and Richard Turner. Efficient low rank gaussian vari- ational inference for neural networks. Advances in Neural Information Processing Systems, 33:4610-4622. 2020. 3,
[56] Tuan Truong, Hoang-Phi Nguyen, Tung Pham, Minh-Tuan Tran, Mehrtash Harandi, Dinh Phung, and Trung Le. Rsam: Learning on manifolds with riemannian sharpness- aware minimization, 2023. 3
[57] Colin Wei, Sham Kakade, and Tengyu Ma. The implicit and explicit regularization effects of dropout. In International conference on machine learning, pages 10181-10192. PMLR, 2020.
[58] Max Welling and Yee Teh. Bayesian learning via stochastic gradient langevin dynamics. In Proceedings of the 28th international conference on machine learning (ICML-II), pages 681-688, 2011. 2.7
[59] Guodong Zhang. Shengyang Sun, David Duvenaud, and Roger Grosse. Noisy natural gradient as variational inference. In International Conference on Machine Learning, pages 5852-5861.
PMLR, 2018. 3
[60] Linfeng Zhang, Jiebo Song, Anni Gao, Jingwei Chen, Chenglong Bao, and Kaisheng Ma. Be your own teacher: Improve the performance of convolutional neural networks via self distillation. In Proceedings of the IEFF/CVF International Conference on Computer Vision, pages 3712-3722, 2019.
13

---

[61] Ying Zhang, Tao Xiang, Timothy M. Hospedales, and Huchuan Lu. Deep mutual learning. 2018 IEEE/CVF Conference on Computer Vision and Pattern Recognition, pages 4320-44 3220-4328, 2018. 2,3
[62] Juntang Zhuang, Boqing Gong, Liangzhe Yuan, Yin Cui, Hartwig Adam, Nicha Dvorekk, Sekhar Tatikonda, James Duncan, and Ting Liu, Surrogate gap minimization improves sharpness-aware training. arXiv preprint arXiv:2203.0088 28060, 2022. 3
14


Reviews:
[{"review_id": "6rrkuDDtob", "reviewer": "Reviewer_P465", "summary": "This paper introduce theories in Bayesian settings and  propose variational inference for the sharpness aware posterior in the context of Bayesian Neural Network. The proposed approach is incorporated with existing state of the art Bayesian Neural Networks and experiments were conducted to show the effectiveness of sharpness aware posterior. The results indicate that incorporating sharpness aware posterior methodolody outperforms baseline in terms of ensemble accuracy , expected callibration error and negative likelihood. Also, experimensts show that models from poroposed approach are less sensitive to noise  and have improved generalization ability", "questions": "Q1) What is the computational head of the proposed approach? ", "limitations": "It would be great to add compared to baseline approaches , what is the computational overhead of the proposed approach.", "rating": 6, "confidence": 2, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "1) The paper proposes to approximate posterior of Bayesian Neural Networks with sharpness aware posterior. \n2) The proposed approach can easily be incorporated to already existing Bayesian Neural Networks\n3) The experiments use state-of-the-art settings\n4) The paper outperforms the baseline methods", "weaknesses": "1) The experments show that proposed approach most of teh times outperforms baseline methods , however, improvement is not that significant. \n2) The experimenst are only for image-classification task\n\n", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "\nWe sincerely appreciate your constructive comments. We carefully address your question to the best of our capabilities and improve our paper based on that.\n\n**Computational head**\n\nWe further report the computational cost of several experiments in Tables 1 and 4 in the attached pdf. Similar to SAM, our flat-seeking method involves the computation of gradients twice. The initial computation is to derive the perturbed model $\\theta'$ and the subsequent one is to update the model. As a result, the training time is nearly double in comparison to non-flat counterparts. We will explicitly discuss this in the limitation section."}, {"review_id": "ZzucmpBlXy", "reviewer": "Reviewer_VfFA", "summary": "The paper proposes a new approach to posterior inference for Bayesian neural networks that takes into account the sharpness/flatness of deep learning models, leading to better generalisation ability. The authors introduce the Sharpness-Aware Posterior (SA-Posterior), which allows the sampling of a set of flat models that improve model generalisation. The paper presents a theoretical framework for the SA-Posterior, including a Bayesian framework and a variational inference approach. The authors demonstrate the effectiveness of their approach through experiments on various datasets, showing that it outperforms existing methods on all metrics of interest.", "questions": "- How does the SA-Posterior approach compare to other recent approaches to improve generalisation in deep learning models?\n- How does the computational complexity of the SA-Posterior approach compare with other approaches to posterior inference for Bayesian neural networks?", "limitations": "- Investigate the potential of the SA posterior approach for other types of neural networks beyond Bayesian neural networks.\n- Investigate the computational complexity of the SA-Posterior approach and develop methods to reduce it.\n- Investigate the potential of the SA-Posterior approach for transfer learning and domain adaptation.", "rating": 5, "confidence": 3, "soundness": 2, "presentation": 2, "contribution": 2, "strengths": "- The paper presents a novel approach to posterior inference for Bayesian neural networks.\n- The theoretical framework for the SA-Posterior is developed, including a Bayesian setting and a variational inference approach.\n- The effectiveness of the SA-Posterior is demonstrated through experiments on various datasets.\n- The paper provides insights into the importance of sharpness/flatness in deep learning models.\n- The paper contributes to the growing body of research on improving the generalisability of deep learning models.", "weaknesses": "- The experiments could be extended to include more datasets and models.\n- The paper does not discuss the computational complexity of the SA-Posterior approach.\n- The paper does not compare the SA-Posterior approach with other recent approaches to improve the generalisation ability of deep learning models.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We sincerely appreciate your constructive comments. We carefully address all the questions listed below to the best of our capabilities and improve our paper based on that.\n\n**The experiments could be extended to include more datasets and models.**\n\nThanks for this comment. In this paper, we conducted the experiments on three datasets including Cifar10, Cifar100, and ImageNet, and compared with 9 baselines including MC-Dropout, Deep-ens, SGLD, SWAG, SWAG-Diag, SGVB, SGVB-LRT, SAM, and bSAM. We believe this current amount of experiments satisfies the standard of our field. \n\nEncouraged by your suggestion, we conduct more experiments on ImageNet for more baselines as reported in Table 1 in the attached pdf. Additionally, we report the largest eigenvalue of the Hessian matrix over several methods applying to WideResNet28x10 on CIFAR100, and the ratio of the largest and fifth eigenvalue (Table 4 in the attached pdf) as evidence that our method updates models to minima having lower curvature.\n\n**The paper does not discuss the computational complexity of the SA-Posterior approach**\n\nWe further report the computational cost of several experiments in Tables 1 and 4 in the attached pdf. Similar to SAM, our flat-seeking method involves the computation of gradients twice. The initial computation is to derive the perturbed model $\\theta'$ and the subsequent one is to update the model. As a result, the training time is nearly double in comparison to non-flat counterparts. We will explicitly discuss this in the limitation section.\n\n**The paper does not compare the SA-Posterior approach with other recent approaches to improve the generalisation ability of deep learning models**\n\nWe provide the comparison of our proposed methods with _bSAM_ and _SAM_ (cf. Table 6 in supplementary), which are recent works using sharpness-aware minimization for improving the generalization ability of deep nets. The results demonstrate that our flat version of BNN outperforms bSAM in the most of metric scores."}, {"review_id": "oD9G0doXM4", "reviewer": "Reviewer_yGhH", "summary": "This paper extends SAM -- sharpness aware minimization framework of Foret et al, who seeks parameters of Neural Networks in the regimes of flat loss landscape. The current understanding is that a model that exhibits a flat loss landscape exhibits better generalization performance. Particularly in this paper, such methods are extended to variational inference for neural networks, where the posterior (and hence neural network training) is obtained through an objective function that seeks flat local minima. Experiments are conducted, which were designed to incorporate their objective function into the current Bayesian Neural Networks. The results show consistent improvements over the vanilla counterparts.", "questions": "--", "limitations": "I could not find the limitation section in the main body of the paper.", "rating": 5, "confidence": 2, "soundness": 3, "presentation": 1, "contribution": 3, "strengths": "The contribution of the paper seems interesting. The framework is utilizing PAC-Bayes like theory, where an upper bound to the true loss function is obtained through a careful balance between empirical loss and certain KL divergence. The later regularizes the Bayesian neural network overfitting. It is interesting to see, how variational inference fits the true posterior, which is in turn modified according to the flat local minima seeking objectives. As SAM is one of the popular methods to understand generalization (while being simple), such extensions to Bayesian Neural Networks might be meaningful to the community.", "weaknesses": "I am not an expert in this specific area of combining geometric insights into variational inference. I lean to accept the paper due to my inability to fully grasp this concept. I hope that other reviewers can provide more meaningful feedback for improvements here.\n\nFor the weakness:\n\n- presentation could perhaps improve a bit more.\n\n1. in the abstract, I do not see the logical connection between generic explanations of Bayesian Neural Networks, and generalization. The switch was to me, a bit disconnected and sounded a bit sudden. In the introduction, I also did not see much the role of the 2nd paragraph. Indeed, there have been many variation inference approaches, but I could not see the connection between the expressivity of the posterior, and the approaches of seeking flatness in variational inference.\n\n2. in the theoretic part, I see several equations that repeat also in other papers, e.g., the upper bound of true loss with empirical and KL divergence, how the true posterior provides tight generalization bound, etc. While I appreciate the author's efforts, especially in the appendix for many derivations, it may also make sense to provide citations more in the corresponding texts, so that the readers may be able to learn also through other papers.\n\n3. lines 170-180 should be extended a bit more. I could not understand the connection between each question on page 5 and also missed insights between each derivation.\n\n4. It might be helpful to include an algorithm part, so that the readers may better comprehend how the method works (and hopefully also why).\n\n- empirical improvements seem rather minor, and it was not clear to me, how these experiments demonstrate generalization.\n\nWhile I appreciate the ample empirical evidence provided by the paper, most of the improvements in terms of accuracy are less than 1 percent. This may raise a question if the experimental design here is adequate to show the generalization performance. Table 3, 4, and 5 need error bars, which is requirements in Bayesian machine learning works.\n\nOn how to test generalization, I think the paper could adapt few-shot settings or experiments from meta-learning. Therein, little training data is provided, and one may be able to test better how the model still generalizes. I am not sure, how testing on vanilla cifar10, cifar100 or imagenet classification score is an indication of generalization performance.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We sincerely appreciate your constructive comments. We carefully address all the questions listed below to the best of our capabilities and improve our paper based on that.\n\n**Presentation could perhaps improve a bit more.**\n\nThank you for pointing this out. We will enhance the clarity and motivation behind the fusion of Bayesian Neural Networks and the concept of generalization in the revised version.\n\n**Provide citations more in the corresponding texts, so that the readers may be able to learn also through other papers.**\n\nThank you for recognizing our efforts in the theoretical derivations. We will definitely provide citations for equations that have also been explored in other works. This will offer readers a broader context for understanding our extensions.\n\n**I could not understand the connection between each question on page 5 and also missed insights between each derivation.**\n\nWe apologize for any confusion caused. We will begin by presenting the proof of Theorem 3.2, which extends the concept of sharpness-aware minimization to general and empirical losses over a distribution Q.\n\nIn the following derivation, we modify the upper-bound of the inequality to explicitly represent _the sharpness over the distribution Q in the model space_. This modification aims to offer readers a clearer understanding of the concept of sharpness-aware minimization over Q.\nNotably, the upper-bound presented in Equation (5) is directly influenced by the insights from Theorem 3.2.\nIn the revised version, we will ensure that these connections are articulated more comprehensively to enhance clarity.\n\n**Include an algorithm part**\n\nWe appreciate your suggestion. While we view our approach as a general framework applicable to various algorithms and settings, we recognize the importance of including specific algorithms for better comprehension. In the revised version, we will present algorithms for different settings to enhance clarity.\n\n\n**Empirical improvements seem rather minor, and it was not clear to me, how these experiments demonstrate generalization. Table 3, 4, and 5 need error bars, which is requirements in Bayesian machine learning works.**\n\nWe understand the importance of the error bar for each experiment in the BNN field and try our best to provide as much as possible in the rebuttal time. Due to the limited rebuttal time, we can only report the error bar for Gaussian posterior experiments as shown in Table 3 in the attached pdf. We will try our best to report the error bars for our approaches and the baselines in the revised version. Meanwhile, we are able to conduct one more setting of SGLD and F-SGLD on ImageNet (Table 1 in attached pdf) and also report the largest eigenvalue of the Hessian matrix over several methods, and the ratio of the largest and fifth eigenvalue (Table 4 in attached pdf) as additional evidence that our method updates models to minima having lower curvature (based on the connection between the geometry of the loss landscape\nand generalization, which has been studied extensively [1]).\n\n[1] Gintare Karolina Dziugaite and Daniel M Roy. Computing non vacuous generalization bounds for deep (stochastic) neural networks with many more parameters than training data. arXiv preprint arXiv:1703.11008, 2017. \n\n**Adapt few-shot settings or experiments from meta-learning.**\n\nThank you for your suggestion on the evaluation of generation under alternate settings. They are directly related to generalization ability, hence potentially applying our approaches, however, we leave it for future works. Currently, our paper mainly focuses on experiments on the Bayesian Neural Network and the effectiveness of sharpness aware posterior. The effectiveness is demonstrated through the consistence improvement of accuracy, which is a reasonable measurement of generalization ability. In Table 5 in the main paper, we also evaluate our proposed approach against the non-flat versions in an out-of-distribution setting. Particularly, the original images added different types of corruption and noise. The consistent improvement in both overall accuracy and ECE is additional evidence of the generalization ability of our proposed methods."}, {"review_id": "NiJaskvoaE", "reviewer": "Reviewer_hAuQ", "summary": "This paper introduces a method called Sharpness-Aware Bayesian Neural Networks (SABNN) that aims to improve generalization performance on test datasets. The key idea is to replace the negative empirical loss function with the negative SAM (Sharpness-Aware Minimization) [1] loss function, enabling consideration of the flatness of the loss surface. The authors provide both theoretical analysis and empirical evidence to demonstrate that SABNN achieves better generalization compared to existing methods when evaluated on test datasets.\n\nReferences\n\n[1] Foret, P., A. Kleiner, H. Mobahi, et al. Sharpness-aware minimization for efficiently improving generalization. In International Conference on Learning Representations. 2021.", "questions": "Experiment\n- It is requested to provide the sharpness scores and the largest Hessian eigenvalues for the entire dataset, as well as for the SWAG and F-SWAG models. I think the eigenvalues can be calculated for the mean $\\mu$ of each method.\n- There is a desire for a discussion regarding potential additional computational costs associated with the proposed method in comparison to the baselines.\n- The results of the Deep Ensemble, which is an ensemble of independently trained models without KL in Bayesian Deep Ensemble and with L2 regularization loss, trained with SAM loss, are also requested. The aim is to assess whether the proposed SABNN method is practically beneficial or not.\n\nI am willing to increase my score after considering the experiment results provided in the rebuttal.", "limitations": "Limitation\n- This paper did not address the limitation of the proposed method.\n- One possible limitation can be additional computational costs.", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "Originality\n- The paper presents a theoretical demonstration of the robustness of sharpness-aware posterior inference for the entire distribution of the true dataset.\n- The authors propose a straightforward variational approach that utilizes a Gaussian approximate posterior for efficient and straightforward inference.\n\nClarity\n- The paper is well-written and provides clear explanations, making it easy to understand and follow.", "weaknesses": "Method\n- The method described in the paper involves a straightforward and nearly direct modification of the likelihood function by replacing the negative empirical loss with the negative SAM loss function.\n\nExperiment\n- In section 4.3 of the paper, the authors mention that the sharpness scores are provided in the supplementary material. However, I was unable to locate the results in the supplementary material.\n- The paper does not discuss any potential additional computational costs associated with the proposed method in comparison to the baselines.\n", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thanks for your constructive comments. Based on your suggestion, we conduct some more experiments (report in the attached pdf file) and hope that we can address some of your points as presented below.\n\n**Largest Hessian eigenvalue**   \nWe report the log scale of the largest eigenvalue of the Hessian matrix over several methods applying to WideResNet28x10 on CIFAR100, and the ratio of the largest and fifth eigenvalue as shown in Table 4 in the pdf attached, which evidently indicates that our method updates models to minima having lower curvature.\n\n**The sharpness scores**      \nWe report the sharpness scores of the PreResNet-164 network for SWAG, and F-SWAG training on CIFAR100 in Figure 1 of the attached pdf. This shows that our proposed approach gains better sharpness scores than the baseline.\n\n**Computational cost**    \n\nWe further report the computational cost of several experiments in Tables 1 and 4 in the attached pdf. Similar to SAM, our flat-seeking method involves the computation of gradients twice. The initial computation is to derive the perturbed model $\\theta'$ and the subsequent one is to update the model. As a result, the training time is nearly double in comparison to non-flat counterparts. We will explicitly discuss this in the limitation section.\n\n**Additional experiments on without the regularization term for the deep ensemble**     \nWe present the result of training Deep-ensemble with SAM following the formula for the flat version (Section 2 supplementary) _without KL (or L2 regularisation)_ in Table 2 of the attached pdf. Each experiment is performed three times and the mean and standard deviation are reported. Based on the result, the flat version *with KL divergence* performs better than the one *without KL divergence*. "}, {"review_id": "J209UKW2EF", "reviewer": "Reviewer_FrCB", "summary": "This paper proposes modifying the loss used for Bayesian neural networks (BNNs) to take into account the sharpness/flatness of the loss with respect to the model parameters. Theory, based on that of sharpness-aware minimization (SAM), is developed to propose this loss modification. Making BNNs sharpness-aware led to increased accuracy, decreased negative log-likelihood (NLL), and decreased expected calibration error (ECE) for large neural network models across CIFAR-10, CIFAR-100, and ImageNet.", "questions": "1. Are different optimization problems used for some of the studied methods? If so, please add those details to Section 3.\n\n2. Is it possible to have results on ImageNet for more of the studied methods?\n\n3. I'd recommend putting the number of repetitions used to create the confidence intervals in the captions of the tables. I see that, in the supplementary material, it was three, but I think that this information, along with the number of samples used for the methods, should be in the main text.\n\n4. How do you think that sharpness-aware methods would interact with Laplace approximation methods, another popular BNN method? (This is not a necessary addition to the paper, in my opinion. However, I am interested in your answer to this.)", "limitations": "Limitations and societal impact are not discussed. One limitation that could be added to the conclusion is that the proposed Gaussian variational approach only uses a diagonal covariance for the posterior approximation. While this is standard, it is a limitation that, by mentioning it, could show areas of potential improvement for the proposed variational method.", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "This paper does an excellent job of combining sharpness-aware and BNN methods. The theory for this is given and the theory is applied to several BNN method, namely stochastic weight averaging Gaussian (SWAG), stochastic gradient Langevin dynamics (SGLD), Monte Carlo (MC) dropout, and deep ensembles, for CIFAR-10 and CIFAR-100. These experiments show reasonable improvements. ", "weaknesses": "1. The difference between the optimization problems used in the different Bayesian methods was not clear enough. The main optimization problem discussed was for variational inference. To my knowledge, however, SWAG and SGLD typically use a procedure that is based on the gradient of the log of the posterior.\n\n2. The incomplete ImageNet experiment. Only SWAG and FSWAG are evaluated for ImageNet, versus all of the methods that were applied to CIFAR-10 and CIFAR-100.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "\nWe sincerely appreciate your constructive comments. We are dedicated to addressing all the questions listed below to the best of our capabilities.\n\n**Is it possible to have results on ImageNet for more of the studied methods?**\n\nWe acknowledged the limited number of experiments on ImageNetWe and would like to report the additional results of SGLD and F-SGLD in Table 1 in the attached pdf. It's important to emphasize that the process of training models on the ImageNet dataset consumes significantly more time compared to other datasets, making it challenging to expand experiments in a limited rebuttal period. However, we wish to maximize our efforts and do as much as possible to add to the revised version.\n\n**I’d recommend putting the number of repetitions used to create the confidence intervals in the captions of the tables**\n\nWe will carefully revise the paper incorporating your suggestions to enhance its quality and comprehensiveness.\n\n**How do you think that sharpness-aware methods would interact with Laplace approximation methods, another popular BNN method?**\n\nThanks for this interesting but challenging question. A possible solution is to start from our proposed sharpness-aware posterior $p^{SA}$ and then perform second-order Taylor expansion for $\\text{ln } p^{SA}(\\theta \\mid \\mathcal{D})$. The $\\theta_{MAP}$ is the SAM solution and we can resort to some techniques mentioned in the paper [1] to approximate the Hessian matrix at $\\theta_{MAP}$.\n\nWe can even go further with more complicated formulations by taking account of the derivative of $s(\\theta) = \\text{argmax}_{\\theta' \\in B\\_{\\rho(\\theta)}} \\mathcal{L}\\_{S(\\theta')}$ in the derivations, which potentially leads to a Hessian matrix.\n\n[1] Daxberger, E. et al. : Laplace redux-effortless bayesian deep learning, NeurIPS21."}]

Author rebuttal:
We appreciate the reviewers' constructive comments. We would like to report additional experiments on both ImageNet and CIFAR datasets, the computation cost, the sharpness scores, and the eigenvalues of the Hessian matrix in the attached pdf.

Meta-review: The paper introduces an approach to posterior inference for Bayesian Neural Networks (BNNs) that considers the sharpness/flatness of deep learning models. The primary contribution is the Sharpness-Aware Posterior (SA-Posterior), which targets flat models, aiming to boost model generalization. The proposed method is theoretically grounded and demonstrates improvements over existing techniques via experiments conducted on datasets like CIFAR-10, CIFAR-100, and ImageNet.

**Strengths of the paper:**
- Proposes a novel combination of sharpness-aware and Bayesian neural network methodologies.
- The theoretical foundation is robust, relying on the principles of sharpness-aware minimization.
- Can be easily integrated into existing Bayesian Neural Network setups.
- Demonstrates empirically that it surpasses baseline methodologies across several datasets.

**Weaknesses of the paper:**
- The improvement over baseline methods, although consistent, is incremental or not statistically significant in many cases.
- Experimental setup is primarily limited to image classification tasks and is not shown to scale up to ImageNet. Fortunately, this is partially mitigated by new results in the rebuttal.
- Computational complexity and overhead when compared to baseline methods aren't addressed.
- Some reviewers highlighted issues in presentation, stating that parts of the paper were not clear enough or seemed disconnected.

The paper makes a meaningful contribution to the field and provides valuable insights. It could benefit from more thorough experiments and clearer presentation of some concepts. I recommend the authors incorporate their new results from the rebuttal into a revision and polish the writing in order to improve clarity.
Decision: Accept (poster)