SlowGuess/NeurIPS_2023_OCR huggingface dataset:
The optimization problem in (15) can be viewed as a probabilistic extension of SAM. Specifically, for each  $\theta \sim \mathbb{Q}_{\phi}$ , the inner max:  $\max_{\tilde{\theta}}\left\{\mathcal{L}_S(\tilde{\theta}) - \lambda d(\tilde{\theta},\theta)\right\}$  seeks a model  $\tilde{\theta}$  maximizing the loss  $\mathcal{L}_S(\tilde{\theta})$  on a soft ball around  $\theta$  controlled by  $\lambda d(\theta ,\tilde{\theta})$ . Particularly, a higher value of  $\lambda$  seeks the optimal  $\tilde{\theta}$  in a smaller ball. Moreover, in the outer min, the term  $\lambda \rho$  trades off between the value of  $\lambda$  and the radius of the soft ball, aiming to find out an optimal  $\lambda^{*}$  and the optimal model  $\tilde{\theta}^*$  maximizing the loss function over an appropriate soft ball. Here we note that SAM also seeks maximizing the loss function but over the ball with the radius  $\rho$  around the model  $\theta$ .
Interestingly, by appointing a particular distance metric between two models, we can exactly recover the SAM formulation as shown in the following theorem.
Theorem 4.3. With the distance metric  $d$  defined as
$$
d (\theta , \tilde {\theta}) = \left\{ \begin{array}{l l} \| \tilde {\theta} - \theta \| _ {2} & \| \tilde {\theta} - \theta \| _ {2} \leq \rho \\ + \infty & \text {o t h e r w i s e} \end{array} , \right. \tag {16}
$$
the OPs in (13), (14) with  $\mathbb{Q}_{\phi} = \delta_{\theta}$ , and (15) equivalently reduce to the OP of SAM as
$$
\min_{\theta}\max_{\substack{\tilde{\theta}:\| \tilde{\theta} -\theta \|_{2}\leq \rho}}\mathcal{L}_{S}\left(\tilde{\theta}\right).
$$
For simplicity, we absorb the biases to the weight matrices
Table 1: Classification accuracy on the CIFAR datasets of the single model setting with one particle. All experiments are trained three times with different random seeds.
Dataset	Method	WideResnet28x10	Pyramid101	Densenet121
CIFAR-10	SAM	96.72 ± 0.007	96.20 ± 0.134	91.16 ± 0.240
OT-MDR (Ours)	96.97 ± 0.009	96.61 ± 0.063	91.44 ± 0.113
CIFAR-100	SAM	82.69 ± 0.035	81.26 ± 0.636	68.09 ± 0.403
OT-MDR (Ours)	84.14 ± 0.172	82.28 ± 0.183	69.84 ± 0.176
Table 2: Classification score on Resnet18. The results of baselines are taken from [46]
Method	CIFAR-10	CIFAR-100
ACC ↑	AUROC ↑	ACC ↑	AUROC ↑
SGD	94.76 ± 0.11	0.926 ± 0.006	76.54 ± 0.26	0.869 ± 0.003
SAM	95.72 ± 0.14	0.949 ± 0.003	78.74 ± 0.19	0.887 ± 0.003
bSAM	96.15 ± 0.08	0.954 ± 0.001	80.22 ± 0.28	0.892 ± 0.003
OT-MDR (Ours)	96.59 ± 0.07	0.992 ± 0.004	81.23 ± 0.13	0.991 ± 0.001
Theorem 4.3 suggests that the OT-based model distributional robustness on the model space is a probabilistic relaxation of SAM in general, while SAM is also a specific crisp case of the OT-based model distributional robustness on the model space. This connection between the OT-based model distributional robustness and SAM is intriguing and may open doors to propose other sharpness-aware training approaches and leverage adversarial training [23, 42] to improve model robustness. Moreover, [46] has shown a connection between SAM and the standard Bayesian inference with a normal posterior. Using the Gaussian approximate posterior  $N(\omega, \nu I)$ , it is demonstrated that the maximum of the likelihood loss  $L(q_{\mu})$  (i.e.,  $\mu$  is the expectation parameter of the Gaussian  $N(\omega, \nu I)$ ) can be lower-bounded by a relaxed-Bayes objective which is relevant to the SAM loss. Our findings are supplementary but independent of the above finding. Furthermore, by linking SAM to the OT-based model distributional robustness on the model space, we can expect to leverage the rich body theory of distributional robustness for new discoveries about improving the model generalization ability.
In this section, we present the results of various experiments<sup>2</sup> to evaluate the effectiveness of our proposed method in achieving distribution robustness. These experiments are conducted in three main settings: a single model, ensemble models, and Bayesian Neural Networks. To ensure the reliability and generalizability of our findings, we employ multiple architectures and evaluate their performance using the CIFAR-10 and CIFAR-100 datasets. For each experiment, we report specific metrics that capture the performance of each model in its respective setting.
5.1 Experiments on a Single Model
To evaluate the performance of our proposed method for one-particle training, we conducted experiments using three different architectures: WideResNet28x10, Pyramid101, and Densenet121. We compared our approach's results against models trained with SAM optimizer as our baseline. For consistency with the original SAM paper, we adopted their setting, using  $\rho = 0.05$  for CIFAR-10 experiments and  $\rho = 0.1$  for CIFAR-100 and report the result in Table 1. In our OT-MDR method, we chose different values of  $\rho$  for each half of the mini-batch  $B$ , and denoted  $\rho_{1}$  for  $B_{1}$  and  $\rho_{2}$  for  $B_{2}$ , where  $\rho_{2} = 2\rho_{1}$  to be simplified (this simplified setting is used in all experiments). To ensure a fair comparison with SAM, we also set  $\rho_{1} = 0.05$  for CIFAR-10 experiments and  $\rho_{1} = 0.1$  for CIFAR-100. Our approach outperformed the baseline with significant gaps, as indicated in Table 1. On average, our method achieved a  $0.73\%$  improvement on CIFAR-10 and a  $1.68\%$  improvement on CIFAR-100 compared to SAM. These results demonstrate the effectiveness of our proposed method for achieving higher accuracy in one-particle model training.
2The implementation is provided in https://github.com/anh-ntv/OT_MDR.git
Table 3: Evaluation of the ensemble Accuracy (%) on the CIFAR-10/100 datasets. We reproduce all baselines with the same hyperparameter for a fair comparison.
Method	CIFAR-10	CIFAR-100
ACC ↑	Brier ↓	NLL ↓	ECE ↓	AAC ↓	ACC ↑	Brier ↓	NLL ↓	ECE ↓	AAC ↓
Ensemble of five Resnet10 models
Deep Ensemble	92.7	0.091	0.272	0.072	0.108	73.7	0.329	0.87	0.145	0.162
Fast Geometric	92.5	0.251	0.531	0.121	0.144	63.2	0.606	1.723	0.149	0.162
Snapshot	93.6	0.083	0.249	0.065	0.107	72.8	0.338	0.929	0.153	0.338
EDST	92.0	0.122	0.301	0.078	0.112	68.4	0.427	1.151	0.155	0.427
DST	93.2	0.102	0.261	0.067	0.108	70.8	0.396	1.076	0.150	0.396
SGD	95.1	0.078	0.264	-	0.108	75.9	0.346	1.001	-	0.346
SAM	95.4	0.073	0.268	0.050	0.107	77.7	0.321	0.892	0.136	0.321
OT-MDR (Ours)	95.4	0.069	0.145	0.021	0.004	79.1	0.059	0.745	0.043	0.054
Ensemble of three Resnet18 models
Deep Ensemble	93.7	0.079	0.273	0.064	0.107	75.4	0.308	0.822	0.14	0.155
Fast Geometric	93.3	0.087	0.261	0.068	0.108	72.3	0.344	0.95	0.15	0.169
Snapshot	94.8	0.071	0.27	0.054	0.108	75.7	0.311	0.903	0.147	0.153
EDST	92.8	0.113	0.281	0.074	0.11	69.6	0.412	1.123	0.151	0.197
DST	94.7	0.083	0.253	0.057	0.107	70.4	0.405	1.153	0.155	0.194
SGD	95.2	0.076	0.282	-	0.108	78.9	0.304	0.919	-	0.156
SAM	95.8	0.067	0.261	0.044	0.107	80.1	0.285	0.808	0.127	0.151
OT-MDR (Ours)	96.2	0.059	0.134	0.018	0.005	81.0	0.268	0.693	0.045	0.045
Ensemble of ResNet18, MobileNet and EfficientNet
Deep Ensemble	89.0	0.153	0.395	0.111	0.126	62.7	0.433	1.267	0.176	0.209
DST	93.4	0.102	0.282	0.070	0.109	71.7	0.393	1.066	0.148	0.187
SGD	92.6	0.113	0.317	-	0.112	72.6	0.403	1.192	-	0.201
SAM	93.8	0.094	0.280	0.060	0.110	76.4	0.347	1.005	0.142	0.177
OT-MDR (Ours)	94.8	0.078	0.176	0.021	0.007	78.3	0.310	0.828	0.047	0.063
We conduct experiments to compare our OT-MDR with bSAM [46], SGD, and SAM [18] on Resnet18. The results, shown in Table 2, demonstrate that the OT-MDR approach consistently outperforms all baselines by a substantial margin. Here we note that we cannot evaluate bSAM on the architectures used in Table 1 because the authors did not release the code. Instead, we run our OT-MDR with the setting mentioned in the original bSAM paper.
5.2 Experiments on Ensemble Models
To investigate the effectiveness of our approach in the context of a uniform distribution over a model space, we examine the ensemble inference of multiple base models trained independently. The ensemble prediction is obtained by averaging the prediction probabilities of all base models, following the standard process of ensemble methods. We compare our approach against several state-of-the-art ensemble methods, including Deep Ensembles [37], Snapshot Ensembles [25], Fast Geometric Ensemble (FGE) [22], and sparse ensembles EDST and DST [40]. In addition, we compare our approach with another ensemble method that utilizes SAM as an optimizer to improve the generalization ability, as discussed in Section 4.3. The value of  $\rho$  for SAM and  $\rho_{1}$ ,  $\rho_{2}$  for OT-MDR is the same as in the single model setting.
To evaluate the performance of each method, we measure five metrics over the average prediction, which represent both predictive performance (Accuracy - ACC) and uncertainty estimation (Brier score, Negative Log-Likelihood - NLL, Expected Calibration Error - ECE, and Average across all calibrated confidence - AAC) on the CIFAR dataset, as shown in Table 3.
Notably, our OT-MDR approach consistently outperforms all baselines across all metrics, demonstrating the benefits of incorporating diversity across base models to achieve distributional robustness. Remarkably, OT-MDR even surpasses SAM, the runner-up baseline, by a significant margin, indicating a better generalization capability.
5.3 Experiment on Bayesian Neural Networks
We now assess the effectiveness of OT-MDR in the context of variational inference, where model parameters are sampled from a Gaussian distribution. Specifically, we apply our proposed method to
the widely-used variational technique as SGVB [34], and compare its performance with the original approach. We conduct experiments on two different architectures, Resnet10 and Resnet18 using the CIFAR dataset, and report the results in Table 4. It is clear that our approach outperforms the original SGVB method in all metrics, showcasing significant improvements. These findings underscore OT-MDR ability to increase accuracy, better calibration, and improve uncertainty estimation.
Table 4: Classification scores of approximate the Gaussian posterior on the CIFAR datasets. All experiments are trained three times with different random seeds.
Dataset	Method	Resnet10	Resnet18
ACC ↑	NLL ↓	ECE ↓	ACC ↑	NLL ↓	ECE ↓
CIFAR-10	SGVB	80.52 ± 2.10	0.78 ± 0.23	0.23 ± 0.06	86.74 ± 1.25	0.54 ± 0.01	0.18 ± 0.02
OT-MDR (Ours)	81.26 ± 0.06	0.81 ± 0.12	0.26 ± 0.08	87.55 ± 0.14	0.52 ± 0.01	0.17 ± 0.01
CIFAR-100	SGVB	54.40 ± 0.98	1.96 ± 0.05	0.21 ± 0.00	60.91 ± 2.31	1.74 ± 0.15	0.24 ± 0.03
OT-MDR (Ours)	55.33 ± 0.11	1.85 ± 0.06	0.18 ± 0.03	63.17 ± 0.04	1.55 ± 0.05	0.20 ± 0.03
Effect of Number of Particle Models. For multiple-particles setting on a single model as mentioned in Section 4.2.1, we investigate the effectiveness of diversity in achieving distributional robustness. We conduct experiments using the WideResnet28x10 model on the CIFAR datasets, training with  $K \in \{1,2,3,4\}$  particles. The results are presented in Figure 1. Note that in this setup, we utilize the same hyper-parameters as the one-particle setting, but only train for 100 epochs to save time. Interestingly, we observe that using two-particles achieved higher accuracy compared to one particle. However, as we increase the number of particles, the difference between them also increases, resulting in worse performance. These results suggest that while diversity can be beneficial in achieving distributional robustness, increasing the number of particles beyond a certain threshold may have diminishing returns and potentially lead to performance deterioration.
Figure 1: Multiple particle classification accuracies on the CIFAR datasets with WideResnet28x10. Note that we train 100 epochs for each experiment and report the average accuracy
Loss landscape. We depict the loss landscape for ensemble inference using different architectures on the CIFAR 100 dataset and make a comparison with the SAM method, which serves as our runner-up. As demonstrated in Figure 2, our approach guides the model towards a lower and flatter loss region compared to SAM, which improves the model's performance. This is important because a lower loss signifies better optimization, while a flatter region indicates improved generalization and robustness. By attaining both these characteristics, our approach enhances the model's ability to achieve high accuracy and stability (Table 3).
Figure 2: Comparing loss landscape of (left) Ensemble of 5 Resnet10, (middle) Ensemble of 3 Resnet18, and (right) Ensemble of ResNet18, MobileNet and EfficientNet on CIFAR-100 dataset training with SAM and OT-MDR. Evidently, OT-MDR leads the model to a flatter and lower loss area
In this paper, we explore the relationship between OT-based distributional robustness and sharpness-aware minimization (SAM), and show that SAM is a special case of our framework when a Dirac delta distribution is used over a single model. Our proposed framework can be seen as a probabilistic extension of SAM. Additionally, we extend the OT-based distributional robustness framework to propose a practical method that can be applied to (i) a Dirac delta distribution over a single model, (ii) a uniform distribution over several models, and (iii) a general distribution over the model space (i.e., a Bayesian Neural Network). To demonstrate the effectiveness of our approach, we conduct experiments that show significant improvements over the baselines. We believe that the theoretical connection between the OT-based distributional robustness and SAM could be valuable for future research, such as exploring the dual form in Eq. (15) to adapt the perturbed radius  $\rho$ .
[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] Maksym Andriushchenko and Nicolas Flammarion. Towards understanding sharpness-aware minimization. In International Conference on Machine Learning, pages 639–668. PMLR, 2022. 1
[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] Aharon Ben-Tal, Dick Den Hertog, Anja De Waegenaere, Bertrand Mellenberg, and Gijs Rennen. Robust solutions of optimization problems affected by uncertain probabilities. Management Science, 59(2):341-357, 2013. 1, 2
[5] Jose Blanchet and Yang Kang. Semi-supervised learning based on distributionally robust optimization. Data Analysis and Applications 3: Computational, Classification, Financial, Statistical and Stochastic Methods, 5:1-33, 2020. 1, 2
[6] Jose Blanchet, Yang Kang, and Karthyek Murthy. Robust Wasserstein profile inference and applications to machine learning. Journal of Applied Probability, 56(3):830-857, 2019. 1, 2
[7] Jose Blanchet and Karthyek Murthy. Quantifying distributional model risk via optimal transport. Mathematics of Operations Research, 44(2):565-600, 2019. 2, 3
[8] Tuan Anh Bui, Trung Le, Quan Tran, He Zhao, and Dinh Phung. A unified wasserstein distributional robustness framework for adversarial training. arXiv preprint arXiv:2202.13437, 2022. 1, 2
[9] 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
[10] Pratik Chaudhari, Anna Choromanska, Stefano Soatto, Yann LeCun, Carlo Baldassi, Christian Borgs, Jennifer T. Chayas, Levent Sagun, and Riccardo Zecchina. Entropy-sgd: biasing gradient descent into wide valleys. Journal of Statistical Mechanics: Theory and Experiment, 2019, 2017. 2
[62] 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. 1
[63] Cédric Villani. Optimal transport: Old and new. 2008. 3
[64] 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. 2
[65] 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. 4
[66] Insoon Yang. Wasserstein distributionally robust stochastic control: A data-driven approach. IEEE Transactions on Automatic Control, 2020. 1, 2
[67] 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. 2
[68] Ying Zhang, Tao Xiang, Timothy M. Hospedales, and Hutchuan Lu. Deep mutual learning. 2018 IEEE/CVF Conference on Computer Vision and Pattern Recognition, pages 4320-4328, 2018. 2
[69] Han Zhao, Remi Tachet Des Combes, Kun Zhang, and Geoffrey Gordon. On learning invariant representations for domain adaptation. In International Conference on Machine Learning, pages 7523-7532. PMLR, 2019. 1
[70] Long Zhao, Ting Liu, Xi Peng, and Dimitris Metaxas. Maximum-entropy adversarial data augmentation for improved generalization and robustness. arXiv preprint arXiv:2010.08001, 2020.1,2
[71] 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