Title: Flat Seeking Bayesian Neural Networks
Venue: NeurIPS 2023
Accepted: yes
Keywords: Bayesian

sharpness-aware

posterior
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

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.

1 Introduction

Bayesian Neural Networks (BNNs) provide a way to interpret deep learning models probabilistically. This is done by setting a prior distribution over model parameters and then inferring a posterior distribution over model parameters based on observed data. This allows us to not only make predictions, but also quantify prediction uncertainty, which is useful for many real-world applications. To sample deep learning models from complex and complicated posterior distributions, advanced particle-sampling approaches such as Hamiltonian Monte Carlo (HMC) [41], Stochastic Gradient HMC(SGHMC) [10], Stochastic Gradient Langevin dynamics (SGLD) [58], and Stein Variational Gradient Descent (SVGD) [36] are often used. However, these methods can be computationally expensive, particularly when many models need to be sampled for better ensembles.

To alleviate this computational burden and enable the sampling of multiple deep learning models from posterior distributions, variational inference approaches employ approximate posteriors to estimate the true posterior. These methods utilize approximate posteriors that belong to sufficiently rich families, which are both economical and convenient to sample from. However, the pioneering works in variational inference, such as [21, 5, 33], assume approximate posteriors to be fully factorized distributions, also known as mean-field variational inference. This approach fails to account for the strong statistical dependencies among random weights of neural networks, limiting its ability to capture the complex structure of the true posterior and estimate the true model uncertainty. To overcome this issue, latter works have attempted to provide posterior approximations with richer

Thanh-Toan Do 1

expressiveness [61, 52, 53, 54, 20, 45, 55, 30, 48]. These approaches aim to improve the accuracy of the posterior approximation and enable more effective uncertainty quantification.

In the context of standard deep network training, it has been observed that flat minimizers can enhance the generalization capability of models. This is achieved by enabling them to locate wider local minima that are more robust to shifts between train and test sets. Several studies, including [27, 47, 15], have shown evidence to support this principle. However, the posteriors used in existing Bayesian neural networks (BNNs) do not account for the sharpness/flatness of the models derived from them in terms of model formulation. As a result, the sampled models can be located in regions of high sharpness and low flatness, leading to poor generalization ability. Moreover, in variational inference methods, using approximate posteriors to estimate these non-sharpness-aware posteriors can result in sampled models from the corresponding optimal approximate posterior lacking awareness of sharpness/flatness, hence causing them to suffer from poor generalization ability.

In this paper, our objective is to propose a sharpness-aware posterior for learning BNNs, which samples models with high flatness for better generalization ability. To achieve this, we devise both a Bayesian setting and a variational inference approach for the proposed posterior. By estimating the optimal approximate posteriors, we can generate flatter models that improve the generalization ability. Our approach is as follows: In Theorem 3.1, we show that the standard posterior is the optimal solution to an optimization problem that balances the empirical loss induced by models sampled from an approximate posterior for fitting a training set with a Kullback-Leibler (KL) divergence, which encourages a simple approximate posterior. Based on this insight, we replace the empirical loss induced by the approximate posterior with the general loss over the entire data-label distribution in Theorem 3.2 to improve the generalization ability. Inspired by sharpness-aware minimization [16], we develop an upper-bound of the general loss in Theorem 3.2, leading us to formulate the sharpnessaware posterior in Theorem 3.3. Finally, we devise the Bayesian setting and variational approach for the sharpness-aware posterior. Overall, our contributions in this paper can be summarized as follows:

Wepropose and develop theories, the Bayesian setting, and the variational inference approach for the sharpness-aware posterior. This posterior enables us to sample a set of flat models that improve the model generalization ability. We note that SAM [16] only considers the sharpness for a single model, while ours is the first work studying the concept and theory of the sharpness for a distribution 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] , MCDropout [18] , Bayesian deep ensemble [35] , and SWAG [39] to demonstrate that the flatseeking 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 a space of solutions. Alternatively, stochastic gradient Langevin dynamics (SGLD) [58] employs first-order Langevin dynamics in the stochastic gradient setting. Additionally, Stein Variational Gradient Descent (SVGD) [36] maintains a set of particles to gradually approach a posterior distribution. Theoretically, all SGHMC, SGLD, and SVGD asymptotically sample from the posterior in the limit of infinitely small step sizes.

Variational Inference : This approach uses an approximate posterior distribution in a family to estimate the true posterior distribution by maximizing a variational lower bound. [21] suggests fitting

a Gaussian variational posterior approximation over the weights of neural networks, which was generalized in [32, 33, 5], using the reparameterization trick for training deep latent variable models. To provide posterior approximations with richer expressiveness, many extensive studies have been proposed. Notably, [38] treats the weight matrix as a whole via a matrix variate Gaussian [22] and approximates the posterior based on this parameterization. Several later works have inspected this distribution to examine different structured representations for the variational Gaussian posterior, such as Kronecker-factored [59, 52, 53], k-tied distribution [54], non-centered or rank-1 parameterization [20, 14]. Another recipe to represent the true covariance matrix of Gaussian posterior is through the low-rank approximation [45, 55, 30, 39].

Dropout Variational Inference: This approach utilizes dropout to characterize approximate posteriors. Typically, [18] and [33] use this principle to propose Bayesian Dropout inference methods such as MC Dropout and Variational Dropout. Concrete dropout [19] extends this idea to optimize the dropout probabilities. Variational Structured Dropout [43] employs Householder transformation to learn a structured representation for multiplicative Gaussian noise in the Variational Dropout method.

2.2 Flat Minima

Flat minimizers have been found to improve the generalization ability of neural networks. This is because they enable models to find wider local minima, which makes them more robust against shifts between train and test sets [27, 47, 15, 44]. The relationship between generalization ability and the width of minima has been investigated theoretically and empirically in many studies, notably [23, 42, 12, 17]. Moreover, various methods seeking flat minima have been proposed in [46, 9, 29, 25, 16, 44]. Typically, [29, 26, 57] investigate the impacts of different training factors such as batch size, learning rate, covariance of gradient, and dropout on the flatness of found minima. Additionally, several approaches pursue wide local minima by adding regularization terms to the loss function [46, 61, 60, 9]. Examples of such regularization terms include softmax output's low entropy penalty [46] and distillation losses [61, 60].

SAM, a method that aims to minimize the worst-case loss around the current model by seeking flat regions, has recently gained attention due to its scalability and effectiveness compared to previous methods [16, 56]. SAM has been widely applied in various domains and tasks, such as metalearning bi-level optimization [1], federated learning [51], multi-task learning [50], where it achieved tighter convergence rates and proposed generalization bounds. SAM has also demonstrated its generalization ability in vision models [11], language models [3], domain generalization [8], and multi-task learning [50]. Some researchers have attempted to improve SAM by exploiting its geometry [34, 31], additionally minimizing the surrogate gap [62], and speeding up its training time [13, 37]. Regarding the behavior of SAM, [28] empirically studied the difference in sharpness obtained by SAM [16] and SWA [24], [40] showed that SAM is an optimal Bayes relaxation of the standard Bayesian inference with a normal posterior, while [44] proved that distribution robustness [4, 49] is a probabilistic extension of SAM.

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.

3.1 Problem Setting and Motivation

We aim to develop Sharpness-Aware Bayesian Neural Networks (SA-BNN). Consider a family of neural networks f θ ( x ) with θ ∈ Θ and a training set S = { ( x 1 , y 1 ) , ..., ( x n , y n ) } where ( x i , y i ) ∼ D . We wish to learn a posterior distribution Q SA S with the density function q SA ( θ |S ) such that any model θ ∼ Q SA S is aware of the sharpness when predicting over the training set S .

We depart with the standard posterior

where the prior distribution P has the density function p ( θ ) and the likelihood has the form

with the loss function ℓ . The standard posterior Q S has the density function defined as

where λ ≥ 0 is a regularization parameter.

We define the general and empirical losses as follows:

Basically, the general loss is defined as 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

Given a distribution Q with the density function q ( θ ) over the model parameters θ ∈ Θ , we define the empirical and general losses over this model distribution Q as

Specifically, the general loss over the model distribution 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 Q is defined as the expectation of the empirical losses incurred by the models sampled from this distribution.

3.2 Our Theory Development

We now present the theory development for the sharpness-aware posterior whose proofs can be found in the supplementary material. Inspired by the Gibbs form of the standard posterior Q S in Eq. (2), we establish the following theorem to connect the standard posterior Q S with the density q ( θ | S ) and the empirical loss L S ( Q ) [7, 2].

Theorem 3.1. Consider the following optimization problem

where we search over Q absolutely continuous w.r.t. P and KL ( · , · ) is the Kullback-Leibler divergence. This optimization has a closed-form optimal solution Q ∗ with the density

which is exactly the standard posterior Q S with the density q ( θ | S ) .

Theorem 3.1 reveals that we need to find the posterior Q S balancing between optimizing its empirical loss L S ( Q ) and simplicity via KL ( Q , P ) . However, minimizing the empirical loss L S ( 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):

Notably, solving the optimization problem (OP) in (4) is generally intractable. To make it tractable, we find its upper-bound which is relevant to the sharpness of a distribution Q over models as shown in the following theorem.

Theorem 3.2. Assume that Θ is a compact set. Under some mild conditions, given any δ ∈ [0; 1] , with the probability at least 1 -δ over the choice of S ∼ D n , for any distribution Q , we have

where f is a non-decreasing function w.r.t. the first variable and approaches 0 when the training size n approaches ∞ .

We note that the proof of Theorem 3.2 is not a trivial extension of sharpness-aware minimization because we need to tackle the general and empirical losses over a distribution Q . To make explicit our sharpness over a distribution Q on models, we rewrite the upper-bound of the inequality as

where the first term E θ ∼ Q [ max θ ′ : ∥ θ ′ -θ ∥≤ ρ L S ( θ ′ ) -L S ( θ ) ] can be regarded as the sharpness over the distribution Q on the model space and the last term f ( max θ ∈ Θ ∥ θ ∥ 2 , n ) is a constant.

Moreover, inspired by Theorem 3.2, we propose solving the following OP which forms an upperbound of the desirable OP in (4)

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 SA S with the density function q SA ( θ |S ) :

where we have defined s ( θ ) = argmax θ ′ : ∥ θ ′ -θ ∥ ≤ ρ L S ( θ ′ ) .

Theorem 3.3 describes the close form of the sharpness-aware posterior distribution Q SA S with the density function q SA ( θ |S ) . Based on this characterization, in what follows, we introduce the SA Bayesian setting that sheds lights on its variational approach.

3.3 Sharpness-Aware Bayesian Setting and Its Variational Approach

Bayesian Setting: To promote the Bayesian setting for sharpness-aware posterior distribution Q SA S , we examine the sharpness-aware likelihood

With this predefined sharpness-aware likelihood, we can recover the sharpness-aware posterior distribution Q SA S with the density function q SA ( θ |S ) :

Variational inference for the sharpness-aware posterior distribution: We now develop the variational inference for the sharpness-aware posterior distribution. Let denote X = [ x 1 , ..., x n ] and

Y = [ y 1 , ..., y n ] . Considering an approximate posterior family { q ϕ ( θ ) : ϕ ∈ Φ } , we have

It is obvious that we need to maximize the following lower bound for maximally reducing the gap KL ( q ϕ , q SA ) :

which can be equivalently rewritten as

Derivation for Variational Approach with A Gaussian Approximate Posterior: Inspired by the geometry-based SAM approaches [34, 31], we incorporate the geometry to the SA variational approach via the distance to define the ball for the sharpness as ∥ θ ′ -θ ∥ diag ( T θ ) =

To further clarify, we consider our SA posterior distribution to Bayesian NNs, wherein we impose the Gaussian distributions to its weight matrices W i ∼ N ( µ i , σ 2 i I ) , i = 1 , . . . , L 1 . The parameter ϕ consists of µ i , σ i , i = 1 , . . . , L . For θ = W 1: L ∼ q ϕ , using the reparameterization trick W i = µ i + diag ( σ i ) ϵ i , ϵ i ∼ N (0 , I ) and by searching θ ′ = W ′ 1: L with W ′ i = µ ′ i + diag ( σ i ) ϵ i , ϵ i ∼ N (0 , I ) , the constraint ∥ θ -θ ′ ∥ diag ( T θ ) = ∥ µ -µ ′ ∥ diag ( T θ ) with µ = µ 1: L and µ ′ = µ ′ 1: L . Thus, the OP in (6) reads

where σ = σ 1: L , ϵ = ϵ 1: L , and we define diag ( T θ ) = diag ( T µ,σ ) in the distance of the geometry.

To solve the OP in (7), we sample ϵ = ϵ 1: L from the standard Gaussian distributions, employ an one-step gradient ascent to find µ ′ , and use the gradient at µ ′ to update µ . Specifically, we find µ ′ [6] (Chapter 9) as

The diagnose of diag ( T µ,σ ) specifies the importance level of the model weights, i.e., the weight with a higher importance level is encouraged to have a higher sharpness via a smaller absolute partial derivative of the loss w.r.t. this weight. We consider diag ( T µ,σ ) = I (i.e., the standard SA BNN ) and diag ( T µ,σ ) = diag ( | µ | σ ) (i.e., the geometry SA BNN ). Here we note that · · represents the element-wise division.

Finally, the objective function in (6) indicates that we aim to 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 sharpnessaware posterior can be integrated into MCMC-based and variational inference-based Bayesian Neural Networks. The supplementary material contains the details on how to derive variational approaches and incorporate the sharpness-awareness into the BNNs used in our experiments 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].

4 Experiments

In this section, we conduct various experiments to demonstrate the effectiveness of the sharpnessaware 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 reparameterization trick), BNNs with stochastic gradient Langevin dynamics (SGLD) [58], MCDropout [18], Bayesian deep ensemble [35], and SWAG [39]. 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. The details of the dataset and implementation are described in the supplementary material 2 .

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 a notable improvement across all experiments. It is worth noting that there is a trade-off between accuracy, negative log-likelihood, and expected calibration error. Nonetheless, our approach obtains a 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

F-SWAG

SWAG

6.4e-01

0.4

0.2

3.5e-03

0.3

0.15

0.2

0.05

0.10

Confidence-

0.1

Confidence

I

0.00

Confidence-

0.05

0.0

0.05

0.00

0.1

F-SWA

SWA

0.10

F-SWAG-Diag

SWAG-Diag

0.05

SWAG

F-SWAG

0.0

0.2

0.4

0.6

0.8

1.0

0.0

0.2

0.4

0.6

0.8

1.0

0.0

0.2

0.4

0.6

0.8

1.0

Confidence(maxprob)

Confidence(maxprob)

Confidence (max prob)

and confidence of the model. To further clarify it, we display the Reliability Diagrams of PreResNet164 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-aware 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-C dataset. This is an extension of the CIFAR-10 designed to 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 sharpness-aware 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-aware one. Particularly, we compare two methods F-SWAG and SWAGbyselecting 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.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 a model θ , the sharpness is evaluated as max || ϵ || 2 ≤ ρ L S ( θ + ϵ ) -L S ( θ ) to measure the change of loss value around θ . We calculate the sharpness score of PreResNet-164 network for SWAG, and F-SWAG training on 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.

5 Conclusion

In this paper, we introduce theories in the Bayesian setting and discuss variational inference for the sharpness-aware 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 a higher flatness. We conducted extensive experiments that leveraged the sharpness-aware posterior with state-of-the-art Bayesian Neural Networks. Our main results show that the models sampled from the proposed posterior outperform their baselines in terms of ensemble accuracy, expected calibration error (ECE), and negative log-likelihood (NLL). This indicates that the flat-seeking 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 out-of-distribution prediction. Overall, the sharpness-aware posterior presents a promising approach for improving the generalization performance of Bayesian neural networks.

Acknowledgements. This work was partly supported by ARC DP23 grant DP230101176 and by the Air Force Office of Scientific Research under award number FA2386-23-1-4044.

References

Momin Abbas, Quan Xiao, Lisha Chen, Pin-Yu Chen, and Tianyi Chen. Sharp-maml: Sharpnessaware model-agnostic meta learning. arXiv preprint arXiv:2206.03996 , 2022. 3

Pierre Alquier, James Ridgway, and Nicolas Chopin. On the properties of variational approximations of gibbs posteriors. The Journal of Machine Learning Research , 17(1):8374-8414, 2016. 4

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 1: Long Papers) , pages 7360-7371, Dublin, Ireland, May 2022. Association for Computational Linguistics. 3

Jose Blanchet and Karthyek Murthy. Quantifying distributional model risk via optimal transport. Mathematics of Operations Research , 44(2):565-600, 2019. 3

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. 1, 3

Stephen Boyd and Lieven Vandenberghe. Convex Optimization . Cambridge University Press, March. 6

Olivier Catoni. Pac-bayesian supervised classification: the thermodynamics of statistical learning. arXiv preprint arXiv:0712.0248 , 2007. 4

Junbum Cha, Sanghyuk Chun, Kyungjae Lee, Han-Cheol Cho, Seunghyun Park, Yunsung Lee, and Sungrae Park. Swad: Domain generalization by seeking flat minima. Advances in Neural Information Processing Systems , 34:22405-22418, 2021. 3

Pratik Chaudhari, Anna Choroma´ nska, 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. 3

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

Xiangning Chen, Cho-Jui Hsieh, and Boqing Gong. When vision transformers outperform resnets without pre-training or strong data augmentations. arXiv preprint arXiv:2106.01548 , 2021. 3

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

Jiawei Du, Daquan Zhou, Jiashi Feng, Vincent YF Tan, and Joey Tianyi Zhou. Sharpness-aware training for free. arXiv preprint arXiv:2205.14083 , 2022. 3

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 on machine learning , pages 2782-2792. PMLR, 2020. 3

Gintare Karolina Dziugaite and Daniel M. Roy. Computing nonvacuous generalization bounds for deep (stochastic) neural networks with many more parameters than training data. In UAI . AUAI Press, 2017. 2, 3

Pierre Foret, Ariel Kleiner, Hossein Mobahi, and Behnam Neyshabur. Sharpness-aware minimization for efficiently improving generalization. In International Conference on Learning Representations , 2021. 2, 3

Stanislav Fort and Surya Ganguli. Emergent properties of the local geometry of neural loss landscapes. arXiv preprint arXiv:1910.05929 , 2019. 3

Yarin Gal and Zoubin Ghahramani. Dropout as a 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, 3, 7

Yarin Gal, Jiri Hron, and Alex Kendall. Concrete dropout. Advances in neural information processing systems , 30, 2017. 3

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, 3

Alex Graves. Practical variational inference for neural networks. Advances in neural information processing systems , 24, 2011. 1, 2

Arjun K Gupta and Daya K Nagar. Matrix variate distributions . Chapman and Hall/CRC, 2018. 3

Sepp Hochreiter and J¨ urgen Schmidhuber. Simplifying neural nets by discovering flat minima. In NIPS , pages 529-536. MIT Press, 1994. 3

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

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

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

Yiding Jiang, Behnam Neyshabur, Hossein Mobahi, Dilip Krishnan, and Samy Bengio. Fantastic generalization measures and where to find them. In ICLR . OpenReview.net, 2020. 2, 3

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

Nitish Shirish Keskar, Dheevatsa Mudigere, Jorge Nocedal, Mikhail Smelyanskiy, and Ping Tak Peter Tang. On large-batch training for deep learning: Generalization gap and sharp minima. In ICLR . OpenReview.net, 2017. 3

Mohammad Khan, Didrik Nielsen, Voot Tangkaratt, Wu Lin, Yarin Gal, and Akash Srivastava. Fast and scalable bayesian deep learning by weight-perturbation in adam. In International Conference on Machine Learning , pages 2611-2620. PMLR, 2018. 2, 3

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 11148-11161. PMLR, 17-23 Jul 2022. 3, 6

Diederik P Kingma and Max Welling. Auto-encoding variational bayes. arXiv preprint arXiv:1312.6114 , 2013. 3

Durk P Kingma, Tim Salimans, and Max Welling. Variational dropout and the local reparameterization trick. Advances in neural information processing systems , 28, 2015. 1, 2, 3, 7

Jungmin Kwon, Jeongseop Kim, Hyunseo Park, and In Kwon Choi. Asam: Adaptive sharpnessaware minimization for scale-invariant learning of deep neural networks. In International Conference on Machine Learning , pages 5905-5914. PMLR, 2021. 3, 6

Balaji Lakshminarayanan, Alexander Pritzel, and Charles Blundell. Simple and scalable predictive uncertainty estimation using deep ensembles. Advances in neural information processing systems , 30, 2017. 2, 7

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

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

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

Wesley J. Maddox, Timur Garipov, Pavel Izmailov, Dmitry Vetrov, and Andrew Gordon Wilson. A Simple Baseline for Bayesian Uncertainty in Deep Learning . Curran Associates Inc., Red Hook, NY, USA, 2019. 2, 3, 7

Thomas M¨ ollenhoff and Mohammad Emtiyaz Khan. Sam as an optimal relaxation of bayes. arXiv preprint arXiv:2210.01620 , 2022. 3

Radford M. Neal. Bayesian Learning for Neural Networks . Springer-Verlag, Berlin, Heidelberg, 1996. 1, 2

Behnam Neyshabur, Srinadh Bhojanapalli, David McAllester, and Nati Srebro. Exploring generalization in deep learning. Advances in neural information processing systems , 30, 2017. 3

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-15202, 2021. 3

Van-Anh Nguyen, Trung Le, Anh Bui, Thanh-Toan Do, and Dinh Phung. Optimal transport model distributional robustness. In Advances in Neural Information Processing Systems , 2023. 3

Victor M-H Ong, David J Nott, and Michael S Smith. Gaussian variational approximation with a factor covariance structure. Journal of Computational and Graphical Statistics , 27(3):465-478, 2018. 2, 3

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

Henning Petzka, Michael Kamp, Linara Adilova, Cristian Sminchisescu, and Mario Boley. Relative flatness and generalization. In NeurIPS , pages 18420-18432, 2021. 2, 3

Cuong Pham, C. Cuong Nguyen, Trung Le, Phung Dinh, Gustavo Carneiro, and Thanh-Toan Do. Model and feature diversity for bayesian neural networks in mutual learning. In Advances in Neural Information Processing Systems , 2023. 2

Hoang Phan, Trung Le, Trung Phung, Anh Tuan Bui, Nhat Ho, and Dinh Phung. Globallocal regularization via distributional robustness. In Francisco Ruiz, Jennifer Dy, and JanWillem 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

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

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

Hippolyt Ritter, Aleksandar Botev, and David Barber. A scalable laplace approximation for neural networks. In 6th International Conference on Learning Representations, ICLR 2018Conference Track Proceedings , volume 6. International Conference on Representation Learning, 2018. 2, 3

Simone Rossi, Sebastien Marmin, and Maurizio Filippone. Walsh-hadamard variational inference for bayesian deep learning. Advances in Neural Information Processing Systems , 33:9674-9686, 2020. 2, 3

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 mean field posteriors in bayesian neural networks. In International Conference on Machine Learning , pages 9289-9299. PMLR, 2020. 2, 3

Marcin Tomczak, Siddharth Swaroop, and Richard Turner. Efficient low rank gaussian variational inference for neural networks. Advances in Neural Information Processing Systems , 33:4610-4622, 2020. 2, 3

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

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

Max Welling and Yee W Teh. Bayesian learning via stochastic gradient langevin dynamics. In Proceedings of the 28th international conference on machine learning (ICML-11) , pages 681-688, 2011. 1, 2, 7

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

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 IEEE/CVF International Conference on Computer Vision , pages 3713-3722, 2019. 3

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-4328, 2018. 2, 3

Juntang Zhuang, Boqing Gong, Liangzhe Yuan, Yin Cui, Hartwig Adam, Nicha Dvornek, Sekhar Tatikonda, James Duncan, and Ting Liu. Surrogate gap minimization improves sharpness-aware training. arXiv preprint arXiv:2203.08065 , 2022. 3

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS:\nThis 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.\n\nWEAKNESSES:\n1. 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.\n\nQUESTIONS:\n1. 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.)\n\nLIMITATIONS:\nLimitations 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.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-06-28"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS:\nOriginality\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.\n\nWEAKNESSES:\nMethod\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\nQUESTIONS:\nExperiment\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.\n\nLIMITATIONS:\nLimitation\n- This paper did not address the limitation of the proposed method.\n- One possible limitation can be additional computational costs.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-05"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS:\nThe 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.\n\nWEAKNESSES:\nI 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.\n\nQUESTIONS:\n--\n\nLIMITATIONS:\nI could not find the limitation section in the main body of the paper.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 2, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-06"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe 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.\n\nSTRENGTHS:\n- 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.\n\nWEAKNESSES:\n- 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.\n\nQUESTIONS:\n- 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?\n\nLIMITATIONS:\n- 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.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 2, "DATE": "2023-07-08"}

comments: 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

STRENGTHS:
1) The paper proposes to approximate posterior of Bayesian Neural Networks with sharpness aware posterior. 
2) The proposed approach can easily be incorporated to already existing Bayesian Neural Networks
3) The experiments use state-of-the-art settings
4) 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. 
2) The experimenst are only for image-classification task

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.
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 3
DATE: 2023-07-12