Fast Distributionally Robust Learning with Variance Reduced Min-Max Optimization

TL;DR

This paper introduces scalable stochastic algorithms with variance reduction for Wasserstein distributionally robust supervised learning, significantly improving convergence and efficiency in large-scale problems.

Contribution

The paper develops and analyzes variance-reduced stochastic extra-gradient algorithms for Wasserstein DRSL, enabling faster convergence and practical scalability.

Findings

Algorithms outperform existing methods in synthetic and real data.

Variance reduction accelerates stochastic min-max optimization.

Proven faster convergence rates than prior approaches.

Abstract

Distributionally robust supervised learning (DRSL) is emerging as a key paradigm for building reliable machine learning systems for real-world applications -- reflecting the need for classifiers and predictive models that are robust to the distribution shifts that arise from phenomena such as selection bias or nonstationarity. Existing algorithms for solving Wasserstein DRSL -- one of the most popular DRSL frameworks based around robustness to perturbations in the Wasserstein distance -- have serious limitations that limit their use in large-scale problems -- in particular they involve solving complex subproblems and they fail to make use of stochastic gradients. We revisit Wasserstein DRSL through the lens of min-max optimization and derive scalable and efficiently implementable stochastic extra-gradient algorithms which provably achieve faster convergence rates than existing approaches. We demonstrate their effectiveness on synthetic and real data when compared to existing DRSL approaches. Key to our results is the use of variance reduction and random reshuffling to accelerate stochastic min-max optimization, the analysis of which may be of independent interest.