Towards Universal Wasserstein Barycenters through Flow Matching

2026-09-29

AI Summary

This study addresses the computational difficulty and inefficiency of approximating universal barycenters for probability distributions with arbitrary weights on the Wasserstein simplex. To overcome this limitation, we propose BaryFM, a flow matching model that, for the first time, parameterizes the entire barycenter family across the Wasserstein simplex using a single network. By integrating an ordinary differential equation (ODE) solver, BaryFM efficiently transports marginal distributions to barycenters corresponding to arbitrary weights, thereby transcending the constraints of traditional fixed-weight computation and enabling continuous weight sampling. Extensive evaluations across ten benchmarks in four tasks, including domain adaptation, demonstrate that BaryFM achieves a superior average rank compared to fifteen baseline methods. Furthermore, its performance is comparable to that of specialized, non-universal solvers, highlighting its effectiveness as a general-purpose framework for Wasserstein barycenter computation.

Abstract

Defining a weighted mean over probability measures under probability metrics is a central tool in probabilistic machine learning. Under the Wasserstein metric, these are called \emph{Wasserstein barycenters}. While most approaches compute barycenters for a fixed weight vector, approximating the whole family of barycenters over the simplex, which we call the \emph{Wasserstein simplex}, remains underexplored. We refer to this problem as \emph{Universal Barycenter Approximation}, and propose \texttt{BaryFM}, a flow matching model transporting the marginal measures into any barycenter in the Wasserstein simplex. Once trained, the network can draw samples from measures in the Wasserstein simplex through an ordinary differential equation. We validate our method on 4 downstream tasks: domain adaptation, generalization, Bayesian posterior aggregation and algorithmic fairness. \texttt{BaryFM} achieves the best average rank among 15 competing methods across 10 domain adaptation benchmarks, matching or surpassing non-universal solvers.

Problem

• Wasserstein barycenters
• Universal Barycenter Approximation
• Wasserstein simplex
• probability measures

Innovation

• Wasserstein Barycenters
• Flow Matching
• Universal Barycenter Approximation
• Wasserstein Simplex
• Domain Adaptation