[Paper Review] Continuous Regularized Wasserstein Barycenters

Lingxiao Li, Aude Genevay+2 more|arXiv (Cornell University)|Aug 28, 2020

Geometric Analysis and Curvature Flows8 citations

TL;DR

This paper proposes a stochastic algorithm for computing continuous regularized Wasserstein barycenters by leveraging a novel dual formulation, enabling efficient online approximation of barycenters for continuous probability distributions using sample access. The method uses primal-dual relationships via dual potentials to implicitly parametrize the barycenter, achieving strong duality and enabling stochastic gradient descent for scalable, continuous barycenter estimation.

ABSTRACT

Wasserstein barycenters provide a geometrically meaningful way to aggregate probability distributions, built on the theory of optimal transport. They are difficult to compute in practice, however, leading previous work to restrict their supports to finite sets of points. Leveraging a new dual formulation for the regularized Wasserstein barycenter problem, we introduce a stochastic algorithm that constructs a continuous approximation of the barycenter. We establish strong duality and use the corresponding primal-dual relationship to parametrize the barycenter implicitly using the dual potentials of regularized transport problems. The resulting problem can be solved with stochastic gradient descent, which yields an efficient online algorithm to approximate the barycenter of continuous distributions given sample access. We demonstrate the effectiveness of our approach and compare against previous work on synthetic examples and real-world applications.

Motivation & Objective

• To overcome the computational difficulty of computing Wasserstein barycenters for continuous distributions, which previous methods restricted to finite supports.
• To develop a continuous approximation of the barycenter that maintains geometric meaningfulness while being practically computable.
• To enable online learning of barycenters from streaming or sampled data by leveraging a dual formulation of the regularized problem.
• To establish strong duality in the regularized Wasserstein barycenter problem to support implicit parametrization via dual potentials.

Proposed method

• Derives a new dual formulation for the regularized Wasserstein barycenter problem, enabling a primal-dual relationship that links the barycenter to dual potentials.
• Parametrizes the barycenter implicitly using the dual potentials of regularized optimal transport problems, avoiding explicit support discretization.
• Applies stochastic gradient descent to optimize the dual variables, allowing online and scalable updates from sample access to input distributions.
• Uses strong duality to ensure the primal solution (the barycenter) is consistently recovered from the optimized dual potentials.
• Employs a regularized transport framework to stabilize optimization and ensure smoothness in the barycenter approximation.
• Constructs a continuous barycenter estimate that adapts incrementally to new samples, enabling real-time or streaming computation.

Experimental results

Research questions

• RQ1 Can a continuous approximation of the Wasserstein barycenter be efficiently computed without restricting support to finite point sets?
• RQ2 How can strong duality in the regularized optimal transport problem be leveraged to implicitly parametrize the barycenter?
• RQ3 Can stochastic gradient descent be effectively applied to the dual formulation to enable online learning of the barycenter from samples?
• RQ4 How does the proposed method compare in accuracy and efficiency to prior finite-support approaches on both synthetic and real-world data?

Key findings

• The proposed method achieves strong duality in the regularized Wasserstein barycenter problem, enabling reliable recovery of the primal barycenter from dual variables.
• The implicit parametrization via dual potentials allows for continuous barycenter estimation without explicit support discretization.
• Stochastic gradient descent on the dual formulation enables an efficient online algorithm that scales to large or streaming datasets.
• The method demonstrates superior performance in continuous barycenter approximation compared to finite-support baselines on synthetic and real-world data.