Optimal Transport Barycenter via Nonconvex-Concave Minimax Optimization

2025-01-24

AI Summary

To address the low accuracy and poor scalability of Wasserstein barycenter computation for high-dimensional discrete probability distributions, this paper proposes WDHA—an algorithm that achieves the first exact and efficient solution to the unregularized Wasserstein barycenter problem. Methodologically, WDHA introduces an alternating Wasserstein–Sobolev doubly geometric optimization framework and devises the first nonconvex–concave minimax algorithm with *O*(*m* log *m*) time and *O*(*m*) space complexity. It integrates Wasserstein gradient flow dynamics, Ḣ¹ dual ascent, and primal–dual alternating iterations, grounded in discrete grid-based density function modeling. Theoretically, we prove convergence to a stable stationary point. Empirically, on 1024×1024 image-scale 2D data, WDHA significantly outperforms Sinkhorn-based methods, achieving breakthroughs in accuracy, computational speed, and scalability.

Technology Category

• Search and Optimization: Mixed Discrete/Continuous Search
• Constraint Satisfaction and Optimization: Mixed Discrete/Continuous Optimization
• Reasoning under Uncertainty: Stochastic Optimization

Application Category

• Systems and Infrastructure for Web, Mobile and WoT: Experiences and lessons learnt from Web-based algorithms and system deployments
• 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

Abstract

The optimal transport barycenter (a.k.a. Wasserstein barycenter) is a fundamental notion of averaging that extends from the Euclidean space to the Wasserstein space of probability distributions. Computation of the unregularized barycenter for discretized probability distributions on point clouds is a challenging task when the domain dimension $d>1$. Most practical algorithms for approximating the barycenter problem are based on entropic regularization. In this paper, we introduce a nearly linear time $O(m log{m})$ and linear space complexity $O(m)$ primal-dual algorithm, the Wasserstein-Descent $dot{mathbb{H}}^1$-Ascent (WDHA) algorithm, for computing the exact barycenter when the input probability density functions are discretized on an $m$-point grid. The key success of the WDHA algorithm hinges on alternating between two different yet closely related Wasserstein and Sobolev optimization geometries for the primal barycenter and dual Kantorovich potential subproblems. Under reasonable assumptions, we establish the convergence rate and iteration complexity of WDHA to its stationary point when the step size is appropriately chosen. Superior computational efficacy, scalability, and accuracy over the existing Sinkhorn-type algorithms are demonstrated on high-resolution (e.g., $1024 imes 1024$ images) 2D synthetic and real data.

Problem

• High-dimensional Data
• Optimal Transport Barycenter
• Computational Efficiency

Innovation

• WDHA algorithm
• optimal transport barycenter
• high-dimensional data