Learning Distributionally Robust Models at Scale via Composite Optimization

TL;DR

This paper introduces scalable composite optimization methods for distributionally robust optimization (DRO), enabling efficient training of robust models on large datasets, overcoming previous computational limitations.

Contribution

It reformulates DRO as a finite-sum composite optimization problem and proposes scalable algorithms, demonstrating improved efficiency and robustness on large-scale data.

Findings

Proposed algorithms outperform prior methods in scalability.

Empirical results show effective robust model training on large datasets.

Reformulation simplifies complex DRO problems for practical use.

Abstract

To train machine learning models that are robust to distribution shifts in the data, distributionally robust optimization (DRO) has been proven very effective. However, the existing approaches to learning a distributionally robust model either require solving complex optimization problems such as semidefinite programming or a first-order method whose convergence scales linearly with the number of data samples -- which hinders their scalability to large datasets. In this paper, we show how different variants of DRO are simply instances of a finite-sum composite optimization for which we provide scalable methods. We also provide empirical results that demonstrate the effectiveness of our proposed algorithm with respect to the prior art in order to learn robust models from very large datasets.