Fast Distributionally Robust Learning with Variance Reduced Min-Max Optimization
This work addresses the need for scalable and efficient DRSL algorithms for building reliable machine learning systems robust to distribution shifts, representing an incremental improvement over existing methods.
The paper tackles the problem of scaling distributionally robust supervised learning (DRSL) with Wasserstein distance, which is limited by complex subproblems and lack of stochastic gradients in existing methods. It introduces stochastic extra-gradient algorithms with variance reduction and random reshuffling, achieving faster convergence rates and demonstrating effectiveness on synthetic and real data.
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.