Research map / Optimization / Nonconvex and smooth optimization
MixMax: Distributional Robustness in Function Space via Optimal Data Mixtures
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Reparameterizes group DRO over bounded functions, proves a minimax theorem showing the worst-case optimum is a data mixture, and gives a concave mixture-weight objective for cross-entropy and squared loss, solved by entropic mirror ascent (MixMax).
Problem
Group distributionally robust optimization is hard for non-convex parametric models and absent for non-parametric learners such as gradient boosting; the best data mixture for group DRO is unknown.
Approach
Move group DRO to function space over bounded functions, where a minimax theorem holds so the worst-case solution equals fitting a specific mixture; for cross-entropy and squared loss the optimal mixture weights maximize a concave objective (average Bayes error), which is optimized by stochastic entropic mirror ascent using per-group proxy models.
Key results
- On ACSIncome and CelebA variants with XGBoost, E2MixMax with data reuse improved worst-group accuracy over balancing by 1.6% (ACSIncome sex, 1 feature) and 2.3%/3.2% (CelebA young, 10 and 5 features), with a Pearson correlation of 0.79 between improvement and oracle gap.
- On sequence modeling, E2MixMax mixture weights gave lower worst-group cross-entropy than DoReMi, group DRO and balanced weights, and performed close to E-MixMax with true probabilities.
- E2MixMax with data reuse matched E2MixMax with a separate split while using fewer samples, indicating training samples can be reused.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Use E²MixMax with data reuse weights instead of balanced group weighting | XGBoost · ACSIncome and CelebA annotation variants | worst-group accuracy: 1.6%; 2.3%; 3.2% | Mixed |
| Use E²MixMax with data reuse weights instead of balanced group weighting | XGBoost · CelebA-Young, 10 features | worst-group accuracy: 62.5 (\pm 0.3) -> 64.8 (\pm 1.6) | Helped |
| Use E²MixMax with data reuse weights instead of balanced group weighting | XGBoost · ACSIncome-Sex, 1 feature | worst-group accuracy: 56.5 (\pm 0.2) -> 58.1 (\pm 1.8) | Helped |
| Use E²MixMax with data reuse weights instead of balanced group weighting | XGBoost · CelebA-Pale Skin, 5 features | worst-group accuracy: 59.5 (\pm 0.3) -> 59.4 (\pm 0.3) | No clear effect |
| Retrain a single XGBoost model on the weighted mixture instead of ensembling proxy models (ablation) | E²MixMax with data reuse · ACSIncome and CelebA annotation variants | worst-group accuracy | Helped |
| Reuse the same samples for proxy fitting and mixture-weight estimation instead of splitting data (ablation) | E²MixMax · Markov-chain sequence modeling | MixMax objective | No clear effect |
| Use data reuse rather than separate proxy and MixMax datasets (ablation) | E²MixMax · Markov-chain sequence modeling, low-sample regime | MixMax objective | Helped |
| Tune the number of optimization steps over 20, 40, 60, and 80 (ablation) | E²MixMax with data reuse on XGBoost · Selected ACSIncome and CelebA variants | worst-group accuracy: 53.7 (\pm 0.5\%) -> 54.3 (\pm 0.6\%) | Mixed |
Surprising findings
- Group DRO's mixture weights performed comparably to balanced weights on the sequence task, though ensembling E2MixMax proxy models outperformed the group DRO trained model.
- Ensembling proxy models with MixMax weights was worse than retraining a single model on the MixMax mixture on tabular data.
Limitations
- Theory requires Bayes-optimal functions to be in the bounded function class and bounded loss; empirical versions only approximate this.
- Covariate shift requires modeling covariate densities, which is data-hungry in high-dimensional spaces.
- Experiments are limited to small transformers on synthetic Markov chains and XGBoost on two tabular datasets; no large-scale evaluation.
Experimental setup
Datasets: ACSIncome CelebA annotations synthetic Markov chain sequences
Benchmarks: ACSIncome CelebA annotations synthetic Markov chain sequences
Metrics: worst-group accuracy worst-group cross-entropy worst-group loss MixMax objective value
Baselines: balanced data DoReMi group DRO (gradient descent-ascent) EMixMax
Builds on: group DRO Sion's minimax theorem DoReMi entropic mirror ascent
Closest papers by meaning
- Focus on the Common Good: Group Distributional Robustness Follows (ICLR 2022 · similarity 0.89)
- Examining and Combating Spurious Features under Distribution Shift (ICML 2021 · similarity 0.88)
- Group Distributionally Robust Optimization with Flexible Sample Queries (NeurIPS 2026 · similarity 0.88)
- Beyond Minimax Rates in Group Distributionally Robust Optimization via a Novel Notion of Sparsity (ICML 2025 · similarity 0.87)
- Efficient Algorithms for Empirical Group Distributionally Robust Optimization and Beyond (ICML 2024 · similarity 0.87)
- Structure-informed Risk Minimization for Robust Ensemble Learning (ICML 2025 · similarity 0.87)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.