Variational Wasserstein Barycenters with c-Cyclical Monotonicity
This work addresses a computational bottleneck for researchers and practitioners using optimal transport in machine learning, offering a more efficient method for aggregating probability distributions, though it is incremental as it builds on existing variational and dual formulation techniques.
The paper tackles the computational burden of computing Wasserstein barycenters in high-dimensional continuous settings by introducing a variational distribution to approximate the true barycenter, framing it as an optimization problem solved via stochastic optimization, and demonstrates effectiveness on subset posterior aggregation and synthetic data.
Wasserstein barycenter, built on the theory of optimal transport, provides a powerful framework to aggregate probability distributions, and it has increasingly attracted great attention within the machine learning community. However, it suffers from severe computational burden, especially for high dimensional and continuous settings. To this end, we develop a novel continuous approximation method for the Wasserstein barycenters problem given sample access to the input distributions. The basic idea is to introduce a variational distribution as the approximation of the true continuous barycenter, so as to frame the barycenters computation problem as an optimization problem, where parameters of the variational distribution adjust the proxy distribution to be similar to the barycenter. Leveraging the variational distribution, we construct a tractable dual formulation for the regularized Wasserstein barycenter problem with c-cyclical monotonicity, which can be efficiently solved by stochastic optimization. We provide theoretical analysis on convergence and demonstrate the practical effectiveness of our method on real applications of subset posterior aggregation and synthetic data.