Variational Wasserstein Barycenters with c-Cyclical Monotonicity

TL;DR

This paper introduces a novel variational approach to compute Wasserstein barycenters efficiently in high-dimensional settings, using c-cyclical monotonicity and stochastic optimization.

Contribution

It develops a continuous approximation method with a variational distribution, enabling tractable dual formulation and efficient computation of Wasserstein barycenters.

Findings

Effective in high-dimensional and continuous settings

Theoretical convergence guarantees

Successful application to real data and synthetic examples

Abstract

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.