Computing Wasserstein Barycenters through Gradient Flows

2025-10-06

AI Summary

Existing discrete Wasserstein barycenter algorithms suffer from poor scalability and require full-sample access. This paper proposes a novel modeling framework based on the Wasserstein gradient flow, reformulating barycenter computation as an energy minimization problem in the space of probability measures—naturally incorporating geometric structure and enabling explicit energy-based regularization. The method employs a minibatch sampling scheme, drastically reducing computational and memory overhead. Convergence is theoretically guaranteed via analysis leveraging the Polyak–Łojasiewicz inequality. Experiments on synthetic datasets and domain adaptation tasks demonstrate superior accuracy, efficiency, and robustness compared to state-of-the-art discrete methods and neural-network baselines. Key contributions include: (i) the first scalable gradient-flow paradigm for Wasserstein barycenters; (ii) interpretable, energy-driven regularization; and (iii) establishing a new standard for minibatch-based Wasserstein barycenter computation.

Technology Category

• Machine Learning: Learning with Manifolds
• Search and Optimization: Mixed Discrete/Continuous Search
• Constraint Satisfaction and Optimization: Mixed Discrete/Continuous Optimization

Application Category

• Graph Algorithms and Modeling for the Web: Efficient manipulation of static and dynamic Web-related graphs
• Web Mining and Content Analysis: Robustness and generalizability of Web mining methods
• Systems and Infrastructure for Web, Mobile and WoT: Experiences and lessons learnt from Web-based algorithms and system deployments

Abstract

Wasserstein barycenters provide a powerful tool for aggregating probability measures, while leveraging the geometry of their ambient space. Existing discrete methods suffer from poor scalability, as they require access to the complete set of samples from input measures. We address this issue by recasting the original barycenter problem as a gradient flow in the Wasserstein space. Our approach offers two advantages. First, we achieve scalability by sampling mini-batches from the input measures. Second, we incorporate functionals over probability measures, which regularize the barycenter problem through internal, potential, and interaction energies. We present two algorithms for empirical and Gaussian mixture measures, providing convergence guarantees under the Polyak-Łojasiewicz inequality. Experimental validation on toy datasets and domain adaptation benchmarks show that our methods outperform previous discrete and neural net-based methods for computing Wasserstein barycenters.

Problem

• Scalable computation of Wasserstein barycenters via gradient flows
• Overcoming limitations of discrete methods requiring full sample sets
• Regularizing barycenters with energy functionals for improved performance

Innovation

• Reformulates barycenter problem as gradient flow
• Uses mini-batches for scalable sampling approach
• Incorporates energy functionals for regularization