MixMax: Distributional Robustness in Function Space via Optimal Data Mixtures
Authors: Anvith Thudi
Chris J. Maddison
Venue: ICLR 2025 Poster (2025)
Keywords: Distributional Robustness
Non-parametric Learning
Function Space
Trustworthy Machine Learning
Primary area: optimization
TL;DR: We show that distributionally robust optimization over the space of bounded functions simplifies to maximum likelihood on a certain, easy-to-find mixture distribution.

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 experiments, 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.

Published as a conference paper at ICLR 2025
MIXMAX: DISTRIBUTIONAL ROBUSTNESS IN FUNCTION SPACE VIA OPTIMAL DATA MIXTURES
Anvith Thudi
Chris J. Maddison
Department of Computer Science University of Toronto and Vector Institute anvi th. thudi @ il, uutorott.cca
Department of Computer Science
University of Toronto and Vector Institute cmadd wtw.eettt.n...e.eeeeeeeeeeeeeee. www.otrr.......... .ddu
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 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, 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, show that the mini- max optimal mixture distribution is the solution of a simple convex optimization problem. Thus, provided one working with model class of universal function approximators, group DRO can be solved by a convex optimization problem fol- lowed by classical risk minimization problem. We call our method MixMax. our experiments, we found that MixMax matched or outperformed the standard group DRO baselines, and in particular, MixMax improved the performance XGBoost over the only baseline, data balancing, for variations of the ACSIncome and CelebA annotations datasets.
INTRODUCTION
Machine learning pipelines often optimize for model that performs well over several different
distributions of data. This can be as the model will be deployed for different groups of users (Strack
ctt al. 2014 Blackard 1998), or in the development of foundation models for a suite of tasks.
Optimizing for the worst-case егтоr over the set of distributions captures this distributional robustness objective, and is called distributionally robust optimization (DRO) (Dupačová 1987 [Shapiro and
Kleywegt 2002 Rahimian and Mehrotra 2019 Oren al. 2019 Sagawa et al. 2019). This is often
called group DRO (Sagawa et al. 2019) when the distribution set is finite.
Despite the importance of obtaining group DRO models, lack effective methods for modern machine learning models. When our loss is convex in our model parameters, methods exist for solving the group DRO optimization (Shapiro and Kleywegt (002). But the case of modern expressive non-linear models, only heuristic methods exist (such as applying methods from the convex literature) (Xie al. 2024 Sagawa et al.) 2011).. Furthermore, these methods are expensive, unstable, and training on simple balanced mixtures of the data distributions sometimes leads to better group DRO solutions (Idrissi et al. (202).. Moreover, no group DRO methods (beyond balancing data) exist for non non-parametric models. We turn to optimizing data mixtures as a possible solution.
Intuitively, mixing the distributions we want robustness over plays an important role in obtaining distributionally robust models. Past work has already shown that how one chooses to weigh sources of data in the loss function has a significant impact on the performance of the final model (Sorscher et al. 2022 Xie et al.) 2024) across tasks, key consideration for large language model training, e.g., Section 3.3 of Bi et al. (024) and the survey by Albalak et al. (2024). Other work has also already shown that in many cases, DRO solutions minimize specific mixture (Arjovsky et al. 2019 Slowik and Bottou| 2021 2022) (existentially). Yet, is not known what the best data mixture
for group DRO is. Minimax theorems provide one way of answering this, they state specific

---

Published as a conference paper at ICLR 2025
mixture distribution whose optima solve the DRO objective. However, the loss functions over the parameters of commonly used model classes (like neural networks) are non-convex, convex minimax theorems do not apply. We address this with classic trick: working over function space.
Modern model classes are able to approximate any function arbitrarily well. Hence, when viewed in function space, these model classes cover all the bounded functions. This subset is also the hypothesis space covering all the Bayes optimal functions (for classification and bounded regression). So we can expect that the group DRO solution is within this subset for modern model classes. Hence, in this paper we study the group DRO problem over the set of bounded functions.
Highlighting the benefits of this, show that the space of bounded functions admits minimax theorem. That is, if we can fit any distribution to Bayes optimal, then fitting to a mixture with the highest Bayes errоr solves the group DRO problem. To obtain this minimax result we bypassed issues with L metrics by using topological dual spaces, which may be of independent interest.
Towards now finding this group DRO solving mixture, we show that these optimal mixture distri- butions maximize a concave objective function over the mixture weights. This result is specific to using cross-entropy and Eg losses, and leverages the structure of their Bayes optimal functions. We further show that one can compute the gradients for this objective given the optimal predictor for each component distribution. Hence, with this gradient oracle, we can use a stochastic entropic mirror ascent algorithm to maximize the objective. We call this selection of mixture weights MixMax. In practice we do not have the optimal predictor for the gradient oracle, but we show experimentally several empirical gradient estimators still lead to near optimal mixture weights.
Beyond the guarantees when fitting near optimal, MixMax has a number of advantages for group DRO. First, given a sufficiently large set of data from each component distribution, finding MixMax weights can be accomplished by fitting a separate model on each source-thh same amount of training compute as training one model on all the data. Furthermore, because the weights be used to ensemble the component models. there is little additional model training overhead. Moreover, unlike previous methods for group DRO, MixMax can be used with non-parametric model classes, like gradient boosting (Friedman, 2001
To illustrate the empirical performance of MixMax when we cannot ensure optimal fitting, we applied it for two real-world model classes. First, we tested MixMax with transformer models on synthetic Markov chain data. Second, we tested MixMax with XGBoost (Chen and Guestrin 2016) on several tabular datasets with different group shifts (Ding et al. 2021: Liu et al. 2015). In all cases, we found that empirical versions MixMax matched or outperformed applicable baseline methods when improvement was possible. In particular, when moderate label shift was present, MixMax yielded relative test accuracy improvements between 2.3 5.9% for XGBoost on variations of ACSIncome (Ding et al. 2021) and CelebA annotations (Liu et al. 2015). Our contributions are:
1. A minimax theorem for DRO over bounded funetions.
2. Showing that applied to cross-entropy and Eg this yields a concave objective to maximize for data mixing (to solve group DRO) which we call MixMax
3. Experiments showing empirical versions of MixMax improved over group DRO alternatives
4. Providing the first group DRO method applicable to non- parameetti learning, and applying it to XGBoost to improve over the baseline of balancing data by upsampling.
2 PRELIMINARIES
In this paper we will consider solving group DRO when we have model classes that are sufficiently expressive to capture the Bayes optimal function for any distribution are dealing with. This will be relaxed in Section 4.2 but for now this formally means we can (and will) reparameterize our optimization to be over generic function space.
We work with functions f X - O from an input domain X' (with a measure dz) to a closed convex output domain O C Rd that have bounded 0 norm in each output coordinate: we denote this
'To the best of our knowledge. there exists no past work on group DRO for non -araamemtt learning algorithms, with the only work on DRO in general being for k-nearest neighbours where the set of distributions forms Wasserstein ball (Chen and Paschalidis 2019.
2

---

Published as a conference paper at ICLR 2025
Simplex centre
Uniform
Components
MixMax mixture (this puper)
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 minimize worst-case cross-entropy loss.
function space by 1 (X,,) For example, when doing classification O could be the probability simplex, with the i'th entry representing the probability of label i. We will specifically consider the set of functions with outputs bounded by some fixed value in every coordinate (a.e.), denoted by (2, (r). We work with data (z. y) where I E X' and y E V, and an associated loss function C(o, y) O x - R+ that is convex in the first argument and continuous in both. An example is cross-entropy where is the label probabilities, and y specific label. With L have the expected loss of a function f E I (X, .) over distribution dp on x ) is Jxxy LJ (x), y ddp(,, y)- Given a set of distributions P, the DRO objective we study is (which is group DRO when P is finite)
inf sup / [(f(x). v)- B (r) dpEP
(DRO)
DRO OVER BOUNDED FUNCTIONS IS SOLVED BY DATA MIXING
optimal data mixture We now show that, letting model class be all bounded functions, fitting
returns a function that solves group DRO (Corollary 3.2). This is done by observing the set of all bounded functions has sufficient amount of regularity under the right topology (i.e., metric). We further show that this optimal mixture is characterized by having the highest Bayes error. Later in Section we demonstrate how to optimize for this mixture in the case of cross-entropy and & loss.
First, Theorem 3.1 formally states that there exists a minimizer to the hardest distribution which solves [RO To show this overcome challenges in applying Sion's minimax theorem bounded functions, in particular, the fact that the set is not compact in any L topology. Our main contribution show that working with the weak*-topologg resolves these issues. This follows from leveraging Mamach-Alaoglu's 's theorem alongside the method of test functions, and introduces mild additional assumptions. For example, we will require that dr is o-finite measure to apply I.P duality theory, which is satisfied if dx is Lebesque measure on R" or counting measure some countable set. We will further require that bounded on x y to use intergral convergence theorems, and discuss the strength of this assumption later. A complete proof is provided in Appendix A.I
Theorem 3.1 (DRO over - DM). Let P be a set of probability distributions dp on the product space X y with y C IR.h for some such that Vdp E P. dp(x.) is absolutely continuous w.r.t given o a-finite measure dx on X. the measure da, and (r=(x,o,(r)
C Th be a closed convex set, and T (X, O) be defined w.rt || Sr).
T (x)
Let the loss function L(o, y) be continuous in both arguments, and convex in the first argument. Furthermore assume C is bounded by some constant M y. If dpx realizes
/ ((f(=),
inf dpECono(P) (ww.oe/BB...eoooeoeee (r(
sup
3

---

Published as a conference paper at ICLR 2025
and there exists f' (x realizing the DRO of the expected loss under dpy, that also realizes the DRO objective
objective, then there exists minimizer fx
/ g)
inf (EBLocNLo dpEP
sup
That is, the DRO objective is solved by fitting a specific distribution in the convex hull of P. If the set of distributions P is finite, then the DRO objective is solved by fitting specific mixture distribution.
While bounded loss may seem strong, can be enforced by choosing the output space carefully, e.g., for cross-entropy loss if we know the minimum probability of any label is >> 0 O accordingly and avoid loss blow-up (from (og(O)). Note, our actual requirement was a condition to enforce that pointwise convergence in the loss at every x implies convergence of the integral for all dp E P. which can be true without loss boundedness depending on P. On assuming aDRO solution over (x, (r) exists, this is true for finite P (Appendix A.1) and hence applies for group DRO.
choose
For the rest of the paper, We will assume there is sufficient regularity (whether imposed by O, P. or both) for pointwise convergence of loss imply the average loss converges and hence the result of Theorem 3.1 applies. We will further assume that for some r < oo we have the bounded functions (z (r) contains the Bayes optimal solutions for all dp E (to leverage structural properties of cross-entropy and E2). Note the last two assumptions are satisfied for cross-entropy if a.e. all labels have 2 probability for some >> 0 (and we choose O accordingly), and for L if the allowed y values are bounded. We leave it to future work to consider applying Theorem without these assumptions, or generalizing Theorem 3. itself.
3.1 GROUP DRO BY MAXIMIZATION
now consider applying Theorem 3.1 for group DRO. For cross-entropy and &g losses, the Bayes optimal functions unique up to measure 0 set over X, and any minimizer of the hardest distribution DRO optimal by Theorem [.1 given DRO): all minimizers have the
(given a DRO solution exists which is true for group average losses dp P. In this case, if the set P is also finite (i.e., group DRO) we further have the max-min optimization over bounded functions reduces to single maximization over a finite dimensional simplex A..
Corollary 3.2 (Group DRO by Maximization). Let be a finite set, and assume J C(J()); y)dp(r; y) uniquely minimized over functions in BL (upto a measure set). Parameterize Conw(P) by X P via Cono(P) = {dpx Edper dn X A Further denote the mini- mizer of I [xf(),, d)dpx(z (x, y) in BI by fx (parameterized by A). Then for À realizing Suржеa" I (()x()), y)dpx (x, v), fx(x) realizes inf ( suPaper [x](x) d)ppp( v). That is, the group DRO optimization over L reduces to a single maximization over A..
Proof. Follows from Theorem 3.1 noting P is finite so DRO solution exists, and rewriting the condition for the mixture with the parameterizations given by uniqueness of solutions.
MIXMAX
We now apply Theorem 3.1 (under the assumptions of the previous section) to cross- cross-entropy and e for finite P. We use p(x) and fp(x) to represent the covariate density and Bayes optimal function of distributions dp P, and px (x) and fx(x)) similarly for the mixture distributions dpx Conu(P).
First note that, by Corollary 3.2 the objective for the optimal mixture weights
/ vl = sup Σ X / ((fxla)..vddb.... XEAP dPEP
sup A
We show that this is concave and that compute its gradients for cross-entropy and L, Thus, can perform entropic mirror ascent Duchi (2018) to solve the constrained optimization. Algorithm describes this approach, which call MixMax ("Mixtures by Maximization").
4

---

Published as a conference paper at ICLR 2025
Algorithm Empirical MixMax
Step size 7,, number of steps n, loss function L (either cross-entropy or (g), and, for each distribution dp in the set P, samples Dp. proxylexact covariate density p(r), and proxy/Bayes optimal prediction function fp-
Require:
Note: If there is no covariate shift then one can set p(æ) ((x) Ydp E P for any fixed ((z); this has no impact due to symmetry in the formula used the algorithm.
Initialize: Xp
for all dp E P
[P|
1: for = 1, do
fa(x) f Lper (pp()))
2:
2per App(2)
3: 1t Edper [D,1 (OwD)D ((fx(x),))
4: g Vi
Ape
5: Ap x Edeep Ape for all dp E P
6: end for
7: return Optapep
Empirical MixMax
Empirical MixMax
Empirical MixMax Reuse
Data
Distributions
Data
Distributions
Data
Distributions
D, p()))
Dp
P(()))
D'
Dr
fx))))(()) ((x)
Figure 2: illustration of the data requirements for the empirical MixMax gradients.
The Objective Concave In the case of cross-entropy, this objective reduces the expected entropy of y conditioned on px (y(z), over ~ Px(x). Therefore it is a concave maximization problem in the mixture weights via the concavity of entropy. In the case of l2 under our assumptions, this objective reduces to the expectation of the conditional variance of given over ~ px which is again concave objective. Appendix A.2 provides full proofs of concavity.
4.1 COMPUTING GRADIENTS OF THE MIXMAX ORJECTIVE
We now demonstrate how to compute the gradient of the objective in Equation 1 i.e., EdpEP VAA)) y) for cross-entropy and & given fp and p(x) and the ability to integrate for all dp E P. We present this cases, and later discuss empirical implementations.
Case 1: No Covariate Shift Suppose there is no covariate shift, i.e., 3 po(x:) such that p(x)
Po(x) Wdp E P. Then, for cross-entropy we have the Bayes optimal function fx(x). pa(y)æ) = Eapep App(y|z) = Edpep (xfp(() and for eg also have fx(x) E (yle)( Edper XpByyplllz)3
Edper )pf((z) Hence given fz(x) Ydp P we compute fx(x) and (x)x(x). Hence, (x (x), 3) follows by product and chain rule, and given the ability to integrate over dp E P we can compute the gradient of the objective.
The more general expression under covariate shift, for cross-entropy and E3. is (x(x) = Enger {p)p())p(2) knowledge of p(z) Vdp E P. Given this, computing (ff, y) follows by product and chain rule, and given the ability integrate over dp E P compute the gradient of the objective.
Case 2: Covariate Shift
Layep (x) Hence computing (x( and Vx)2 would further require
5

---

Published as a conference paper at ICLR 2025
4.2 EMPIRICAL MIXMAX GRADIENTS AND PRACTICAL CONSIDERATIONS
Section 4.] described how to compute the gradients of the MixMax objective given we have the Bayes optimal function (and covariate density function if necessary) for each distribution, and integrate exactly over the domain. We now decribe empirical gradients which remove these assumptions, summarized in Figure 2 These add varying sources of empirical error.
Empirical MixMax (EMixMax)
Computing the MixMax gradients is often intractable due to the integral over dp E P. However, given datasets D, for dy E one can compute empirical MixMax (EMixMax) gradients and update À with stochastic entropic mirror ascent (Algorithm
EMMMM (A)
Σ Σ VxAp((f((E).)). dpE Dpl
(x:y)ED;
We prove that this is an unbiased estimator and discuss its computational complexity in Appendix B
Empirical? MixMax One also usually does not know the Bayes optimal functions (z or the densities p(æ). However can train model on samples from dp P to obtain approximations fp(x)) and p. and then use another set of samples compute EMMizMoz We call this approach Empirical MixMax Max.xMax. More precisely, assuming covariate shift, given i.idd datasets D Dr.D2 vdp E P. the Maxixxax gradients are
MixMaxXX (A) Σ Σ VxlpB
Σ >pfo(o).1) dPEP
dpEP 1Dpl
(z.
where fp(x) arg min
(((()).v)
D
It is challenging to evaluate the expectation of M and we leave this to future work. Instead, we study this estimator empirically and find it gives results close to EMixMax. In the case of covariate shift, we also fit density approximations p(x)) on the second dataset D2 and use the covariate shift definition of fx. Algorithm [ describes the changes for MixMxax. We also consider E.2 MixMax with Data Reuse, i.e., Dp D',, as a more sample efficient alternative.
How to Use MixMax Weights Finally, MixMax mixture weights can be used in two ways. We can use the weights to define a mixture distribution and fit new model on it, which Theorem 3.1 states will result in a group DRO solution if fit optimally. Alternatively, we can use the mixture weights to combine our (approximate) fp models (i.e., the fx formula), which also returns an approximation of the best model for the MixMax mixture distribution. These approaches provide trade-off: the first approach has higher training cost but results in single model and hence smaller inference cost. We employed both approaches in our experiments and found they both improved over baselines.
4.3 ILLUSTRATING MIXMAX
Here we investigated patterns for the group DRO solutions found by EMixMax, running with many samples as to be representative of the true group DRO solution. To understand Binary Classification we considered two cases: one where the two distribution were mirror opposites of each other, and one where this was not the case. Specifically we considered binary classification distributions with
the same covariate probabilities p(x) uni 1). but with pt (1(1x) = 0.5cos (Tæ) + 0.5 and -0.5cos(T(2
(2 (x)) -0. -0.5cos + 0.5 (Figure 3a or (2(1
0.2)) - 0.5 (Figure (
Note the MixMax objective maximizes average entropy within the mixture span, and in the first
case we found EMixMax selected random guessing which is the maximal entropy distribution
Formmaly, we require (xL(f((z)) y) have sufficient regularity (e.g., bounded).

---

Published as a conference paper at ICLR 2025
1.0
1,0
0.8-
0.8
MaaMat
Immopy
0.2
0.2
MaMac
0.0-
Mecmum Entropy
0.6
(a) Maximum Entropy Conditional in Mixture Span (b) Maximum Entropy Conditional not in Mixture Span
Figure 3: EMixMax for cross-entropy maximizes average prediction entropy within the mixture span.
(the true group DRO solution). In the second case where random guessing is not in the span, EMixMax found mixture closer to random guessing than the individual distributions (expected for the true group DRO solution). For regression we considered deterministic predictions and varied which distributions had the extreme predictors. Specifically, we considered a fixed p(z) = =0, 1]. but with P1 (1) xz) 0. O.2cos( Tx) + 0.5, P2 0.1, and pa (1[z) 0.15 setting (s(1(x) 0.8 Appendix [ respectively). We found EMixMax always found the extreme functions and balanced those, which minimizes maximal & error (true group DRO solution).
(Figure 6a) and Figure 6b
5 RELATED WORK
Distributionally Robust Optimization DRO is motivated by several applications not limited to just machine learning, such as resource allocation (Dupačová 1987). Approaches for DRO often focus on "uncertainty" sets P (the set of distributions
minimizing the maximum error over) that admit a duality theory (Rahimian and Mchrotra 2019 Delage and Ye 2010; Bertsimas et al. 2010
Wiesemann et al. 2013 Ben-Tall et al. (013)) such Lagragian duality, conic duality, Fenchel duality, etc. For theorem 3.1 we assume nothing on P, and for some results, that is finite similar to work on
the Cutting-Surface approaches DRO (Rahimian and Mehrotral 2019 Priug and Wozabal 2007
Rahimian et al. 2019; [ansal et al. 2018); the finite P case is widely called group DRO (Sagawa al. 2019) deep learning for its applications where well-specified "groups" exist in the collected data. In particular, for finite P the equivalence to the convex hull of the set of distributions used in Theorem already appears as Lemma 17 Rahimian and Mchrotral (2019) and is wide-spread in the finite P literature. On similarity in theory, our main theorem that DRO is solved by data mixing most closely related to the saddle point approaches stochastic programming (Shapiro) and Kleywegt 2002 Dupačovál 1987 Nemirovski et al. 2009 Zhang al. 2024 Yu et | 2024). These results present duality theory where solving for the hardest distribution is enough when the set of parameters to minimize w.r.t is subset of _=. In the case of Dupacová (1987) focus is placed on sets of all distributions meeting certain moment constraints (e.g., only the mean and variance are known). Shapiro and Kleywegt 2002);; Zhang et al. (2024) YYu et al. (202)) consider the case of an arbitrary finite set of distributions like our work, proposing sample average approaches for optimal solutions. Our method also sample averages, but further leverages the structure of cross-entropy and E loss over bounded functions to collapse the min optimization to single maximization.
Nevertheless, the best of our knowledge, our work is the first to extend these past minimax approaches to DRO over the set of all bounded functions (Theorem 3.I.. This yields a concave objective maximize (Corollary 3.2) which we find yields good results for practical applications involving non non-maramttric learning and expressive non-linear model classes.
Data Mixing The work of Xie et al. (2024). building past work on DRO for deep learn- ing (Sagawa et al., 2219 Oren et al.) 2019), highlighted empirically how optimizing dataset mixtures can lead to better performance due to faster convergence or due to the optima being better suited for the set of downstream tasks. On this, past work has highlighted the role of dataset selection for convergence rates (Sorscher et al..
several downstream tasks. However, it is not clear whether this is
7

---

Published as a conference paper at ICLR 2025
2022; Kolossov al. 2023) alongside increasing sample access (Jain et al. 2024). In this paper asked what the role of dataset selection, in particular data mixing, is for having optima better suited for deployments with uncertain downstream tasks. The invariant risk minimization (Arjovsky et al., 2019) and data bias (Slowik and Botton, 2021 2022) literature have considered the same question, but the analysis there has been limited problems satisfying the KKT conditions or studying local minima in I8", and do not discuss how to find the best data mixture, We also note recent work by Fan al. (2023) proposed data mixing methods that build on Xie al. (2024). but departed away from the DRO objective (and Xie et al. 2024) performed comparably with enough compute).
6 EXPERIMENTS
We only proved the guarantees of MixMax when we can obtain the optimal model over all bounded functions for our distributions. This often not possible in practice. Hence, we tested how empirical implementations of MixMax (described in Section 4.2 performed for real-world model classes with only finite samples from each distribution; we observed MixMax still improved over the baselines even when such empirical errors were introduced.
In our experiments, we ran EMixMax for 10 steps with 1) = 2.0 for all the sequence modeling tasks, and for 20 steps with 7 = 0.1 for the tabular datasets unless otherwise specified; preliminary testing showed that this was enough to have the objective converge within 0.01 between iterates. We used Nvidia RTX 2080 Ti and A100 GPUs to accelerate our experiments involving small transformers, and otherwise used Intel Xeon Silver 4210 CPUs and AMD EPYC 7643 CPUs. We used the GPTNeo architecture (Black et al., 2021) for the transformers (hyperparameters described in Appendix C
6.1 EMIXMAX AND E'MIXMAX VARIANTS MAXIMIZE MIXMAX OBJECTIVE COMPARABLY
Here we investigated how Max with various data splits compared to EMixMax in maximizing the MixMax objective. Specifically, we considered tokens (0. 1,2,3) and sequences generated by Markov chain, and the task was to model the mixed distribution of sequences from lengths 1 to 10 where the probability a sequence was of length was always 1/10. We constructed three Markov chains perform group DRO over by independently sampling transition probabilities from a symmetric Dirichlet distribution with magnitude 1.0. We then constructed training datasets by varying the number of samples per length, and considered either using all the samples to perform EMixMax with the ground truth prediction probabilities, MaM with (75 : 25). (50 : 50), and (25 : 75) split between the proxy model (a small transformer) training set and EMixMax set, and E MiMixax with Data Reuse. This is shown in Figure 11 (in Appendix DJ, where we see that for low samples Mixixa with Data Reuse performs better than alternative E MixMax approaches, but as we increase the samples all methods are comparable to EMixMax.
Now just comparing E MixMax with Data Reuse to E MixMax with a (75 : 25) split, but varying the similarity in the set of distribution (by changing the Dirichlet magnitude) with fixed 800 samples per length, we found in Figure (in Appendix DD they perform comparably. We conclude one can effectively resuse training samples to run MixMax and be more sample efficient.
6.2 MIMXX PERFORMS BETTER THAN GROUP DRO BASELINES FOR SEQUENCE MODELING
Here we investigated how MiMixMax compared other mixture finding methods (balanced and DoReMi (Xie et al.) 2024) which follows the near sample optimal algorithm of Zhang et al. (2024) for convex group DRO) and the original gradient descent and ascent group DRO algorithm (Sagawal ct al. 2019) across sets of distributions with varying similarity. To do so we considered the same sequence modeling task as Section 6.1 with Markov chain transition probabilities samples from symmetric Dirichlet distributions with magnitudes 1.0, 3.0, 5.0, 7.0, 10.0 to represent increasing similarity between the Markov chains. For all methods we took a training set of 800 samples per length and a test set of 200 samples per length from each Markov chain. We applied MiMMXaax given a small transformer trained for next token prediction on 600 of the 800 training samples per length (leaving the other 200 training samples per length to run EMixMax). We further ran EMixMax with the true probabilities on the 200 held-our samples as a reference for best empirical performance. We applied DoReMi and group DRO using the same small transformer architecture and all 800 training
8

---

Published as a conference paper at ICLR 2025
7A
76
74-
74-
12-
722
1.-
IMieMax
Grapg
7.-
EMMishar with Unzakemme Presice
63-
LMO Tnzaforec
68
DoRthli with Trandrmen Utlanned
E'MiMta Trunforsco Erocablic
group 1&k razed Traudomnee
66-
56-
Disihict Magniizds
Disihist Magniuut
(a) Comparing mixtures found by different methods.
(b) Comparing group DRO Models
Figure 4: MixMax found better mixture weights on sequence modeling task, and its ensemble model performed better than the group DRO trained model. The improvement stronger when the distributions are less similar. We present the mean and 0.5 standard deviation of the worst group (i.c., Markov chain) cross- cross-cntropy the Bayes optimal function for the different methods' mixture weights in Figure 4a Figure 4b further compares the performance between using E Mixixx weights to ensemble its proxy functions and the model given by training with group DRO.
samples per length (per Markov chain). Hyperparameters are described in Appendix C and we reported the results for each method with the hyperparameter settings that had the lowest maximum group cross-entropy test set loss over 15 trials of generating sets of Markov chains and samples.
We found Maxxxaax returned better mixtures to to than other methods, being close to the optimal performance (EMixMax), and that ensembling its proxy transformers with its mixture weights performed better than training with group DRO and was still close to optimal. Specifically, in Figure we compared the worst group cross-entropy loss of the Bayes optimal function for the mixture given by M'MixMax. EMixMax using ground truths to represent optimal performance, DoRcMi, the average group DRO mixture weights, and also balanced weights. As seen, fitting to the mixture given by MaM always performed best, doing even better as the distributions became less similar (lower magnitude). We also found the mixture weights returned by group DRO performed comparably to balanced, consistent with previous findings on group DRO performance Idrissi et al.) (2022). However, we note that group DRO's intended use is to apply the model it trained and not its mixture weights. In Figure 4b we observed using Max mixture weights to ensemble its proxy models performed better than the group DRO trained model; note the methods have comparable training compute as the 3 models in the ensemble trained on 1/3 of the training set each. These results were consistent even if we used fewer training samples (Figures 7 7ajin Appendix [D} highlighting that we had sufficient samples for both experiments.
4a
6.3 E'MIXMAX WITH DATA REUSE BEATS BALANCING DATA FOR XGBOOST
We compared E MixMax with Data Reuse weights to data balancing for non-pane con-parametric learning algorithms, in particular XGBoost Chen and Guestrin (2016) which is known to be the state of the art for tabular data. We selected ACSIncome Ding et al. (2021) (released under the MIT license) and CelebA annotations Liu et al. (2015) (released for non-commercial use g test on. For ACSIncome we constructed the dataset from the first 10 American states in alphabetic order, and considered group shifts from race and sex. We further constructed variations of the dataset using all the features, the first 2 features, and the first feature to introduce varying covariate shifts. For CelebA annotations, used attractiveness as the label with Young and Pale Skin as the group shifts, and constructed variations of the dataset using all features, the first 10 features, and the first 5 features 4 We used random 80%- 20% train-test splits in all settings. We applied MaMixMax with Data Reuse by using XBoos models (trained on the group data) as proxies for label probabilities, and modeled covariate probabilities using Gaussian kernel density estimation. We then returned an XGBoost
'The agreement for use is at hhttpss:: ///m// mm/ ab. ie. -uhk edu. edw.eed..kkkkk htt.:hk//ppoo..co 'The different number of features compared to ACSIncome was because the features were binary while Ancome first features take values.
eebA.

---

Published as a conference paper at ICLR 2025
Quadle (ywa)
Pazoa Conclation p-sako 0002
E-MisMa with Dut Resse Balancod
Toppooocaat Oves Bakaseed
Ben Lize
20-
70-
1.0
60-
0.5
15
0.0
-05
Greup Accuray
(a) Comparing non-parametric group DRO methods
(b) Correlation with room for improvement
Figure 5: Mixxxxxx with Data Reuse improved worst group accuracy over balancing data more when there was bigger room for improvement. In Figure 5a we present the mean and standard deviation (over 5 trials) of the worst group accuracy of E Mix Max with Data Reuse and balanced data as a function of the oracle accuracy for that setting. In Figure 5b we plot E'MixMax with Data Reuse's 's improvement in worst group accuracy over balanced data function of the Oracle worst group accuracy (i.e., having a model trained for each distribution) improvement over balanced data; observed a Pearson correlation of 0.792 with p-value 0.002.
model trained on the same training data but with the loss reweighed according weights returned by M'MixMax with Data Reuse. For the balanced baseline, we returned an XGBoost model trained with the data balanced by by normalizing the loss for each group. Lastly, as measure for best possible performance, we also presented the worst oracle accuracy, where oracle accuracy refers to the test accuracy an XGBoost model trained on the same distribution (where we need access to the group identity at test time use this model). Hyperparameters are described in Appendix C and we reported the results for the hyperparameter setting with the best performance over trials with random train-test splits.
As seen in Figure 5a Table in Appendix D. M-MixMax with Data Reuse matched or outperformed the worst group accuracy of data balancing in all settings, and similarly for worst group loss (Table 3 in Appendix D). In particular, as seen Figure 5b we observed the improvement over the worsf group accuracy of data balancing was stronger when more room for improvement (a larger gap between oracle and balanced worst group accuracy): we observed a 0.79 Pearson correlation factor. In the end, we found Maxixx with Data Reuse improved worst group accuracy on ACSIncome with groups by sex and one feature by 1.6% (2.8% relative gain), and CelebA with groups by Young and 10 and 5 features by 2.3% and 3.2% (3. 7% and 5.9% relative gains).
possible
7 CONCLUSION
In this paper we showed that group DRO over bounded functions can be solved by fitting to optimal data mixture, and that maximizing particular concave objective returns the optimal mixture weights for cross-entropy and E loss. We called this method for finding data mixtures MixMax. Our experi- ment on a simple sequence modeling task showed that even empirical versions Mix Max improved previous parametric group DRO baselines. An empirical version of MixMax also shown to improve over the baseline of balancing data for non- parameetettiic learning algorithms, specifically XGBoost, for which no previous group DRO methods were proposed. We leave open the problem of applying our minimax theorem to other losses, and proposing better empirical versions of MixMax.
Limitations The empirical versions of MixMax can immediately scale to provide group DRO solutions to large generative modeling over high-dimensional spaces (e.g.. decoder-only LLMs) where there is no need to model covariate shifts as there are no inputs. However, main technical limitation of MixMax methods is the need to model covariate shifts covariate shifts exist, which often requires large amounts of data in nigh-dimensional covariate spaces. We also acknowledge that DRO may be used to claim model is fair, despite it still carrying societal biases. We hope future uses will take care in considering the claims DRO and cannot make.
10

---

Published as a conference paper at ICLR 2025
ACKNOW
Resources used in preparing this research were provided in part by the Province of Ontario, the Government of Canada through CIFAR, and companies sponsoring the Vector Institute. We acknowl- edge the support of the Natural Sciences and Engineering Research Council of Canada (NSERC), RGPIN-202) 10TIN-2021-34444 100445. Anvith Thudi is supported by a Vanier Fellowship from NSERC. We thank Relu Patrascu for administrating and procuring the compute infrastructure used for the experiments in this paper. We would also like to thank Ayoub El Hanchi, Leo Cotta, Nishkrit Desai, Stephan Rabanser, Sierra Wyllie, Alon Albalak, Nicolas Papernot, Colin Raffel, and many others at the Vector Institute for discussions contributing to this paper.
REFERENCES
A. Albalak, Y. Elazar, S. M. Xie, S. Longpre, N. Lambert, X. Wang, N. Muennighoff, B. Hou, L. Pan, H. Jeong, al. A survey on data selection for language models. arXiv preprint arX22:2402.16827, 2024.
M. Arjovsky, L. Bottou, I. Guirajani, and D. Lopez-Paz. Invariant risk minimization. arXiv preprint arXiv:1907.02893, 2019.
M. Bansal, K.-L. Huang, and S. Mchrotra. Decomposition algorithms for two-stage distributionally robust mixed binary programs. SIAM Journal on Optimization, 28(3):2360-22 2018.
A. Ben- Ben-Tal, D. Den Hertog, A. De Waegenaere, B. Melenberg, and G. Rennen. Robust solutions of optimization problems affected by uncertain probabilities. Management Science, 59(2):341-357, 2013.
D. Bertsimas, X. V. Doan, K. Natarajan, and C.-P. Teo. Models for minimax stochastic linear optimization problems with risk aversion. Mathematics of Operations Research, 35(3):580-605 2010.
X. Bi, Chen, Chen, Chen, D. Dai, C. Deng, Ding, K. Dong, Du, Z. Fu, et al. Deepseek I1m: Scaling open-source language models with longtermism. arXiv preprint arXiv:2401.02954, 2024.
S. Black, L. Gao, P. Wang, C. Leahy, and S. Biderman. GPT-Neo: Large Scale Autoregressive Language Modeling with Mesh-Tensorflow, Mar. 2021. URL / /ttp:////../////0/////8 zenodo. 5297715 If you use this software, please cite it using these metadata.
J. Blackard. Covertype. htpps///dii.rr/11.24432//////
UCI Machine
Learning
Repository, 1998.
DOI:
R. Chen and I. Paschalidis. Selecting optimal decisions via distributionally robust nearest-neighhor regression. Advances in Neural Information Processing Systems, 32, 2019.
T. Chen and C. Guestrin. Xgboost: A scalable tree boosting system. In Proceedings of the 22nd aem sigkdd international conference on knowledge discovery and data mining, pages 785-794, 2016.
E. Delage and Ye. Distributionally robust optimization under moment uncertainty with application to data-driven problems. Operations research, 58(3):595-612, 2010.
Ding, M. Hardt, J. Miller, and L. Schmidt. Retiring adult: New datasets for fair machine learning. Advances in neural information processing systems, 34:6478-6490, 2021.
J. C. Duchi. Introductory lectures stochastic optimization. The mathematics of data, 25:99-186. 2018.
J. Dupačová. The minimax approach to stochastic programming and an illustrative application. Stochasties: International Journal of Probahility and Stochastic Processes, 20(1):73-88, 1987.
S. Fan, M. Pagliardini, and M. Jaggi. Doge: Domain reweighting with generalization estimation. arXiv preprint 200i:22100.53999 2023.

---

Published as a conference paper at ICLR 2025
G. B. Folland. Real analysis: modern techniques and their applications, volume 40, John Wiley & Sons, 1999.
J. H. Friedman. Greedy function approximation: gradient boosting machine. Annals of statistics, pages 119-1132, 2001.
B. Idrissi, M. Arjovsky, M. Pezeshki, and D. Lopez-Paz. Simple data balancing achieves competi- tive worst-group-acoon PMLR, 2022.
In Conference on Causal Learning and Reasoning, pages 336-351.
A. Jain, A. Montanari, and E. Sasoglu. Scaling laws for learning with real and surrogate data. arXiv preprint 43X1v:2402.04376, 2024.
G. Kolossov, A. Montanari, and P. Tandon. Towards statistical theory of data selection under weak supervision. arXiv preprint WVXv:2339.14563, 2023.
Z. Liu, P. Luo, X. Wang, and X. Tang. Deep learning face attributes in the wild. In Proceedings of International Conference on Computer Vision (ICCV), December 2015.
A. Nemirovski, A. Juditsky, G. Lan, and A. Shapiro. Robust stochastic approximation approach to stochastic programming, SIAM Journal on optimization, 19(4):1574-160) 2009.
Y. Oren, S. Sagawa, B. Hashimoto, and P. Liang. Distributionally robust language modeling. arXiv preprint #rXiv:1909.020m0 2019.
G. Pflug and Wozabal. Ambiguity in portfolio selection. Quantitative Finance, 7(4):435 442, 2007.
H. Rahimian and S. Mchrotra. Distributionally robust optimization: A review. arXiv preprint arXiv:1908.0550 05659, 2019.
H. Rahimian, G. Bayraksan, and T. Homem-de Mello. Identifying effective scenarios in distribution- ally robust stochastic programs with total variation distance. Mathematical Programming, 173(1): 393-430, 2019.
Sagawa, P. W. Koh, T. B. Hashimoto, and Liang. Distributionally robust neural networks for group shifts: On the importance of regularization for worst-case generalization. arXiv preprint arXiv:1911.08731, 2019.
D. W. Scott. On optimal and
histograms. Biometrika, 66(3):605-610) 1979.
A. Shapiro and A. Kleywegt. Minimax analysis of stochastic problems. Optimization Methods and Software, 17(3):223-54,, 2002.
M. Sion. On general minimax theorems. 1958.
A. Slowik and L. Bottou. Algorithmic bias and data bias: Understanding the relation between distributionally robust optimization and data curation. arXiv preprint arXiv:2106.0944 2021.
A. Slowik and L. Bottou. On distributionally robust optimization and data rebalancing. In Interna- tional Conference on Artificial Intelligence and Statistics, pages 1283-1297. PMLR, 2022.
B. Sorscher, R. Geirhos, S. Shekhar, S. Ganguli, and A. Morcos. Beyond neural scaling laws: beating power law scaling via data pruning. Advances in Neural Information Processing Systems, 35: 19523-19536, 2022.
B. Strack. J. P. DeShazo, C. Gennings, J. L. Olmo, S. Ventura, K. J. Cios, J. N. Clore, et al. Impact of hbale measurement on hospital readmission rates: analysis of 70,000 clinical database patient records. BioMed research international, 2014, 2014.
W. Wiesemann, D. Kuhn, and B. Rustem. Robust markov decision processes. Mathematics of Operations Research, 3831)1153-183, 2013.
12

---

Published as a conference paper at ICLR 2025
S. M. Xie, H. Pham, X. Dong, N. Du, H. Liu, Y. Lu, P. S. Liang, Q. V. Le, T. Ma, and A. W. Yu. Doremi: Optimizing data mixtures speeds up language model pretraining. Advances in Neural Information Processing Systems, 36, 2024.
D. Yu, Y. Cai, W. Jiang, and Zhang. Efficient algorithms for empirical group distributional robust optimization and beyond. arXiv preprini urXiv:2403. 035622, 2024.
L. Zhang, P. Zhao, Z.-H. Zhuang, T. Yang, and Z.-H. Zhou. Stochastic approximation approaches to group distributionally robust optimization. Advances Neural Information Processing Systems, 36, 2024.
APPENDIX
A PROOFS
A.1 PROOF OF THEOREM 3.1
Proof. First note P sup dppoocu)))) as the objective is linear in dp. More formally we have
/
/ (lf(x).))dp(c 8)
((f(t)
sup dpe
sup ApeConu(P)
by P C Conv(P) (). Furthermore we have:
/ ((f(e). dpECone(P)
/ ((f())))) ApEP
sup
sup xEA
Σ
Apdp(x.
E X (L((())
sup A dpEP
/
I sup dpEP.
where we used the parameterization of the convex hull for the first equality, then linear- ity of the integral for the second, and then the fact the convex combination of a finite set of scalars is less than the maximum scalar in that set. This allows us to conclude sup4per I' [(f(s), u)dp(r,)) = SundnECono(P)) L(f(x), v)dp(z,)) v)- Hence we have, a minimizer exists for supaper J L(f(x), y)dp(z,)) y) then it must also be a minimizer of sUPdpeComu(P) I ((f(x), v)dp(zz y) (as they are equal for all f). Hence we now characterize the minimizer of suPdpECono(P) f [(f(x), y)dp(r, y).
Note we can apply Sion's 's minimax theorem Sion (9588 to interchange the supremum and infimum By (r) is compact in topology that the objective is continuous in. Note (x,,( (r) is compact in the Weak -topoooy theorem and continuity of dilation Folland (1999) Now note (x is closed subset of B, RR) (r) and hence is also compact. The fact (x,,(xxo) (r) closed follows from O being closed and that weak* convergence implies pointwise convergence (a.e.) (L (r) (a set containing (x=((xo) which is proven below.
(X,, R (as the dual of T.2 TR"")) by Manach-Alaogluis
Note convergence in weak* for functions in (L=(XxR="((7) means the functions converge pointwise a.e. This is as, for g L and the limit of fn in the w fa)gdr = f lim (f fn gg) by dominated convergence functions are bounded. Now taking g sign (f fx)]xXE for some finite measure set E, note have g L' as sign simm (f((x)))))) fn (x)) S Xx < oc. Applying this g to our previous limit expression have:
lim, I'f
13

---

Published as a conference paper at ICLR 2025
/
/
0=
lim (f - fulgdz = n=200
lim (f fn)sign( lim (f fu))xedx n-v00
n0000
= =f
lim (f - (f-fn)lxes fn)lxEde 1-+00
By proposition 2.16 in Folland 1999) we have that -im -o is countably coverable by finite measured sets (by the o-finite assumption), i.e., &' = U E, where Vi J Xx, dæ < oo, this implies the functions converge pointwise a.e in X.
a.e on E. As the space
Now C is bounded, and the functions converge pointwise in the weak -topolooo we also have by the dominated convergence theorem objective continuous the weak* topology on - More formally, by continuity of C in its O argument, we have if fn() + f(x) then lim LLf (1), y) [(f(x),)) y). From this and the previously shown fact that if fn weak"
then fz) + (2) we have for any convergent sequence in the weak ttoo-ooo,,ogy, fn. wwekk f. that
lim / ((UA(E))
/ lim ()dp()) -
f [C(f(E)
where we pulled the limit through the integral by the boundedness of & (which allows dominated convergence). Note this is the definition of the objective being continuous in topology, so conclude f [(((x), y)dp(x,)) is continuous over the weak*- -tooology over f.
Thus with the weak* -taoooty on - satisfy the compactness and continuity conditions for Sion's minimax theorem, which implies
/ [((())
/ ((F()). v)
sup inf dpECone(P) (x (r() So if f achieves the inf in the LHS and dp, achieves the sup in the RHS, then note
inf sup feBr=(x.o) dpEConu(P)
/ (L)
/
sup dpECone(P)
/ (lf(x).))dpx(el)
inf (e
to conclude that f (x)) u)dx((z,)) in B C(J(x), v)dpx(r, v). So in particu- lar is minimizer of dpx. Note minimizers of dpx exist by the compactness of (x (r) and continuity of the objective over this function space.
Remark on Existence of DRO Solution The theorem required that there exists a DRO solution in the first place, and this is satisfied if P is finite by the continuity of the supremum over finite set of continuous functions (note continuity of the objective for single dp was proven above). Further care needed when P is not finite, but the paper focuses on the finite case and so we do not discuss this further.
A.2 CONCAVITY OF OBJECTIVES
Fact A.1. If L is cross-entropy, then [(f2(=) 9) is concave in
Proof. We note fa - Px >n (2())(()(z)) where Ap(x) Edv'r Ep [
App(z)
App(z)
px(2)
Note that that- -y log(E (p(=))(pl))) (p)))(())))) Edpe Ap(x)) ly -(0g(p((l/z)))p(ll) by the concavity of entropy. Applying this we have
14

---

Published as a conference paper at ICLR 2025
J 1 -
Σ >x(xx) / -(og(((sl=)))(ula)))) pa(x)
x
dpEp
/ x
Σ App(x) dpEP ps(x)
/ -log(p(elz)) px(z)
Σ dpEP
Xp
log log(((gl=))pp(/a))) (2)
Note this inequality holds for an arbitrary set of distributions P. To now conclude con- cavity the mixture weights, let pi = PA we have I aaaa (z), (z) 2 a C(Ax, (x), y)dpx, (x. + (1 a)
P>2- Then by the above inequality
v). This proves concavity w.r. the mixture weights X, as desired.
Fact A.2. 1f C. is Ey. then [(Ix(x)) v) is concave in À
Proof. We assume fA is the bayes optimal solution, hence we have the loss is the bayes error App(ar)
Sx Var (U Note px (v) = Edper {p(x))(ular) where >p(x)) Edd'c A
Asp(z) and variance is concave w.r.t mixture weights. Thus we have (x(
f ((()())
= -/x Var Varpxtoo)(Vll)p(()
Σ \ pa(x)
dpEr
Xpp(x)
arr (())))))()))) px(x)
px(x)
dpEP
Σ Ap / ((rp)ol))(ll))p()) dpEP
Σ Σ (((s()))))((..)) (3) dpEP
Note this inequality holds for an arbitrary set of distributions P. To now conclude con- cavity to the mixture weights, let Pı PA1. P2 (I,y) 2 Pxa- Then by the above inequality a [(fx, (x). v)dx, (x. v) + (1 a) (U))(()) (z) This proves concavity to mixture weights.
we have [ (z),
B EMIXMAX GRADIENT DISCUSSION
We show that the expectation of EEMixMax (A) satisfies EMixMax is Unbiased:
ED EMixMax (A)
Σ
VxApC(A(Z))
y)dpx))
dpEP
which is the gradient of the MixMax objective (Section
4.1
15

---

Published as a conference paper at ICLR 2025
Fact B.1. V FEMixMax(X) (X) is an unbiased gradient estimator for the MixMax objective (Equation
assuming VAC(fx(2), y) is bounded and L(fx(x). y) is differentiable a.e (W.rt X).
Proof. Note if V ((f(((z)) y) is bounded, then by dominated convergence we have the derivative pulls into the integral for the MixMax objective (by the limit definition of derivatives, and the limit pulling in by dominated convergence)
VA VA E >x (L dpEP
Σ / [AAAA((AAA))) dpEP
Now note the expectation of EMixMax over sampling datasets of size D i.dd from dp E P is:
Σ
Σ
Σ
ApSo(t)ii
Ppll
dpEP
(e,n))Dy
dpEp
Σ
ED
VxApC
Σ
(pFe(z))
1Dpl
dpEP
(
dpEp
1D
Σ / VxAPC
Σ
XpFo(z)i1
dp(x.))
Dpl
dpEP P t=1
dpEp
VxdpC
Σ Apfo(o). dp(x,y)
dpEP
dpEP
Note by the previous equality this is the gradient of the MixMax objective (Equation D
The EMixMax gradient computation in Algorithm currently scales at ((p (D))). This be prohibitive when the datasets are large. One reduce the cost to O(IP))) by doing mini-batch sampling at every step. particular, for every optimization step and y E P one samples Bp Dr uniformly where |Bp = B. One then replaces line 3 in Algorithm) ]
Complexity of EMixMax
with
1
Σ dpEP
Σ
((Aa()).).
B
(T.y)E
C EXPERIMENTAL SETUPS
C.1 TOY EXPERIMENTS
For the binary classification task, We ran EMixMax with D; consisting of 10000 samples from each distribution, step size 0.5, and for 100 steps.
We ran EMixMax with D; again consisting of 10000 samples from each distribution, and = 0.0001 for 100 steps.
C.2 SEQUENCE MODELING
For the MixMax methods, letting y denote sequence, note the task was to model p(y) and so had no covariates and hence had no covariate shift (as one can take X' to be a singleton). Thus MixMax only needed functions for the token probabilities from each distribution, and considered both the true probabilities (as the optimal baseline) and transformers trained on each Markov chain as the proxy optimal functions. The transformer used is GPTNeo Black et al. (202)) with 6 hidden states, 2 hidden
16

---

Published as a conference paper at ICLR 2025
layers, 2 attention heads, 8 intermediate size, and with 12 max position embeddings. In both cases we used 200 samples per length to run EMixMax, keeping the other 600 to train the proxy model for Mixixax.. The proxy model is trained for 20 epochs using Adam W with learning rates 0.01, 0.001 and 0.0001 (and otherwise default Pytorch hyperparameters).
We implemented DoReMi for finding mixture weights using the same transformer architecture as above, with the reference model being trained on a balanced dataset. The DoReMi reference and proxy models were again trained for 20 epochs using Adam V with learning rate 0.01, 0.001 and 0.0001 (and otherwise default Pytorch hyperparameters) The learning rate for the mixture weights was 0.1. We also implemented group DRO with the same architecture, with the same model weights optimizer and mixture weights learning rate 0.1. Furthermore, we tested varying minibatch sizes for the training of all models (50, 100, 200), and number of steps used for group DRO (150, 300, 600). In Figure 4 we reported the results for the hyperparameter settings with the lowest mean error for each method over 15 trials of randomly generating markov chains according to the varying magnitudes.
C.3 TABULAR DATA
For running Gaussian kernel density estimation, used the Scott method for bandwidth selec- tion Scott 1979). For fitting the model on M2MixMax with data reuse mixture weights, the reference models used for Mixixx with data reuse on the Income dataset, and the baseline of balancing the dataset, we implemented XGBoost with depths 6, 6,8, 10, number of trees 100, 200, 300, and learning rates 0.01, 0.1 and reported the results for the hyperparameter setting with the best average accuracy over 5 trials with random train-test splits. For CelebA, the oracle accuracies and the reference models used Max with data reuse were always trained using depth 8, number of trees 200, and learning rate 0.1; we found hyperparameter sweeping the models on individual groups had marginal impact, and anyways would only improve our results (the baseline of balaneing always has hyperparameter sweep). For training the MixMax mixture balanced mixture, the loss for XGBoost using the sample weight parameter: for MixMax it is set to D for every datapoint in group p. and for balanced it is set to 1/ After normalization this gives the desired weighting of each group.
reweighed
For the ensembled model results, we took the hyperparameters used for the retraining experiments and ran them for a separate 5 random train-test splits to evaluate the performance of ensembling the proxy model.
D ADDITIONAL TABLES AND FIGURES
6.7
0.8
08
07
MisMaa
0.6
Mishia
0.5
+04
0.4
0.3-
02-
6.1
0.3
(a) Extreme Functions Example
(b) Extreme Functions Example 2
Figure 6: EMixMax for deterministic regression finds the average of the extreme functions. In general MixMax maximizes expected variance and here we plot the expected y of the EMixMax mixture.
17

---

Published as a conference paper at ICLR 2025
Datasets
Features
Oracle (%)
Balanced (%)
E²MixMax (%)
All
79.1
80.0(#1.2)
80.3(+0.6)
Inc-Race
2
71.2(+00.1
70.7(+0.2)
70.6(+0.3)
60.0(+0.1)
53.7(+0.4)
53.7(+0.5)
All
81.2(+0.1)
81.1(+0.2)
81.1(+0.1)
Inc-Sex
2
72.8(+0.1)
71.4(+0.2)
71.6(+0.1)
62.8(+0.2)
56.5(+0.2)
58.1(#1.8)
All
77.0(=0.1)
75.6(+0.2)
76.0(+1.0)
CelebA-Young
10
68.3(
62.5(+0.3)
64.8(
5
63.4(+0..)
53.8(+0.2)
57.0(
All
78.2(+0.3)
77.((+0.)
77.1(+0.6)
CelebA-Pale Skin
10
67.6 67.6(+0.1)
66.8(+0.2)
67.1(+0.3)
62.4(+0.2)
59.5(+0..)
59.4(+0.3)
Table MaxixMa with data reuse improved worst group accuracy over the baseline of balancing data for tabular dataset using XGBoost, in particular when more significant label shifts were present (see Figure for a visualization). We present the average minimum accuracy over the groups from 5 trials with random training-test splits for E with data reuse, balancing by up-sampling, and the "oracle" accuracy where model was trained for each group and evaluated on the same group (as example of optimal performance). Results are bolded if the better method's average performance was outside one standard deviation of the other method. We modified the datasets to study a variety of shifts by only including the first N features.
Datasets
Features
Oracle (%)
Retrained MaxixMax (%)
Ensemble MaMixMax (%)
All 2
79.1(#1.2)
80.3(+0.6)
76.2(++0.4
Inc-Race
71.2(#0.1)
70.6(
70.4(+0.3)
60.0(10.1)
53.7(10.55)
53.3(+0.33
All
81.2(#0.1)
81.1
76.6(+0.2)
Inc-Sex
2
72.8(+0.1)
71.6(+0..)
69.8(+1.2)
1
62.8(+0.2)
58.1
55.9(+0.11
All
77.0(20.1)
76.0(+1.0))
76.0
CelebA- Young
68.3(+0.3)
64.8(+1.6)
62.6(+2.7)
5
63.4(+0.3)
57.0(+0.6)
+1.6
All
78.2(+0.3)
77.1
75.9(+0.6)
CelebA-Pale Skin
10
67.6(+0.1)
67.1(+0.3)
67.2 +0.2)
5
62.2
59.4(+1..)
58.0(
Table 2:
Using E MixMax with data reuse mixture weights train a new model performs better than ensembling the proxy models (following line 2 in Algorithm] D. We present the settings from T and similarly report the average minimum accuracy over the groups from 5 trials with random training-test splits for both methods of using MixMax mixture weights. Results are bolded if the better method's average performance was outside one standard deviation of the other method.
18

---

Published as a conference paper at ICLR 2025
Datasets
Features
Balanced
M'xixa
All
0.426(=0.013)
0.422(=0.007
Inc-Race
2
0.557(=0.001)
0.559 +0.003)
0.707(=00004)
0.688 1688(=00020)
All
0.403(=00004)
0.405(=000055
Inc-Sex
0.544(=0.002)
0.541 0.410600006
0.616(=00001)
0.611 A.611(10.005)
All
0.495(=0.002)
CelebA-Young
0.644(=00.00)
0.614 (00019
5
0.706(=0.002)
0.670 +0.007)
All
0.465(=0.001)
0.468 =0.006
CelebA-Pale Skin
10
0.594(=0.001)
0.588 #0.005
5
0.647(=0.001)
0.641 (=0.005)
Table Max with data reuse improved worst group loss over the baseline balancing data for tabular dataset using XGBoost. We present the average maximum loss over the groups from 5 trials with random training-test splits for MixMxx with data reuse and balancing by up-sampling. Results bolded if the better method's average performance was outside one standard deviation of the other method. We modified the datasets to study a variety shifts by only including the first N features.
Dataset
Features
Untuned MixMax
Tuned E2MixMax
Inc-Race Inc-Sex
1
53.7(=0.5%
54.3(+0.6%
1
58.1(=1.8%)
58.4(+1.8%)
10
64.8((+1.6%))
66.0((+2.1%))
CelebA-Young
5
57.0%
37.0%(+0.6%)
CelebA-Pale Skin
5
59.4%((+0.3%))
59.4%(+0.33%)
Table 4: MaMixMax with data reuse can sometimes be improved by doing hyperparameter search number of optimization steps (20, 40, 80). We present the settings from Table where the gap between oracle accuracy and the balanced accuracy was more than 2%, and include the average minimum accuracy over groups hyperparameter tuning. We bolded results where the tuned performance was outside one standard deviation of untuned.
trials of Ma with data reuse with further
Word Coopp Tax CavesTainny
Wore Groug Tax Caveo Eatmopy
2.4
7.2
7.4
EMieMat
3.0
MaMixMax Temformer Proxies
group DRO
EMix
6.8
DoRcMi with Transoomees
ElMidaa with Transfmas Ensemble
group
10
Dinichlet Magnitude
Dirichee Magnitude
(a) Optimally Fitting Mixture
(b) Empirically Fitting Mixture
Figure A reproduction of the experimental setup for Figure 7 but with fewer training samples (500 instead of 800). We still see that model (empirically by ensembling transformers optimally) fitted to Mixixa mixture weights performs better than DoReMi and group DRO mixture weights, and the model trained by group DRO. Furthermore, still see that the improvement is stronger when the distributions are less similar.
19

---

Published as a conference paper at ICLR 2025
Income Ract
Income Sex
Crichá-Yori
Aceinsy Gan (5)
0.5
Crich4-Yoing
Polo
0.0
-0.5
20
30
40
50
60
80
Steps
We see that Max with data reuse can benefit from early stopping for several datasets, as seen by the peaks in performance, however these trends are statistically weak given the standard deviations reported in Table 4 Here plot the worst group accuracy gain over running Mixixxax with data reuse for 20 steps when running MaMixMax with data reuse for 20, 40, 60, 80 steps. Further hyperparameter details are described in Appendix C This is done for the datasets observed to have big label shifts, described in Table
Figure 8:
7.5
7.5-
2.4
7.4
73
7.3
22-
7.
Woett Group
2.1
Woest Group
7.1
7.0
6.
MixcMax with Transformers
MMiiMaax with Transformes
MixMa with Dam Reuse and Transformets
MixMax with Duta Reuse und Transfommess
688
10
10
Diniclle Magninde
Drichle Magninde
(a) 800 Training Samples
(b) 500 Training Samples
Figure 9: MaxixMax with data reuse is much worse than Mixixa and can be slightly better when one has fewer training samples. Here we consider the experimental setup from Figure for fitting the various MixMax mixture weights optimally, and present results for when the same or 500 training points are used to fit the proxy models and run EMixMax giving the data reuse version (instead of the 600 200 and 300 200 split for proxy training and MixMax).
20

---

Published as a conference paper at ICLR 2025
E-MisMaa with Dana Resse
Wont Greany Toc Ooo--iikegy
E MisiMaa 50.50
722
Maa
2.0
200
300
400
600
700
800
Triining Set Size
Figure 10: Various empirical MixMax approaches perform comparably for worst group cross-entropy. with MaxixMax with data reuse performing closest to EMixMax
(with ground truths) in the small training set regime. Here we plot the mean and 95% confidence interval (over 45 trials of sampling new sets of Markov-chains from the symmetric Dirichlet distribution with magnitude 1) of the worst group cross-entropp of the Bayes optimal function defined by the methods's 's mixture weights, changing the number of training samples per length.
MixMon Objuative
6.90
6.85
8.50
EMixMaa
6.75
ErMisMaa with
MiMixMaa 25/35
6.70
L'MisMax 30:50
EiMicMax
6,55
Triining Sei Sizee
Figure 11: Various empirical MixMax approaches perform comparably in maximizing the MixMax objective, with M²MixMax with data reuse performing closest to EMixMax (which uses ground truths) in the small training set regime. Here we plot the mean and 95% confidence interval (over 45 trials of sampling new sets of Markov-chain from the symmetric Dirichlet distribution with magnitude 1) of the MixMax objective value for the methods's 's mixture weights, changing the number of training samples per length.
21


Reviews:
[{"review_id": "oNJS7CT4wu", "reviewer": "Reviewer_13Yx", "summary": "This paper provides a reparameterization of the group DRO objective from the parameter space to function space under the assumption that the hypothesis class for the predictor contains the set of the bayes optimal classifiers. Under this reparameterization, they introduce an algorithm to optimize this objective. Empirical results show the proposed algorithm outperforms groupDRO.", "questions": "1)Under the assumption that the bayes optimal predictor is a subset of the set of hypothesis, wouldn't the same f minimize the objective for any dp in the set P? If this is the case, this weakens the theoretical contribution.", "limitations": "", "rating": 8, "confidence": 2, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "1)The problem is important and relevant to OOD generalization.\n\n2)The proposed method is novel in its approach and theoretically justified.\n\n3)Experimental results show that the proposed method outperforms groupDRO.\n\n4)The paper is well written.", "weaknesses": "See questions.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "oJHzUt8HtQ", "reviewer": "Reviewer_8SGa", "summary": "The paper formalizes the DRO problem in the case of bounded functions and shows that a solution is an optimal data mixture. In case of cross-entropy and L2 losses the problem can be further simplified to a single maximization of a concave function. An empirical version of the problem is further presented even in the case where the Bayes optimal functions are unknown. Several experiments on both synthetic and real data demonstrate that the empirical version of the MixMax improves over previous group DRO baselines in parametric cases as well as over non parametric data balancing algorithms.", "questions": "- For the experiment described in 6.3, which strategy was used to upsample the minority class ? Was the training loss reweighed in order to de-bias the loss estimation ?\n\n- Two different ways of using the MixMax weights were described, either retraining a model on the weighted dataset or reweighing models at inference time. Do you have a sense of how those 2 approaches compare in practice under different settings ?", "limitations": "", "rating": 8, "confidence": 4, "soundness": 4, "presentation": 3, "contribution": 3, "strengths": "- The major theoretical result that under certain conditions the DRO problem is equivalent to a single maximization over a concave objective function and that a solution can be achieved by fitting an optimal data mixture.\n\n- The authors also provide a practical algorithm that computed the empirical MixMax solution.", "weaknesses": "- $Emperical^2$ $MixMax$ relies on an accurate estimation of $\\hat{f}_p(x)$, which in practice requires to optimally fit a model on every distribution, including performing a HP search by cross validation. This might be computationally prohibitive in some practical cases. Could you discuss computational trade-offs or potential approximations that could be used in resource-constrained settings? How those approximations would impact the  $Emperical^2$ $MixMax$ estimate?\n\n- In some of the experiments where a Mixmax method was compared to an alternative method, the MixMax requires more compute which makes the comparison less fair. (e.g. experiment 6,3). It would be helpful to have experiments that control for computational budget across methods.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "iMQeJbWSwk", "reviewer": "Reviewer_DjaM", "summary": "This paper aims to provide a solution to the problem of group distributionally robust optimization (DRO). It considers the class of all bounded functions to address this issue. The group DRO problem is a specific case of the broader DRO problem, where the set of distributions is finite and consists of predefined distributions. \n\nThe paper first demonstrates that group DRO over the space of bounded functions satisfies a minimax theorem. Additionally, it claims that when the loss function is either cross-entropy or mean squared error, the problem can be transformed into a convex optimization problem, which subsequently leads to a classical risk minimization problem.", "questions": "1- I believe the VC dimension of the class of bounded functions is unbounded. Could this present a challenge for solving real-world problems?\n\n2- It is claimed that in Equation 1 the objective is concave with respect to $\\lambda$. However, the convexity of $f_\\lambda$ with respect to $\\lambda$ is not clear. Although it is true that, if the class is rich enough, it can approximate the error of $p_\\lambda(y|x)$, it is not exactly equal to this function. \n\n3-To prove that the function $f_\\lambda$ is concave with respect to $\\lambda$, it must be shown that $f_{\\alpha \\lambda_1 + (1 - \\alpha) \\lambda_2} \\geq \\alpha f_{\\lambda_1} + (1 - \\alpha) f_{\\lambda_2}$. However, in Equation 3, this is not demonstrated. Instead, concavity with respect to $p$ is proven in that equation.", "limitations": "", "rating": 6, "confidence": 2, "soundness": 2, "presentation": 2, "contribution": 2, "strengths": "The issue is both intriguing and significant in the literature.", "weaknesses": "The theoretical proofs are quite brief and unclear. It would be beneficial to discuss and clarify the proofs step by step.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "stJmNJg3l5", "reviewer": "Reviewer_MsQo", "summary": "The paper proposes a reparameterization technique for the group DRO problem to address the non-convexity of the parameters, under the most popular machine learning backgrounds. This reparameterization from parameter space to bounded function space admits a generalized minimax theorem. Moreover, the paper proves that group DRO can be decomposed into a convex optimization problem followed by a risk minimization problem for the frequently-used cross-entropy and mean squared error. Comprehensive experimental results align closely with the theoretical guarantee.", "questions": "I would appreciate it if the authors were willing to answer some of the questions I raised in the weaknesses above.", "limitations": "", "rating": 5, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "1. The idea of reparameterization seems novel in the context of group DRO.\n2. The experiments are comprehensive in terms of the choice of datasets and the diverse ways of presentation.\n3. The paper is written clearly, especially for the theoretical results.", "weaknesses": "There are multiple disadvantages, mainly regarding the methodology itself, together with its relevance to the existing literature on (empirical) group DRO.\n\n1. The main contribution of this paper (Theorem 3.1) seems a direct generalization of Sion's minimax theorem by transforming the topology in the function space. \n2. The key discovery is that group DRO can be solved by fitting a specific mixture distribution. However, this seems like a straightforward idea in the optimization community (Nemirovski et al., 2009, Section 3.2) and recent works on group DRO (Zhang et al., 2023; Yu et al., 2024) also use this property to reformulate the optimization problem. I understand that the second step in this paper, i.e., transforming the minimax object into a concave optimization problem of the distributional weights, is different from the ways in the papers I cited, but the first step (supported by the main theorem), has already been exploited in a similar way\n3. The reparameterization enables line 2 of Algorithm 1, which allows the possibility of iterative training and the whole MixMax schema. However, I am a little skeptical about parameterizing $f_\\lambda(\\cdot)$ in this form. It seems that MixMax exerts a constraint on the function class in the group DRO objective, which is not satisfactory. What happens if $f_\\lambda(\\cdot)$ can not be expressed via the closed form shown in the paper?\n4. As an optimization paper in ICLR, the authors should consider adding complexity analysis of the proposed algorithm, either the computation complexity of Empirical MixMax, or the sample complexity of Population MixMax. From my understanding, Algorithm 1 needs to compute the empirical risk among each data group per iteration. This seems too expensive for modern machine learning or large-scale optimization tasks. Although the experiments verify the effectiveness of MixMax over other group DRO approaches, theoretical justifications still need to be presented in an optimization paper. \n5. The original form of group DRO involves risk function in the form of expectation, which is intractable in reality. That's why stochastic approximation algorithms prevail. However, as it requires computing the risk function for each group in Algorithm 1, I was wondering if MixMax can be directly applied to group DRO (at population level). If not, then the main contribution of this paper should be the reparameterization of empirical group DRO, rather than group DRO.\n6. The baselines used in the experiments could be improved. The algorithm proposed by Sagawa et al. (2020) suffers from a suboptimal complexity, which may cause unfair comparisons in Figure 4. The authors could try adopting the near-optimal algorithms of Zhang et al. (2023, Algorithm 1 and 2), whose complexity is $|P|$ times lower than that of Sagawa et al. (2020). \n7. Both the algorithm and the experiments target empirical group DRO problems. However, only group DRO algorithms are compared. It would be more persuasive to compare with empirical group DRO algorithms (Yu et al., 2024) since the latter could use the problem structure to further improve the complexity.\n\n[1] Nemirovski, A., Juditsky, A., Lan, G., and Shapiro, A. Robust stochastic approximation approach to stochastic programming. SIAM Journal on Optimization, 2009.\n\n[2] Zhang, L., Zhao, P., Yang, T., and Zhou, Z.-H. Stochastic approximation approaches to group distributionally robust optimization. In NeurIPS, 2023.\n\n[3] Yu, D., Cai, Y., Jiang, W., and Zhang, L. Efficient Algorithms for Empirical Group Distributionally Robust Optimization and Beyond. In ICML, 2024.\n\n[4] Sagawa, S., Koh, P. W., Hashimoto, T. B., and Liang, P. Distributionally robust neural networks for group shifts: On the importance of regularization for worst-case generalization. In ICLR, 2020.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)