SlowGuess/NeurIPS_2023_OCR huggingface dataset:
$$
\left. \min  _ {\mu , \sigma} \left\{\lambda \mathbb {E} _ {\epsilon} \left[ \max  _ {\left\| \mu^ {\prime} - \mu \right\| _ {\operatorname {d i a g} \left(T _ {\mu , \sigma}\right)} \leq \rho} \mathcal {L} _ {\mathcal {S}} \left(\left[ \mu_ {i} ^ {\prime} + \operatorname {d i a g} (\sigma_ {i}) \epsilon_ {i} \right] _ {i = 1} ^ {L}\right) \right] \right\}, \right. \tag {7}
$$
where  $\sigma = \sigma_{1:L}$ ,  $\epsilon = \epsilon_{1:L}$ , and we define  $\mathrm{diag}(T_{\theta}) = \mathrm{diag}(T_{\mu,\sigma})$  in the distance of the geometry.
To solve the OP in (7), we sample  $\epsilon = \epsilon_{1:L}$  from the standard Gaussian distributions, employ an one-step gradient ascent to find  $\mu'$ , and use the gradient at  $\mu'$  to update  $\mu$ . Specifically, we find  $\mu'$  [6] (Chapter 9) as
$$
\mu^ {\prime} = \mu + \rho \frac {\operatorname {d i a g} (T _ {\mu , \sigma}) \nabla_ {\mu} \mathcal {L} _ {\mathcal {S}} \left(\left[ \mu_ {i} + \operatorname {d i a g} (\sigma_ {i}) \epsilon_ {i} \right] _ {i = 1} ^ {L}\right)}{\left\| \operatorname {d i a g} (T _ {\mu , \sigma}) \nabla_ {\mu} \mathcal {L} _ {\mathcal {S}} \left(\left[ \mu_ {i} + \operatorname {d i a g} (\sigma_ {i}) \epsilon_ {i} \right] _ {i = 1} ^ {L}\right) \right\|}.
$$
The diagnose of  $\mathrm{diag}(T_{\mu, \sigma})$  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  $\mathrm{diag}(T_{\mu, \sigma}) = \mathbb{I}$  (i.e., the standard SA BNN) and  $\mathrm{diag}(T_{\mu, \sigma}) = \mathrm{diag}\left(\frac{|\mu|}{\sigma}\right)$  (i.e., the geometry SA BNN). Here we note that  $\ddagger$  represents the element-wise division.
<sup>1</sup>We absorb the biases to 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.
Method	PreResNet-164	WideResNet28x10
ACC↑	NLL↓	ECE↓	ACC↑	NLL↓	ECE↓
Variational inference
MC-Dropout	79.50 ± 0.37	0.9162 ± 0.0103	0.0993 ± 0.0033	82.30 ± 0.19	0.6500 ± 0.0049	0.0574 ± 0.0028
F-MC-Dropout	81.06 ± 0.44	0.7027 ± 0.0049	0.0514 ± 0.0047	83.24 ± 0.11	0.6144 ± 0.0068	0.0250 ± 0.0027
Deep-ens	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 ± 0.0186	0.0239 ± 0.0047	82.40 ± 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 ± 0.50	0.6595 ± 0.0019	0.0587 ± 0.0048	82.23 ± 0.19	0.6078 ± 0.0006	0.0113 ± 0.0020
F-SWAG	80.93 ± 0.27	0.6704 ± 0.0049	0.0350 ± 0.0025	83.57 ± 0.26	0.5757 ± 0.0136	0.0196 ± 0.0015
Table 2: Classification score on CIFAR-10 dataset.Each experiment is repeated three times with different random seeds and reports the mean and standard deviation.
Method	PreResNet-164	WideResNet28x10
ACC ↑	NLL ↓	ECE ↓	ACC ↑	NLL ↓	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
Deep-ens	96.39 ± 0.09	0.1277 ± 0.0030	0.0108 ± 0.0015	96.96 ± 0.10	0.1031 ± 0.0076	0.0087 ± 0.0018
F-Deep-ens	96.70 ± 0.04	0.1031 ± 0.0016	0.0057 ± 0.0031	97.11 ± 0.10	0.0851 ± 0.0011	0.0059 ± 0.0012
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 ± 0.0002	96.43 ± 0.05	0.1336 ± 0.004	0.0385 ± 0.0003
Sample
SWAG-Diag	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.0043 ± 0.0005	97.05 ± 0.08	0.0888 ± 0.0052	0.0043 ± 0.0004
SWAG	96.03 ± 0.02	0.1232 ± 0.0022	0.0053 ± 0.0004	96.32 ± 0.08	0.1122 ± 0.0009	0.0088 ± 0.0006
F-SWAG	96.25 ± 0.03	0.11062 ± 0.0014	0.0056 ± 0.0002	97.09 ± 0.14	0.0883 ± 0.0004	0.0036 ± 0.0008
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 sharpness-aware 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].
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 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-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<sup>2</sup>.
2The implementation is provided in https://github.com/anh-ntv/flat_bnn.git
Table 3: Classification scores of approximate the Gaussian posterior on the CIFAR datasets. Each experiment is repeated three times with different random seeds and reports the mean and standard deviation.
Method	Resnet10	Resnet18
ACC ↑	NLL ↓	ECE ↓	ACC ↑	NLL ↓	ECE ↓
Experiments on 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 ± 0.67	1.3704 ± 0.01	0.0301 ± 0.03	70.12 ± 1.02	1.121 ± 0.23	0.036 ± 0.06
SGVB	54.40 ± 0.98	1.968 ± 0.05	0.214 ± 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 ± 0.65	1.906 ± 0.02	0.207 ± 0.00	62.58 ± 0.53	1.612 ± 0.03	0.224 ± 0.00
Experiments on Cifar-10 dataset
SGVB-LRT	84.98 ± 1.87	0.422 ± 0.10	0.043 ± 0.04	89.10 ± 1.32	0.344 ± 0.02	0.033 ± 0.02
F-SGVB-LRT	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.014 ± 0.02
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 score on ImageNet dataset
Model	Densenet-161	ResNet-152
ACC ↑	NLL ↓	ECE ↓	ACC ↑	NLL ↓	ECE ↓
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	0.0167
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
Figure 1: Comparing loss landscape of PreResNet-164 on CIFAR-100 dataset training with SWAG and F-SWAG method. For visualization purposes, we sample two models for each SWAG and F-SWAG and then plot the loss landscapes. It can be observed that the loss landscapes of our F-SWAG are flatter, supporting our argument for the flatter sampled models.
Figure 2: Reliability diagrams for PreResNet164 on CIFAR-100. 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 score on CIFAR-10-C on PreResNet-164 model when training with CIFAR-10. The full result on each type of corruption is displayed in the supplementary material.
Corruption	ECE↓	Accuracy ↑
SWAG-D	F-SWAG-D	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.35	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-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.
In Figure 1, we plot the loss-landscape of the models sampled from our proposal of sharpness-aware posterior against the non-sharpness-aware one. Particularly, we compare two methods F-SWAG and SWAG by selecting four random models sampled from the posterior distribution of each method under the same hyper-parameter settings. As observed, our method not only improves the generalization of ensemble inference, demonstrated by classification results in Section 4.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  $\theta$ , the sharpness is evaluated as  $\max_{||\epsilon ||_2\leq \rho}\mathcal{L}_S(\theta +\epsilon) - \mathcal{L}_S(\theta)$  to measure the change of loss value around  $\theta$ . 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.
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.
[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.03996, 2022. 3
[2] 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
[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 1: 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. 1, 3
[6] Stephen Boyd and Lieven Vandenberghe. Convex Optimization. Cambridge University Press, March. 6
[7] Olivier Catoni. Pac-bayesian supervised classification: the thermodynamics of statistical learning. arXiv preprint arXiv:0712.0248, 2007. 4
[8] 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
[9] Pratik Chaudhari, Anna Choromanska, 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
[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 Boqing Gong. When vision transformers outperform resnets without pre-training or strong data augmentations. arXiv preprint arXiv:2106.01548, 2021.3
[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-4328, 2018. 2, 3
[62] 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