[Paper Review] Faster Wasserstein Distance Estimation with the Sinkhorn Divergence

Lénaïc Chizat, Pierre Roussillon+3 more|arXiv (Cornell University)|Jun 15, 2020

Geometric Analysis and Curvature Flows57 references54 citations

TL;DR

The paper proposes estimating W2^2 between smooth densities using the Sinkhorn divergence, which enables higher entropic regularization and yields faster computation with comparable sample complexity, plus a Richardson-extrapolated variant for improved efficiency.

ABSTRACT

The squared Wasserstein distance is a natural quantity to compare probability distributions in a non-parametric setting. This quantity is usually estimated with the plug-in estimator, defined via a discrete optimal transport problem which can be solved to $ε$-accuracy by adding an entropic regularization of order $ε$ and using for instance Sinkhorn's algorithm. In this work, we propose instead to estimate it with the Sinkhorn divergence, which is also built on entropic regularization but includes debiasing terms. We show that, for smooth densities, this estimator has a comparable sample complexity but allows higher regularization levels, of order $ε^{1/2}$, which leads to improved computational complexity bounds and a strong speedup in practice. Our theoretical analysis covers the case of both randomly sampled densities and deterministic discretizations on uniform grids. We also propose and analyze an estimator based on Richardson extrapolation of the Sinkhorn divergence which enjoys improved statistical and computational efficiency guarantees, under a condition on the regularity of the approximation error, which is in particular satisfied for Gaussian densities. We finally demonstrate the efficiency of the proposed estimators with numerical experiments.

Motivation & Objective

• Motivate efficient estimation of W2^2 between densities from sampled or discretized data.
• Introduce Sinkhorn divergence as a debiased entropically regularized OT estimator for W2^2.
• Analyze statistical and computational trade-offs of regularization and debiasing.
• Develop and analyze Richardson-extrapolated variants to improve accuracy and efficiency.
• Provide theoretical bounds and numerical validation for the proposed estimators.

Proposed method

• Define the Sinkhorn divergence S_lambda as T_lambda(mu,nu) - 1/2(T_lambda(mu,mu) + T_lambda(nu,nu)).
• Introduce the Richardson-extrapolated estimator R_lambda = 2 S_lambda(mu,nu) - S_{sqrt{2} lambda}(mu,nu).
• Show that |S_lambda(mu,nu) - W2^2(mu,nu)| = O(lambda^2) under regularity, with a second-order expansion.
• Provide sample complexity bounds comparing plug-in, S_lambda, and R_lambda estimators under sampling and discretization.
• Analyze computational complexity of Sinkhorn iterations and provide practical guidelines for lambda choice.
• Discuss Richardson-extrapolated estimator performance under abstract regularity assumptions and in Gaussian cases.

Experimental results

Research questions

• RQ1 How does the Sinkhorn divergence approximate the squared Wasserstein distance for smooth densities?
• RQ2 What are the statistical and computational trade-offs when using entropic regularization and debiasing for estimating W2^2?
• RQ3 Can Richardson extrapolation yield estimators with improved bias-variance trade-offs for estimating W2^2?
• RQ4 How do the proposed estimators perform under random sampling versus grid discretization?
• RQ5 What are practical and theoretical guarantees for computation time and accuracy in high dimensions?

Key findings

• Sinkhorn divergence provides a debiased estimator whose bias decays as lambda^2, enabling larger regularization than T_lambda while preserving accuracy.
• For independent samples, S_lambda achieves comparable sample complexity to the plug-in estimator but with improved computational guarantees due to larger permissible lambda.
• Richardson extrapolation (R_lambda) can further reduce bias and improve efficiency under suitable regularity, with explicit bounds in Gaussian cases.
• When marginals are discretized on grids, entropy regularization improves stability and yields faster error decay than the unregularized problem, under Lipschitz log-densities.
• Empirical results confirm faster computation and favorable accuracy for S_lambda and R_lambda compared to the standard plug-in and T_lambda approaches.