Training on mixtures of data distributions is now common in many modern machine learning pipelines, useful for performing well on several downstream tasks. Group distributionally robust optimization (group DRO) is one popular way to learn mixture weights for training a specific model class, but group DRO methods suffer for non-linear models due to non-convex loss functions and when the models are non-parametric. We address these challenges by proposing to solve a more general DRO problem, giving a method we call MixMax. MixMax selects mixture weights by maximizing a particular concave objective with entropic mirror ascent, and, crucially, we prove that optimally fitting this mixture distribution over the set of bounded predictors returns a group DRO optimal model. Experimentally, we tested MixMax on a sequence modeling task with transformers and on a variety of non-parametric learning problems. In all instances MixMax matched or outperformed the standard data mixing and 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.

We're looking to hire an experienced data engineer to maintain, improve and continue to add to Mixmax's data infrastructure and data culture. You'll be a key communicator, working cross functionally in our entirely remote organization. You'll have the opportunity to work with multiple stakeholders, technologies, and data sources. As a data engineer you'll be a part of a team specifically focused on ensuring the company runs on accurate and repeatable data. This means being a member of a team that values continuous and collective learning, culture over process, data driven development and always asking tons of questions.