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.

翻译：分布鲁棒性是一个有前景的框架，用于训练对对抗样本和数据分布偏移不太敏感的深度学习模型。以往的工作主要集中在数据空间中利用分布鲁棒性。在这项工作中，我们探索了模型空间中基于最优传输的分布鲁棒性框架。具体而言，我们考察了一个以给定模型分布为中心的Wasserstein球内的模型分布，该分布旨在最大化损失。我们发展了理论，使我们能够学习最优鲁棒中心模型分布。有趣的是，我们发展的理论允许我们灵活地将锐度感知概念融入训练中，无论是单个模型、集成模型还是贝叶斯神经网络，只需考虑中心模型分布的特定形式。这些形式包括单个模型上的狄拉克δ分布、多个模型上的均匀分布以及通用的贝叶斯神经网络。此外，我们证明了当使用单个模型上的狄拉克δ分布时，锐度感知最小化（SAM）是我们框架的一个特例，而我们的框架可被视为SAM的概率扩展。为了验证我们框架在上述设定中的有效性，我们进行了大量实验，结果表明与基线相比有显著改善。