Source-linked AI summary
Computational Optimal Transport: Complexity by Accelerated Gradient Descent Is Better Than by Sinkhorn's Algorithm
Pavel Dvurechensky, Alexander Gasnikov, Alexey Kroshnin
TL;DR
The paper studies computational methods for approximating general OT distances, including settings where the transportation plan's sparsity matters. It analyzes Sinkhorn-based and APDAGD-based approaches, reporting complexity bounds better than the state-of-the-art result and broader regularizer applicability for APDAGD.
Problem
Existing Sinkhorn-based OT computation is tied to entropic regularization, while applications may require sparse transportation plans from more general regularizers.
Method
The paper analyzes Sinkhorn's algorithm and a novel Adaptive Primal-Dual Accelerated Gradient method for approximating general OT distances.
Results
Both analyzed algorithms achieve complexity bounds better than the state-of-the-art result.
Takeaways & Limitations
APDAGD is applicable beyond entropic regularization and can address strongly convex problems with linear constraints.
Takeaways & Limitations
Sinkhorn's algorithm and other mentioned algorithms are specific to entropic regularization.
Abstract
from arXiv · showhide
We analyze two algorithms for approximating the general optimal transport (OT) distance between two discrete distributions of size $n$, up to accuracy $\varepsilon$. For the first algorithm, which is based on the celebrated Sinkhorn's algorithm, we prove the complexity bound $\widetilde{O}\left({n^2/\varepsilon^2}\right)$ arithmetic operations. For the second one, which is based on our novel Adaptive Primal-Dual Accelerated Gradient Descent (APDAGD) algorithm, we prove the complexity bound $\widetilde{O}\left(\min\left\{n^{9/4}/\varepsilon, n^{2}/\varepsilon^2 \right\}\right)$ arithmetic operations. Both bounds have better dependence on $\varepsilon$ than the state-of-the-art result given by $\widetilde{O}\left({n^2/\varepsilon^3}\right)$. Our second algorithm not only has better dependence on $\varepsilon$ in the complexity bound, but also is not specific to entropic regularization and can solve the OT problem with different regularizers.
1. Introduction
The paper targets more accurate and flexible computation of discrete optimal transport distances, improving complexity bounds beyond prior methods. It analyzes Sinkhorn’s algorithm and introduces APDAGD for strongly convex regularizers beyond entropy.
- OT distances support applications including unsupervised and semi-supervised learning, clustering, text classification, image retrieval, and statistics.
- Motivation: The state-of-the-art approach applies Sinkhorn’s algorithm to entropy-regularized OT, but its cubic dependence on ε prevents accurate approximation.
- Motivation: General strongly convex regularization can produce sparse transportation plans, but Sinkhorn’s algorithm is specific to entropic regularization.
- Contributions: The paper seeks better complexity bounds for OT distance approximation and a flexible solver supporting different regularizers.
- Contributions: The contributions include improved Sinkhorn analysis, APDAGD complexity results, and numerical illustrations of practical algorithm performance.
- Contributions: APDAGD uses linesearch and accelerated convergence for the primal objective and constraint feasibility with a general strongly convex regularizer.
2. Sinkhorn’s Algorithm
The section refines Sinkhorn’s convergence analysis and applies it to approximate the OT distance, obtaining a new complexity bound with improved dependence on ε.
- Improved iteration analysis: The refined analysis links the regularized dual objective to marginal feasibility at each Sinkhorn iteration.This relationship is used to sharpen the iteration estimate.
- Improved iteration analysis: The main proof innovation adds a first component to the iteration bound, whereas earlier analyses used only the second component.The added component follows from the objective–feasibility relation in Lemma 2.
- Complexity: Theorem 2 establishes the resulting approximation guarantee for the output plan, and the authors report better dependence on ε and ∥C∥∞ than prior work.The analysis also theoretically justifies replacing zero measure values with small positive values.
- OT approximation procedure: The procedure regularizes with γ = ε/(4 ln n), enforces positive perturbed marginals, applies Sinkhorn, and projects the result onto U(r, c).The projection step produces a feasible transport plan satisfying the required approximation guarantee.
- Complexity: Each Sinkhorn iteration requires O(n^2) arithmetic operations, so the total complexity follows by multiplying the iteration bound by n^2.Finding the perturbed marginals costs O(n) time.
3. Accelerated Gradient Descent
This section develops APDAGD for linearly constrained strongly convex problems and applies it to regularized OT, yielding a flexible accelerated approach with improved complexity dependence on ε.
- General method: APDAGD solves linearly constrained minimization problems through adaptive primal-dual accelerated gradient updates.Its convergence analysis separately controls primal objective residual and constraint feasibility.
- General method: The algorithm uses linesearch and an online stopping criterion based on duality gap and constraint infeasibility.It does not require a known Lipschitz constant or an estimate of the dual-radius parameter R.
- General method: Theorem 3 guarantees convergence under strong convexity and a bounded minimum-norm dual solution, while validating the stopping criterion.The algorithm outputs primal and dual iterates when both objective and feasibility tolerances are met.
- Application to OT: APDAGD applies to entropy, squared Euclidean, and other strongly convex regularizers.Entropy regularization uses the ℓ1 norm, while Euclidean regularization uses the Euclidean norm for strong convexity.
- Application to OT: The OT procedure applies APDAGD to the entropy-regularized problem, projects iterates onto U(r, c), and stops when the objective and projection criteria are satisfied.The cost matrix is normalized and the regularization is set to γ = ε/(3 ln n).
- Implementation: The APDAGD procedure can be parallelized and efficiently implemented when the Sinkhorn kernel exp(−C/γ) is easy to apply.The paper gives regular-grid squared-Euclidean costs as an example.
- Complexity: Each APDAGD iteration and projection costs O(n^2) arithmetic operations, producing the stated total complexity bound.The paper compares this bound with Sinkhorn and reports improved dependence on ε and ∥C∥∞ at an n^1/4 cost.
4. Experiments
The experiments compare Sinkhorn’s algorithm and APDAGD on resized MNIST image pairs, measuring working time across accuracy ε and problem size n.
- Experimental setup: MNIST images were resized from 28 × 28 to 28·s × 28·s pixels to vary the number of pixels n.Zero-valued measure elements were replaced with 10^-6 and the measures were normalized.
- Experimental setup: For varying ε, the experiments used s = 1, several ε values from 0.025 to 0.12, and averages over 10 randomly selected image pairs.Both algorithms were run until their respective stopping criteria were fulfilled.
- Dependence on problem size: For varying n, the experiments fixed ε = 0.1, used s values from 1 to 8, and averaged results over 5 randomly selected image pairs.Figure 1 (right) illustrates working time for different n.
- Comparison: Figure 1 compares the working time of Algorithm 2, Sinkhorn’s algorithm, and Algorithm 4, APDAGD.
5. Conclusion
The conclusion summarizes two algorithms for approximating general OT distances between discrete distributions and reports their complexity bounds. It also emphasizes APDAGD’s broader applicability beyond entropic regularization.
- Conclusion: The paper analyzes two algorithms for approximating general OT distances between two discrete distributions.
- Conclusion: The first algorithm uses entropic regularization and Sinkhorn’s algorithm.
- Conclusion: The second algorithm uses entropic regularization and the novel Adaptive Primal-Dual Accelerated Gradient method.
- Conclusion: Both complexity bounds are reported as better than the state-of-the-art result.
- Conclusion: APDAGD is not specific to entropic regularization, incorporates line search, and has an accelerated convergence rate.
Supplementary Material
The supplementary material gives additional derivations, implementation details, and proofs supporting the paper’s Sinkhorn- and APDAGD-based analyses.
- Supplementary Material: The supplement provides details on the dual problem for regularized OT used in the Sinkhorn analysis.
- Supplementary Material: It describes an efficient APDAGD implementation when the Sinkhorn kernel is easy to apply, such as on regular grids with squared Euclidean costs.
- Supplementary Material: The supplement includes missing proofs for the APDAGD-based approach.
1 Details for the Sinkhorn's Algorithm Approach
This section derives the dual problem for regularized OT and explains the variable transformation used to obtain the formulation analyzed in Section 2.
- Dual derivation: The dual problem for the regularized OT problem is derived for use in Section 2.
- Variable transformation: The derivation changes variables using u = −y/γ − 1/2 and v = −z/γ − 1/2, then disregards a constant and rescales the objective.These transformations produce the dual problem considered in Section 2.
2.1 E cient Implementation for Entropic Regularization
APDAGD can reuse the Sinkhorn kernel K through matrix-vector products, making its entropic-regularization steps efficient and parallelizable under the same favorable kernel structure.
- APDAGD computes its steps using multiplication of the Sinkhorn kernel K = exp(−C/γ) by vectors.This applies to the entropic-regularized OT specialization.
- Kernel application can be accelerated when the measures lie on regular grids and C is squared Euclidean distance.
- The gradient calculation applies K to ey/γ and ez/γ, while evaluating ϕ(λ) also uses K.Both operations are described as easy when the kernel is easy to apply.
- APDAGD uses the same set of operations as Sinkhorn’s algorithm and can therefore be implemented in a parallel framework.The remaining operations involve dual vectors of size 2n.
2.2 Adaptive Primal-Dual Accelerated Gradient Descent (APDAGD) Convergence Analysis
The paper supplies convergence-rate proofs for APDAGD, viewing it as accelerated gradient descent on a dual problem with primal-iterate reconstruction.
- The appendix provides missing convergence-rate proofs for APDAGD applied to constrained convex optimization.
- The underlying constrained problem uses a γ-strongly convex objective, while its Lagrange dual is written as a minimization problem.
- The Lagrange-multiplier space Λ is generally unbounded.
- APDAGD applies Adaptive Accelerated Gradient Descent to the dual problem and reconstructs a primal iterate.
2.3 Adaptive Accelerated Gradient Descent for Convex Optimization
This section develops Adaptive Accelerated Gradient Descent with adaptive smoothness handling, proximal geometry, and convergence analysis designed to support primal-dual applications.
- AAGD is analyzed for convex functions ϕ with L-Lipschitz-continuous gradients over a generally unbounded closed convex set Λ.
- The method uses a norm, a prox-function d, and the associated Bregman divergence V[ζ](λ).The prox-function is required to be 1-strongly convex with respect to the chosen norm.
- The framework accommodates standard Euclidean, entropy, ℓ1/ℓ2, simplex, nuclear-norm, and spectahedron proximal setups.
- AAGD avoids requiring the Lipschitz constant L as an explicit algorithm input and proves convergence in a primal-dual-friendly form.This supports application to dual problems with primal-iterate reconstruction.
- The convergence rate is established for the iterates ηk, with inner-cycle work bounded in terms of k and log2 factors.Each inner check requires two oracle calls.
- The algorithm adaptively doubles its estimate Mk until the smoothness check succeeds, then updates Lk+1 = Mk/2.The inner cycle is finite because doubling eventually reaches the gradient’s Lipschitz scale.
- The convergence analysis bounds weighted function values using linearizations of ϕ and Bregman-divergence terms.
2.4 Adaptive Primal-Dual Accelerated Gradient Descent for Constrained Convex Optimization
APDAGD extends the adaptive dual method to strongly convex constrained optimization, reconstructing primal iterates and stopping when objective and feasibility errors meet prescribed tolerances.
- APDAGD applies AAGD to a constrained convex problem with γ-strongly convex objective f and linear constraint operator A.
- Strong convexity makes the dual objective ϕ smooth, with its gradient determined through the unique solution x(λ) of a strongly convex subproblem.
- The method reconstructs an approximate primal solution while optimizing the dual problem, using a Euclidean proximal setup in the multiplier space.
- The algorithm stops when the primal-dual objective residual and equality-constraint violation satisfy εf and εeq thresholds.
- Under a dual-solution bound ∥λ∗∥2 ≤ R, Theorem 2 provides guarantees for the reconstructed primal points and dual iterates.
- The inner cycle uses at most 4k + 4 + 2 log2 iterations after iteration k.