Optimal transport methods improve robustness of adversarial learning
Brenier Meets Adversarial Training: Optimal Transport Geometry for Robust Learning
Machine Learning
Summary
When computers learn from data, they can sometimes get tripped up if the data changes a bit. The authors study a way to train computers to be tougher against such changes, using ideas from math about moving things optimally (called optimal transport). They show that the usual ways people try to 'fool' these learning machines don’t work as efficiently because they break certain mathematical rules. The authors propose new techniques that respect these rules and lead to better defenses against tricky data, which helps computers learn more reliably in changing situations.
What this means in practice
- •For machine learning engineers: Build models more resilient to changes in data by incorporating optimal transport-based adversarial training to improve robustness and generalization.
- •For robotics control teams: Enhance control systems' resistance to uncertain environment changes by applying robust learning methods grounded in optimal transport geometry.
Authors
Alireza Abdollahpoorrostam, Ehsan Sharifian, Buse Şen, Marco Cuturi, Daniel Kuhn
Abstract
Distributionally robust optimization (DRO) provides a principled framework for learning under distribution shift, but its practical use is hindered by the difficulty of evaluating worst-case risks for nonconvex loss functions. We study a penalized DRO formulation in which the adversary may choose any distribution but incurs a Wasserstein penalty for deviating from the empirical distribution. We show that the adversary's problem can be reformulated as an optimization problem over transport maps that push empirical samples to adversarial ones, and we prove that optimal maps are cyclically monotone. We also show that standard adversarial training---based on per-sample local optimization---violates cyclical monotonicity and wastes transport costs unless the adversary is severely restricted. We propose two remedies. First, we introduce multi-start particle ascent, which alternates parallel gradient ascent with reassignment to enforce cyclical monotonicity across samples. Second, we parameterize adversarial maps as gradients of input-convex neural networks, which guarantees cyclical monotonicity by construction. Experiments on robust regression, image classification, and robust control show that our methods consistently outperform standard adversarial training and state-of-the-art baselines, achieving improved robustness and better generalization under distribution shift.