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

TL;DR

This paper presents a scalable Sobolev gradient ascent algorithm for Wasserstein barycenter computation, offering convergence guarantees without complex projections and demonstrating superior empirical performance.

Contribution

It introduces a novel constraint-free dual formulation and a simplified, scalable gradient ascent method tailored to Sobolev geometry for barycenter optimization.

Findings

SGA achieves convergence rates comparable to classical subgradient methods.

The algorithm eliminates the need for costly $c$-concavity projections.

Numerical experiments show SGA outperforms existing barycenter solvers.

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 discretized over 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.