Source-linked AI summary
Near-linear time approximation algorithms for optimal transport via Sinkhorn iteration
Jason Altschuler, Jonathan Weed, Philippe Rigollet
TL;DR
General optimal transport lacked a near-linear approximation guarantee despite practical algorithms for large datasets. This paper analyzes entropically regularized Sinkhorn iteration, derives a greedy Greenkhorn variant, and proves near-linear guarantees for both approaches, with Greenkhorn faster in practice.
Problem
General optimal transport distances had practical algorithms but no known near-linear time approximation guarantee for arbitrary cost matrices.
Method
The paper analyzes Sinkhorn iteration for entropic regularization, adds feasible rounding, and derives Greenkhorn by updating one row or column at a time.
Results
The paper proves the first near-linear time approximation guarantee for general nonnegative cost matrices, and Greenkhorn significantly outperforms Sinkhorn in practice.
Takeaways & Limitations
Sinkhorn Distances provide a theoretically justified near-linear approach to general discrete optimal transport, while Greenkhorn offers the same guarantee with better practical performance.
Takeaways & Limitations
The analysis requires η of order log n/ε, and runtime dependence on marginal imbalance can grow to Ω(n) when a marginal has a constant-sized entry.
Abstract
from arXiv · showhide
Computing optimal transport distances such as the earth mover's distance is a fundamental problem in machine learning, statistics, and computer vision. Despite the recent introduction of several algorithms with good empirical performance, it is unknown whether general optimal transport distances can be approximated in near-linear time. This paper demonstrates that this ambitious goal is in fact achieved by Cuturi's Sinkhorn Distances. This result relies on a new analysis of Sinkhorn iteration, which also directly suggests a new greedy coordinate descent algorithm, Greenkhorn, with the same theoretical guarantees. Numerical simulations illustrate that Greenkhorn significantly outperforms the classical Sinkhorn algorithm in practice.
1. INTRODUCTION
Optimal transport distances are important across machine learning, statistics, and computer vision, but general OT remained computationally expensive. This work addresses the open question of whether general OT can be approximated in near-linear time.
- Motivation: Optimal transport includes earth mover’s and Wasserstein distances used to compare probability measures and point clouds.These applications span machine learning, statistics, and computer vision.
- Motivation: Recent Sinkhorn Distances enabled large-scale applications by using entropic penalization, but lacked theoretical guarantees for general OT.The unresolved issue was whether this approach could approximate general OT in time nearly linear in the input size n^2.
- Related work: Prior general-purpose methods required eO(n^2.5) time, while faster alternatives relied on metric or low-dimensional ℓp assumptions.Orlin’s algorithm had O(n^3 log n) complexity, and earlier faster methods restricted the cost matrix structure.
- Contribution: The paper provides the first near-linear time guarantee for approximating OT with general nonnegative cost matrices.The result applies without requiring the cost matrix to arise from an underlying metric.
- Contribution: The contribution includes a new Sinkhorn analysis, parameter-tuning guidance, a greedy Greenkhorn variant, and a parallelizable rounding algorithm.The analysis yields iteration bounds independent of the matrix dimension n, while Greenkhorn significantly outperforms classical Sinkhorn in practice.
2. OPTIMAL TRANSPORT IN NEAR-LINEAR TIME
The proposed OT algorithm approximately projects an entropically transformed cost matrix onto the target marginal constraints, then rounds the result to feasibility. Its runtime inherits the near-linear complexity of the projection subroutine.
- Algorithm: Algorithm 1 forms A = exp(−ηC), approximately projects it toward the target marginal set, and rounds the result to a feasible transport plan.The approximate projection may violate feasibility, so rounding is required before output.
- Algorithm: The projection and rounding stages are implemented by separate subroutines, Proj and round.Sinkhorn and Greenkhorn provide alternative implementations of Proj, while round uses a simple procedure analyzed separately.
- Rounding: The rounding procedure rescales rows and columns before adding a correction that restores the prescribed marginals.Its correction uses the residual row and column errors.
- Guarantee: Algorithm 1 runs in O(n^2+S) time, where S is the runtime of the projection subroutine.When ∥C∥∞≤L, the stated bound is O(n^2L^3(log n)ε^−3).
- Guarantee: Accounting for bit complexity gives O(n^2L^4(log n)^2ε^−4), which remains truly near-linear.The arithmetic-only theorem ignores exponential bit complexity, whose maximum is bounded during the iterations.
3. LINEAR-TIME APPROXIMATE SINKHORN PROJECTION
The paper analyzes approximate Sinkhorn projection in total variation, obtaining stronger iteration bounds and extending the analysis to the greedy Greenkhorn algorithm.
- Approximate Sinkhorn projection: The Sinkhorn projection is the unique minimizer Pη = XAY of the entropic program, with diagonal scaling matrices X and Y unique up to a constant factor.The scaling representation follows from first-order optimality conditions.
- Approximate Sinkhorn projection: The entropic-penalized OT algorithm computes an approximate projection of A = exp(−ηC), then rounds it into the feasible set Ur,c.The projection may violate the target marginals by at most ε′ in ℓ1 before rounding.
- Sinkhorn analysis: Sinkhorn alternately rescales all rows and columns to match r and c, but prior analyses controlled convergence in ℓ2, which is inadequate for probability distributions.Distributions can have small ℓ2 distance while having disjoint support, whereas ℓ1 distance is more natural for statistical applications.
- Sinkhorn analysis: Theorem 2 gives an ℓ1 approximation guarantee for Sinkhorn, with a dimension-independent iteration bound beyond the uniform case and a runtime stronger than the prior bound by up to a factor of n.The extra factor n in the earlier result arose from converting an ℓ2 bound to an ℓ1 bound.
- Sinkhorn analysis: The new analysis uses a Kullback–Leibler potential whose decrease equals the divergence to the current marginals, and Pinsker’s inequality converts this control into an ℓ1 guarantee.This yields an iteration bound k* ≤ 4ε′−2 log(s/ℓ) for Sinkhorn.
- Greenkhorn: Greenkhorn updates one row or column greedily with an exact line search, using O(n) entry updates per iteration instead of Sinkhorn’s O(n2).Its guarantee is O(n(ε′)−2 log(s/ℓ)) iterations, and experiments report better practical performance than Sinkhorn.
- Greenkhorn: The new Sinkhorn analysis transfers to Greenkhorn with only trivial modifications, providing an explicit near-linear-time guarantee unavailable from earlier Sinkhorn analyses.The greedy update selects the best row or column at each step.
4. PROOF OF THEOREM 1
The proof bounds the error from approximate projection and rounding, then derives the algorithm’s runtime by analyzing projection and rounding costs.
- The approximate projection output B is rounded into the transport polytope, with the rounding error controlled by marginal violations.The rounding procedure outputs a feasible matrix and bounds its ℓ1 distance from B using the input and target marginals.
- The proof compares B with an optimal transport solution by viewing B as the entropy-regularized optimum for its own approximate marginals.Because B = XAY with positive diagonal scalings, it solves the regularized problem over the marginals induced by B.
- Rounding and projection therefore combine to produce the approximation guarantee after substituting Algorithm 1’s prescribed η and ε′ values.The proof concludes by plugging those parameter choices into the preceding error bounds.
- The cost gap is bounded by 2η^-1 log n plus a term proportional to the marginal violations and ∥C∥∞.The entropy terms are bounded by 2 log n, while Hölder’s inequality controls the cost contribution from marginal discrepancies.
- Both Sinkhorn and Greenkhorn require O(n2(ε′)^-2(log n + η∥C∥∞)) time for approximate projection, while rounding costs O(n2).The bound uses s ≤ n2 and log(1/ℓ) = η∥C∥∞ for A = exp(−ηC).
5. EMPIRICAL RESULTS
Experiments evaluate Sinkhorn and Greenkhorn for earth mover’s distance on MNIST and tunable synthetic images. Greenkhorn is reported to outperform Sinkhorn, especially for sparse data and small η.
- 5. EMPIRICAL RESULTS: The experiments compute earth mover’s distance between normalized grayscale images using pixel-location ℓ1 distances as transport costs.The target marginals represent image pairs, and the cost matrix contains pairwise pixel-location distances.
- 5.1 MNIST: On 10 random MNIST image pairs, Greenkhorn significantly outperforms Sinkhorn in both the short and long term.Figure 2 compares their convergence behavior using the approximate projection step.
- 5.2 Random images: The synthetic framework varies foreground-square area across 20%, 50%, and 80% to test both algorithms in different regimes.Background intensities are drawn from [0, 1], while foreground intensities are drawn from [0, 50].
- 5.2 Random images: Greenkhorn performs better as the number of background pixels increases, consistent with updating only salient rows and columns.The authors identify sparse data as a setting where Greenkhorn is especially advantageous because Sinkhorn updates background-associated rows and columns.
- 5.2 Random images: For small constant η, Greenkhorn runs faster than Sinkhorn with no loss in accuracy.Figure 3 also compares median competitive ratios across synthetic images with different foreground proportions.
A.1 Proof of Lemma 3
The proof establishes the first inequality using a divergence argument between successive iterates.
- The first inequality is proved by an argument analogous to the proof of Lemma 2.
- The divergence K(A(1)∥A(0)) treats the matrices as elements of the simplex ∆n2.
- For the second claim, A(0) has unit ℓ1 norm and smallest entry ℓ/s, so its Sinkhorn projection exists uniquely.The proof invokes positivity and Sinkhorn’s uniqueness result before introducing the corresponding scaling factors.
- The argument then uses the probability-vector properties of r and c to derive the required bound.
A.2 Proof of Lemma 5
The proof reduces the row-update case to changes localized to one row and one scaling coordinate; the column case is identical.
- Only the row-update case needs to be proved because the column-update case is exactly symmetric.
- After the update, A′ and A′′ differ only in row I, while x′′ and x′ differ only in entry I and y′′ equals y′.
- The proof then uses the fact that the updated row marginal satisfies rI(A′′) = rI.
A.3 Proof of Lemma 6
The proof combines nonnegativity, a bounded scalar inequality, and Pinsker’s inequality to derive an ℓ1 bound, then concludes the claim from the resulting displays.
- s −1 −log s and K(α∥¯β) are nonnegative, and when ρ(α, β) ≤1, s −1 −log s ≥(s −1)2/5.
- Pinsker’s inequality converts the preceding scalar and divergence terms into a bound involving (s−1)2 and ∥α −¯β∥2 1.
- The triangle inequality and convexity bound the discrepancy by 7/5(s−1)2+ 7/2∥α−¯β∥2 1.
- The lemma’s claim follows from the two derived displays.
A.4 Proof of Lemma 7
The proof analyzes mass removed during row-then-column rescaling, bounds the resulting error using marginal discrepancies, and establishes an O(n2) runtime bound.
- The rounded matrix G is nonnegative because F′′ and the row and column errors are nonnegative, with both error norms equal to 1−∥F′′∥1.
- The total removed mass ∆ is tracked across row and column rescaling steps to bound the ℓ1 difference between F and G.
- The proof separately bounds the two sums in equation (5), using that c(F) is entrywise larger than c(F′) for the second sum.
- The resulting relation is c ∥1/∥errr∥1 = 2∆+ 1 −∥F∥1 ≤∥r(F) −r∥1 + 2∥c(F) −c∥1.
- O(n2) runtime follows by observing the cost of each rescaling and computing the matrix errrerr⊤.
A.5 Randomized variant of rounding algorithm (Algorithm 2)
The rounding analysis is symmetrized by randomly choosing row-then-column or column-then-row removal, producing an improved bound; the section also compares Greenkhorn with Stochastic Sinkhorn.
- Randomized rounding: The original rounding guarantee bounds the ℓ1 difference between F and G, while the order of row and column removal creates an asymmetry between marginal errors.
- Randomized rounding: With probability 1/2, the randomized variant performs Algorithm 2; otherwise, it removes column mass before row mass to obtain an improved bound.
- Empirical comparison: The empirical comparison evaluates Greenkhorn against Stochastic Sinkhorn, which uses stochastic averaged gradient optimization of the dual entropic-penalty problem.
- Empirical comparison: Greenkhorn greedily selects a row or column and performs an exact closed-form line search, whereas Stochastic Sinkhorn uses a different coordinate-descent method.
- Empirical comparison: The experiment uses random images with 20% foreground pixels, initializes both algorithms with the same primal solution, and rounds iterates to the feasible polytope.
- Empirical comparison: Stochastic Sinkhorn is tested with C = 1/(Ln), 3/(Ln), and 5/(Ln), while Greenkhorn is tested with two values of η.