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 data space. In this work, we explore an optimal transport-based distributional robustness framework on model spaces. Specifically, we examine a model distribution in a Wasserstein ball of a given center model distribution that maximizes the loss. We have developed theories that allow us to learn the optimal robust center model distribution. Interestingly, through our developed theories, we can flexibly incorporate the concept of sharpness awareness into training a single model, ensemble models, and Bayesian Neural Networks by considering specific forms of the center model distribution, such as 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 viewed as a probabilistic extension of SAM. We conduct extensive experiments to demonstrate the usefulness of our framework in the aforementioned settings, and the results show remarkable improvements in our approaches to the baselines.

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