Multi-distribution Learning: From Worst-Case Optimality to Lexicographic Min-Max Optimality

Guanghui Wang, Umar Syed, Robert E. Schapire, Jacob Abernethy

Keywords: Multi-distribution Learning, Lexicographic Min-Max Optimization, Game Theory

TL;DR: Obtaining an (approximate) lex-min-max solution can be as easy as achieving (approximate) min-max optimality

Abstract: We study multi-distribution learning (MDL), where the goal is to train a model that performs well across a set of underlying groups or distributions. The predominant performance metric in MDL is min-max optimality, which guarantees minimum error on the ``hardest'' distribution presented. Optimizing this metric has the unfortunate side effect of producing models that potentially sacrifice performance gains on non-worst-case groups. In the present work we propose a natural alternative, the lexicographic min-max (lex-min-max) objective, which promotes balanced performance by sequentially minimizing the worst, second-worst, and subsequent group losses. Despite its non-convex nature, we show that obtaining an (approximate) lex-min-max solution can be as easy as achieving (approximate) min-max optimality. We develop an efficient algorithm that directly approximates lex-min-max optimality via implementing stochastic no-regret dynamics on a regularized variant of the classical min-max objective. Our method is efficient and easy to implement, and it advances the frontier of multi-distribution learning by providing stronger, hierarchy-aware performance guarantees.

Submission Number: 106