PaperlayerView exact revision on arXiv

Source-linked AI summary

Faster Wasserstein Distance Estimation with the Sinkhorn Divergence

Lenaic Chizat, Pierre Roussillon, Flavien Léger, François-Xavier Vialard, Gabriel Peyré

arXiv:2006.08172v2math.OCmath.STstat.ML

TL;DR

Estimating squared Wasserstein distance from empirical or discretized measures is statistically and computationally difficult. The paper replaces the biased plug-in estimator with debiased Sinkhorn divergence and Richardson extrapolation, showing comparable sample complexity with improved computational guarantees under smoothness and regularity assumptions. Numerical experiments support the efficiency of these estimators.

  • Problem

    The plug-in estimator for squared Wasserstein distance has computational difficulties and suffers from n^-2/d estimation error for densities.

  • Method

    The paper estimates squared Wasserstein distance using Sinkhorn divergence, analyzes random samples and uniform-grid discretizations, and studies Richardson extrapolation.

  • Results

    The proposed estimators have comparable or improved sample-complexity guarantees and improved computational guarantees relative to the plug-in estimator, with numerical experiments confirming their benefits.

  • Takeaways & Limitations

    Debiased entropic regularization can increase accuracy and efficiency for estimating squared Wasserstein distance when the problem has a smooth nature.

  • Takeaways & Limitations

    The Sinkhorn divergence estimator’s theoretical statistical bound is slightly weaker than the plug-in estimator’s, although experiments suggest at least comparable statistical efficiency.

Abstract

from arXiv · show

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.

1 Introduction

The paper studies efficient estimation of squared Wasserstein distance from empirical or discretized measures, replacing biased plug-in estimation with debiased entropic estimators. Sinkhorn divergence and Richardson extrapolation offer improved computational guarantees while retaining comparable or better supported statistical performance.

  • Squared Wasserstein distance compares probability densities in applications including generative modeling, two-sample testing, shape matching, and computer graphics.
  • The plug-in estimator suffers from n^-2/d estimation error for densities and computational cost scaling as n^2 log(n)/ε^2 with Sinkhorn’s algorithm.
  • Sinkhorn divergence debiases entropic optimal transport and approximates W 2 with error of order λ^2, versus λ log(1/λ) for the uncorrected quantity.
  • The plug-in estimator provides a baseline rate of n^-2/d, while Richardson-extrapolated estimators can achieve better computational and sample-complexity bounds under stronger regularity assumptions.The assumption is shown to hold for Gaussian densities.
  • For independent samples, the Sinkhorn divergence estimator has comparable sample complexity and improved computational guarantees when λ is properly chosen, with analogous analysis for uniform-grid discretizations.
  • The paper targets computational rather than statistical optimality, and numerical experiments confirm benefits of the proposed estimators.

2 Refined approximation bound for the Sinkhorn divergence

The section analyzes how debiased entropic optimal transport approximates squared Wasserstein distance. Under boundedness and regularity conditions, its approximation error is controlled at second order in the regularization parameter through Fisher information.

  • The dynamical formulation expresses entropic optimal transport through interpolating densities and velocity fields satisfying a continuity equation.Its objective combines kinetic energy with integrated Fisher information.
  • The Sinkhorn divergence approximates W 2 with error O(λ^2) for probability measures with bounded densities and supports.
  • The approximation bound depends on Fisher information along the Wasserstein geodesic connecting the two measures.The endpoint Fisher informations alone do not make this geodesic quantity straightforward to bound.
  • Bounds on the geodesic Fisher information follow from assumptions on the Brenier potential, including bounded, Lipschitz Hessian conditions; quadratic potentials give a simpler bound.
  • The Hessian-Lipschitz assumption is demanding and potentially non-sharp, although it can be avoided in one dimension and Gaussian cases admit an explicit expression.

3 Performance analysis of the Sinkhorn divergence estimator

The section compares the plug-in estimator with Sinkhorn-divergence estimators for random samples and grid-discretized densities. Sinkhorn divergence permits larger regularization, improving computational guarantees while retaining comparable statistical performance under stated assumptions.

  • 3.1 Analysis of the plug-in estimator: Sinkhorn’s algorithm estimates the regularized plug-in cost in O(n^2||c||_∞^2/(λε)) operations for ε-accuracy.Approximating the unregularized cost requires λ of order ε/log(n/α), yielding O(n^2 log(n/α)||c||_∞^2/ε^2) operations.
  • 3.2 Performance of the Sinkhorn divergence estimator given random samples: The Sinkhorn-divergence estimator uses a population approximation assumption |Sλ(µ,ν)−W2^2(µ,ν)| ≤ λ^2I and chooses λ according to the sample size.Its statistical bound uses d′ = 2⌊d/2⌋ and achieves accuracy ε ≍ n^-2/(d′+4) with probability 1−δ.
  • 3.2 Performance of the Sinkhorn divergence estimator given random samples: Sinkhorn divergence needs regularization of order ε^1/2 for ε-accuracy, compared with order ε for the plug-in estimator.This larger admissible regularization level is identified as the source of a computational advantage, although the available statistical bound is insufficient for the strongest theoretical comparison.
  • 3.2 Performance of the Sinkhorn divergence estimator given random samples: For random samples, both estimators achieve ε ≍ n^-2/(d′+4), but Sinkhorn divergence requires ˜O(n^2ε^-1.5) time versus ˜O(n^2ε^-2) for the regularized plug-in cost.The paper notes that the statistical exponent approaches the plug-in rate as dimension grows, while numerical experiments show greater efficiency in practice.
  • 3.3 Performance of the Sinkhorn divergence estimator given densities discretized on grids: The grid analysis obtains discretization error of order h^2 for the entropy-regularized problem, whereas order h is suggested by a naive analysis at λ = 0.The resulting bounds demonstrate the potential usefulness of the debiasing terms beyond random sampling.
  • 3.3 Performance of the Sinkhorn divergence estimator given densities discretized on grids: For Lipschitz continuous log-densities on grids, Sinkhorn divergence achieves ˜O(ε^-(3d/2+3/2)) time with h ≍ ε^3/4 and λ ≍ ε^1/2.The corresponding regularized plug-in guarantee is ˜O(ε^-(2d+2)) with h ≍ ε and λ ≍ ε/log(1/ε).
  • 3.3 Performance of the Sinkhorn divergence estimator given densities discretized on grids: The grid complexity comparisons are upper bounds rather than optimal guarantees and do not exploit convolution-based reductions in Sinkhorn iteration cost.A separate analysis cited in the paper gives ˜O(ε^-(2d+1)) for the regularized plug-in estimator in a similar setting.

4 Towards faster estimation with Richardson extrapolation

Richardson extrapolation cancels the leading λ² bias of the Sinkhorn divergence, yielding improved statistical and computational guarantees under a stronger regularity assumption, satisfied for Gaussian densities.

  • Richardson extrapolation: Richardson extrapolation removes the leading Fisher-information bias by combining Sinkhorn divergences at different regularization levels.When Sλ = S0 + λ²I + o(λ²), the extrapolated estimator satisfies Rλ = S0 + o(λ²).
  • Efficiency guarantees: Under the abstract assumption Sλ − W2² = λ²I + O(λ⁴), the extrapolated estimator achieves improved statistical and computational guarantees.The analysis uses λ ≍ n^−1/(d′+8) and bounds the approximation error by λ⁴.
  • Efficiency guarantees: With probability 1 − δ, the estimator reaches ε-accuracy using ˜O(ε^−(d′+11)/2) Sinkhorn computations.The ˜O notation hides poly-logarithmic factors in 1/δ.
  • Efficiency guarantees: The improvement is statistical for d > 8 and computational for d > 6 relative to the plug-in estimator.The stated thresholds follow from the abstract regularity assumption.
  • Gaussian case: For Gaussian densities, a Taylor expansion verifies the required regularity and gives an explicit Fisher-information expression.The second-order term is consistent with the general expansion of the Sinkhorn divergence.

5 Numerical experiments

Experiments evaluate dual-potential estimation under random sampling and grid discretization, showing that debiased estimators generally achieve lower errors with larger optimal regularization and favorable computational times.

  • Experimental setup: The experiments assess proposed estimators on synthetic problems, with reproducible code made available by the authors.The reported figures focus on dual-potential errors rather than only scalar Wasserstein errors.
  • Experimental setup: The experiments use L1 error on the estimated dual potential because scalar Wasserstein error can vanish spuriously as n or λ varies.The dual potential is the Lagrange multiplier associated with the first marginal constraint.
  • Random sampling: For random sampling in d = 5, Sλ and Rλ have smaller bias than Tλ and require fewer samples with larger λ for a given accuracy.Figure 2 varies n and λ, while Figure 3 optimizes n and λ for computational time.
  • Random sampling: For small target accuracies, Sλ and Rλ compare favorably to Tλ in computational time after optimizing over n and λ.The comparison selects the largest λ and smallest n that achieve the target L1 potential error.
  • Grid discretization: On one-dimensional grids, Sλ and Rλ have lower errors than Tλ at fixed (h, λ) and after selecting each method’s optimal λ⋆(h).Their optimal regularization parameters are systematically larger than Tλ’s.

6 Conclusion and open questions

The paper concludes that entropic regularization with debiasing can improve accuracy and efficiency for smooth problems, while noting gaps between theory and observed performance.

  • Conclusion: For smooth problems, entropic regularization with debiasing may increase both accuracy and computational efficiency in squared Wasserstein estimation.The conclusion concerns the proposed estimators’ supported smooth setting.
  • Open questions: Numerical experiments suggest that the Sinkhorn divergence estimator is at least as statistically efficient as the plug-in estimator, although the theoretical bound is slightly weaker.Estimation guarantees for Kantorovich potentials appear similar but are not covered by the theory.

A Bounds on the approximation error

The approximation error from entropic regularization is controlled by Fisher information along the Wasserstein geodesic. The paper derives bounds for this quantity, including an explicit Gaussian formula.

  • Bounds on Fisher information: Under Hessian bounds and Lipschitz regularity of the Brenier potential, I0(µ, ν) is bounded by endpoint Fisher information and a regularity term.For quadratic Brenier potentials, the bound simplifies to I0(µ, ν) ≤ 2κ^-1I0(µ, µ).
  • Gaussian case: For centered Gaussian measures, I0(µ, ν) = tr S^-1 with S = (A^1/2BA^1/2)^1/2.The resulting expansion contains 1/8(2 tr S^-1 − tr A^-1 − tr B^-1).

B Computational complexity of Sinkhorn’s algorithm

The paper analyzes Sinkhorn’s alternating updates and derives complexity bounds for approximating entropic and unregularized transport. Smaller regularization requires more iterations, while discrete updates cost O(n^2) each.

  • Algorithm: Sinkhorn’s algorithm alternates updates of the dual potentials u and v, with each update enforcing one marginal constraint.
  • Convergence analysis: The convergence analysis controls the dual-objective gap of Sinkhorn iterates through successive-update inequalities.
  • Interpretation: The complexity discussion uses the dependence on λ because smaller regularization makes the computational problem harder.
  • Complexity bounds: For discrete measures with n atoms, an ε-accurate estimate of Tλ takes time O(n^2∥c∥∞^2/(λε)).
  • Complexity bounds: Choosing λ = ε/[4(log(n) + log(1/α))] yields an ε-accurate estimate of unregularized transport in O(n^2 log(n)∥c∥∞^2/ε).

C Properties of the plug-in estimator

The plug-in estimator is analyzed through empirical-process and concentration bounds for discrete optimal transport. Its statistical and computational rates depend strongly on dimension, with a distinct regime at d = 4.

  • Statistical analysis: The expected plug-in estimation error is bounded using Rademacher complexity, Dudley’s chaining, and covering numbers for convex Lipschitz functions.
  • Concentration: The concentration bound applies to both regularized and unregularized transport when measures are supported on a set of diameter D.
  • Dimension-dependent rates: For d > 4, choosing n ≳ log(2/δ)^(d/4)ε^-d/2 leads to computational complexity ˜O(ε^-d−2).
  • Dimension-dependent rates: At d = 4, ε ≍ n^-1/2(log(n) + log(2/δ)) and the resulting complexity is O(ε^-6(log n + ...)).

D Analysis of the Sinkhorn divergence estimator given samples

For empirical measures, the Sinkhorn divergence matches the plug-in estimator’s statistical rate up to constants while permitting substantially larger regularization. This changes the resulting computational complexity in its favor.

  • Statistical rate: The Sinkhorn divergence estimator has the same statistical bound as the regularized transport estimator up to a constant.
  • Analysis: The analysis combines prior bounds for regularized transport with an error decomposition separating approximation, sampling, concentration, and debiasing terms.
  • Regularization choice: Optimizing the regularization gives λ ≍ n^-1/(d′+4) and estimation error n^-2/(d′+4).
  • Computational complexity: For the Sinkhorn divergence estimator, the resulting computational complexity is ˜O(ε^-d′−5.5).
  • Regularization choice: With target accuracy ε, the Sinkhorn divergence requires λ ≍ ε^-1/2, whereas the plug-in estimator requires λ ≲ ε/log(n).

E Analysis of deterministic discretization

The deterministic discretization analysis bounds how replacing smooth densities by uniform-grid measures affects entropic transport and Sinkhorn-divergence estimators. Optimizing the discretization and regularization scales yields an accuracy-dependent computational bound.

  • Discretization construction: Uniform-grid discretization assigns each cell’s mass to its closest grid point, producing a discrete measure with n = h^-d grid points.The construction uses finite-volume cell masses and labels the resulting grid points.
  • Discretization stability: For M-Lipschitz log-densities, the entropic transport-cost discretization error is bounded between -h^2(1 + M) and min{h, h^2(λ^-1 + M + 1)}.The constants hidden by ≲ depend only on d and the cost parameter C.
  • Discretization stability: The direct discretization argument gives Tλ(µh, νh) − Tλ(µ, ν) ≲ h for all λ ≥ 0, while regularity of the optimal plan yields a sharper h^2 λ^-1 + λ^2 bound.The sharper estimate combines first- and second-order cell errors with regularity of the regularized optimal density.
  • Domain and asymptotics: The torus analysis requires a modified heat-kernel cost to connect the compact-domain formulation with the Euclidean Wasserstein setting, with asymptotically exponentially small replacement error.The corresponding constants in one comparison may depend on µ and ν because the statement is purely asymptotic.
  • Sinkhorn-divergence estimator: For the Sinkhorn divergence, minimizing h^2 λ^-1 + λ^2 gives λ = h^2/3 and an O(h^4/3) error.In terms of target accuracy ε, the corresponding scales are h ≍ ε^3/4 and λ ≍ ε^1/2.
  • Computational complexity: The resulting computational complexity is O(ε^-3d/2−3/2), using n = h^-d and the entropic solver cost O(n^2 λ^-1 ε^-1).The uncorrected analysis instead uses λ ≍ ε/log(n) after balancing an O(λ log(n)) regularization term with O(h) discretization error.

F Analysis of the Gaussian case

The Gaussian analysis expands the entropic transport and Sinkhorn-divergence quantities using explicit covariance formulas. It identifies the unregularized limit with the squared Bures covariance distance and characterizes the regularization-error expansion.

  • Explicit Gaussian formulas: For Gaussian measures N(a, A) and N(b, B), the entropic transport cost admits an explicit formula in terms of means, covariances, matrix square roots, and log determinants.The regularized covariance term uses D_AB,λ = (A^1/2 B A^1/2 + λ^2 I/4)^1/2.
  • Unregularized limit: At λ = 0, the squared Wasserstein distance is ||a − b||^2 + tr(A) + tr(B) − 2tr(S), where S = (A^1/2BA^1/2)^1/2.The covariance contribution is the squared Bures distance between positive definite matrices.
  • Approximation error: The Sinkhorn-divergence approximation error can be expanded by combining fourth-order expansions of the Gaussian matrix and log-determinant terms.The resulting expression is simplified in terms of Bures distances after cancellation of log-determinant terms.
  • Approximation error: The expansion terms are nonzero unless A = B, and their signs are determined by the covariance relationship.This conclusion follows from the resulting Bures-distance representation.

G.1 Sampling method

The experiments use centered elliptically contoured distributions with randomly generated covariances and compactly supported radial sampling. They compare estimator errors and visualize discretized potentials, while noting that scalar cost-error minima can reflect sign changes.

  • Distributional setting: The random-sampling experiments use centered elliptically contoured distributions whose optimal transport cost and dual potential have closed forms determined by the covariances.The distributions are chosen from a common elliptically contoured family so the Brenier potential is quadratic and Fisher-information conditions apply.
  • Covariance generation: Covariances are generated from normalized Gaussian random matrices, producing trace-one positive definite matrices with controlled anisotropy at α = 1/3.The eigenvalue ratio is controlled with high probability for large d.
  • Sampling method: Samples are generated from a uniform sphere direction and a compactly supported radial variable R = α|arctan(Z/β)|^1/d, with β = 2 producing bell-shaped densities.The resulting law is centered, elliptically contoured, compactly supported, has covariance A, and has finite Fisher information when A is nonsingular.
  • Random-sampling results: In d = 10, the plug-in estimator follows the n^-2/d convergence rate, while Sλ and Rλ have smaller estimated-cost errors than the plug-in estimator in the reported experiment.Tλ has substantially larger bias than Sλ and Rλ in this comparison.
  • Interpretation of errors: The scalar cost-error minimum observed for Sλ and Rλ results from the error changing sign as n grows, so it should not be interpreted as uniform improvement.The paper therefore also reports L1 error on the potential, which is not subject to this sign-change phenomenon.
  • Discretization experiment: For large λ, the discretization experiment shows systematically smaller errors when debiasing terms are included.Figure 7 also tracks convergence of estimated dual potentials toward the unregularized potentials as h or λ decreases.
Loading 2006.08172v2…