MixMax: Distributional Robustness in Function Space via Optimal Data Mixtures
This addresses robustness across user groups in machine learning, offering a more efficient solution for non-convex or non-parametric settings, though it is incremental as it builds on existing group DRO methods.
The paper tackles the challenge of group distributionally robust optimization (group DRO) when losses are non-convex or models are non-parametric by reparameterizing it from parameter space to function space, showing that MixMax matches or outperforms standard baselines, improving XGBoost performance on datasets like ACSIncome and CelebA.
Machine learning models are often required to perform well across several pre-defined settings, such as a set of user groups. Worst-case performance is a common metric to capture this requirement, and is the objective of group distributionally robust optimization (group DRO). Unfortunately, these methods struggle when the loss is non-convex in the parameters, or the model class is non-parametric. Here, we make a classical move to address this: we reparameterize group DRO from parameter space to function space, which results in a number of advantages. First, we show that group DRO over the space of bounded functions admits a minimax theorem. Second, for cross-entropy and mean squared error, we show that the minimax optimal mixture distribution is the solution of a simple convex optimization problem. Thus, provided one is working with a model class of universal function approximators, group DRO can be solved by a convex optimization problem followed by a classical risk minimization problem. We call our method MixMax. In our experi ments, we found that MixMax matched or outperformed the standard group DRO baselines, and in particular, MixMax improved the performance of XGBoost over the only baseline, data balancing, for variations of the ACSIncome and CelebA annotations datasets.