MixMax: Distributional Robustness in Function Space via Optimal Data Mixtures

TL;DR

MixMax introduces a novel approach to group distributionally robust optimization by reparameterizing it in function space, enabling convex optimization solutions for non-convex models, and demonstrating improved performance on real datasets.

Contribution

This paper proposes MixMax, a new method that reformulates group DRO in function space, allowing convex optimization solutions for complex models and improving robustness.

Findings

MixMax matches or outperforms standard group DRO baselines.

MixMax improves XGBoost performance over data balancing.

The method is effective on ACSIncome and CelebA datasets.

Abstract

Machine learning models are often required to perform well across several pre-defined settings, such as a set of user groups. Worst-case performance is a common metric to capture this requirement, and is the objective of group distributionally robust optimization (group DRO). Unfortunately, these methods struggle when the loss is non-convex in the parameters, or the model class is non-parametric. Here, we make a classical move to address this: we reparameterize group DRO from parameter space to function space, which results in a number of advantages. First, we show that group DRO over the space of bounded functions admits a minimax theorem. Second, for cross-entropy and mean squared error, we show that the minimax optimal mixture distribution is the solution of a simple convex optimization problem. Thus, provided one is working with a model class of universal function approximators, group DRO can be solved by a convex optimization problem followed by a classical risk minimization problem. We call our method MixMax. In our experi ments, we found that MixMax matched or outperformed the standard group DRO baselines, and in particular, MixMax improved the performance of XGBoost over the only baseline, data balancing, for variations of the ACSIncome and CelebA annotations datasets.