Towards Universal Wasserstein Barycenters through Flow Matching
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.
Related
- Generalized Wasserstein Flow Matching: Nested Couplings for Point Cloud Generation
- Diffusion Flow Matching: Dimension-Improved KL Bounds and Wasserstein Guarantees
- The p-Bures-Wasserstein distance via 3×3 block matrices (with application to p-Bures-Wasserstein barycenters)
- Sliced Wasserstein Barycenters: The Analysis Approach within the Barycentric Coding Model in the Wasserstein Space
- Capacity-Constrained Wasserstein Barycenters: Existence, Duality, and Entropic Regularization
BibTeX
@misc{montesuma2026towards,
title = {{Towards Universal Wasserstein Barycenters through Flow Matching}},
author = {Eduardo Fernandes Montesuma},
howpublished = {arXiv},
year = {2026},
month = sep,
doi = {10.48550/arxiv.2609.38547},
eprint = {2609.38547},
archivePrefix = {arXiv},
url = {https://arxiv.org/abs/2609.38547},
note = {CC0},
}
RIS
TY - UNPB
TI - Towards Universal Wasserstein Barycenters through Flow Matching
AU - Eduardo Fernandes Montesuma
PY - 2026
DA - 2026/09/29
PB - arXiv
DO - 10.48550/arxiv.2609.38547
UR - https://arxiv.org/abs/2609.38547
AB - 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.
ER -