MixMax: Distributional Robustness in Function Space via Optimal Data Mixtures

We thank the reviewer for their time and feedback! We discuss the questions raised by the reviewer below, and have incorporated a lot of their feedback, strengthening our paper. In particular, we believe DoReMi implements a variant of Zhang et al’s algorithm. We now state this in the paper. Because we already compare MixMax to DoReMi, we believe this addresses the question of comparisons to near-optimal convex optimization algorithms, but we are happy to discuss this further with the reviewer.

We have also added detail on the computational complexity of EMixMax. Furthermore, in the discussion below we clarify the differences with past parametric convex-concave optimization methods and our setting, and now cite the mentioned papers. We hope the reviewer considers raising their score given these changes.

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.

To clarify how we proved Theorem 3.1, note Sion’s minimax theorem requires a topology that is compact over the space we are minimizing. In the case of optimization over $\mathbb{R}^n$, this is satisfied by restricting the set to be bounded and closed. However, for $L^{\infty}$, a bounded and closed subset is not compact in the respective $L^{\infty}$ norm; in fact, the subset would not be compact in any $L^p$ norm (as mentioned in the paper). Our first contribution in the paper is making the connection that the weak*-topology over bounded functions is compact (by Banach-Alaoglu’s theorem), and furthermore, the group DRO objective is continuous in the weak*-topology (shown by the method of test functions and $L^p$ space duality). Hence, Sion’s minimax theorem can apply over the set of bounded functions by the existence of a suitably topology.

To the best of our knowledge, no prior minimax theorem over bounded functions existed. Furthermore, our paper provides several other contributions to the theory and practice of group DRO over bounded functions which builds on this first result of saddle points existing (see discussion to questions below).

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

Thank you for the references, we have added them to the related work section. To clarify the differences in setup, all of the above papers consider convex-concave optimization over a subset of $\mathbb{R}^n$ where saddle-points were already known to exist, and devise algorithms to jointly converge to the saddle point (in both the minimization parameter and maximization parameter). The key difference is our optimization is over bounded functions, where apriori saddle-points were not known to exist. Our first step was to prove saddle-points exist for our problem.

To briefly clarify the “second step” described by the reviewer, another difference between our setup and the cited papers is that even if a saddle-point exists for our problem, we cannot readily evaluate the function space gradients of our objective (e.g., the literature on gradient boosting to approximate these gradients). Hence, we can not directly apply the gradient based methodology provided in previous papers to obtain our saddle-points. We bypass this by showing that for two common loss functions, the mixture for the saddle-point satisfies a concave maximization over a finite-dimensional simplex. We elaborate on how we show this in response to the next question.

We thank the reviewer for their time and feedback! We discuss the weaknesses and questions raised by the reviewer below, and believe we have addressed the reviewer’s primary concern regarding the proof of concavity by adding a final step to the proof which was implicit before. We ask the reviewer to consider raising their score.

The theoretical proofs are quite brief and unclear. It would be beneficial to discuss and clarify the proofs step by step.

We clarify the proofs of concavity in our response to the questions below, and have added more detail to those proofs in the paper.

I believe the VC dimension of the class of bounded functions is unbounded. Could this present a challenge for solving real-world problems?

Yes, in theory this is true, by using a classical analysis the VC dimension of bounded functions would show up as an error in our ability to learn the $\hat{f_p}$. However, our experiments suggest that MixMax is an improvement on baselines in real-world settings when we learn over highly expressive model classes (e.g., XGBoost).

It is claimed that in Equation 1 the objective is concave with respect to λ. However, the convexity of $f_λ$ 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λ(y|x), it is not exactly equal to this function.

To clarify, we showed the MixMax objective is concave in $\lambda$: see our response to the next question for an elaboration on the proof. This is distinct from the Bayes optimal function $f_{\lambda}$ being concave or convex in $\lambda$.

To prove that the function $f_{\lambda}$ is concave with respect to $\lambda$, it must be shown that $f_{αλ_1+(1−α)λ_2}≥ αf_{\lambda_1}+(1−α)f_{\lambda_2}$. However, in Equation 3, this is not demonstrated. Instead, concavity with respect to p is proven in that equation.

Our proof of concavity to the distribution $dp$ implies the concavity condition you stated for the MixMax objective, but this was not explicit in the writing. We have now added this step to the proofs in the appendix.

Briefly, we had shown $MixMax(\sum \lambda_i p_{i}) \geq \sum \lambda_i MixMax(p_i)$. Defining $p_1 = p_{\lambda_1}$ and $p_2 = p_{\lambda_2}$, this implies $MixMax(p_{\alpha \lambda_1 + (1-\alpha)\lambda_2}) = MixMax(\alpha p_{\lambda_1} + (1-\alpha)p_{\lambda_2}) \geq \alpha MixMax(p_{\lambda_1}) + (1-\alpha)MixMax(p_{\lambda_2})$. Hence the Mixmax objective is concave in the mixture weights.

Thank you for your response.

Dear authors, I appreciate your answer and your clarification regarding the concavity of the objective function when $f_\lambda$ is equal to the Bayes optimal solution.

I would like to clarify that my comment about the proofs being short and unclear is not limited to the proof of concavity that you mentioned. For a paper that primarily contributes theoretical insights, I find your mathematical proofs to be brief and compressed, lacking sufficient detail (for example, in the proof of Theorem 3.1).

Additionally, your explanation regarding the VC dimension of the class of bounded functions did not convince me. If the VC dimension is unbounded, it suggests that the model may not generalize well; specifically, the empirical minimax may not converge around its expected value uniformly. The reason you may be obtaining good experimental results could be because you have parameterized the function class in a way that makes it more restricted, which in turn reduces the VC dimension. However, this implies that your class may not be general enough to match the performance of the Bayes optimal classifier, as you assumed in your proofs. Consequently, as I mentioned in my second question, this undermines your proof of the concavity of the objective function. While you have demonstrated that the objective function is concave when $f_\lambda$ equals the Bayes optimal solution, it must also be shown that your new restricted class can achieve this solution.

We thank the reviewer for their time and feedback! Below we answered the question raised by the reviewer. If there are no additional concerns we ask the reviewer to consider raising their score.

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.

The Bayes optimal function is not the same for every distribution, hence a single f won’t minimize every distribution. For a specific example, consider binary classification over a single input. Take $P = \{p_1,p_2\}$ where $p_1(y=1) = 0.9$, $p_2(y=1) = 0.2$. Then the Bayes optimal function are $f_1(1) = 0.9$ and $f_2(1) = 0.2$ and MixMax would give $f_{\lambda}(1) = 0.5$ with optimal MixMax weights being $\lambda = [3/7, 4/7]$. Notice $f_{\lambda}$ is not optimal for $p_1$ or $p_2$.

We look forward to your thoughts on our response. Let us know if there is anything more we can do to address your comments!

Thanks again for you time,

Authors

We thank the reviewer for their time and interesting questions! We discuss the weaknesses and questions raised by the reviewer below. Your comments have helped us improve the paper. In particular, we have now added experiments directly comparing the two ways of using MixMax weights.

Emperical2MixMax 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 Emperical2 MixMax estimate?

We agree that understanding how different approximations of $f_p(x)$ impact mixture finding is interesting and important to study. Towards this, in figure 10 (in the Appendix) we studied how obtaining the approximations $\hat{f_p}(x)$ by training with fewer samples (and hence with less compute) affected the final mixture performance. We found performance converged quickly with samples: training with ¼ of the dataset gave mixtures comparable to using the whole dataset. We hope future work characterizes more broadly under what approximation errors to $f_p$ we can recover close to optimal mixture weights. For example, what happens if we use smaller models instead of less samples to save compute?

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

This is a great point. The only baseline applicable in Section 6.3 (non-parametric models) is data balancing. It is not fair to compare compute in this setting, as data balancing uses 0 compute to return its mixture weight, but it can be arbitrarily suboptimal.

For the parametric problems in Section 6.2, there are algorithmic baselines available and MixMax uses comparable compute to them. For example, both DoReMi and Group DRO require training one model on all the data, whereas we require training three models on one-third of the data each: up to hyperparameter tuning the number of training steps (e.g., for early-stopping), this is the same cost.

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?

Thanks for raising this question, this is in fact a typo. We reweighed the training loss for XGBoost according to the MixMax weights or balanced weights. This is now described in more detail in Appendix C.3.

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 ?

This is an interesting question! We have added results comparing the two ways of using the MixMax weights for the tabular experiments, shown in Table 2 in the Appendix. We observed training a new model performs better or as good as ensembling the proxies across all the tabular experiments. Retraining comes with additional training cost, while ensembling comes with more inference cost, so we leave it to practitioners to decide if the improved performance offsets the varying costs.

We look forward to your thoughts on our response. Let us know if there is anything more we can do to address your comments!

Thanks again for you time,

Authors