Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Sobolev Gradient Ascent for Optimal Transport: Barycenter Optimization and Convergence Analysis

Kaheon Kim, Bohan Zhou, Changbo Zhu, Xiaohui Chen

TL;DR

This work tackles the exact Wasserstein barycenter computation for grid-supported distributions by formulating a constraint-free concave dual and developing a Sobolev gradient ascent in the $dot{H}^1$ geometry. By avoiding $c$-concavity projection steps, the proposed SGA method achieves global convergence at rates matching Euclidean subgradient methods while maintaining tractable per-iteration costs on grids. The authors establish strong duality, provide a coordinate-wise gradient scheme, and prove convergence guarantees for the barycenter optimization, with empirical results showing superior accuracy and efficiency over existing exact barycenter solvers in 2D, 3D, and real-data video scenarios. The approach is especially advantageous for exact barycenter computation on moderate-dimensional grids, offering a simpler and more stable alternative to entropic regularization or MMOT-based methods, though scalability to very high dimensions remains an open challenge.

Abstract

This paper introduces a new constraint-free concave dual formulation for the Wasserstein barycenter. Tailoring the vanilla dual gradient ascent algorithm to the Sobolev geometry, we derive a scalable Sobolev gradient ascent (SGA) algorithm to compute the barycenter for input distributions supported on a regular grid. Despite the algorithmic simplicity, we provide a global convergence analysis that achieves the same rate as the classical subgradient descent methods for minimizing nonsmooth convex functions in the Euclidean space. A central feature of our SGA algorithm is that the computationally expensive $c$-concavity projection operator enforced on the Kantorovich dual potentials is unnecessary to guarantee convergence, leading to significant algorithmic and theoretical simplifications over all existing primal and dual methods for computing the exact barycenter. Our numerical experiments demonstrate the superior empirical performance of SGA over the existing optimal transport barycenter solvers.

Sobolev Gradient Ascent for Optimal Transport: Barycenter Optimization and Convergence Analysis

TL;DR

This work tackles the exact Wasserstein barycenter computation for grid-supported distributions by formulating a constraint-free concave dual and developing a Sobolev gradient ascent in the geometry. By avoiding -concavity projection steps, the proposed SGA method achieves global convergence at rates matching Euclidean subgradient methods while maintaining tractable per-iteration costs on grids. The authors establish strong duality, provide a coordinate-wise gradient scheme, and prove convergence guarantees for the barycenter optimization, with empirical results showing superior accuracy and efficiency over existing exact barycenter solvers in 2D, 3D, and real-data video scenarios. The approach is especially advantageous for exact barycenter computation on moderate-dimensional grids, offering a simpler and more stable alternative to entropic regularization or MMOT-based methods, though scalability to very high dimensions remains an open challenge.

Abstract

This paper introduces a new constraint-free concave dual formulation for the Wasserstein barycenter. Tailoring the vanilla dual gradient ascent algorithm to the Sobolev geometry, we derive a scalable Sobolev gradient ascent (SGA) algorithm to compute the barycenter for input distributions supported on a regular grid. Despite the algorithmic simplicity, we provide a global convergence analysis that achieves the same rate as the classical subgradient descent methods for minimizing nonsmooth convex functions in the Euclidean space. A central feature of our SGA algorithm is that the computationally expensive -concavity projection operator enforced on the Kantorovich dual potentials is unnecessary to guarantee convergence, leading to significant algorithmic and theoretical simplifications over all existing primal and dual methods for computing the exact barycenter. Our numerical experiments demonstrate the superior empirical performance of SGA over the existing optimal transport barycenter solvers.
Paper Structure (24 sections, 10 theorems, 31 equations, 5 figures, 2 tables, 5 algorithms)

This paper contains 24 sections, 10 theorems, 31 equations, 5 figures, 2 tables, 5 algorithms.

Table of Contents

  1. Introduction
  2. Literature review
  3. Contribution and limitation
  4. Notations
  5. Preliminary: optimal transport computation
  6. Constraint-free Sobolev gradient ascent
  7. Wasserstein barycenter computation
  8. Sobolev gradient ascent for barycenter computation
  9. Empirical study
  10. 2D synthetic distribution
  11. 3D interpolation
  12. Real data example
  13. Technical Details
  14. Technical Lemmas
  15. Proof to \ref{['thm:onestep']}
  16. ...and 9 more sections

Key Result

Theorem 1

Let $\{ f^{(t)} \}_{t=1}^T$ be the sequence computed from Algorithm alg:onestep. Assuming that $\mu$ and $\nu$ are absolutely continuous with respect to the Lebesgue measure, and $\{ f^{(t)} \}_{t=1}^T$ are continuous on $\Omega$. Then, where $M = \max_{t=1}^T \| \mu - (T_{(f^{(t-1)})^c})_{\#} \nu \|_{\dot{\mathbb{H}}^{-1}}$, $\tilde{f}$ is any $c$-concave maximizer and $f^{(\mathrm{best})}$ is

Figures (5)

  • Figure 1: Empirical convergence rates of parallel (Algorithm \ref{['alg:barycenter']}), sequential (Algorithm \ref{['alg:barycenter(sequential)']}), and random (Algorithm \ref{['alg:barycenter(random)']}) SGA algorithms.
  • Figure 2: Comparison of weighted Waserstein barycenter of densities supported on different shapes.
  • Figure 3: Interpolation using SGA between a 3D ball and a 3D cube.
  • Figure 4: Left half displays 16 surveillance video frames. Right half displays their baryceters computed by CWB (top left), DSB (top right), WDHA (bottom left) and SGA (bottom right).
  • Figure 5: Comparison of weighted Waserstein barycenter of densities supported on different shapes.

Theorems & Definitions (12)

  • Theorem 1: Convergence rate for SGA
  • Remark 2
  • Lemma 3
  • Theorem 4: Strong duality and barycenter characterization
  • Lemma 5: Dual gradient
  • Theorem 6: Main theorem: convergence rate for barycenter optimization with SGA
  • Remark 7
  • Lemma 8
  • Lemma 9
  • Lemma 10
  • ...and 2 more