Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

MixMax: Distributional Robustness in Function Space via Optimal Data Mixtures

Anvith Thudi, Chris J. Maddison

TL;DR

MixMax reframes group DRO by optimizing over bounded functions in function space, proving a minimax result that the worst-case distribution can be realized as an optimal data mixture. It then shows that, for cross-entropy and $\ell_2^2$ losses, the MixMax objective is concave in the mixture weights, enabling a convex optimization over mixture weights followed by standard risk minimization. The paper introduces EMixMax and E$^2$MixMax to compute gradients empirically, including data-reuse variants, and demonstrates improved worst-group performance over baselines on sequence modeling and tabular tasks with non-parametric models like XGBoost. These findings offer a scalable, non-parametric DRO approach with practical gradient estimators and actionable prescription for mixture-based robustness in real-world deployments. Overall, MixMax provides a principled, efficient path to distributional robustness when model classes are expressive and data come from multiple sources.

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.

MixMax: Distributional Robustness in Function Space via Optimal Data Mixtures

TL;DR

MixMax reframes group DRO by optimizing over bounded functions in function space, proving a minimax result that the worst-case distribution can be realized as an optimal data mixture. It then shows that, for cross-entropy and losses, the MixMax objective is concave in the mixture weights, enabling a convex optimization over mixture weights followed by standard risk minimization. The paper introduces EMixMax and EMixMax to compute gradients empirically, including data-reuse variants, and demonstrates improved worst-group performance over baselines on sequence modeling and tabular tasks with non-parametric models like XGBoost. These findings offer a scalable, non-parametric DRO approach with practical gradient estimators and actionable prescription for mixture-based robustness in real-world deployments. Overall, MixMax provides a principled, efficient path to distributional robustness when model classes are expressive and data come from multiple sources.

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.
Paper Structure (36 sections, 3 theorems, 14 equations, 11 figures, 4 tables, 1 algorithm)

This paper contains 36 sections, 3 theorems, 14 equations, 11 figures, 4 tables, 1 algorithm.

Table of Contents

  1. Introduction
  2. Preliminaries
  3. DRO Over Bounded Functions is Solved by Data Mixing
  4. Group DRO by Maximization
  5. MixMax
  6. The Objective is Concave
  7. Computing Gradients of the MixMax Objective
  8. Case 1: No Covariate Shift
  9. Case 2: Covariate Shift
  10. Empirical MixMax Gradients and Practical Considerations
  11. Empirical MixMax (EMixMax)
  12. Empirical$\mathbf{^2}$ MixMax
  13. How to Use MixMax Weights
  14. Illustrating MixMax
  15. Related Work
  16. ...and 21 more sections

Key Result

Theorem 3.1

Let $P$ be a finite set of probability distributions $dp$ on the product space $\mathcal{X} \times \mathcal{Y}$ with $\mathcal{Y} \subset \mathbb{R}^n$ for some $n$, such that $\forall dp \in P$, $dp(x)$ is absolutely continuous w.r.t a given $\sigma$-finite measure $dx$ on $\mathcal{X}$. Let $\math then there exists a minimizer $f_{\lambda}$ of the expected loss under $dp_{\lambda}$ that also rea

Figures (11)

  • Figure 1: MixMax for classification picks the label probability that maximizes entropy (is closest to the centre of the simplex) in the convex hull of the distributions. We illustrate the label probabilities given by MixMax compared to balancing the distributions when there is only one input and the objective is to minimize worst-case cross-entropy loss.
  • Figure 2: An illustration of the data requirements for the empirical MixMax gradients.
  • Figure 3: EMixMax for cross-entropy maximizes average prediction entropy within the mixture span.
  • Figure 4: E$^2$MixMax found better mixture weights on a sequence modeling task, and its ensemble model performed better than the group DRO trained model. The improvement is stronger when the distributions are less similar. We present the mean and $0.5$ standard deviation of the worst group (i.e., Markov chain) cross-entropy of the Bayes optimal function for the different methods' mixture weights in Figure \ref{['fig:autoregressive_opt']}. Figure \ref{['fig:autoregressive_emp']} further compares the performance between using E$^2$MixMax weights to ensemble its proxy functions and the model given by training with group DRO.
  • Figure 5: E$^2$MixMax with Data Reuse improved worst group accuracy over balancing data more when there was bigger room for improvement. In Figure \ref{['fig:tab_exp']} we present the mean and $1$ standard deviation (over $5$ trials) of the worst group accuracy of E$^2$MixMax with Data Reuse and balanced data as a function of the oracle accuracy for that setting. In Figure \ref{['fig:MixMax_corr']} we plot E$^2$MixMax with Data Reuse's improvement in worst group accuracy over balanced data as a function of the Oracle worst group accuracy (i.e., having a model trained for each distribution) improvement over balanced data; we observed a Pearson correlation of $0.792$ with p-value $0.002$.
  • ...and 6 more figures

Theorems & Definitions (8)

  • Theorem 3.1: DRO over $L^{\infty}$ = DM
  • Corollary 3.2: Group DRO by Maximization
  • proof
  • Theorem A.1: sion1958general
  • proof
  • proof
  • proof
  • proof