An Efficient Algorithm For Computing Optimal Wasserstein Ball Center

Yuntao Wang, Yuxuan Li, Qingyuan Yang, Hu Ding

Keywords: Wasserstein barycenter, model ensembling, fairness

TL;DR: We introduce and solve a new type of probability center that enhences the fairness performance of WB.

Abstract: Wasserstein Barycenter (WB) is a fundamental problem in machine learning, whose objective is to find a representative probability measure that minimizes the sum of its Wasserstein distance to given distributions. WB has a number of applications in various areas. However, in some applications like model ensembling, where it aggregates predictions of different models on the label space, WB may lead to unfair outcome towards underrepresented groups (e.g., a "minority'' distribution may be far away from the obtained WB under Wasserstein distance). To address this issue, we propose an alternative objective called ``Wasserstein Ball Center (WBC)''. Specifically, WBC is a distribution that encompasses all input distributions within the minimum Wasserstein distance, which can be formulated as a minmax optimization problem. We show that the WBC problem with fixed support is equivalent to solving a large-scale linear programming (LP) instance, which is quite different from the previous LP model for WB. By incorporating some novel observations on the induced normal equation, we propose an efficient algorithm that accelerates the interior point method by $O(Nm)$ times ($N$ is the number of distributions and $m$ is the support size). Finally, we conduct a set of experiments on both synthetic and real-world datasets. We demonstrate the computational efficiency of our algorithm, and showcase its better accuracy on model ensembling under heterogeneous data distributions.

Primary Area: optimization

Code Of Ethics: I acknowledge that I and all co-authors of this work have read and commit to adhering to the ICLR Code of Ethics.

Submission Guidelines: I certify that this submission complies with the submission instructions as described on https://iclr.cc/Conferences/2025/AuthorGuide.

Reciprocal Reviewing: I understand the reciprocal reviewing requirement as described on https://iclr.cc/Conferences/2025/CallForPapers. If none of the authors are registered as a reviewer, it may result in a desk rejection at the discretion of the program chairs. To request an exception, please complete this form at https://forms.gle/Huojr6VjkFxiQsUp6.

Anonymous Url: I certify that there is no URL (e.g., github page) that could be used to find authors’ identity.

No Acknowledgement Section: I certify that there is no acknowledgement section in this submission for double blind review.

Submission Number: 9372