Distributionally Robust Optimization via Iterative Algorithms in Continuous Probability Spaces

TL;DR

This paper introduces an iterative algorithm with convergence guarantees for distributionally robust optimization in continuous probability spaces, addressing computational challenges and enabling practical solutions.

Contribution

It develops a globally convergent iterative method for continuous DRO problems using convex analysis and transport maps, bridging theory and practical algorithms.

Findings

Achieves global convergence under mild assumptions.

Efficient computation of worst-case distributions via modified JKO scheme.

Provides convergence rate and iteration complexity analysis.

Abstract

We consider a minimax problem motivated by distributionally robust optimization (DRO) when the worst-case distribution is continuous, leading to significant computational challenges due to the infinite-dimensional nature of the optimization problem. Recent research has explored learning the worst-case distribution using neural network-based generative models to address these computational challenges but lacks algorithmic convergence guarantees. This paper bridges this theoretical gap by presenting an iterative algorithm to solve such a minimax problem, achieving global convergence under mild assumptions and leveraging technical tools from vector space minimax optimization and convex analysis in the space of continuous probability densities. In particular, leveraging Brenier's theorem, we represent the worst-case distribution as a transport map applied to a continuous reference measure and reformulate the regularized discrepancy-based DRO as a minimax problem in the Wasserstein space. Furthermore, we demonstrate that the worst-case distribution can be efficiently computed using a modified Jordan-Kinderlehrer-Otto (JKO) scheme with sufficiently large regularization parameters for commonly used discrepancy functions, linked to the radius of the ambiguity set. Additionally, we derive the global convergence rate and quantify the total number of subgradient and inexact modified JKO iterations required to obtain approximate stationary points. These results are potentially applicable to nonconvex and nonsmooth scenarios, with broad relevance to modern machine learning applications.