Optimal Transport Model Distributional Robustness
Authors: Van-Anh Nguyen
Trung Le
Anh Tuan Bui
Thanh-Toan Do
Dinh Phung
Venue: NeurIPS 2023 poster (2023)
Keywords: Distributional Robustness
Sharpness-aware
SAM

Abstract: Distributional robustness is a promising framework for training deep learning models that are less vulnerable to adversarial examples and data distribution shifts. Previous works have mainly focused on exploiting distributional robustness in the data space. In this work, we explore an optimal transport-based distributional robustness framework in model spaces. Specifically, we examine a model distribution within a Wasserstein ball centered on a given model distribution that maximizes the loss. We have developed theories that enable us to learn the optimal robust center model distribution. Interestingly, our developed theories allow us to flexibly incorporate the concept of sharpness awareness into training, whether it's a single model, ensemble models, or Bayesian Neural Networks, by considering specific forms of the center model distribution. These forms include a Dirac delta distribution over a single model, a uniform distribution over several models, and a general Bayesian Neural Network. Furthermore, we demonstrate that Sharpness-Aware Minimization (SAM) is a specific case of our framework when using a Dirac delta distribution over a single model, while our framework can be seen as a probabilistic extension of SAM. To validate the effectiveness of our framework in the aforementioned settings, we conducted extensive experiments, and the results reveal remarkable improvements compared to the baselines.

Optimal Transport Model Distributional Robustness
Van-Anh Nguyen'
Trung Le
Anh Tuan Bui'
Thanh- Do
Dinh Phung 1.2
Depaatment of Data Science and AI, Monash University, Australia VinAl, Vietnam
(van-anh. nguyen, trunglm, tuan. bui, toan.do, dinh Shung)hunggg @monaah.
Abstract
Distributional robustness promising framework for training deep learning models that are less vulnerable adversarial examples and data distribution shifts. Previous works have mainly focused on exploiting distributional robustness the data space. In this work, we explore optimal transport- based distributional robustness framework model spaces. Specifically, we examine a model dis- tribution within Wasserstein ball centered maximizes the loss. We have developed theories that enable us to learn the optimal robust center model distribution. Interestingly, our developed theories allow us flexibly incorporate the concept of sharpness awareness into training, whether it's single model, ensemble models, Bayesian Neural Networks, by considering specific forms of the center model distribution. These forms include a Dirac delta distribution
given model distribution that
single model, a uniform distribution over several models, and general Bayesian Neural Network. Furthermore, we demonstrate that Sharpness- Aware Minimization (SAM) is specific case of our framework when using a Dirac delta distribution over single model, while our framework can be seen as a probabilistic extension of SAM. validate the effectiveness of our framework the aforementioned settings, we conducted extensive experiments, and the results reveal remarkable improvements compared the baselines.
Introduction
Distributional robustness (DR) is promising framework for learning and decision-making under uncertainty, which has gained increasing attention in recent years [4, 15, 16, 6]. The primary objective of DR is to identify the worst-case data distribution within the vicinity of the ground-truth data distribution, thereby challenging the model's robustness to distributional shifts. DR has been widely applied to various fields, including semi- remi-supervised learning [5, 11, 66]. transfer learning and domain adaptation [39. 16, 69, 50, 51, 38, 52, 57], domain generalization [57, 70]. and improving model robustness [57, 8, 61]. Although the principle of DR can applied to either data space or model space, the majority of previous works DR have primarily concentrated exploring its applications in data space.
Sharpness-awan aware minimization (SAM) [18] has emerged as effective technique for enhancing the generalization ability of deep learning models. SAM aims to find a perturbed model within the vicinity of a current model that maximizes the loss over training set. The success of SAM and variants [36, 33, 62] has inspired further investigation into its formulation and behavior, as evidenced by recent works such [31, 45, 2]. While [30] empirically studied the difference in sharpness obtained by SAM [18] and SWA [26], and [46] demonstrated that SAM is an optimal Bayes relaxation of standard Bayesian inference with normal posterior, none of the existing works have explored the connection between SAM and distributional robustness.
37th Conference on Neural Information Processing Systems (NcurIPS 2023).

---

In this work, we study the theoretical connection between distributional robustness in model space and sharpness- aware minimization (SAM), as they share conceptual similarity. We examine Optimal Transport-basee distributional robustness in model space by considering Wasserstein ball centered around a model distribution and searching for the worst-case distribution that maximizes the empirical loss on training set. By controlling the worst-case performance, it is expected to have smaller generalization error, as demonstrated by a smaller empirical loss and sharpness. We then develop rigorous theories that suggest the strategy to learn the center model distribution. We demonstrate the effectiveness of our framework by devising the practical methods for three cases of model distribution: (i) Dirac delta distribution over single model, (ii) uniform distribution over several models, and (iii) a general model distribution (i.e., a Bayesian Neural Network [48, 56]). Furthermore, show that SAM is specific case of our framework when using Dirac delta distribution over single model, and our framework can be regarded as a probabilistic extension of SAM.
In summary, our contributions in this work are as follows:
We propose a framework for enhancing model generalization by introducing an Optimal Transport (OT) model distribution robustness approach (named OT-MDR). To the best of our knowledge, this is the first work that considers distributional robustness within the model space.
We have devised three practical methods tailored to different types of model distributions. Through extensive experiments, we demonstrate that our practical methods effectively im- prove the generalization of models, resulting in higher natural accuracy and better uncertainty estimation.
Our theoretical findings reveal that our framework can be considered as a probabilistic extension of the widely-used sharpness-aware minimization (SAM) technique. In fact, SAM can be viewed as a specific case within our comprehensive framework. This observation not only explains the outperformance of our practical method over SAM but also sheds light on future research directions for improving model generalization through the perspective of distributional robustness.
2 Related Work
2.1 Distributional Robustness
Distributional Robustness (DR) is a promising framework for enhancing machine learning models in terms of robustness and generalization. The main idea behind DR is to identify the most challenging distribution that is in close proximity to a given distribution and then evaluate the model's 's performance on this distribution. The proximity between two distributions can be measured using either a f- divergence [4, 15, 16, 43, 47] or Wasserstein distance [6, 21, 35, 44, 60]. In addition, some studies [7, 61] have developed a dual form for DR, which allows it to be integrated into the training process of deep learning models. DR has been applied to wide range of domains, including semi-supervised learning [5, 11. 66], domain adaptation [57, 50]. domain generalization [57, 70], and improving model robustness [57, 8, 61].
2.2 Flat Minima
The generalization ability of neural networks can be improved by finding flat minimizers, which allow models to find wider local minima and increase their robustness to shifts between training and test sets [29, 55, 17]. The relationship between the width of minima and generalization ability has been studied theoretically and empirically in many works, including [24, 49, 13, 19, 53J. Various methods have been proposed to seek flat minima, such as those presented in [54, 10, 32, 27, 18]. Studies such as [32, 28, 64] have investigated the effects of different training factors, including batch-size, learning rate, gradient covariance, and dropout, on the flatness of found minima. In addition, some approaches introduce regularization terms into the loss function to pursue wide local minima, such as low-entropy penalties for softmax outputs [54] and distillation losses [68, 67, 10].
SAM [18] is method that seeks flat regions by explicitly minimizing the worst-case loss around the current model. has received significant attention recently for its effectiveness and scalability compared to previous methods. SAM has been successfully applied to various tasks and domains, in-

---

cluding meta-learning bi-level optimization [1], federated learning [59], vision models [12], language models [3], domain generalization [9]. Bayesian Neural Networks [53], and multi-task learning [58]. For instance, SAM has demonstrated its capability to improve meta-learning bi-level optimization in [1], while in federated learning, SAM achieved tighter convergence rates than existing works and proposed generalization bound for the global model [59]. Additionally, SAM has shown its ability to generalize well across different applications such as vision models, language models, domain generalization, and multi-task learning. Some recent works have attempted to enhance SAM's performance by exploiting its geometry [36, 33], minimizing surrogate gap [71], and speeding up its training time [14, 41]. Moreover, [30] empirically studied the difference in sharpness obtained by SAM and SWA [26], while [46] demonstrated that SAM is an optimal Bayes relaxation of the standard Bayesian inference with normal posterior.
3 Distributional Robustness
In this section, we present the background on the OT-based distributional robustness that serves our theory development in the sequel. Distributional robustness (DR) is emerging framework for learning and decision maaaiinnmmmiingggg under uncertainty. which seeks the worst-case expected loss among a ball of distributions, containing all distributions that are close the empirical distribution [20].
Here we consider generic Polish space endowed with distribution Q. Let f S - RR be real-valued (risk) function and c S x S + I__ be cost function. Distributional robustness setting aims to find the distribution Q in the vicinity of Q and maximizes the risk in the E form [61, 7]:
max (Q.0)ce
E (z)),
where >> O and We denotes the optimal transport (OT) or a Wasserstein distance [63] for a metric e, which is defined as:
We
② Q @+Q)===
inf vEr(G.0)
cddy.
where r @ is the set of couplings whose marginals are Q and Q
With the assumption that f E L' (Q) is upper semi-continuous and the cost c is non-negative lower semi-continuous satisfying c(z, z') = = z', [7] shows that the dual form for Eq. (1) is:
min Ac + Ez~q mmax {f (z') (c
[61] further employs a Lagrangian for Wasserstein uncertainty sets to arrive at relaxed version with A > 0:
max [=)l(=)] - AW. (Q. @): Bswolmax (x)) Ac (=
(4)
Proposed Framework
4.1 OT based Sharpness-aware Distribution Robustness
Given family deep nets fa where O, let Qo with the density function go where @ be family of distributions over the parameter space O. To improve the generalization ability of the optimal Qa+, we propose the following distributional robustness formulation on the model space:
min Φ
max (
Ls
i
where S - (x11111) .... (NN, yN) is a training set, d (0,0) 110 0|L (p 1) is a distance on the model space, and we have defined
Ls ②
Ean [Cs (0)) = Eggg
Σ ==1
(So (2))
3

---

with the loss function E.
The OP in (5) seeks the most challenging model distribution Q in the WS ball around Qo and then finds Qø, & E & which minimizes the worst loss. To derive a solution for the OP in (5), we define
{
[ (0)00)
1/8 spj.
T po
Moreover, the following theorem characterizes the solution of the OP in (5).
Theorem 4.1. The OP in (5) is equivalent to the following OP:
min max Ls (7),
(6)
GE pEL
where Ls (B)
(0,0)~
$ N Exer l (Jg (In)
We now need solve the OP in (6). To make solvable, we add the entropic regularization term as
min max @C TET
(7)
where HI (7) returns the entropy of the distribution y with the trade-off parameter 1/A. We note that when À approaches +o0, the OP in (7) becomes equivalent to the OP in (6). The following theorem
indicates the solution of the OP in (7).
Theorem 4.2. When p = +CG, the inner max in the OP in (7) has the solution which a distribution with the density function
7" (0,0)= - 2s ( 7" (810).
where 7" (010) JB (P) 96 (0) is the density function of the distribution Qá and (0) = (0' 110' - 0112 S p} is the p-ball around O.
exp A (C)(())
Referring Theorem 4.2, the OP in (7) hence becomes:
Ls (o)]-
min E 6E O~
(8)
The OP in (8) implies that given a model distribution Qo. we sample models ~ Qs- For each individual model 0, we further sample the particle models 0 (0(0 0) where 7"
oc worth noting that the particle models ~ (ō () x exp (ACs (0)} seek the modes of the distribution (0 0), aiming to obtain high and highest likelihoods (i.e., exp {ACs (0) (0)}) the implementation, we employ stochastic gradient Langevin dynamics (SGLD) [65] to sample the particle models 0.
exp { ACCs { Subsequently, we update Qa to minimize the average of Ls 0
Additionally, in
In what follows, we consider three cases where Qs is (i) Dirac delta distribution over single model, (ii) a uniform distribution over several models, and (iii) a general distribution over the model space (i.e., a Bayesian Neural Network (BNN)) and further devise the practical methods for them.
4.2 Practical Methods
4.2.1
Single-Model OT-based Distributional Robustness
We first examine the case where Qa is Dirac delta distribution over single model i.e., Q de where the Dirac delta distribution. Given the current model 0 we sample K particles 0 using SGLD with only two-step sampling. To diversify the particles, in addition to adding small Gaussian

---

noise to each particle, given mini-batche B. for each particle dj. we randomly split B = [B) B into two equal halves and update the particle models as follows:
VOLB! Bt (0)
0 -
+ c
V (0) 2
VoL Bt
0
0% = 0 +
+ ck1
(9)
VoLB
(0h
where ll the identity matrix and c: o (0, pl).
Furthermore, we base on the particle models to update the next model as follows:
K
0 = 0
V VOLB
0
(10)
K
where 1) > 0 is a learning rate.
It is worth noting that in the update formula (9), different random splits B [ B ....., to encourage the diversity of the particles 0..k. Moreover, we can benefit from the parallel computing to estimate these particles in parallel, which costs |B| (i.e., the batch size) gradient
operations. Additionally, similar to SAM [18], when computing the gradient VeL (8(() we set the corresponding Hessian matrix to the identity one.
Particularly, the gradient VoLgg 0% is evaluated on the second half of the current batch so that the entire batch is used for the update. Again, we can take advantage of the parallel computing to evaluate VaLng 0% k .,, K all once, which costs [B] (i.e., the batch size) gradient operations. Eventually, with the aid of the parallel computing, the total gradient operations in our approach is B], which similar SAM.
4.2.2 Ensemble OT-based Distributional Robustness
We now examine the case where Qo is uniform distribution over several models, i.c., Qo = M Σ Σm=1 6o..- In the context of ensemble learning, for each base learner 0m, , M, 1,---
seek particle models mmn: k = 1,..., as in the case of single model.
VeLBI (0m)
P 0
+ c
IV likk (0m) l2
VoL Ba (Omnk)
0 Øl mk p
+ +
(11)
V
(mink) ll2
where ll is the identity matrix and e 2 Furthermore, we base on the particles to update the next base learners follows:
Om 0m
VoLn
(Pms))
(12)
K
worth noting that the random splits B [
mk mk m = 1 M,K M 1,..,, K of the
B2
current batch and the added Gaussian noise constitute the diversity of the base learners 0I.M- In our developed ensemble model, we do not invoke any term to explicitly encourage the model diversity.
4.2.3 BNN OT-based Distributional Robustness
We finally examine the where Qo is an approximate posterior BNN. To simplify the context, assume that Q6 consists of Gaussian distributions M(4, diag(o?)), L, over
5

---

the weight matrices 0 = WI:z'- : Given 0 = WIZL ~ Qs. the reparameterization trick reads Wi i diag(o))K) with Ką (O, T). We next sample K particle models 02. [Wis]u,, and k -,,..., K from 70 0) as in Theorem 4.2. To sample 0g in the ball around 0, for each layer I, we indeed sample jits the ball around ii and then form Wix = juue + diag(o)) We randomly split B = [B B k 1,..., K and update the approximate Gaussian distribution as
1, ,...., L.
Vn Lnz (lur + diag (ot) xil1)
pi 14 + p
nu C mg (1u + diag (0,) ki],)]12
V B3 (lit + diag (01) sa])) Fh = phu + P || μu B (1 + diag (at) kalp) ll2
+ cx.
02 = L + diag (o,) Kil k l....,,.. K,
"
n LB (xa) and -
p P
K
where c & N pT,
4.3 Connection of SAM and OT-based Model Distribution Robustness
In what follows, we show a connection between OT-based model distributional robustness and SAM. Specifically, prove that SAM is a specific case of OT-based distributional robustness the model space with a particular distance metric. To depart, we first recap the formulation of the OT-based distributional robustness the model space in (5): min SE
min max S (
Ls
E (As (0))
(13)
(.W.
By linking to the dual form in (3), we reach the following equivalent OP:
max { Ls (0) Ad (o. (0.0)
min min
Ap Ea-Qs
(14)
σ AS0
Considering the simple case wherein Qs cquivalently entails
ôg is a Dirac delta distribution. The OPs in (14)
{ Ap + max (0s(0) (d))>a(0.0)) (1x00)
min min XSO
(15)
The optimization problem in (15) can be viewed as probabilistic extension of SAM. Specifically, for each 0 ~ Qo. the inner max: maxq Lss (0) Ad (0,0) seeks a model 0 maximizing the loss (()
soft ball around 9 controlled by Ad(O, Ø). Particularly, higher value of À seeks the optimal 0 smaller ball. Moreover, the outer min, the term Ap trades off between the value of and the radius of the soft ball, aiming find out optimal X and the optimal model Ô 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 p around the model 6.
Interestingly, by appointing a particular distance metric between two models, we can exactly recover the SAM formulation shown in the following theorem.
Theorem 4.3, With the distance metric d defined as
d (0.00
[ 110 01l2
11-00125
(16)
+00
otherwise
the OPs in (13), (14) with Qs = dg. and (15) equivalently reduce to the OP of SAM as
Ls (0)
min
max 0:110-01a25
For simplicity, absorb the biases to the weight matrices

---

Table 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
P1ramid101
Densenet121
CIFAR-10
SAM OT-MDR (Ours)
96.72 + 0.007
96.20 + 0.134
91.16 0.240
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 Resnet18. The results of baselines are taken from [46]
CIFAR-10
CIFAR-100
Method
ACC
AUROCT
ACC1 1
AUROC 1
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
1 0.001
80.22
1 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 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 intriguing and may open doors propose other sharpness- aware training approaches and leverage adversarial training [23, 42] to improve model robustness. Moreover, [46] has shown connection between SAM and the standard Bayesian inference with a normal posterior. Using the Gaussian approximate posterior ((w, v/), it is demonstrated that the maximum of the likelihood loss L(qu) (i.c., μ is the expectation parameter of the Gaussian N (w, vf)) can be lower-bounded by relaxed-Bay -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.
5 Experiments
In this section, we present the results of various experiments 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 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 experi- ments using three different architectures: WedeRessNe22x 10, Pyramid101, and Densenet121, We compared our approach's 's results against models trained with SAM optimizer as our baseline. For consistency with the original SAM paper, we adopted their setting, using p = 0.05 for CIFAR-10 experiments and
0.1 for CIFAR- and report the result in Table 1. In our OT-MDR method, we chose different values of p for each half of the mini-batch B, and denoted #i for B, and P2 for By, where P2 = 2p1 to be simplified (this simplified setting is used in all experiments). To ensure a fair comparison with SAM, we also set #1 = 0.05 for CIFAR- 10 experiments and P1 = 0.1 for CIFAR- Our approach outperformed the baseline with significant gaps. indicated in Table 1. On average, our method achieved a 0.73% improvement on CIFAR- 10 and a 1.68% improvement on CIFAR- compared to SAM. These results demonstrate the effectiveness of our proposed method for achieving higher accuracy in one- one-particle model training.
Thh implementation is provided in https: ptttps://githue http//nn-ntv/O1... MDR. git
7

---

Table Evaluation the ensemble Accuracy (%) on the CIFAR-10/100 datasets. We reproduce all baselines with the same hyperparameter for a fair comparison.
CIEAR-10
CIFAR-100
Method
ACC 1
Brier
NLL
BCE
AAC L
ACC
Brier
NLL
ECE
AAC
Ensemble of five Resnet 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
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.074
0.11
69.6
0.412
1.123
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 EfficientNeet
Deep Ensemble
89.0
0.153
0.395
0.111
0.126
62.7
0.433
1.267
0.176
0.209
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.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 18. 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 because the authors did not release the code. Instead, we OT-MDR with the setting mentioned 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, 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 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, 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 p for SAM and P1. P2 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 Lig-Liielihood. NLL,, Expected Calibration Error ECE, and Average across all calibrated confidence AAC) on the CIFAR dataset, as shown in Table
Notably, our OT-MDR approach consistently outperforms all baselines across all metrics, demonstrat- ing 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, indicat- ing a better generalization capability.
5.3 Experiment on Bayesian Neural Networks
We now assess the effectiveness of OT-MDR the context variational inference, where model parameters are sampled from a Gaussian distribution. Specifically, apply our proposed method to

---

the widely-used variational technique SGVB [34], and compare its performance with the original approach. We conduct experiments on two different architectures, Resnet10 and Resnet18 18 using the CIFAR dataset, and report the results in Table 4. It is clear that our approach outperforms the original SGVB method metries, 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 the CIFAR datasets. All experiments are trained three times with different random seeds.
Resnet10
Resnet18
Dataset
Method
ACC
NLI.
BCE
ACC1
NLI.
BCB
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.35 + 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
5.4
Ablation Study
Effect of Number of Particle Models. For multiple- particles setting
single model as mentioned in Sec- tion 4.2.1, we investigate the effectiveness of diversity in achieving distributional robustness. We conduct experi- ments using the WideResnet28x 10 model on the CIFAR datasets, training with K {1, 2, 3, 4} particles. The re- sults are presented in Figure 1. Note that in this setup, utilize the same hyper-parameters as the one-particle setting, but only train for 100 epochs to save time. In- terestingly, we observe that using too-particles achieved higher accuracy compared to one particle. However, as we increase the number of particles, the difference be- tween them also increases, resulting in worse performance. These results suggest that while diversity can be benefi- cial in achieving distributional robustness, increasing the number particles beyond certain threshold may have diminishing returns and potentially lead to performance deterioration.
00.
Accuracy
700.2
63.3
56
85.1
CIAALDD
ORARID
Number of particies
Figure Multiple particle classification accuracies on the CIFAR datasets with WideResmet28x10. Note that train 100 epochs for each experiment and re- port the average accuracy
Loss landscape. We depict the loss landscape for ensemble inference using different architectures on the CIFAR 100 dataset and make 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 important because a lower loss signifies better optimization, while flatter region indicates improved generalization and robustness. By attaining both these characteristics, our approach enhances the model's ability achieve high accuracy and stability (Table 3).
SAM
OT-MDR
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 flatter and lower loss area

---

Conclusion
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 Dirac delta distribution is used over a single model. Our proposed framework can be seen as a probabilistic extension of SAM. Additionally, propose a practical method that can be applied to (i) a Dirac delta distribution over a single model, (ii) uniform distribution
extend the OT-based distributional robustness framework to
several models, and (iii) 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) adapt the perturbed radius p.
This work 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-agnostie meta learning. arXiv preprint arXiv:2206.03996, 2022.
[2] Maksym Andriushchenko and Nicolas Flammarion. Towards understanding sharpness-aware minimization. International Conference on Machine Learning, pages 639-668. PMLR, 2022.
[3] Dara Bahri, Hossein Mobahi, and Yi Tay. Sharpness-awan 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 Melenberg, and Gijs Rennen. Robust solutions of optimization problems affected by uncertain probabilities. Management Science, 59(2):3413357, 2013. 1,
[5] Jose Blanchet and Yang Kang. Semi-supervised learning based distributionally robust optimization. Data Analysis and Applications 3: Computational, Classification, Financial, Statistical and Stochastic Methods, 5:1-33, 2020. 1,
[6] Jose Blanchet, Yang Kang, and Karthyek Murthy. Robust wasserstein profile inference and applications to machine learning. Journal of Applied Probability, 56(3):830-88) 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,
[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, 01-222055-22111 2021. 3
[10] 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. 2
10

---

[11] Ruidi Chen and Ioannis C Paschalidis. A robust learning approach for regression models based on distributionally robust optimization. Journal of Machine Learning Research, 19(13), 2018.
1,2
[12] Xiangning Chen, Cho-Jui Hsieh, and Boqing Gong. When vision transformers outperform resnets without pre-training or strong data augmentations. arXiv preprini arXiv:2106.01548,
2021.
[13] Laurent Dinh, Razvan Pascanu, Samy Bengio, and Yoshua Bengio. Sharp minima can generalize for deep nets. In International Conference Machine Learning, pages 1019-1028. PMLR, 2017.
[14] 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
[15] John C Duchi, Peter W Glynn, and Hongseok Namkoong. Statistics of robust optimization A generalized empirical likelihood approach. Mathematics of Operations Research, 2021. 1, 2
[16] John Duchi, Tatsunori Hashimoto, and Hongseok Namkoong. Distributionally robust losses
against mixture covariate shifts. Under review, 2019. 1,2
[17] 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.
[18] Pierre Foret, Ariel Kleiner, Hossein Mobahi, and Behnam Neyshabur, Sharpness-awan aware min- imization for efficiently improving generalization. In International Conference on Learning Representations, 2021. 1, 8
[19] Stanislav Fort and Surya Ganguli. Emergent properties of the local geometry of neural loss landscapes. arXiv preprint 19X00099,, 2019. 2
[20] Rui Gao, Xi Chen, and Anton J Kleywegt. Wasserstein distributional robustness and regulariza- tion statistical learning. arXiv e-prints, pages arXiv-1712, 2017. 3
[21] Rui Gao and Anton J Kleywegt. Distributionally robust stochastic optimization with wasserstein distance. arXiv preprint 164.v.160.002/999 2016.
[22] Timur Garipov, Pavel Izmailov, Dmitrii Podoprikhin, Dmitry P Vetrov, and Andrew G Wilson. Loss surfaces, mode connectivity, and fast ensembling of dnns. Advances in neural information processing systems, 31, 2018.
[23] lan J Goodfellow, Jonathon Shlens, and Christian Szegedy. Explaining and harnessing adversar- ial examples. arXiv preprint arXiv:1412.657. 2014. 7
[24] Sepp Hochreiter and Jürgen Schmidhuber. Simplifying neural nets by discovering flat minima. In NIPS, pages 529-536. MIT Press, 1994.
[25] Gao Huang, Yixuan Li, Geoff Pleiss, Zhuang Liu, John E. Hoperoft, and Kilian Q. Wein- berger. Snapshot ensembles: Train 1, get m for free. In International Conference on Learning Representations, 2017. 8
(26) Pavel Izmailov, Dmitrii Podoprikhin, Timur Garipov, Dmitry Vetrov, and Andrew Gordon Wilson. Averaging weights leads to wider optima and better generalization. arXiv preprint arXiv:1803055007, 2018. 3.3
[27] 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.
[28] Stanislaw Jastrzebski, Zachary Kenton, Devansh Arpit, Nicolas Ballas, Asja Fischer, Yoshua Bengio, and Amos J. Storkey. Three factors influencing minima sgd. ArXiv, 10.0422,, 2017.

---

[29] Yiding Jiang, Behnam Neyshabur, Hossein Mobahi, Dilip Krishnan, and Samy Bengio. Fantastic generalization measures and where find them. In ICLR. OpenReview.net, 2020. 2
[30] 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. 1,3
[31] Jean Kaddour, Linqing Liu, Ricardo Silva, and Matt Kusner. fair comparison of two popular flat minima optimizers: Stochastic weight averaging vs. sharpness-aware minimization. arXiv preprint arXiv:2202.0066 1, 2022.
[32] 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.
[33] 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 Machine Learning. volume 162 of Proceedings of Machine Learning Research, pages 11148-14 11188-111 111. PMLR, 17-23 Jul 2022. 1,3
[34] Durk P Kingma, Tim Salimans, and Max Welling. Variational dropout and the local reparame- terization trick. Advances in neural information processing systems, 28, 2015. 9
[35] Daniel Kuhn, Peyman Mohajerin Esfahani, Viet Anh Nguyen, and Soroosh Shafieezadeh- Abadeh. Wasserstein distributionally robust optimization: Theory and applications in machine learning. In Operations Research & Management Science in the Age of Analytics, pages 130-166. INFORMS, 2019. 2
[36] 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. 1,
[37] Balaji Lakshminarayanan, Alexander Pritzel, and Charles Blundell. Simple and scalable predictive uncertainty estimation using deep ensembles. Advances in neural information processing systems, 30, 2017.
[38] Trung Le, Tuan Nguyen, Nhat Ho, Hung Bui, and Dinh Phung. Lamda: Label matching deep domain adaptation. In Marina Meila and Tong Zhang, editors, Proceedings of the 38th Inter- national Conference on Machine Learning, volume 139 of Proceedings of Machine Learning Research, pages 6043-6054. PMLR, 18-24 Jul 2021. 1
[39] Jacho Lee and Maxim Raginsky. Minimax statistical learning with wasserstein distances. In NeurIPS, pages 2692-2701, 2018.
[40] Shiwei Liu, Tianlong Chen, Zahra Atashgahi, Xiaohan Chen, Ghada Sokar, Elena Mocanu, Mykola Pechenizkiy, Zhangyang Wang, and Decebal Constantin Mocanu. Deep ensembling with no overhead for either training or testing: The all-round blessings of dynamic sparsity. In International Conference on Learning Representations, 2022. 8
[41] 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.
[42] Aleksander Madry, Aleksandar Makelov, Ludwig Schmidt, Dimitris Tsipras, and Adrian Vladu. Towards deep learning models resistant to adversarial attacks. In International Conference on Learning Representations, 2018.7
[43] Takeru Miyato, Shin-ichi Maeda, Masanori Koyama, Ken Nakae, and Shin Ishii. Distributional smoothing with virtual adversarial training. arXiv preprint
2015. 2
[44] Peyman Mohajerin Esfahani and Daniel Kuhn. Data-driven distributionally robust optimization using the wasserstein metrie: Performance guarantees and tractable reformulations. arXiv e-prints, pages arXiv-1505, 2015. 2
12

---

[45] Thomas Möllenhoff and Mohammad Emtiyaz Khan. Sam as an optimal relaxation of bayes. arXiv preprint arXiv:2210.016200 2022. 1
[46] Thomas Möllenhoff and Mohammad Emtiyaz Khan. SAM as an optimal relaxation of bayes. In
The Eleventh International Conference on Learning Representations, 2023. 1, 7,
[47] Hongseok Namkoong and John C Duchi. Stochastic gradient methods for distributionally robust optimization with f-divergences. In NIPS, volume 29, pages
2016. 2
[48] Radford M Neal. Bayesian learning for neural networks, volume 1 118, Springer Science & Business Media, 2012. 2
[49] Behnam Neyshabur, Srinadh Bhojanapalli, David McAllester, and Nati Srebro. Exploring generalization in deep learning. Advances in neural information processing systems, 30, 2017.
2
[50] Tuan Nguyen, Trung Le, Nhan Dam, Quan Hung Tran, Truyen Nguyen, and Dinh Q Phung. Tidot: A teacher imitation learning approach for domain adaptation with optimal transport. In IJCAI, pages 2862-2868, 2021. 1
[51] Tuan Nguyen, Trung Le, He Zhao, Quan Hung Tran, Truyen Nguyen, and Dinh Phung. Most: multi-source domain adaptation via optimal transport for student-teache learning. In Cassio de Campos and Marloes H. Maathuis, editors, Proceedings of the Thirty-Sevnnth Conference on Uncertainty in Artificial Intelligence, volume 161 of Proceedings of Machine Learning Research, pages 225 235. PMLR, 27-30 Jul 2021.
[52] Tuan Nguyen, Van Nguyen, Trung Le, He Zhao, Quan Hung Tran, and Dinh Phung. Cycle class consistency with distributional optimal transport and knowledge distillation for unsupervised domain adaptation. In James Cussens and Kun Zhang, editors, Proceedings of the Thirty-Ei Eighth Conference on Uncertainty in Artificial Intelligence, volume 180 of Proceedings of Machine Learning Research, pages 1519-1529. PMLR, 01-05 Aug 2022.
[53] Van-Anh Nguyen, Tung-Long Vuong, Hoang Phan, Thanh-Toan Do, Dinh Phung, and Trung Le. Flat seeking bayesian neural network. Advances in Neural Information Processing Systems, 2023. 3,3
[54] 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. 2
[55] Henning Petzka, Michael Kamp, Linara Adilova, Cristian Sminchisescu, and Mario Boley. Relative flatness and generalization. In NeurlPS, pages 18420-18432, 2021. 2
[56] Cuong Pham, C. Cuong Nguyen, Trung Le, Phung Dinh, Gustavo Carneiro, and Thanh- Do. Model and feature diversity for bayesian neural networks in mutual learning. In Advances in Neural Information Processing Systems, 2023. 2
[57] 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 Artificial Intelligence and Statistics, volume 206 of Proceedings of Machine Learning Research, pages 7644-7664. PMLR, 25-27 Apr 2023. 1,2
[58] Hoang Phan, Ngoc Tran, Trung Le. Toan Tran, Nhat Ho, and Dinh Phung. Stochastic multiple target sampling gradient descent. Advances in neural information processing systems, 2022.
[59] Zhe Qu, Xingyu Li, Rui Duan, Yao Liu, Bo Tang, and Zhuo Lu. Generalized federated learning via sharpness aware minimization. arXiv preprini arXiv:2206.0221 02618, 2022. 3
[60] Soroosh saatfiezzade--Abaade Peyman Mohajerin Esfahani, and Daniel Kuhn. Distributionally robust logistic regression. arXiv preprint arXiv: 1509.09259 2015. 2
[61] Aman Sinha, Hongseok Namkoong, and John Duchi. Certifying some distributional robustness
with principled adversarial training. In International Conference on Learning Representations, 2018. 1, 2,
13

---

[62] Tuan Truong, Hoang- Phi Nguyen, Tung Pham, Minh-Tuam Tran, Mehrtash Harandi, Dinh Phung. and Trung Le. Rsam: Learning on manifolds with riemanmian sharpness-award aware minimization, 2023.
[63] Cédric Villani. Optimal transport: Old and new. 2008.
[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.
[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.
[66] Insoon Yang. Wasserstein distributionally robust stochastic control: A data-driven approach. IEFE Transactions on Automatic Control, 2020. 1,
[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 Huchuan Lu. Deep mutual learning. 2018 IEEESCVF Conference on Computer Vision and Puttern Recognition, pages 4320-4328, 2018.
[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-entrop 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 22Xx:22033080653 2022. 3
14


Reviews:
[{"review_id": "xUzC2Pj0DI", "reviewer": "Reviewer_P1BN", "summary": "This paper proposes a novel distributional robustness optimization problem on the space of model parameters (OT-MDR). By formulating the uncertainty set with Wasserstein distance and introducing an entopy regularization, the explicit solution to the primal problem is derived. A practical algorithm is proposed based on sampling techniques from SGLD. The algorithm is extended beyond point estimation to accomodate emsembling and bayesian learning. The dual of OT-MDR turns out as a stochastic relaxation of the objective of SAM, bridging sharpness aware optimization and distributional robustness. Empirical results on Cifar validates the advantage of proposed method over baseline optimizers.", "questions": "1. The practical algorithm is based on SGLD which shows that stochastic gradient descent could converge to the posterior distribution of model parameters. However, this paper adopts two-step sampling such that the sampled parameters might deviate much from the distribution in Theorem 4.2. Is this a significant gap between the empirical algorithm and theory?\n", "limitations": "na", "rating": 7, "confidence": 3, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "1. I might be unfamiliar with literature on sharpness aware optimization, but from my perspective this paper is the first to reveal an intrinsic connection between distributional robustness and local sharpness. Theorem 3.4 is intriguing since it might inspire further research to apply abundant literature on distributional robustness (and adversarial robustness which is highly relative) to solving sharpness aware optimization. \n1. The proposed method is a stochatstic extension of SAM but it introduces limited computational burden. On the other hand, the algorithm well extends to bayesian learning because of its stochastic formulation of model parameters. ", "weaknesses": "This paper proves OT-MDR to be more general than SAM but the advantage of OT-MDR in the case of single model and emsembling is theoretically vague. From Theorem 3.4, it seems that OT-MDR is equivalent to SAM for single model except that OT-MDR softly expands the uncertainty ball around the centering model, by allowing parameters to go beyond the radius $\\rho$. The gap is insufficient to account for the empiricial success of OT-MDR for single model. Additionally, is it possible that enlarging the radius $\\rho$ for SAM would make its loss landscape approach OT-MDR in Figure 2?", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We sincerely appreciate your constructive comments. We are dedicated to addressing all the questions listed below to the best of our capabilities.\n\n**Advantage of OT-MDR in the case of a single model and ensembling**  \nIn Theorem 4.3, we show a connection of our OT-MDR developed by applying distributional robustness (DR) to model distribution and SAM. This theorem indicates that for a specific distance $d(\\theta, \\tilde{\\theta})$ as in Theorem 4.3, our OT-MDR is reduced exactly to SAM. However, in general, from the modeling perspective, using a more general distance, our OT-MDR can be regarded as a probabilistic relaxation of SAM (cf. Theorem 4.2).\n\nDifferent from SAM, we cast the problem to sample the particles from a target distribution (i.e., $\\gamma^*(\\theta, \\tilde{\\theta})$, which has a closed-form) instead of maximizing the loss using one-step gradient ascent update. Although we observe that the peak of the performance at two or three particles, these multiple-particles versions consume multiplicatively more memory and take multiplicatively longer training time. For a fair comparison, we use the practical version with one single particle to compare to SAM. However,  our particles are sampled in multiple steps using stochastic gradient Langevin dynamics (SGLD). We observe that multiple-step sampling with small Gaussian noises to diversify the particles really helpful to boost the performance. Moreover, randomly splitting a mini-batch into multiple smaller sub-batches, each of which is used for one sampling step, is also beneficial to diversify the particles.  \n\nIn practice, to balance between performance and computation complexity, we propose using two-step sampling, which was empirically demonstrated to outperform SAM and other baselines as in Tables 1, 2, and 3.\n\n**Enlarging the radius for SAM would make its loss landscape approach OT-MDR in Figure 2**   \nIn our experiment, for SAM, we use the optimal perturbation radius $\\rho$ reported in the original paper. \n\n Additionally, SAM is sensitive to the perturbation radius $\\rho$ as shown in [1]. Please refer to Figure 21 of Appendix E.6 in [1], which shows that the testing accuracies get decreasing when we increase $\\rho$ from its optimal values. Moreover, Figure 2 in [1] further indicates that the $m$-sharpness increases if we enlarge $\\rho$.  \n\n\n**Adopting two-step sampling such that the sampled parameters might deviate much from the distribution in Theorem 4.2: Is this a significant gap between the empirical algorithm and theory?**\n\nThe situation here is quite similar to SAM. In the theory of SAM, it suggests minimizing the sharpness \n$\n\\max_{\\Vert\\delta\\Vert\\leq\\rho} \\{\\mathcal{L}\\_{S}(\\theta+\\delta)-\\mathcal{L}\\_{S}(\\theta)\\}\n$\n, which requires multiple gradient ascent steps to find an optimal perturbation. However, in the practical version, it proposes using only one gradient ascent step to find $\\delta$. Moreover, Figure 4 in [1] shows the effect of different numbers of projected gradient ascent steps for SAM, indicating that using more projected gradient ascent steps does not necessarily improve the performance.  \n\nFor our OT-MDR, we adopt two-step sampling using SGLD with two random sub-batches to diversify the particles. We observe that it is sufficient to gain good empirical results, while still balancing the training times. We empirically find that the more-step sampling can slightly boost the performance, but in return more significantly increase the training times. Therefore, we opt the two-step sampling strategy as a practical method.\n\nWe conjecture that two-step sampling is good due to an accumulative effect of gradually moving the models $\\theta$ or the distribution of models $\\theta$ to flatter regions. The two-step sampling particles assist us in moving the models $\\theta$ to a flatter region, while the models $\\theta$ in a flatter region support the two-step sampling in finding more accurate particles.    \n\n[1] Andriushchenko, M. and Flammarion, N. (2022, June). Towards understanding sharpness-aware minimization. In International Conference on Machine Learning (pp. 639-668). PMLR.\n\n[2] Foret, P., Kleiner, A., Mobahi, H., and Neyshabur, B. (2020). Sharpness-aware minimization for efficiently improving generalization. arXiv preprint arXiv:2010.01412."}, {"review_id": "VOYNMZa2IX", "reviewer": "Reviewer_NFva", "summary": "This paper explores the optimal transport distributional robustness framework in the model parameters spaces, and provides practical algorithms to solve the corresponding optimization problems. This flexible framework can be used to train a single model, ensemble models and Bayesian neural networks in a sharpness-aware fashion. The authors additionally establish that the sharpness-aware minimization serves as a special instance within their proposed framework. The experimental results demonstrate the usefulness of this framework.", "questions": "\nQ1. In SAM [1], their experiments show that using 2-Wasserstein distance is better than using \\infty-Wasserstein distance. For the current paper, in theorem 4.3, it shows that choosing center distribution as a Dirac delta in OT-MDR with 2-Wasserstein distance is equivalent to SAM (using 2-Wasserstein distance). But based on theorem 4.2, the algorithm solves the minimization problem of OT-MDR with \\infty-Wasserstein distance. And the algorithm of OT-MDR still outperforms the SAM (with 2-Wasserstein distance). Does this probably imply the performance of OT-MDR can be further improved when using algorithms that solves the minimization problem of OT-MDR with 2-Wasserstein norm? Here are a few more delicate methods [2,3] to solve the general distributional robustness learning problem.\n\n  [1] Foret, P., Kleiner, A., Mobahi, H., & Neyshabur, B. (2020). Sharpness-aware minimization for efficiently improving generalization. \n\n  [2] Wang, G., & Chizat, L. (2022). An exponentially converging particle method for the mixed nash equilibrium of continuous games. \n\n  [3] Trillos, C. G., & Trillos, N. G. (2023). On adversarial robustness and the use of Wasserstein ascent-descent dynamics to enforce it. \n\n\nQ2. In figure 1, the difference of accuracy for using different numbers of particles is actually small. Since this figure doesn’t provide the information standard deviation, it might be difficult to conclude that using 2 particles gives the best results. It would be better if the author could provide different random seeds results of this experiment. It will probably provide more insight.\n\nQ3. As mentioned in the weakness, the experimental results in table 1, 3 and 4 don't contain standard deviation. Could the author provide more results on those experiments?\n", "limitations": "In general, I like the idea and the presentation of this paper. The authors could probably further improve this paper through more experiments and add some analysis of the observations from comparing SAM and OT-MDR.", "rating": 6, "confidence": 3, "soundness": 2, "presentation": 3, "contribution": 3, "strengths": "This paper considers an important question of neural networks training. It is clearly written and easy to follow. The idea of extending the form of SAM to the DRO framework is natural but very meaningful, since it provides more well-developed tools to understand and improve the generalization of training models.\n", "weaknesses": "1. In the theory, this paper shows that SAM is equivalent to the optimization problem in the OT-MDR framework when using a particular distance metric and picking the center distribution as a Dirac delta distribution. In the experiments, the results show (1) OT-MDR performs better than SAM for both single model and ensemble models cases; (2) The loss landscape found by OT-MDR is more flatten than the one found by SAM. This probably implies the training algorithm of OT-MDR is better than directly using SGD to solve SAM. But the paper doesn’t provide much intuitive understanding or analysis of these observations. \n\n2. Most of the experimental results don’t provide standard deviation, which makes the results less solid.\n", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We sincerely appreciate your constructive comments. We are dedicated to addressing all the questions listed below to the best of our capabilities.\n\n\n**Q1. In SAM [1], their experiments show that using 2-Wasserstein distance is better. But based on theorem 4.2, the algorithm solves the minimization problem of OT-MDR82 with $\\infty$-Wasserstein distance still outperforms the SAM. Does this probably imply the performance of OT-MDR can be further improved when using algorithms that solve the minimization problem of OT-MDR with 2-Wasserstein norm?**\n\n\nWe develop our theories for the distance $d(\\theta, \\tilde{\\theta}) = \\Vert \\theta - \\tilde{\\theta} \\Vert_2^p$. In Theorem 4.2, we propose a tractable solution when $p \\rightarrow +\\infty$ in which the ball $B_\\rho$ and the geometry use the L2 (i.e., Euclidian) distance as in SAM. However, as you mention, the OT-based distributional robustness on a model space and its connection to SAM leave open doors to develop other theoretically-guaranteed and practical methods in the future by considering a more general family of distances between two models or customizing the OT-based distributional robustness itself. \n\nWe really appreciate your recommended papers [2,3],  as we believe they are strong tools to develop a more generalized version of this work.\n\n\n**Q2. In figure 1, the difference of accuracy for using different numbers of particles is actually small and doesn’t provide the information standard deviation.**\n\nWe would like to update the result of multiple particles as below. Experiments for each number of particles are run three times and report the mean and standard deviation. We agree that the gap between one and two-particle is small but accuracy on two-particle still yields the highest performance consistently.\n\nTable 5. Multiple particle classification accuracies on the CIFAR datasets with WideResnet28x10\n| K particles        |      1       |       2     |         3     |       \n| :---                    |    :----:   |    :----:   |     :----:   | \nAccuracy | 83.22 $\\pm$ 0.18 | 83.31 $\\pm$ 0.16 | 82.95 $\\pm$ 0.08\n\n**Q3. The experimental results in Tables 1, 3, and 4 don’t contain standard deviation**\n\nWe acknowledge the importance of mean and standard deviation when dealing with distribution robustness problems. Meanwhile, we are able to provide additional results for single models (Table 1 in the main paper), Ensemble models (Table 3 in the main paper), and Gaussian posterior models (Table 4 in the main paper) within the attached PDF. Each experiment is run at least three times with different random seeds. It is worth noting that we managed to re-produce the result of both baselines and our methods for a fair comparison. Due to the time limitation of the rebuttal period, we focus on re-producing experiments with our OT-MDR and will report the remaining baselines in the revised version. Based on the results, our proposed methods consistently demonstrate remarkable performance in most of the experiments. On the other hand, the baseline models, particularly the Gaussian posterior experiment, yield a notably high standard deviation."}, {"review_id": "1ljs5iC9oj", "reviewer": "Reviewer_WK2o", "summary": "This paper proposes an optimal-transport-based distribution robustness framework on model space. The proposed theory considers the worst-case loss w.r.t a Wasserstein distance from the model distribution, which can be considered a probabilistic extension of sharpness-aware minimization (SAM).", "questions": "I suggest the author fully discuss the aforementioned related works and improve their methods or add comparisons to corresponding algorithms PAC Bayes community.", "limitations": "No, I suggest the authors add a section to discuss the limitation of the proposed algorithm, e.g., the increment of computation complexity of BNN and the multiple particle setting.", "rating": 5, "confidence": 4, "soundness": 2, "presentation": 3, "contribution": 2, "strengths": "- The paper is well-written and easy to follow.\n- The proposed approach shows empirical improvements compared to SGD, SAM, DGVB, etc.", "weaknesses": "- The theoretical and empirical results are not surprising to me. Moreover, they do not conduct generalization analysis, which is important in learning theory.\n- I do not agree with the claim that this paper is the first work that considers distribution robustness within the model space. The sharpness-related learning problem has already been well addressed in PAC Bayes theory. The counterpart of the worst-case awareness studied in this paper can be found in [1]. And other Wasserstein or more general theories can be referred to [2] and [3].\n- The single-model, ensemble, and BNN cases are also common algorithm realizations in the PAC Bayes community.\n\n>[1] Jiang, Yiding, Behnam Neyshabur, Hossein Mobahi, Dilip Krishnan, and Samy Bengio. \"Fantastic Generalization Measures and Where to Find Them.\" In *International Conference on Learning Representations*.\n\n>[2] Amit, R., Epstein, B., Moran, S. and Meir, R., Integral Probability Metrics PAC-Bayes Bounds. In *Advances in Neural Information Processing Systems*.\n\n>[3] Haddouche, Maxime, and Benjamin Guedj. \"Wasserstein PAC-Bayes Learning: Exploiting Optimisation Guarantees to Explain Generalisation.\" (2023).", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We sincerely appreciate your constructive comments. We are dedicated to addressing all the questions listed below to the best of our capabilities.\n\n**Do not conduct generalization analysis, which is important in learning theory**\n\nThanks for this comment. Generalization bound is an interesting theoretical development that strengthens our work. We will definitely consider it in future development and the papers [1,2] suggested by you possibly shed light for us to develop the bounds. \n\nIn the scope of this research, we aim to find out theoretical connections between OT-based distributional robustness on a model space and sharpness-aware minimization, which are further harvested to propose practical methods on single models, ensemble models, and Bayesian neural networks.\n\n**Do not agree with the claim that this paper is the first work that considers distribution robustness within the model space.**\n\n\nThe work [1] trained more than 10,000 models over two image classification datasets, namely, CIFAR-10 and Street View House Numbers (SVHN). In\norder to create diverge generalization behaviors, they carefully varied hyper-parameters that influence generalization and selected multiple optimization algorithms with different stopping criteria for training convergence. They finally investigated 40 complexity\nmeasures to reach conclusions. Although the empirical conclusions from this paper are interesting and beneficial, we cannot see any distributional robustness applied to model space. \n\nThe work [2] is a solid study of PAC-Bayes generalization bound for the gap between general and empirical losses. Specifically, this work developed the total variation metric and the Wasserstein (WS) distance bounds. The work [3] further developed WS-style PAC-Bayes generalization bound for the data-free and data-dependent priors. \n\nTo summarize, we agree that these works are related to ours, hence we will definitely mention and discuss them in the related work. However, we respectfully disagree with the reviewer the point: _these works studied WS-based distributional robustness on the model space, hence affecting the novelty of our work_.\n\n[1] Jiang, Y. et al.,  \"Fantastic Generalization Measures and Where to Find Them.\" In International Conference on Learning Representations.\n\n[2] Amit, R. et al., Integral Probability Metrics PAC-Bayes Bounds. In Advances in Neural Information Processing Systems.\n\n[3] Haddouche, M. et al.. \"Wasserstein PAC-Bayes Learning: Exploiting Optimisation Guarantees to Explain Generalisation.\" (2023).\n\n\n**Computation complexity**\n\nWe would like to report the training time of our method and baselines (Table 4 of the attached pdf). It is with noting that our method takes longer time for training since it involves calculating the gradient three times sequentially. However, for the first two times, we only need to calculate the gradients for half of the data batch. In total, the number of gradients needs to calculate is equal to the SAM method."}, {"review_id": "IIURJPQzyw", "reviewer": "Reviewer_X3P4", "summary": "This paper explores the relationship between OT-based distributional robustness and sharpness-aware minimization (SAM), and shows that SAM is a special case of the proposed framework when a Dirac delta distribution is used over a single model. To demonstrate the effectiveness of the framework, the authors perform several experiments and show improvements over the baselines.", "questions": "1. Theorem 4.2 holds for $p=+\\infty$, can the results be extended to more general settings?\n2. To make the OP problem solvable, the authors resort to the entropy-regularized version of OP. However, when establishing the connection between SAM and OT-based model distribution robustness, the authors use the original version of OP. Can the authors elaborate more on this?", "limitations": "The practical methods stated in Section 4.2 are not easy to follow.", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "1. This paper is theoretically sound.\n2. The experiment verifies the effectiveness of the proposed framework.", "weaknesses": "The practical methods stated in Section 4.2 are not easy to follow. It would be better if the authors can provide a pseudocode to help the readers understand the methods easily.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We sincerely appreciate your constructive comments. We are dedicated to addressing all the questions listed below to the best of our capabilities.\n\n\n**Theorem 4.2 holds for $p=+\\infty$, can the results be extended to more general settings?**\n\nOur distance of interest is $d(\\theta,\\tilde{\\theta})=\\Vert\\theta-\\tilde{\\theta}\\Vert_{2}^{p}$ and we develop a tractable solution when $p \\rightarrow +\\infty$ in this paper. For the case of a general $p$ or even a general distance between two models as surveyed in this paper [1], we possibly need to resort to the dual form in Eq. (3) mentioned in the main paper.\n\n\n**To make the OP problem solvable, the authors resort to the entropy-regularized version of OP. However, when establishing the connection between SAM and OT-based model distribution robustness, the authors use the original version of OP. Can the authors elaborate more on this?**\n\nSimilar to other OT-based practical methods such as Sinkhorn, we need to resort to the entropy-regularized term with the trade-off $1/\\lambda$ in order to develop a tractable solution for our OT-based distributional robustness on the model space. It is obvious that when $\\lambda$ is sufficiently large the estimating solution can tightly approximate the ideal solution. \n\nMoreover, to establish the theoretical connection of OT-based distributional robustness on the model space and SAM, we use the original form of the proposed OT-based distributional robustness but employ a specific distance. This purpose is to demonstrate that our proposed OT-based distributional robustness is richer than SAM in terms of modelling and can be viewed as a probabilistic extension of SAM.\n\nPutting all together, from the modeling perspective, the current proposed practical solution is a relaxation of the original OT-based distributional robustness, while SAM is another specific case in the spectrum of our OT-based distributional robustness. However, we believe that this viewpoint is interesting and opens doors to develop future efficient practical methods. For instance, we can investigate the dual form in Eq. (3) using more advanced and meaningful distances between two models.  \n\n\n[1] Klabunde, M., Schumacher, T., Strohmaier, M., and Lemmerich, F. (2023). Similarity of Neural Network Models: A Survey of Functional and Representational Measures. arXiv preprint arXiv:2305.06329.\n\n**The practical methods stated in Section 4.2 are not easy to follow. It would be better if the authors can provide a pseudocode to help the readers understand the methods easily.**\n\nThanks for this advice. We will definitely provide pseudocode in the appendix of the revised version. "}]

Author rebuttal:
We appreciate the reviewers' valuable comments and would like to report additional experiments in the attached pdf.

Meta-review: The strengths of the paper lie in its theoretical foundation and innovative insights into the connection between distributional robustness and local sharpness. The empirical evidence provided, showing improvements over existing methods, adds further credibility to the work.

While there were initial concerns raised by reviewers about the novelty, presentation clarity, and the absence of standard deviations in the results, the authors addressed these effectively in their rebuttal. They clarified misunderstandings, updated experimental data, and engaged actively in discussions, boosting confidence in their work.

There may still be room for debate on specific aspects of novelty and methodology, but the overall quality and impact of the paper are substantial.

Considering all of these factors, I recommend accepting this paper for publication. Its contributions to the field are significant and warrant its inclusion in the conference.
Decision: Accept (poster)