MixMax: Distributional Robustness in Function Space via Optimal Data Mixtures

## Abstract

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 M ixMax. In our experiments, we found that M ixMax matched or outperformed the standard group DRO baselines, and in particular, M ixMax improved the performance of XGBoost over the only baseline, data balancing, for variations of the ACSIncome and CelebA annotations datasets.

## 1 Introduction

Machine learning models are often required to perform well over different distributions of data. For example, models may be deployed for different groups of users [, ], or in the development of foundation models for a suite of tasks. Worst-case error is one way to summarize the performance of a model over a set of test-time distributions. The algorithms that optimize this worst case error are called distributionally robust optimization (DRO) methods [, , , , ], and group DRO methods [] when the distirbution set is finite.

Despite the importance of robust models, we lack effective methods for group DRO over modern model classes. When our loss is convex in the model parameters, methods exist for solving the group DRO objective []. But in the case of modern expressive non-linear models, only heuristic methods exist [, ]. Furthermore, these methods are expensive, unstable, and training on simple balanced mixtures of the data distributions sometimes leads to better group DRO solutions [].

In this paper we show that group DRO for expressive models can be broken up into two sub-problems: a maximization to find a “hardest” training mixture, and fitting a model to that hardest mixture. Although DRO solutions are risk minimizers with respect to a specific mixture [, , ], it is not known how to find these mixtures. A minimax theorem would give a stationary condition for these mixtures, but it not clear whether minimax theorems exist for modern model classes.

We derive a general minimax theorem for group DRO, which gives a stationary condition for this “hardest" mixture, by making a classical move: reparameterize the group DRO objective to be over the model’s function space as opposed to parameter space. We make a key assumption, which is that the model’s function space covers all bounded functions. Although this assumption may exclude some settings of interest, as the community scales up model sizes the space of bounded functions may be a good proxy of the hypothesis space we are exploring. Lastly, to obtain this minimax theorem over function space we bypassed issues with $L^{p}$ metrics by using topological dual spaces, which may be of independent interest.

Crucially, we show that the stationary condition for this “hardest" mixture is given by a concave maximization problem for cross entropy and $\ell_{2}^{2}$, which means it can be found efficiently in principle in these cases. Our result leverages the structure of the Bayes optimal functions for cross entropy and $\ell_{2}^{2}$. We derive a variety of practical, empirical methods for this objective, which we collectively call *MixMax*. In principle, the mixture weights returned by MixMax methods can be used to obtain the group DRO optimal model by training on the corresponding MixMax mixture.

MixMax has a number of advantages for group DRO, even under practical constraints where we do not necessarily have the proven guarantees. 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—the same amount of training compute as training one model on all the data. Furthermore, because the weights can 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 [] (^1^11To the best of our knowledge, there exists no past work on group DRO for non-parametric learning algorithms, with the only work on DRO in general being for k-nearest neighbours where the set of distributions forms a Wasserstein ball [].).

To illustrate the empirical performance of MixMax , we applied it for two real-world model classes (^2^22The code is at [https://github.com/Anvith-Thudi/MixMax](https://github.com/Anvith-Thudi/MixMax)). First, we tested MixMax with transformer models on synthetic Markov chain data. Second, we tested MixMax with XGBoost [] on several tabular datasets with different group shifts [, ]. In all cases, we found that empirical versions of MixMax matched or outperformed applicable baseline methods when improvement was possible. In particular, when a moderate label shift was present, MixMax yielded relative test accuracy improvements between $2.3-5.9\%$ for XGBoost on variations of ACSIncome [] and CelebA annotations []. Our contributions are:

- 1.
A minimax theorem for DRO over bounded functions.
- 2.
Showing that applied to cross-entropy and $\ell_{2}^{2}$, 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-parametric learning, and applying it to XGBoost to improve over the baseline of balancing data by upsampling.

Figure: 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.
Refer to caption: x1.png

## 2 Preliminaries

In this paper we will consider solving group DRO when we have model classes that are sufficiently expressive to cover all bounded functions. This will be relaxed in Section 4.2, but for now this formally means we can (and will) reparameterize our optimization to be over all bounded functions. We can restrict to this subset of functions without loss of generality as we assume the output space is bounded and hence the group DRO and Bayes optimal functions are bounded.

Specifically, we work with functions $f:\mathcal{X}\to\mathcal{O}$ from an input domain $\mathcal{X}$ (with a measure $dx$) to a closed convex output domain $\mathcal{O}\subset\mathbb{R}^{n}$ that have bounded $L^{\infty}$ norm in each output coordinate: we denote this function space by $L^{\infty}(\mathcal{X},\mathcal{O})$. For example, when doing classification $\mathcal{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 $r$ in every coordinate (a.e.), denoted by $B_{L^{\infty}(\mathcal{X},\mathcal{O})}(r)$. We work with data $(x,y)$ where $x\in\mathcal{X}$ and $y\in\mathcal{Y}$, and an associated loss function $\mathcal{L}(o,y):\mathcal{O}\times\mathcal{Y}\rightarrow\mathbb{R}^{+}$ that is convex in the first argument and continuous in both. An example is cross-entropy where $o$ is the label probabilities, and $y$ is a specific label. With $\mathcal{L}$ we have the expected loss of a function $f\in L^{\infty}(\mathcal{X},\mathcal{Y})$ over a distribution $dp$ on $\mathcal{X}\times\mathcal{Y}$ is $\int_{\mathcal{X}\times\mathcal{Y}}\mathcal{L}(f(x),y)dp(x,y)$. Given a set of distributions $P$, the DRO objective we study is (which is group DRO when $P$ is finite)

$$ $\inf_{f\in B_{L^{\infty}}(r)}\sup_{dp\in P}\int\mathcal{L}(f(x),y)dp(x,y).$ (DRO) $$

## 3 DRO Over Bounded Functions is Solved by Data Mixing

We now show that, letting our model class be all bounded functions, fitting to an optimal data mixture returns a function that solves group DRO (Corollary 3.2). This is done by observing the set of all bounded functions has a 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 4 we demonstrate how to optimize for this mixture in the case of cross-entropy and $\ell_{2}^{2}$ loss.

First, Theorem 3.1 formally states that there exists a minimizer to the hardest distribution which solves DRO. To show this we overcome challenges in applying Sion’s minimax theorem to bounded functions, in particular, the fact that the set is not compact in any $L^{p}$ topology. Our main contribution is to show that working with the weak*-topology resolves these issues. This follows from leveraging Banach-Alaoglu’s theorem alongside the method of test functions, and introduces mild additional assumptions. For example, we will require that $dx$ is a $\sigma$-finite measure to apply $L^{p}$ duality theory, which is satisfied if $dx$ is Lebesque measure on $R^{n}$ or counting measure on some countable set. We will further require that $\mathcal{L}$ is bounded on $\mathcal{O}\times\mathcal{Y}$ to use intergral convergence theorems, and discuss the strength of this assumption later. A complete proof is provided in Appendix A.1

###### Theorem 3.1 (DRO over $L^{\infty}$ = DM) .

Let $P$ be a 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 $\mathcal{O}\subset\mathbb{R}^{m}$ be a closed convex set, and $L^{\infty}(\mathcal{X},\mathcal{O})$ be defined w.r.t the measure $dx$, and $B_{L^{\infty}(\mathcal{X},\mathcal{O})}(r)=\{f\in L^{\infty}(\mathcal{X},
\mathcal{O}):||f||_{\infty}\leq r\}$.

Let the loss function $\mathcal{L}(o,y)$ be continuous in both arguments, and convex in the first argument. Furthermore assume $\mathcal{L}$ is bounded by some constant $M$ on $\mathcal{O}\times\mathcal{Y}$. If $dp_{\lambda}$ realizes

$$ $\sup_{dp\in Conv(P)}\inf_{f\in B_{L^{\infty}(\mathcal{X},\mathcal{O})}(r)}\int \mathcal{L}(f(x),y)dp(x,y),$ $$

and there exists an $f^{*}\in B_{L^{\infty}(\mathcal{X},\mathcal{O})}(r)$ realizing the DRO objective, then there exists a minimizer $f_{\lambda}$ of the expected loss under $dp_{\lambda}$ that also realizes the DRO objective

$$ $\inf_{f\in B_{L^{\infty}(\mathcal{X},\mathcal{O})}(r)}\sup_{dp\in P}\int \mathcal{L}(f(x),y)dp(x,y).$ $$

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 a specific mixture distribution.

While bounded loss may seem strong, it can be enforced by choosing the output space $\mathcal{O}$ carefully, e.g., for cross-entropy loss if we know the minimum probability of any label is $\epsilon>0$ we can choose $\mathcal{O}$ accordingly and avoid loss blow-up (from $log(0)$). 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\in P$, which can be true without loss boundedness depending on $P$. On assuming a DRO solution over $B_{L^{\infty}(\mathcal{X},\mathcal{O})}(r)$ exists, this is true for finite $P$ (Appendix A.1) and hence applies for group DRO.

For the rest of the paper, We will assume there is sufficient regularity (whether imposed by $\mathcal{O}$, $P$, or both) for pointwise convergence of loss to imply the average loss converges and hence the result of Theorem 3.1 applies. We will further assume that for some $r<\infty$ we have the bounded functions $B_{L^{\infty}(\mathcal{X},\mathcal{O})}(r)$ contains the Bayes optimal solutions for all $dp\in P$ (to leverage structural properties of cross-entropy and $\ell_{2}^{2}$). Note the last two assumptions are satisfied for cross-entropy if a.e. all labels have $\geq\epsilon$ probability for some $\epsilon>0$ (and we choose $\mathcal{O}$ accordingly), and for $\ell_{2}^{2}$ if the allowed $y$ values are bounded. We leave it to future work to consider applying Theorem 3.1 without these assumptions, or generalizing Theorem 3.1 itself.

###### Theorem 3.1 (DRO over L ∞ superscript 𝐿 L^{\infty} italic_L start_POSTSUPERSCRIPT ∞ end_POSTSUPERSCRIPT = DM) .

### 3.1 Group DRO by Maximization

We now consider applying Theorem 3.1 for group DRO. For cross-entropy and $\ell_{2}^{2}$ losses, the Bayes optimal functions are unique up to a measure 0 set over $\mathcal{X}$, and so any minimizer of the hardest distribution is DRO optimal by Theorem 3.1 (given a DRO solution exists which is true for group DRO): all minimizers have the same average losses on $dp\in 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 $\Delta^{P}$.

###### Corollary 3.2 (Group DRO by Maximization) .

Let $P$ be a finite set, and assume $\int\mathcal{L}(f(x),y)dp(x,y)$ is uniquely minimized over functions in $B_{L^{\infty}(\mathcal{X},\mathcal{O})}(r)$ (upto a measure 0 0 set). Parameterize $Conv(P)$ by ${\lambda\in\Delta^{P}}$ via ${Conv(P)=\{dp_{\lambda}\coloneqq\sum_{dp\in P}\lambda_{p}dp:\lambda\in\Delta^{
P}\}}$. Further denote the minimizer of $\int\mathcal{L}(f(x),y)dp_{\lambda}(x,y)$ in $B_{L^{\infty}(\mathcal{X},\mathcal{O})}(r)$ by $f_{\lambda}$ (parameterized by $\lambda$). Then for $\lambda$ realizing $\sup_{\lambda\in\Delta^{P}}\int\mathcal{L}(f_{\lambda}(x),y)dp_{\lambda}(x,y)$, $f_{\lambda}(x)$ realizes $\inf_{f\in B_{L^{\infty}(\mathcal{X},\mathcal{O})}(r)}\sup_{dp\in P}\int
\mathcal{L}(f(x),y)dp(x,y)$.
That is, the group DRO optimization over $L^{\infty}$ reduces to a single maximization over $\Delta^{P}$.

###### Proof.

Follows from Theorem 3.1, noting $P$ is finite so a DRO solution exists, and rewriting the condition for the mixture with the parameterizations given by uniqueness of solutions.
∎

###### Corollary 3.2 (Group DRO by Maximization) .

###### Proof.

## 4 M ixMax

Figure: Algorithm 1 Empirical MixMax

We now apply Theorem 3.1 (under the assumptions of the previous section) to cross-entropy and $\ell_{2}^{2}$ for finite $P$.
We use $p(x)$ and $f_{p}(x)$ to represent the covariate density and Bayes optimal function of distributions $dp\in P$, and $p_{\lambda}(x)$ and $f_{\lambda}(x)$ similarly for the mixture distributions $dp_{\lambda}\in Conv(P)$.

First note that, by Corollary 3.2, the objective for the optimal mixture weights is

$$ $\sup_{\lambda\in\Delta^{P}}\int\mathcal{L}(f_{\lambda}(x),y)dp_{\lambda}(x,y)= \sup_{\lambda\in\Delta^{P}}\sum_{dp\in P}\lambda_{p}\int\mathcal{L}(f_{\lambda }(x),y)dp(x,y)$ (1) $$

We show that this is concave and that we can compute its gradients for cross-entropy and $\ell_{2}^{2}$. Thus, we can perform entropic mirror ascent [] to solve the constrained optimization. Algorithm 1 describes this approach, which we call MixMax ("Mixtures by Maximization").

#### The Objective is Concave

In the case of cross-entropy, this objective reduces to the expected entropy of $y$ conditioned on $x$, $p_{\lambda}(y|x)$, over $x\sim p_{\lambda}(x)$. Therefore it is a concave maximization problem in the mixture weights via the concavity of entropy. In the case of $\ell_{2}^{2}$, under our assumptions, this objective reduces to the expectation of the conditional variance of $y$ given $x$ over $x\sim p_{\lambda}(x)$, which is again a concave objective. Appendix A.2 provides full proofs of concavity.

### 4.1 Computing Gradients of the M ixMax Objective

We now demonstrate how to compute the gradient of the objective in Equation 1, i.e., $\sum_{dp\in P}\int\nabla_{\lambda}\lambda_{p}\mathcal{L}(f_{\lambda}(x),y)dp(x
,y)$ (^3^33Formally, we require $\nabla_{\lambda}\mathcal{L}(f_{\lambda}(x),y)$ to have sufficient regularity (e.g., is bounded).), for cross-entropy and $\ell_{2}^{2}$ given $f_{p}$ and $p(x)$ and the ability to integrate for all $dp\in P$. We present this in cases, and later discuss empirical implementations.

#### Case 1: No Covariate Shift

Suppose there is no covariate shift, *i.e.*, $\exists~{}p_{0}(x)$ such that $p(x)=p_{0}(x)~{}\forall dp\in P$. Then, for cross-entropy we have the Bayes optimal function $f_{\lambda}(x)=p_{\lambda}(y|x)=\sum_{dp\in P}\lambda_{p}p(y|x)=\sum_{dp\in P}
\lambda_{p}f_{p}(x)$ and for $\ell_{2}^{2}$ we also have $f_{\lambda}(x)=\mathbb{E}_{y\sim p_{\lambda}(y|x)}y=\sum_{dp\in P}\lambda_{p}
\mathbb{E}_{y\sim p(y|x)}y=\sum_{dp\in P}\lambda_{p}f_{p}(x)$. Hence given $f_{p}(x)~{}\forall dp\in P$ we can compute $f_{\lambda}(x)$ and $\nabla_{\lambda}f_{\lambda}(x)$. Hence, $\nabla_{\lambda}\lambda_{p}\mathcal{L}(f_{\lambda}(x),y)$ follows by product and chain rule, and given the ability to integrate over $dp\in P$ we can compute the gradient of the objective.

#### Case 2: Covariate Shift

The more general expression under covariate shift, for cross-entropy and $\ell_{2}^{2}$, is $f_{\lambda}(x)=\frac{\sum_{dp\in P}\lambda_{p}f_{p}(x)p(x)}{\sum_{dp^{\prime}
\in P}\lambda_{p^{\prime}}p^{\prime}(x)}$. Hence computing $f_{\lambda}(x)$ and $\nabla_{\lambda}f_{\lambda}(x)$ would further require knowledge of $p(x)~{}\forall dp\in P$. Given this, computing $\nabla_{\lambda}\lambda_{p}\mathcal{L}(f_{\lambda}(x),y)$ follows by product and chain rule, and given the ability to integrate over $dp\in P$ we can compute the gradient of the objective.

### 4.2 Empirical M ixMax and Practical Considerations

Figure: Figure 2: An illustration of the data requirements for the empirical MixMax approaches.
Refer to caption: x2.png

Section 4.1 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 can integrate exactly over the domain. We now decribe empirical implementations which remove these assumptions, summarized in Figure 2. These add varying sources of empirical error.

#### Empirical MixMax (EMixMax)

Computing MixMax weights is often intractable due to the integral over $dp\in P$. However, given datasets $D_{p}$ for $dp\in P$ one can compute empirical MixMax (EMixMax) weights,

$$ $\displaystyle\hat{\lambda}^{*}_{\mathrm{EMixMax}}=\arg\max_{\lambda\in\Delta^{ P}}\sum_{dp\in P}\frac{1}{|D_{p}|}\sum_{(x,y)\in D_{p}}\lambda_{p}\mathcal{L}( f_{\lambda}(x),y),$ $$

via a stochastic entropic mirror ascent using the gradient estimator $\nabla_{\lambda}\lambda_{p}\mathcal{L}(f_{\lambda}(x),y)$. EMixMax weights are empirical risk minimizers and should converge to the MixMax weights with enough samples. We use EMixMax for all our experiments (see Algorithm 1).

#### Empirical 2 MixMax

One also usually does not know the Bayes optimal functions $f_{p}(x)$ or the densities $p(x)$. However one can train a model on samples from $dp\in P$ to obtain approximation $\hat{f_{p}}(x)$ and $\hat{p}$, and then use another set of samples to run EMixMax. We call this approach Empirical^2 MixMax or E^2MixMax. More precisely, in the case of no covariate shift, given independent datasets $D_{p}$ and $D_{p}^{\prime}$ for $dp\in P$, E^2MixMax returns

$$ $\displaystyle\hat{\lambda}^{*}_{\mathrm{E}^{2}\mathrm{MixMax}}=$ $\displaystyle\arg\max_{\lambda\in\Delta^{P}}\sum_{dp\in P}\frac{1}{|D_{p}|} \sum_{(x,y)\in D_{p}}\lambda_{p}\mathcal{L}\left(\sum_{dp\in P}\lambda_{p}\hat {f_{p}}(x),y\right)$ $\displaystyle\text{where }\hat{f_{p}}(x)=\arg\min_{f}\frac{1}{|D_{p}^{\prime}| }\sum_{(x,y)\in D_{p}^{\prime}}\mathcal{L}\left(f(x),y\right).$ $$

If $\hat{f_{p}}$ are close to the true $f_{p}(x)$, then the mixture weights are still useful for the underlying distributions. In the case of covariate shift, we also fit density approximations $\hat{p}(x)$ on the second dataset $D_{p}^{\prime}$ and use the covariate shift definition of $f_{\lambda}$. Algorithm 1 describes the changes for E^2MixMax . We also consider *E^2MixMax with Data Reuse*, i.e., $D_{p}=D_{p}^{\prime}$, as a more sample efficient alternative.

#### How to Use M ixMax Weights

Finally, MixMax mixture weights can be used in two ways. We can use the weights to define a mixture distribution and fit a new model on it, which Theorem 3.1 states will result in a group DRO solution if we fit optimally. Alternatively, we can use the mixture weights to combine our (approximate) $f_{p}$ models (i.e., the $f_{\lambda}$ formula), which also returns an approximation of the best model for the MixMax mixture distribution. These approaches provide a trade-off: the first approach has higher training cost but results in a single model and hence smaller inference cost. We employed both approaches in our experiments and found they both improved over baselines.

### 4.3 Illustrating M ixMax

Figure: (a) Maximum Entropy Conditional in Mixture Span
Refer to caption: x3.png

Here we investigated patterns for the group DRO solutions found by EMixMax, running it 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)=unif[0,1]$, but with $p_{1}(1|x)=0.5cos(\pi x)+0.5$ and $p_{2}(1|x)=-0.5cos(\pi x)+0.5$ (Figure 3(a)) or $p_{2}(1|x)=-0.5cos(\pi(x-0.2))+0.5$ (Figure 3(b)). 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 (the true group DRO solution). In the second case where random guessing is not in the span, EMixMax found a 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(x)=unif[0,1]$, but with $p_{1}(1|x)=0.2cos(\pi x)+0.5$, $p_{2}(1|x)=0.1$, and $p_{3}(1|x)=0.15$ or setting $p_{3}(1|x)=0.8$ (Figure 6(a)) and Figure 6(b) in Appendix C respectively). We found EMixMax always found the extreme functions and balanced those, which minimizes maximal $\ell_{2}^{2}$ error (true group DRO solution).

## 5 Related Work

#### Distributionally Robust Optimization

DRO is motivated by several applications not limited to just machine learning, such as resource allocation []. Approaches for DRO often focus on “uncertainty" sets $P$ (the set of distributions we are minimizing the maximum error over) that admit a duality theory [, , , , ], such as Lagragian duality, conic duality, Fenchel duality, etc. For theorem 3.1 we assume nothing on $P$, and for some results, that it is finite similar to work on the Cutting-Surface approaches to DRO [, , , ]; the finite $P$ case is widely called group DRO [] in 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 3.1 already appears as Lemma 17 in Rahimian and Mehrotra [] and is wide-spread in the finite $P$ literature. On similarity in theory, our main theorem that DRO is solved by data mixing is most closely related to the saddle point approaches to stochastic programming [, ]. These results present duality theory where solving for the hardest distribution is enough when the set of parameters to minimize w.r.t is a subset of $\mathbb{R}^{n}$. In the case of Dupačová [] a focus is placed on sets of all distributions meeting certain moment constraints (e.g., only the mean and variance are known). Shapiro and Kleywegt [] instead considers the case of an arbitrary finite set of distributions like our work, proposing sample average approaches for optimal solutions. Our method also uses sample averages, but further leverages the structure of cross-entropy and $\ell_{2}^{2}$ loss to collapse the max-min optimization to a single maximization.

Nevertheless, to 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.1). This yields a concave objective to maximize (Corollary 3.2) which we find yields good results for practical applications involving non-parametric learning and expressive non-linear model classes.

#### Data Mixing

The work of Xie et al. [], building on past work on DRO for deep learning [, ], highlighted empirically how optimizing dataset mixtures can lead to better performance over several downstream tasks. However, it is not clear whether this is 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 [, ], alongside increasing sample access []. In this paper we 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 [] and data bias [, ] literature have considered the same question, but the analysis there has been limited to problems satisfying the KKT conditions or to studying local minima in $\mathbb{R}^{n}$, and do not discuss how to find the best data mixture. We also note recent work by Fan et al. [] proposed data mixing methods that build on Xie et al. [], but departed away from the DRO objective (and Xie et al. [] 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 is 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 $\eta=2.0$ for all the sequence modeling tasks, and for $20$ steps with $\eta=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 [] for the transformers (hyperparameters described in Appendix B).

### 6.1 EMixMax and E 2 MixMax variants Maximize MixMax Objective Comparably

Here we investigated how E^2MixMax with various data splits compared to EMixMax in maximizing the MixMax objective. Specifically, we considered 4 tokens $\{0,1,2,3\}$ and sequences generated by a 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 $i$ was always $1/10$. We constructed three Markov chains to 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, E^2MixMax with $(75:25),(50:50),$ and $(25:75)$ split between the proxy model (a small transformer) training set and EMixMax set, and E^2MixMax with Data Reuse. This is shown in Figure 11 (in Appendix C), where we see that for low samples E^2MixMax with Data Reuse performs better than alternative E^2MixMax approaches, but as we increase the samples all methods are comparable to EMixMax .

Now just comparing E^2MixMax with Data Reuse to E^2MixMax with a $(75:25)$ split, but varying the similarity in the set of distribution (by changing the Dirichlet magnitude) with a fixed $800$ samples per length, we found in Figure 9 (in Appendix C) they perform comparably. We conclude one can effectively resuse training samples to run MixMax and be more sample efficient.

### 6.2 E 2 MixMax Performs Better than Group DRO Baselines for Sequence Modeling

Figure: (a) Comparing mixtures found by different methods.
Refer to caption: x5.png

Here we investigated how E^2MixMax compared to other mixture finding methods (balanced and DoReMi []) and the original gradient descent and ascent group DRO algorithm [] 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 E^2MixMax 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-out samples as a reference for best empirical performance. We applied DoReMi and group DRO using the same small transformer architecture and all $800$ training samples per length (per Markov chain). Hyperparameters are described in Appendix B, 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 E^2MixMax returned better mixtures to fit 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 4(a) we compared the worst group cross-entropy loss of the Bayes optimal function for the mixture given by E^2MixMax , EMixMax using ground truths to represent optimal performance, DoReMi, the average group DRO mixture weights, and also balanced weights. As seen, fitting to the mixture given by E^2MixMax 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 []. However, we note that group DRO’s intended use is to apply the model it trained and not its mixture weights. In Figure 4(b) we observed using E^2MixMax 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(b) 7(a) in Appendix C) highlighting that we had sufficient samples for both experiments.

### 6.3 E 2 M ixMax with Data Reuse Beats Balancing Data for XGBoost

Figure: (a) Comparing non-parametric group DRO methods
Refer to caption: x7.png

We compared E^2MixMax with Data Reuse weights to data balancing for non-parametric learning algorithms, in particular XGBoost [] which is known to be the state of the art for tabular data. We selected ACSIncome [] (released under the MIT license) and CelebA annotations [] (released for non-commercial use (^4^44The agreement for use is at [https://mmlab.ie.cuhk.edu.hk/projects/CelebA.html](https://mmlab.ie.cuhk.edu.hk/projects/CelebA.html).)) to 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, we 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 (^5^55The different number of features compared to ACSIncome was because the features were binary while ACSIncome first two features take on more values.). We used random $80\%-20\%$ train-test splits in all settings. We applied E^2MixMax with Data Reuse by using XGBoost 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 model trained on the same training data but up-sampled according to weights returned by E^2MixMax with Data Reuse. For the balanced baseline, we returned an XGBoost model trained with the data balanced by up-sampling the smaller groups. Lastly, as a measure for best possible performance, we also presented the worst oracle accuracy, where oracle accuracy refers to the test accuracy of an XGBoost model trained on the same distribution (where we need access to the group identity at test time to use this model). Hyperparameters are described in Appendix B, and we reported the results for the hyperparameter setting with the best performance over $5$ trials with random train-test splits.

As seen in Figure 5(a) (and Table 1 in Appendix C), E^2MixMax with Data Reuse matched or outperformed the worst group accuracy of data balancing in all settings, and similarly for worst group loss (Table 2 in Appendix C). In particular, as seen in Figure 5(b) we observed the improvement over the worst group accuracy of data balancing was stronger when more room for improvement was possible (a larger gap between oracle and balanced worst group accuracy): we observed a $0.79$ Pearson correlation factor. In the end, we found E^2MixMax 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).

## 7 Conclusion

In this paper we showed that group DRO over bounded functions can be solved by fitting to an optimal data mixture, and that maximizing a particular concave objective returns the optimal mixture weights for cross-entropy and $\ell_{2}^{2}$ loss. We called this method for finding data mixtures MixMax. Our experiment on a simple sequence modeling task showed that even empirical versions MixMax improved over previous parametric group DRO baselines. An empirical version of MixMax was also shown to improve over the baseline of balancing data for non-parametric 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, a main technical limitation of MixMax methods is the need to model covariate shifts if covariate shifts exist, which often requires large amounts of data in high-dimensional covariate spaces. We also acknowledge that DRO may be used to claim a model is fair, despite it still carrying societal biases. We hope future uses will take care in considering the claims DRO can and cannot make.

## Acknowledgements

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 acknowledge the support of the Natural Sciences and Engineering Research Council of Canada (NSERC), RGPIN-2021-03445. 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.

## Appendix

## Appendix A Proofs

### A.1 Proof of Theorem 3.1

###### Proof.

First note $\sup_{dp\in P}=\sup_{dp\in Conv(P)}$ as the objective is linear in $dp$. Now note we can apply Sion’s minimax theorem [] if $B_{L^{\infty}(\mathcal{X},\mathcal{O})}(r)$ is compact in a topology that the objective is continuous in. Note $B_{L^{\infty}(\mathcal{X},\mathbb{R}^{m})}(r)$ is compact in the Weak*-topology on $L^{\infty}(\mathcal{X},\mathbb{R}^{m})$ (as the dual of $L^{1}(\mathcal{X},\mathbb{R}^{m})$) by Banach-Alaoglu’s theorem and continuity of dilation []. Now note $B_{L^{\infty}(\mathcal{X},\mathcal{O})}(r)$ is a closed subset of $B_{L^{\infty}(\mathcal{X},\mathbb{R}^{m})}(r)$ and hence is also compact. The fact $B_{L^{\infty}(\mathcal{X},\mathcal{O})}(r)$ is closed follows from $\mathcal{O}$ being closed and that weak* convergence implies pointwise convergence (a.e.) in $B_{L^{\infty}(\mathcal{X},\mathbb{R}^{m})}(r)$ (a set containing $B_{L^{\infty}(\mathcal{X},\mathcal{O})}(r)$), which is proven below.

Note convergence in weak* for functions in $B_{L^{\infty}(\mathcal{X},\mathbb{R}^{m})}(r)$ means the functions converge pointwise a.e. This is as, for $g\in L^{1}$ and $f$ the limit of $f_{n}$ in the weak*-topology, $0=\lim_{n\rightarrow\infty}\int(f-f_{n})gdx=\int\lim_{n\rightarrow\infty}(f-f_
{n})gdx$ by dominated convergence as functions are bounded. Now taking $g=sign(\lim_{n\rightarrow\infty}(f-f_{n}))\chi_{E}$ implies by proposition 2.16 in [] that $\lim_{n\rightarrow\infty}(f-f_{n})(x)=0$ a.e on $E$ for finite measure sets $E$. As the space is countably coverable by finite measured sets (by the $\sigma$-finite assumption), this implies the functions converge pointwise a.e.

Now as $\mathcal{L}$ is bounded, and the functions converge pointwise in the weak*-topology, we also have by the dominated convergence theorem the objective is continuous in the weak* topology on $L^{\infty}$.

Thus we satisfy the conditions for Sion’s minimax theorem, which implies

$$ $\sup_{dp\in Conv(P)}\inf_{f\in B_{L^{\infty}(\mathcal{X},\mathcal{O})}(r)}\int \mathcal{L}(f(x),y)p(x,y)=\inf_{f\in B_{L^{\infty}(\mathcal{X},\mathcal{O})}(r )}\sup_{dp\in Conv(P)}\int\mathcal{L}(f(x),y)p(x,y)$ $$

So if $f^{*}$ achieves the inf in the RHS and $dp_{\lambda}$ achieves the sup in the RHS, then note

$$ $\int\mathcal{L}(f^{*}(x),y)dp_{\lambda}(x,y)\leq\sup_{dp\in Conv(P)}\int \mathcal{L}(f^{*}(x),y)dp(x,y)=\inf_{f\in B_{L^{\infty}(\mathcal{X},\mathcal{O })}(r)}\int\mathcal{L}(f(x),y)dp_{\lambda}(x,y)$ $$

conclude that $\int\mathcal{L}(f^{*}(x),y)dp_{\lambda}(x,y)=\inf_{f\in B_{L^{\infty}(\mathcal
{X},\mathcal{O})}(r)}\int\mathcal{L}(f(x),y)dp_{\lambda}(x,y)$. So in particular $f^{*}$ is a minimizer of $dp_{\lambda}$. Note minimizers of $dp_{\lambda}$ exist by the compactness of $B_{L^{\infty}(\mathcal{X},\mathbb{R}^{m})}(r)$ and continuity of the objective over this function space.

∎

###### Proof.

#### 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 a finite set of continuous functions (note continuity of the objective for a single $dp$ was proven above). Further care is needed when $P$ is not finite, but the paper focuses on the finite $P$ case and so we do not discuss this further.

### A.2 Concavity of Objectives

###### Fact A.1 .

If $\mathcal{L}$ is cross-entropy, then $\int\mathcal{L}(f_{\lambda}(x),y)dp_{\lambda}(x,y)$ is concave in $\lambda$.

###### Proof.

We note ${f_{\lambda}=p_{\lambda}(y|x)=\sum_{dp\in P}\lambda_{p}(x)p(y|x)}$ where ${\lambda_{p}(x)=\frac{\lambda_{p}p(x)}{\sum_{dp^{\prime}\in P}\lambda_{p^{
\prime}}p^{\prime}(x)}=\frac{\lambda_{p}p(x)}{p_{\lambda}(x)}}$. Note that ${-\int_{\mathcal{Y}}\log(\sum_{dp\in P}\lambda_{p}(x)p(y|x))\sum_{dp\in P}
\lambda_{p}(x)p(y|x)\geq\sum_{dp\in{}}\lambda_{p}(x)\int_{\mathcal{Y}}-log(p(y
|x))p(y|x)}$ by the concavity of entropy. Applying this we have

$$ $\int_{\mathcal{X}}\int_{\mathcal{Y}}-log(p_{\lambda}(y|x))p_{\lambda}(y|x)p_{ \lambda}(x)\\ \geq\int_{\mathcal{X}}\left(\sum_{dp\in P}\lambda_{p}(x)\int_{\mathcal{Y}}-log (p(y|x))p(y|x)\right)p_{\lambda}(x)\\ =\int_{\mathcal{X}}\left(\sum_{dp\in P}\frac{\lambda_{p}p(x)}{p_{\lambda}(x)} \int_{\mathcal{Y}}-log(p(y|x))p(y|x)\right)p_{\lambda}(x)\\ =\sum_{dp\in P}\lambda_{p}\int_{\mathcal{X}}\int_{\mathcal{Y}}-log(p(y|x))p(y| x)p(x)$ (2) $$

this proves concavity w.r.t $\lambda$, as desired.

∎

###### Fact A.2 .

If $\mathcal{L}$ is $\ell_{2}^{2}$, then $\int\mathcal{L}(f_{\lambda}(x),y)p_{\lambda}(x,y)$ is concave in $\lambda$

###### Proof.

We assume $f_{\lambda}$ is the bayes optimal solution, hence we have the loss is the bayes error $\int_{\mathcal{X}}Var_{p_{\lambda}(y|x)}(y|x)p_{\lambda}(x)$. Note $p_{\lambda}(y|x)=\sum_{dp\in P}\lambda_{p}(x)p(y|x)$ where $\lambda_{p}(x)=\frac{\lambda_{p}p(x)}{\sum_{dp^{\prime}\in P}\lambda_{p^{
\prime}}p^{\prime}(x)}=\frac{\lambda_{p}p(x)}{p_{\lambda}(x)}$, and variance is concave w.r.t mixture weights. Thus we have

$$ $\int\mathcal{L}(f_{\lambda}(x),y)p_{\lambda}(x,y)\\ =\int_{\mathcal{X}}Var_{p_{\lambda}(y|x)}(y|x)p_{\lambda}(x)\\ \geq\int_{\mathcal{X}}\left(\sum_{dp\in P}\lambda_{p}(x)Var_{p(y|x)}(y|x) \right)p_{\lambda}(x)\\ =\int_{\mathcal{X}}\left(\sum_{dp\in P}\frac{\lambda_{p}p(x)}{p_{\lambda}(x)} Var_{p(y|x)}(y|x)\right)p_{\lambda}(x)\\ =\sum_{dp\in P}\lambda_{p}\int_{\mathcal{X}}Var_{p(y|x)}(y|x)p(x)\\ =\sum_{dp\in P}\lambda_{p}\int\mathcal{L}(f_{p}(x),y)p(x,y)$ (3) $$

This proves concavity
∎

###### Fact A.1 .

###### Proof.

###### Fact A.2 .

###### Proof.

## Appendix B Experimental Setups

### B.1 Toy Experiments

For the binary classification task, We ran EMixMax with $D_{i}$ consisting of $10000$ samples from each distribution, step size $\eta=0.5$, and for $100$ steps.

We ran EMixMax with $D_{i}$ again consisting of $10000$ samples from each distribution, and $\eta=0.0001$ for $100$ steps.

### B.2 Sequence Modeling

For the MixMax methods, letting $y$ denote a sequence, note the task was to model $p(y)$ and so had no covariates and hence had no covariate shift (as one can take $\mathcal{X}$ to be a singleton). Thus MixMax only needed functions for the token probabilities from each distribution, and we 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 [] with $6$ hidden states, $2$ hidden 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 E^2MixMax . The proxy model is trained for $20$ epochs using AdamW 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 AdamW 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.

### B.3 Tabular Data

For running Gaussian kernel density estimation, we used the Scott method for bandwidth selection []. For fitting the model on E^2MixMax with data reuse mixture weights, the reference models used for E^2MixMax with data reuse on the Income dataset, and the baseline of balancing the dataset, we implemented XGBoost with depths $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 in E^2MixMax with data reuse were always trained used 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 balancing always has a hyperparameter sweep).

## Appendix C Additional Tables and Figures

**Table 1: E^2MixMax 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 5(b) for a visualization). We present the average minimum accuracy over the groups from $5$ trials with random training-test splits for E^2MixMax with data reuse, balancing by up-sampling, and the "oracle" accuracy where a model was trained for each group and is evaluated on the same group (as an 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 ($\%$) | Balanced ($\%$) | E^2MixMax ($\%$) |
| --- | --- | --- | --- | --- |
| Inc-Race | All | $79.1(\pm 1.2)$ | $80.0(\pm 1.2)$ | $80.3(\pm 0.6)$ |
| $2$ | $71.2(\pm 0.1)$ | $70.7(\pm 0.2)$ | $70.6(\pm 0.3)$ |  |
| $1$ | $60.0(\pm 0.1)$ | $53.7(\pm 0.4)$ | $53.7(\pm 0.5)$ |  |
| Inc-Sex | All | $81.2(\pm 0.1)$ | $81.1(\pm 0.2)$ | $81.1(\pm 0.1)$ |
| $2$ | $72.8(\pm 0.1)$ | $71.4(\pm 0.2)$ | $71.6(\pm 0.1)$ |  |
| $1$ | $62.8(\pm 0.2)$ | $56.5(\pm 0.2)$ | $\mathbf{58.1(\pm 1.8)}$ |  |
| CelebA-Young | All | $77.0(\pm 0.1)$ | $75.6(\pm 0.2)$ | $\mathbf{76.0(\pm 1.0)}$ |
| $10$ | $68.3(\pm 0.3)$ | $62.5(\pm 0.3)$ | $\mathbf{64.8(\pm 1.6)}$ |  |
| $5$ | $63.4(\pm 0.3)$ | $53.8(\pm 0.2)$ | $\mathbf{57.0(\pm 0.6)}$ |  |
| CelebA-Pale Skin | All | $78.2(\pm 0.3)$ | $77.3(\pm 0.2)$ | $77.1(\pm 0.6)$ |
| $10$ | $67.6(\pm 0.1)$ | $66.8(\pm 0.2)$ | $\mathbf{67.1(\pm 0.3)}$ |  |
| $5$ | $62.4(\pm 0.2)$ | $59.5(\pm 0.3)$ | $59.4(\pm 0.3)$ |  |

**Table 2: E^2MixMax with data reuse improved worst group loss over the baseline of 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 E^2MixMax with data reuse and balancing by up-sampling. 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 | Balanced | E^2MixMax |
| --- | --- | --- | --- |
| Inc-Race | All | $0.426(\pm 0.013)$ | $0.422(\pm 0.007)$ |
| $2$ | $0.557(\pm 0.001)$ | $0.559(\pm 0.003)$ |  |
| $1$ | $0.707(\pm 0.004)$ | $\mathbf{0.688(\pm 0.020)}$ |  |
| Inc-Sex | All | $0.403(\pm 0.004)$ | $0.405(\pm 0.005)$ |
| $2$ | $0.544(\pm 0.002)$ | $\mathbf{0.541(\pm 0.006)}$ |  |
| $1$ | $0.616(\pm 0.001)$ | $\mathbf{0.611(\pm 0.003)}$ |  |
| CelebA-Young | All | $0.495(\pm 0.002)$ | $\mathbf{0.488(\pm 0.015)}$ |
| $10$ | $0.644(\pm 0.002)$ | $\mathbf{0.614(\pm 0.019)}$ |  |
| $5$ | $0.706(\pm 0.002)$ | $\mathbf{0.670(\pm 0.007)}$ |  |
| CelebA-Pale Skin | All | $0.465(\pm 0.001)$ | $0.468(\pm 0.006)$ |
| $10$ | $0.594(\pm 0.001)$ | $\mathbf{0.588(\pm 0.005)}$ |  |
| $5$ | $0.647(\pm 0.001)$ | $\mathbf{0.641(\pm 0.005)}$ |  |

**Table 3: E^2MixMax with data reuse can sometimes be improved by doing a hyperparameter search over number of optimization steps ($20,40,60,80$). We present the settings from Table 1 where the gap between oracle accuracy and the balanced accuracy was more than $2\%$, and include the average minimum accuracy over groups over $5$ trials of E^2MixMax with data reuse with further hyperparameter tuning. We bolded results where the tuned performance was outside one standard deviation of untuned.**
| Dataset | Features | Untuned E^2MixMax | Tuned E^2MixMax |
| --- | --- | --- | --- |
| Inc-Race | $1$ | $53.7(\pm 0.5\%)$ | $\mathbf{54.3(\pm 0.6\%)}$ |
| Inc-Sex | $1$ | $58.1(\pm 1.8\%)$ | $58.4(\pm 1.8\%)$ |
| CelebA-Young | $10$ | $64.8\%(\pm 1.6\%)$ | $66.0\%(\pm 2.1\%)$ |
| $5$ | $57.0\%(\pm 0.6\%)$ | $57.0\%(\pm 0.6\%)$ |  |
| CelebA-Pale Skin | $5$ | $59.4\%(\pm 0.3\%)$ | $59.4\%(\pm 0.3\%)$ |

Figure: (a) Extreme Functions Example 1
Refer to caption: x9.png

Figure: (a) Optimally Fitting Mixture
Refer to caption: x11.png

Figure: Figure 8: We see that E^2MixMax 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 3. Here we plot the worst group accuracy gain over running E^2MixMax with data reuse for $20$ steps when running E^2MixMax with data reuse for $20,40,60,80$ steps. Further hyperparameter details are described in Appendix B. This is done for the datasets observed to have big label shifts, described in Table 1
Refer to caption: x13.png

Figure: (a) $800$ Training Samples
Refer to caption: x14.png

Figure: Figure 10: Various empirical MixMax approaches perform comparably for worst group cross-entropy, with E^2MixMax 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-entropy of the Bayes optimal function defined by the methods’s mixture weights, changing the number of training samples per length.
Refer to caption: x16.png

Figure: Figure 11: Various empirical MixMax approaches perform comparably in maximizing the MixMax objective, with E^2MixMax 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-chains from the symmetric Dirichlet distribution with magnitude $1$) of the MixMax objective value for the methods’s mixture weights, changing the number of training samples per length.
Refer to caption: x17.png

## Peer Reviews

confidence: 2
contribution: 3
initial_rating: 6
presentation: 4
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.
soundness: 3
strengths: 1)The problem is important and relevant to OOD generalization.

2)The proposed method is novel in its approach and theoretically justified.

3)Experimental results show that the proposed method outperforms groupDRO.

4)The paper is well written.
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.
weaknesses: See questions.

---

confidence: 4
contribution: 3
initial_rating: 8
presentation: 3
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 ?

- 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 ?
soundness: 4
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.

- The authors also provide a practical algorithm that computed the empirical MixMax solution.
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.
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?

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

---

confidence: 2
contribution: 2
initial_rating: 6
presentation: 2
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?

2- 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. 

3-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.
soundness: 2
strengths: The issue is both intriguing and significant in the literature.
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. 

The 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.
weaknesses: The theoretical proofs are quite brief and unclear. It would be beneficial to discuss and clarify the proofs step by step.

---

confidence: 4
contribution: 2
initial_rating: 3
presentation: 3
questions: I would appreciate it if the authors were willing to answer some of the questions I raised in the weaknesses above.
soundness: 3
strengths: 1. The idea of reparameterization seems novel in the context of group DRO.
2. The experiments are comprehensive in terms of the choice of datasets and the diverse ways of presentation.
3. The paper is written clearly, especially for the theoretical results.
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.
weaknesses: There are multiple disadvantages, mainly regarding the methodology itself, together with its relevance to the existing literature on (empirical) group DRO.

1. 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. 
2. 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
3. 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?
4. 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. 
5. 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.
6. 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). 
7. 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.

[1] Nemirovski, A., Juditsky, A., Lan, G., and Shapiro, A. Robust stochastic approximation approach to stochastic programming. SIAM Journal on Optimization, 2009.

[2] Zhang, L., Zhao, P., Yang, T., and Zhou, Z.-H. Stochastic approximation approaches to group distributionally robust optimization. In NeurIPS, 2023.

[3] Yu, D., Cai, Y., Jiang, W., and Zhang, L. Efficient Algorithms for Empirical Group Distributionally Robust Optimization and Beyond. In ICML, 2024.

[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.

Meta-review: The paper studies group DRO settings with non-convex losses and/or non-parametric model classes. The paper proposes reparameterizing the problem to function space and prove that group DRO over the space of bounded functions admits a minimax theorem. They further show that when the loss is squared or cross-entropy, the problem can be reformulated as a concave maximization problem. Finally, the paper provides a suite of experiments that demonstrate the effectiveness of the proposed approach. While there were initial concerns regarding clarity and novelty of the results, all were resolved in the rebuttal phase and the paper was judged well above the acceptance threshold.

Decision: Accept (Poster)