# MixMax: Distributional Robustness in Function Space via Optimal Data Mixtures

arXiv: https://arxiv.org/abs/2406.01477 · HTML: https://mlsanity.com/paper/2406.01477/

## Key facts

- Cited by 1 papers in the map (data through 14 Oct 2026).
- Builds on 16 of its 36 references that are in the map.
- First posted on arXiv 3 Jun 2024.

## Builds on (most cited first, 16 of 16)

- [XGBoost: A Scalable Tree Boosting System](https://mlsanity.com/paper/1603.02754/) (2016)
- [Deep Learning Face Attributes in the Wild](https://mlsanity.com/paper/1411.7766/) (2014)
- [Invariant Risk Minimization](https://mlsanity.com/paper/1907.02893/) (2019)
- [Distributionally Robust Neural Networks for Group Shifts: On the Importance of Regularization for Worst-Case Generalization](https://mlsanity.com/paper/1911.08731/) (2019)
- [DeepSeek LLM: Scaling Open-Source Language Models with Longtermism](https://mlsanity.com/paper/2401.02954/) (2024)
- [Beyond neural scaling laws: beating power law scaling via data pruning](https://mlsanity.com/paper/2206.14486/) (2022)
- [Retiring Adult: New Datasets for Fair Machine Learning](https://mlsanity.com/paper/2108.04884/) (2021)
- [DoReMi: Optimizing Data Mixtures Speeds Up Language Model Pretraining](https://mlsanity.com/paper/2305.10429/) (2023)
- [A Survey on Data Selection for Language Models](https://mlsanity.com/paper/2402.16827/) (2024)
- [Simple data balancing achieves competitive worst-group-accuracy](https://mlsanity.com/paper/2110.14503/) (2021)
- [Distributionally Robust Language Modeling](https://mlsanity.com/paper/1909.02060/) (2019)
- [DoGE: Domain Reweighting with Generalization Estimation](https://mlsanity.com/paper/2310.15393/) (2023)
- [Towards a statistical theory of data selection under weak supervision](https://mlsanity.com/paper/2309.14563/) (2023)
- [Scaling laws for learning with real and surrogate data](https://mlsanity.com/paper/2402.04376/) (2024)
- [Algorithmic Bias and Data Bias: Understanding the Relation between Distributionally Robust Optimization and Data Curation](https://mlsanity.com/paper/2106.09467/) (2021)
- [Efficient Algorithms for Empirical Group Distributionally Robust Optimization and Beyond](https://mlsanity.com/paper/2403.03562/) (2024)

## Built on by (most cited first, 1 of 1)

- [MaD-Mix: Multi-Modal Data Mixtures via Latent Space Coupling for Vision-Language Model Training](https://mlsanity.com/paper/2602.07790/) (2026)

## Abstract (arXiv)

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.