Gradient Flow Sampler-based Distributionally Robust Optimization

We propose a mathematically principled PDE gradient flow framework for distributionally robust optimization (DRO). Exploiting the recent advances in the intersection of Monte Carlo sampling and statistical optimal transport, we show that our theoretical framework can be implemented as practical algorithms for sampling from worst-case distributions and, consequently, DRO. While numerous previous works have relied on dual reformulation techniques, we contribute a sound and complete gradient flow view based on SDEs or PDEs that can be used to construct new algorithms for general, potentially non-convex, losses. Without loss of generality, we solve a class of Wasserstein and entropy-regularized DRO problems using the recently-discovered Wasserstein Fisher-Rao and Stein variational gradient flows. Notably, we also show some simple reductions of our framework recover exactly previously proposed popular DRO methods, and provide new insights into their theoretical limits and optimization dynamics of DRO. Numerical studies based on stochastic gradient descent on machine learning tasks provide empirical backing for our theoretical findings.

Machine learning models learn from the data we give them, but they often fail when reality looks a little different. A photo might be tweaked to fool them, or the training data might have left out whole groups. Researchers guard against this by training models against the "worst-case" data they might face, but the usual recipes only work for simple problems, not the large neural networks people actually use. What if finding that worst-case data is really just a sampling problem? So we let a cloud of "particles," each one a stand-in for a possible data point, drift toward the most troublesome examples, guided by rules borrowed from physics. We even gave the cloud a "birth-and-death" rule, so it could abandon dead ends and gather where the real difficulty lies. This works where older methods don't, and it offers many different worst-case scenarios instead of just one. In our experiments, models trained this way resisted attacks noticeably better.