SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Van-Anh Nguyen<sup>1</sup> Trung Le<sup>1</sup> Anh Tuan Bui<sup>1</sup> Thanh-Toan Do<sup>1</sup> Dinh Phung<sup>1,2</sup>
2VinAI, Vietnam
{van-anh.nguyen, trunglm, tuan.bui, toan.do, dinh.phung}@monash.edu
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.
Distributional robustness (DR) is a 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-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 be applied to either data space or model space, the majority of previous works on DR have primarily concentrated on exploring its applications in data space.
Sharpness-aware minimization (SAM) [18] has emerged as an 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 a training set. The success of SAM and its variants [36, 33, 62] has inspired further investigation into its formulation and behavior, as evidenced by recent works such as [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 a normal posterior, none of the existing works have explored the connection between SAM and distributional robustness.
In this work, we study the theoretical connection between distributional robustness in model space and sharpness-aware minimization (SAM), as they share a conceptual similarity. We examine Optimal Transport-based distributional robustness in model space by considering a Wasserstein ball centered around a model distribution and searching for the worst-case distribution that maximizes the empirical loss on a training set. By controlling the worst-case performance, it is expected to have a smaller generalization error, as demonstrated by a smaller empirical loss and sharpness. We then develop rigorous theories that suggest us 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) a Dirac delta distribution over a single model, (ii) a uniform distribution over several models, and (iii) a general model distribution (i.e., a Bayesian Neural Network [48, 56]). Furthermore, we show that SAM is a specific case of our framework when using a Dirac delta distribution over a 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)-based 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 improve 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.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 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 a 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].
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, 53]. 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 a method that seeks flat regions by explicitly minimizing the worst-case loss around the current model. It 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 a 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 a 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 an emerging framework for learning and decision-making under uncertainty, which seeks the worst-case expected loss among a ball of distributions, containing all distributions that are close to the empirical distribution [20].
Here we consider a generic Polish space  $S$  endowed with a distribution  $\mathbb{Q}$ . Let  $f: S \to \mathbb{R}$  be a real-valued (risk) function and  $c: S \times S \to \mathbb{R}_+$  be a cost function. Distributional robustness setting aims to find the distribution  $\tilde{\mathbb{Q}}$  in the vicinity of  $\mathbb{Q}$  and maximizes the risk in the  $\mathbb{E}$  form [61, 7]:
$$
\max  _ {\tilde {\mathbb {Q}}: \mathcal {W} _ {c} (\tilde {\mathbb {Q}}, \mathbb {Q}) <   \epsilon} \mathbb {E} _ {\tilde {\mathbb {Q}}} [ f (z) ], \tag {1}
$$
where  $\epsilon > 0$  and  $\mathcal{W}_c$  denotes the optimal transport (OT) or a Wasserstein distance [63] for a metric  $c$ , which is defined as:
$$
\mathcal {W} _ {c} (\tilde {\mathbb {Q}}, \mathbb {Q}) := \inf  _ {\gamma \in \Gamma (\tilde {\mathbb {Q}}, \mathbb {Q})} \int c d \gamma , \tag {2}
$$
where  $\Gamma\left(\tilde{\mathbb{Q}},\mathbb{Q}\right)$  is the set of couplings whose marginals are  $\tilde{\mathbb{Q}}$  and  $\mathbb{Q}$ .
With the assumption that  $f \in L^{1}(\mathbb{Q})$  is upper semi-continuous and the cost  $c$  is a non-negative lower semi-continuous satisfying  $c(z, z') = 0$  iff  $z = z'$ , [7] shows that the dual form for Eq. (1) is:
$$
\left. \min  _ {\lambda \geq 0} \left\{\lambda \epsilon + \mathbb {E} _ {z \sim \mathbb {Q}} \left[ \max  _ {z ^ {\prime}} \left\{f \left(z ^ {\prime}\right) - \lambda c \left(z ^ {\prime}, z\right) \right\} \right] \right\}. \right. \tag {3}
$$
[61] further employs a Lagrangian for Wasserstein-based uncertainty sets to arrive at a relaxed version with  $\lambda \geq 0$ :
$$
\max  _ {\tilde {\mathbb {Q}}} \left\{\mathbb {E} _ {\tilde {\mathbb {Q}}} [ f (z) ] - \lambda \mathcal {W} _ {c} (\tilde {\mathbb {Q}}, \mathbb {Q}) \right\} = \mathbb {E} _ {z \sim \mathbb {Q}} [ \max  _ {z ^ {\prime}} \{f \left(z ^ {\prime}\right) - \lambda c \left(z ^ {\prime}, z\right) \} ]. \tag {4}
$$
4.1 OT based Sharpness-aware Distribution Robustness
Given a family of deep nets  $f_{\theta}$  where  $\theta \in \Theta$ , let  $\mathbb{Q}_{\phi}$  with the density function  $q_{\phi}(\theta)$  where  $\phi \in \Phi$  be a family of distributions over the parameter space  $\Theta$ . To improve the generalization ability of the optimal  $\mathbb{Q}_{\phi^*}$ , we propose the following distributional robustness formulation on the model space:
$$
\min  _ {\phi \in \Phi} \max  _ {\tilde {\mathbb {Q}}: \mathcal {W} _ {d} (\tilde {\mathbb {Q}}, \mathbb {Q} _ {\phi}) \leq \rho} \mathcal {L} _ {S} (\tilde {\mathbb {Q}}), \tag {5}
$$
where  $S = \{(x_1, y_1), \ldots, (x_N, y_N)\}$  is a training set,  $d(\theta, \tilde{\theta}) = \| \theta - \tilde{\theta} \|_2^p$  ( $p \geq 1$ ) is a distance on the model space, and we have defined
$$
\mathcal {L} _ {S} \left(\tilde {\mathbb {Q}}\right) = \mathbb {E} _ {\theta \sim \tilde {\mathbb {Q}}} \left[ \mathcal {L} _ {S} \left(\theta\right) \right] = \mathbb {E} _ {\theta \sim \tilde {\mathbb {Q}}} \left[ \frac {1}{N} \sum_ {n = 1} ^ {N} \ell \left(f _ {\theta} \left(x _ {n}\right), y _ {n}\right) \right]
$$
with the loss function  $\ell$
The OP in (5) seeks the most challenging model distribution  $\tilde{\mathbb{Q}}$  in the WS ball around  $\mathbb{Q}_{\phi}$  and then finds  $\mathbb{Q}_{\phi}, \phi \in \Phi$  which minimizes the worst loss. To derive a solution for the OP in (5), we define
$$
\Gamma_ {\rho , \phi} = \left\{\gamma : \gamma \in \cup_ {\tilde {\mathbb {Q}}} \Gamma \left(\tilde {\mathbb {Q}}, \mathbb {Q} _ {\phi}\right), \mathbb {E} _ {(\theta , \tilde {\theta}) \sim \gamma} \left[ d (\theta , \tilde {\theta}) \right] ^ {1 / p} \leq \rho \right\}.
$$
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  _ {\phi \in \Phi} \max  _ {\gamma \in \Gamma_ {\rho , \phi}} \mathcal {L} _ {S} (\gamma), \tag {6}
$$
where  $\mathcal{L}_S(\gamma) = \mathbb{E}_{(\theta ,\tilde{\theta})\sim \gamma}\left[\frac{1}{N}\sum_{n = 1}^N\ell \left(f_{\tilde{\theta}}(x_n),y_n\right)\right].$
We now need to solve the OP in (6). To make it solvable, we add the entropic regularization term as
$$
\min  _ {\phi \in \Phi} \max  _ {\gamma \in \Gamma_ {\rho , \phi}} \left\{\mathcal {L} _ {S} (\gamma) + \frac {1}{\lambda} \mathbb {H} (\gamma) \right\}, \tag {7}
$$
where  $\mathbb{H}(\gamma)$  returns the entropy of the distribution  $\gamma$  with the trade-off parameter  $1 / \lambda$ . We note that when  $\lambda$  approaches  $+\infty$ , 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 = +\infty$ , the inner max in the OP in (7) has the solution which is a distribution with the density function
$$
\gamma^ {*} (\theta , \tilde {\theta}) = q _ {\phi} (\theta) \gamma^ {*} (\tilde {\theta} \mid \theta),
$$
where  $\gamma^{*}\left(\tilde{\theta} \mid \theta\right) = \frac{\exp\left\{\lambda\mathcal{L}_{S}\left(\tilde{\theta}\right)\right\}}{\int_{B_{\rho(\theta)}}\exp\left\{\lambda\mathcal{L}_{S}\left(\theta^{\prime}\right)\right\}d\theta^{\prime}}$ ,  $q_{\phi}(\theta)$  is the density function of the distribution  $\mathbb{Q}_{\phi}$ , and  $B_{\rho}(\theta) = \{\theta^{\prime}:\| \theta^{\prime} - \theta \|_{2}\leq \rho \}$  is the  $\rho$ -ball around  $\theta$ .
Referring to Theorem 4.2, the OP in (7) hence becomes:
$$
\min  _ {\phi \in \Phi} \mathbb {E} _ {\theta \sim \mathbb {Q} _ {\phi}, \tilde {\theta} \sim \gamma^ {*} (\tilde {\theta} | \theta)} \left[ \mathcal {L} _ {S} (\tilde {\theta}) \right]. \tag {8}
$$
The OP in (8) implies that given a model distribution  $\mathbb{Q}_{\phi}$ , we sample models  $\theta \sim \mathbb{Q}_{\phi}$ . For each individual model  $\theta$ , we further sample the particle models  $\tilde{\theta} \sim \gamma^{*}(\tilde{\theta} \mid \theta)$  where  $\gamma^{*}\left(\tilde{\theta} \mid \theta\right) \propto \exp\left\{\lambda \mathcal{L}_{S}\left(\tilde{\theta}\right)\right\}$ . Subsequently, we update  $\mathbb{Q}_{\phi}$  to minimize the average of  $\mathcal{L}_{S}\left(\tilde{\theta}\right)$ . It is worth noting that the particle models  $\tilde{\theta} \sim \gamma^{*}(\tilde{\theta} \mid \theta) \propto \exp\left\{\lambda \mathcal{L}_{S}\left(\tilde{\theta}\right)\right\}$  seek the modes of the distribution  $\gamma^{*}(\tilde{\theta} \mid \theta)$ , aiming to obtain high and highest likelihoods (i.e.,  $\exp\left\{\lambda \mathcal{L}_{S}\left(\tilde{\theta}\right)\right\}$ ). Additionally, in the implementation, we employ stochastic gradient Langevin dynamics (SGLD) [65] to sample the particle models  $\tilde{\theta}$ .
In what follows, we consider three cases where  $\mathbb{Q}_{\phi}$  is (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 (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  $\mathbb{Q}_{\phi}$  is a Dirac delta distribution over a single model  $\theta$ , i.e.,  $\mathbb{Q}_{\phi} = \delta_{\theta}$  where  $\delta$  is the Dirac delta distribution. Given the current model  $\theta$ , we sample  $K$  particles  $\tilde{\theta}_{1:K}$  using SGLD with only two-step sampling. To diversify the particles, in addition to adding a small Gaussian
noise to each particle, given a mini-batch  $B$ , for each particle  $\tilde{\theta}_k$ , we randomly split  $B = [B_k^1, B_k^2]$  into two equal halves and update the particle models as follows:
$$
\tilde {\theta} _ {k} ^ {1} = \theta + \rho \frac {\nabla_ {\theta} \mathcal {L} _ {B _ {k} ^ {1}} (\theta)}{\| \nabla_ {\theta} \mathcal {L} _ {B _ {k} ^ {1}} (\theta) \| _ {2}} + \epsilon_ {k} ^ {1},
$$
$$
\tilde {\theta} _ {k} ^ {2} = \tilde {\theta} _ {k} ^ {1} + \rho \frac {\nabla_ {\theta} \mathcal {L} _ {B _ {k} ^ {2}} \left(\tilde {\theta} _ {k} ^ {1}\right)}{\| \nabla_ {\theta} \mathcal {L} _ {B _ {k} ^ {2}} \left(\tilde {\theta} _ {k} ^ {1}\right) \| _ {2}} + \epsilon_ {k} ^ {2}, \tag {9}
$$
where  $\mathbb{I}$  is the identity matrix and  $\epsilon_k^1,\epsilon_k^2\sim \mathcal{N}(0,\rho \mathbb{I})$
Furthermore, we base on the particle models to update the next model as follows:
$$
\theta = \theta - \frac {\eta}{K} \sum_ {k = 1} ^ {K} \nabla_ {\theta} \mathcal {L} _ {B} \left(\tilde {\theta} _ {k} ^ {2}\right), \tag {10}
$$
It is worth noting that in the update formula (9), we use different random splits  $B = [B_k^1, B_k^2]$ ,  $k = 1, \dots, K$  to encourage the diversity of the particles  $\tilde{\theta}_{1: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  $\nabla_{\theta} \mathcal{L}_{B_k^2}(\tilde{\theta}_k^2(\theta))$ , we set the corresponding Hessian matrix to the identity one.
Particularly, the gradient  $\nabla_{\theta}\mathcal{L}_{B_k^2}\left(\tilde{\theta}_k^1\right)$  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  $\nabla_{\theta}\mathcal{L}_{B_k^2}\left(\tilde{\theta}_k^1\right), k = 1,\dots ,K$  all at 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  $2|B|$ , which is similar to SAM.
4.2.2 Ensemble OT-based Distributional Robustness
We now examine the case where  $\mathbb{Q}_{\phi}$  is a uniform distribution over several models, i.e.,  $\mathbb{Q}_{\phi} = \frac{1}{M}\sum_{m=1}^{M}\delta_{\theta_m}$ . In the context of ensemble learning, for each base learner  $\theta_m, m = 1,\dots,M$ , we seek  $K$  particle models  $\tilde{\theta}_{mk}, k = 1,\dots,K$  as in the case of single model.
$$
\tilde {\theta} _ {m k} ^ {1} = \theta + \rho \frac {\nabla_ {\theta} \mathcal {L} _ {B _ {m k} ^ {1}} (\theta_ {m})}{\| \nabla_ {\theta} \mathcal {L} _ {B _ {m k} ^ {1}} (\theta_ {m}) \| _ {2}} + \epsilon_ {m k} ^ {1},
$$
$$
\tilde {\theta} _ {m k} ^ {2} = \tilde {\theta} _ {m k} ^ {1} + \rho \frac {\nabla_ {\theta} \mathcal {L} _ {B _ {m k} ^ {2}} \left(\tilde {\theta} _ {m k} ^ {1}\right)}{\| \nabla_ {\theta} \mathcal {L} _ {B _ {m k} ^ {2}} \left(\tilde {\theta} _ {m k} ^ {1}\right) \| _ {2}} + \epsilon_ {m k} ^ {2}, \tag {11}
$$
where  $\mathbb{I}$  is the identity matrix and  $\epsilon_{mk}^{1},\epsilon_{mk}^{2}\sim \mathcal{N}(0,\rho \mathbb{I})$
Furthermore, we base on the particles to update the next base learners as follows:
$$
\theta_ {m} = \theta_ {m} - \frac {\eta}{K} \sum_ {k = 1} ^ {K} \nabla_ {\theta} \mathcal {L} _ {B} \left(\tilde {\theta} _ {m k} ^ {2}\right). \tag {12}
$$
It is worth noting that the random splits  $B = [B_{mk}^{1}, B_{mk}^{2}]$ ,  $m = 1, \dots, M$ ,  $k = 1, \dots, K$  of the current batch and the added Gaussian noise constitute the diversity of the base learners  $\theta_{1: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 case where  $\mathbb{Q}_{\phi}$  is an approximate posterior or a BNN. To simplify the context, we assume that  $\mathbb{Q}_{\phi}$  consists of Gaussian distributions  $\mathcal{N}(\mu_l,\mathrm{diag}(\sigma_l^2)),l = 1,\ldots ,L$  over
the weight matrices  $\theta = W_{1:L}^1$ . Given  $\theta = W_{1:L} \sim \mathbb{Q}_{\phi}$ , the reparameterization trick reads  $W_{l} = \mu_{l} + \mathrm{diag}(\sigma_{l})\kappa_{l}$  with  $\kappa_{l} \sim \mathcal{N}(0,\mathbb{I})$ . We next sample  $K$  particle models  $\tilde{\theta}_k = [\tilde{W}_{lk}]_{lk}, l = 1, \ldots, L$  and  $k = 1, \ldots, K$  from  $\gamma^{*}(\tilde{\theta} \mid \theta)$  as in Theorem 4.2. To sample  $\tilde{\theta}_k$  in the ball around  $\theta$ , for each layer  $l$ , we indeed sample  $\tilde{\mu}_{lk}$  in the ball around  $\mu_l$  and then form  $\tilde{W}_{lk} = \tilde{\mu}_{lk} + \mathrm{diag}(\sigma_l)\kappa_l$ . We randomly split  $B = [B_k^1, B_k^2]$ ,  $k = 1, \ldots, K$  and update the approximate Gaussian distribution as
$$
\tilde {\mu} _ {l k} ^ {1} = \mu_ {l} + \rho \frac {\nabla_ {\mu_ {l}} \mathcal {L} _ {B _ {k} ^ {1}} \left(\left[ \mu_ {l} + \operatorname {d i a g} (\sigma_ {l}) \kappa_ {l} \right] _ {l}\right)}{\| \nabla_ {\mu_ {l}} \mathcal {L} _ {B _ {k} ^ {1}} \left(\left[ \mu_ {l} + \operatorname {d i a g} (\sigma_ {l}) \kappa_ {l} \right] _ {l}\right) \| _ {2}} + \epsilon_ {k} ^ {1},
$$
$$
\tilde {\mu} _ {l k} ^ {2} = \tilde {\mu} _ {l k} ^ {1} + \rho \frac {\nabla_ {\mu_ {l}} \mathcal {L} _ {B _ {k} ^ {2}} \left(\left[ \tilde {\mu} _ {l k} ^ {1} + \operatorname {d i a g} (\sigma_ {l}) \kappa_ {l} \right] _ {l}\right)}{\| \nabla_ {\mu_ {l}} \mathcal {L} _ {B _ {k} ^ {2}} \left(\left[ \tilde {\mu} _ {l k} ^ {1} + \operatorname {d i a g} (\sigma_ {l}) \kappa_ {l} \right] _ {l}\right) \| _ {2}} + \epsilon_ {k} ^ {2},
$$
$$
\tilde {\theta} _ {k} = \left[ \tilde {\mu} _ {l k} ^ {2} + \operatorname {d i a g} (\sigma_ {l}) \kappa_ {l} \right] _ {l}, k = 1, \dots , K,
$$
$$
\mu = \mu - \frac {\eta}{K} \sum_ {k = 1} ^ {K} \nabla_ {\mu} \mathcal {L} _ {B} \left(\tilde {\theta} _ {k}\right) \text {a n d} \sigma = \sigma - \eta \nabla_ {\sigma} \mathcal {L} _ {B} (\theta),
$$
where  $\epsilon_k^1,\epsilon_k^2\sim \mathcal{N}(0,\rho \mathbb{I})$
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, we prove that SAM is a specific case of OT-based distributional robustness on the model space with a particular distance metric. To depart, we first recap the formulation of the OT-based distributional robustness on the model space in (5):
$$
\min  _ {\phi \in \Phi} \max  _ {\tilde {\mathbb {Q}}: \mathcal {W} _ {d} (\tilde {\mathbb {Q}}, \mathbb {Q} _ {\phi}) \leq \rho} \mathcal {L} _ {S} (\tilde {\mathbb {Q}}) = \min  _ {\phi \in \Phi} \max  _ {\tilde {\mathbb {Q}}: \mathcal {W} _ {d} (\tilde {\mathbb {Q}}, \mathbb {Q} _ {\phi}) \leq \rho} \mathbb {E} _ {\theta \sim \tilde {\mathbb {Q}}} [ \mathcal {L} _ {S} (\theta) ]. \tag {13}
$$
By linking to the dual form in (3), we reach the following equivalent OP:
$$
\left. \min  _ {\phi \in \Phi} \min  _ {\lambda > 0} \left\{\lambda \rho + \mathbb {E} _ {\theta \sim \mathbb {Q} _ {\phi}} \left[ \max  _ {\tilde {\theta}} \left\{\mathcal {L} _ {S} (\tilde {\theta}) - \lambda d (\tilde {\theta}, \theta) \right\} \right] \right\}. \right. \tag {14}
$$
Considering the simple case wherein  $\mathbb{Q}_{\phi} = \delta_{\theta}$  is a Dirac delta distribution. The OPs in (14) equivalently entails
$$
\left. \min  _ {\theta} \min  _ {\lambda > 0} \left\{\lambda \rho + \max  _ {\tilde {\theta}} \left\{\mathcal {L} _ {S} (\tilde {\theta}) - \lambda d (\tilde {\theta}, \theta) \right\} \right\}. \right. \tag {15}
$$