MixMax: Distributional Robustness in Function Space via Optimal Data Mixtures

2024-06-03

AI Summary

To address the convergence and solvability challenges of group distributionally robust optimization (group DRO) under non-convex losses and non-parametric models, this paper proposes MixMax—a function-space perspective framework. First, it establishes the minimax theorem for group DRO in function space—the first such result. Second, it proves that under cross-entropy and mean squared error losses, the optimal mixture distribution over groups can be recovered exactly via convex optimization, ensuring both theoretical tractability and computational scalability. Third, it enables efficient optimization through function-space reparameterization and universal function approximators (e.g., XGBoost). Empirical evaluation on ACSIncome and CelebA demonstrates that MixMax matches or surpasses state-of-the-art group DRO methods: worst-group accuracy improves by 3.2% over data-balanced baselines, significantly enhancing fairness and robustness across heterogeneous subpopulations.

Technology Category

• Search and Optimization: Mixed Discrete/Continuous Search
• Constraint Satisfaction and Optimization: Mixed Discrete/Continuous Optimization
• Machine Learning: Distributed Machine Learning & Federated Learning

Application Category

• User Modeling, Personalization and Recommendation: Fairness-aware retrieval and ranking
• Economics, Online Markets and Human Computation: Fairness, privacy, and diversity in economic environments
• Search and Retrieval-Augmented AI: Web learning to rank, online learning, and counterfactual learning for ranking

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.

Problem

• Address non-convex loss in group DRO
• Reparameterize group DRO to function space
• Solve group DRO via convex optimization

Innovation

• Reparameterize group DRO
• Minimax theorem application
• Convex optimization solution