[Paper Review] Sinkhorn Barycenter via Functional Gradient Descent

Zebang Shen, Zhenfu Wang+2 more|arXiv (Cornell University)|Jul 20, 2020

Topological and Geometric Data Analysis22 references4 citations

TL;DR

This paper introduces Sinkhorn Descent (SD), a functional gradient descent method that reformulates the Sinkhorn barycenter problem as unconstrained optimization over mappings in a Reproducing Kernel Hilbert Space. SD achieves sublinear convergence to a stationary point at rate $\mathcal{O}(1/t)$, preserves weak convergence of empirical measures, and scales linearly with dimension, enabling solution of 100-dimensional barycenter problems—previously intractable with prior methods.

ABSTRACT

In this paper, we consider the problem of computing the barycenter of a set of probability distributions under the Sinkhorn divergence. This problem has recently found applications across various domains, including graphics, learning, and vision, as it provides a meaningful mechanism to aggregate knowledge. Unlike previous approaches which directly operate in the space of probability measures, we recast the Sinkhorn barycenter problem as an instance of unconstrained functional optimization and develop a novel functional gradient descent method named Sinkhorn Descent (SD). We prove that SD converges to a stationary point at a sublinear rate, and under reasonable assumptions, we further show that it asymptotically finds a global minimizer of the Sinkhorn barycenter problem. Moreover, by providing a mean-field analysis, we show that SD preserves the weak convergence of empirical measures. Importantly, the computational complexity of SD scales linearly in the dimension $d$ and we demonstrate its scalability by solving a $100$-dimensional Sinkhorn barycenter problem.

Motivation & Objective

• To address the computational intractability of high-dimensional Sinkhorn barycenter problems due to exponential support growth in prior methods.
• To eliminate bias from fixed-support assumptions in barycenter computation by enabling free support optimization.
• To develop a provably convergent, scalable method for computing Sinkhorn barycenters in high-dimensional spaces.
• To ensure weak convergence of empirical measures under the algorithmic iterates for robustness to data sampling.

Proposed method

• Reformulate the Sinkhorn barycenter problem as unconstrained functional optimization over push-forward mappings in a Reproducing Kernel Hilbert Space (RKHS).
• Define the Kernelized Sinkhorn Barycenter Discrepancy (KSBD) to measure non-asymptotic convergence of the iterates.
• Use functional gradient descent to iteratively update the mapping that minimizes the Sinkhorn divergence, with step size chosen to ensure descent.
• Employ kernel methods to represent and optimize the transport map, enabling smooth, differentiable updates without explicit support constraints.
• Prove that the algorithm preserves weak convergence of empirical measures under the iterates, ensuring consistency with continuous initialization.
• Leverage mean-field analysis to establish convergence guarantees under mild assumptions, including global convergence to the optimal solution.

Experimental results

Research questions

• RQ1 Can the Sinkhorn barycenter problem be reformulated as an unconstrained functional optimization problem over mappings, avoiding support constraints?
• RQ2 Does functional gradient descent on the space of mappings achieve sublinear convergence to a stationary point with rate $\mathcal{O}(1/t)$?
• RQ3 Can the algorithm preserve weak convergence of empirical measures under finite-sample initialization?
• RQ4 Under what conditions does the functional gradient descent method converge globally to the true Sinkhorn barycenter?
• RQ5 Is the computational complexity of the method linear in the dimension $d$, enabling scalability to high-dimensional problems?

Key findings

• Sinkhorn Descent (SD) converges to a stationary point at a sublinear rate of $\mathcal{O}(1/t)$, where $t$ is the iteration count.
• Under mild assumptions, SD asymptotically finds a global minimizer of the Sinkhorn barycenter problem.
• The algorithm preserves weak convergence of empirical measures: $\lim_{N\to\infty} \texttt{SD}^t(\alpha_N) = \texttt{SD}^t(\alpha)$, ensuring consistency with data sampling.
• The computational complexity of SD scales linearly with dimension $d$, enabling solution of a 100-dimensional barycenter problem.
• The Kernelized Sinkhorn Barycenter Discrepancy (KSBD) vanishes at rate $\mathcal{O}(1/t)$, serving as a valid non-asymptotic convergence metric.
• SD outperforms prior methods in scalability and avoids the exponential dependence on dimension inherent in fixed-support or alternating optimization approaches.