# Optimal Transport Model Distributional Robustness

arXiv: https://arxiv.org/abs/2306.04178 · HTML: https://mlsanity.com/paper/2306.04178/

## Key facts

- Cited by 5 papers in the map (data through 14 Oct 2026).
- Builds on 51 of its 72 references that are in the map.
- First posted on arXiv 7 Jun 2023.

## Builds on (most cited first, 30 of 51)

- [Explaining and Harnessing Adversarial Examples](https://mlsanity.com/paper/1412.6572/) (2014)
- [Towards Deep Learning Models Resistant to Adversarial Attacks](https://mlsanity.com/paper/1706.06083/) (2017)
- [Simple and Scalable Predictive Uncertainty Estimation using Deep Ensembles](https://mlsanity.com/paper/1612.01474/) (2016)
- [Sharpness-Aware Minimization for Efficiently Improving Generalization](https://mlsanity.com/paper/2010.01412/) (2020)
- [On Large-Batch Training for Deep Learning: Generalization Gap and Sharp Minima](https://mlsanity.com/paper/1609.04836/) (2016)
- [Averaging Weights Leads to Wider Optima and Better Generalization](https://mlsanity.com/paper/1803.05407/) (2018)
- [Deep Mutual Learning](https://mlsanity.com/paper/1706.00384/) (2017)
- [Exploring Generalization in Deep Learning](https://mlsanity.com/paper/1706.08947/) (2017)
- [Loss Surfaces, Mode Connectivity, and Fast Ensembling of DNNs](https://mlsanity.com/paper/1802.10026/) (2018)
- [A Simple Baseline for Bayesian Uncertainty in Deep Learning](https://mlsanity.com/paper/1902.02476/) (2019)
- [Computing Nonvacuous Generalization Bounds for Deep (Stochastic) Neural Networks with Many More Parameters than Training Data](https://mlsanity.com/paper/1703.11008/) (2017)
- [Fantastic Generalization Measures and Where to Find Them](https://mlsanity.com/paper/1912.02178/) (2019)
- [Be Your Own Teacher: Improve the Performance of Convolutional Neural Networks via Self Distillation](https://mlsanity.com/paper/1905.08094/) (2019)
- [Variational Dropout and the Local Reparameterization Trick](https://mlsanity.com/paper/1506.02557/) (2015)
- [Sharp Minima Can Generalize For Deep Nets](https://mlsanity.com/paper/1703.04933/) (2017)
- [Regularizing Neural Networks by Penalizing Confident Output Distributions](https://mlsanity.com/paper/1701.06548/) (2017)
- [SWAD: Domain Generalization by Seeking Flat Minima](https://mlsanity.com/paper/2102.08604/) (2021)
- [Entropy-SGD: Biasing Gradient Descent Into Wide Valleys](https://mlsanity.com/paper/1611.01838/) (2016)
- [Snapshot Ensembles: Train 1, get M for free](https://mlsanity.com/paper/1704.00109/) (2017)
- [Certifying Some Distributional Robustness with Principled Adversarial Training](https://mlsanity.com/paper/1710.10571/) (2017)
- [ASAM: Adaptive Sharpness-Aware Minimization for Scale-Invariant Learning of Deep Neural Networks](https://mlsanity.com/paper/2102.11600/) (2021)
- [Data-driven Distributionally Robust Optimization Using the Wasserstein Metric: Performance Guarantees and Tractable Reformulations](https://mlsanity.com/paper/1505.05116/) (2015)
- [Three Factors Influencing Minima in SGD](https://mlsanity.com/paper/1711.04623/) (2017)
- [When Vision Transformers Outperform ResNets without Pre-training or Strong Data Augmentations](https://mlsanity.com/paper/2106.01548/) (2021)
- [Pitfalls of In-Domain Uncertainty Estimation and Ensembling in Deep Learning](https://mlsanity.com/paper/2002.06470/) (2020)
- [Distributionally Robust Stochastic Optimization with Wasserstein Distance](https://mlsanity.com/paper/1604.02199/) (2016)
- [Statistics of Robust Optimization: A Generalized Empirical Likelihood Approach](https://mlsanity.com/paper/1610.03425/) (2016)
- [Surrogate Gap Minimization Improves Sharpness-Aware Training](https://mlsanity.com/paper/2203.08065/) (2022)
- [Wasserstein Distributionally Robust Optimization: Theory and Applications in Machine Learning](https://mlsanity.com/paper/1908.08729/) (2019)
- [Towards Understanding Sharpness-Aware Minimization](https://mlsanity.com/paper/2206.06232/) (2022)

## Built on by (most cited first, 5 of 5)

- [Antibody: Strengthening Defense Against Harmful Fine-Tuning for Large Language Models via Attenuating Harmful Gradient Influence](https://mlsanity.com/paper/2603.00498/) (2026)
- [Promoting Ensemble Diversity with Interactive Bayesian Distributional Robustness for Fine-tuning Foundation Models](https://mlsanity.com/paper/2506.07247/) (2025)
- [Unveiling m-Sharpness Through the Structure of Stochastic Gradient Noise](https://mlsanity.com/paper/2509.18001/) (2025)
- [Sharpness-Aware Minimization in Logit Space Efficiently Enhances Direct Preference Optimization](https://mlsanity.com/paper/2603.18258/) (2026)
- [Sharpness-Aware Teleportation on Riemannian Manifolds](https://mlsanity.com/paper/2309.17215/) (2023)

## Abstract (arXiv)

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.