Optimal Transport Barycenter via Nonconvex-Concave Minimax Optimization

TL;DR

This work tackles the efficient computation of unregularized Wasserstein barycenters on discretized grids by recasting the problem as a nonconvex-concave minimax optimization. The proposed WDHA algorithm alternates between $\dot{\mathbb{H}}^1$-gradient ascent on Kantorovich potentials and Wasserstein-gradient updates for the barycenter, enabling near-linear time per Kantorovich update and linear space usage. Under mild density-bounds assumptions and careful step-size choices, WDHA provably converges to a stationary point of a regularized barycenter functional $\mathcal F_{\alpha,\beta}$, with an $O(1/T)$ rate for the averaged gradient norm. Empirically, WDHA outperforms Sinkhorn-based methods on high-resolution 2D data, producing sharper barycenters and substantial speedups, and scales to 3D cases with reasonable compute times, highlighting its practical impact for large-scale OT problems.

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 \times 1024$ images) 2D synthetic and real data.

Figures (5)

• Figure 1: Illustration of Wasserstein barycenters computed by different methods. The goal is to compute the barycenter of four uniform densities supported on the square, circle, heart, and cross, respectively, as displayed in the top left image. The blended shape shown in the top middle image is the barycentric density computed using our method. Barycentric densities computed using CWB and DSB with $\text{reg}=0.005$, and their thresholded versions are shown in the top right image and the bottom three images.
• Figure 2: Top row displays three exemplary digit 8 images. Bottom row displays barycenters computed by different methods using 300 iterations.
• Figure 3: Illustration of Wasserstein barycenters computed using WDHA, CWB and DSB.
• Figure 4: Illustration of Wasserstein barycenters computed using WDHA for 3D distributions. The results are smoothed using a Gaussian filter.
• Figure 5: Plot of output from Algorithm \ref{['alg:c3']} to uniform distributions on round disks.

Theorems & Definitions (15)

• Lemma 3.1: Strong concavity and smoothness of $\mathcal{I}_\nu^\mu$
• Lemma 3.2
• Definition 3.3
• Lemma 3.4: Lipschitzness of Wasserstein gradient $\nabla\!\!\!\!\nabla\mathcal{L}^\mu$
• Lemma 3.5: Smoothness of $\mathcal{L}^\mu$
• Theorem 3.6: Convergence rate of WDHA
• Remark 3.7
• Lemma 2.1
• proof
• proof : Proof of Lemma \ref{['lem:1']}