A Particle-Flow Algorithm for Free-Support Wasserstein Barycenters

2025-09-14

AI Summary

This work addresses three key limitations of entropy-regularized methods for computing Wasserstein barycenters: solution blurring, support restriction, and deviation from the formal Riemannian geometry of the Wasserstein space. We propose a free-support particle-based flow algorithm grounded in the intrinsic Riemannian structure of the Wasserstein manifold. Particle evolution is driven by the average optimal transport displacement field, and barycenter projection—rather than Monge mapping—is employed to avoid entropy regularization entirely, thereby preserving solution sharpness and geometric fidelity. The method is theoretically guaranteed to be stable, convergent, and resolution-invariant. Empirical evaluation demonstrates superior accuracy and scalability over state-of-the-art linear programming and regularized solvers across diverse tasks, including probabilistic averaging, Bayesian posterior aggregation, image prototype extraction, and large-scale clustering.

Technology Category

• Machine Learning: Learning with Manifolds
• Search and Optimization: Sampling/Simulation-based Search
• Computer Vision: Other Foundations of Computer Vision

Application Category

• Graph Algorithms and Modeling for the Web: Efficient manipulation of static and dynamic Web-related graphs
• Responsible Web: Human-perceived consequences of algorithmic deployment on the web
• Systems and Infrastructure for Web, Mobile and WoT: Experiences and lessons learnt from Web-based algorithms and system deployments

Abstract

The Wasserstein barycenter extends the Euclidean mean to the space of probability measures by minimizing the weighted sum of squared 2-Wasserstein distances. We develop a free-support algorithm for computing Wasserstein barycenters that avoids entropic regularization and instead follows the formal Riemannian geometry of Wasserstein space. In our approach, barycenter atoms evolve as particles advected by averaged optimal-transport displacements, with barycentric projections of optimal transport plans used in place of Monge maps when the latter do not exist. This yields a geometry-aware particle-flow update that preserves sharp features of the Wasserstein barycenter while remaining computationally tractable. We establish theoretical guarantees, including consistency of barycentric projections, monotone descent and convergence to stationary points, stability with respect to perturbations of the inputs, and resolution consistency as the number of atoms increases. Empirical studies on averaging probability distributions, Bayesian posterior aggregation, image prototypes and classification, and large-scale clustering demonstrate accuracy and scalability of the proposed particle-flow approach, positioning it as a principled alternative to both linear programming and regularized solvers.

Problem

• Computing Wasserstein barycenters without entropic regularization
• Developing geometry-aware particle-flow algorithm for barycenters
• Avoiding Monge maps through barycentric projection technique

Innovation

• Particle-flow algorithm for barycenters
• Avoids entropic regularization approach
• Uses barycentric projections instead maps