Worst-case generation via minimax optimization in Wasserstein space

TL;DR

This paper introduces a Wasserstein space-based minimax framework for worst-case generative modeling, improving scalability and expressiveness over traditional discrete distributionally robust optimization methods.

Contribution

It develops a continuous, transport map-based approach for worst-case generation, with a GDA algorithm and neural network parameterization that ensures convergence and efficiency.

Findings

Effective worst-case generation demonstrated on synthetic data

Neural transport maps enable scalable risk-induced generation

Method outperforms traditional discrete DRO approaches

Abstract

Worst-case generation plays a critical role in evaluating robustness and stress-testing systems under distribution shifts, in applications ranging from machine learning models to power grids and medical prediction systems. We develop a generative modeling framework for worst-case generation for a pre-specified risk, based on min-max optimization over continuous probability distributions, namely the Wasserstein space. Unlike traditional discrete distributionally robust optimization approaches, which often suffer from scalability issues, limited generalization, and costly worst-case inference, our framework exploits the Brenier theorem to characterize the least favorable (worst-case) distribution as the pushforward of a transport map from a continuous reference measure, enabling a continuous and expressive notion of risk-induced generation beyond classical discrete DRO formulations. Based on the min-max formulation, we propose a Gradient Descent Ascent (GDA)-type scheme that updates the decision model and the transport map in a single loop, establishing global convergence guarantees under mild regularity assumptions and possibly without convexity-concavity. We also propose to parameterize the transport map using a neural network that can be trained simultaneously with the GDA iterations by matching the transported training samples, thereby achieving a simulation-free approach. The efficiency of the proposed method as a risk-induced worst-case generator is validated by numerical experiments on synthetic and image data.