Efficient Stochastic Gradient Descent for Learning with Distributionally Robust Optimization

TL;DR

This paper introduces an efficient stochastic gradient descent algorithm for distributionally robust optimization, improving convergence and generalization in machine learning models through a novel subset sampling approach.

Contribution

The paper presents a new stochastic gradient descent method tailored for DRO problems, with theoretical analysis on subset size growth and empirical validation showing improved performance.

Findings

The algorithm converges efficiently with increasing subset sizes.

Empirical results show improved model generalization.

Theoretical analysis balances stochastic error and computational effort.

Abstract

Distributionally robust optimization (DRO) problems are increasingly seen as a viable method to train machine learning models for improved model generalization. These min-max formulations, however, are more difficult to solve. We therefore provide a new stochastic gradient descent algorithm to efficiently solve this DRO formulation. Our approach applies gradient descent to the outer minimization formulation and estimates the gradient of the inner maximization based on a sample average approximation. The latter uses a subset of the data in each iteration, progressively increasing the subset size to ensure convergence. Theoretical results include establishing the optimal manner for growing the support size to balance a fundamental tradeoff between stochastic error and computational effort. Empirical results demonstrate the significant benefits of our approach over previous work, and also illustrate how learning with DRO can improve generalization.