PaperlayerView exact revision on arXiv

Source-linked AI summary

SinkSLOT: Sinkhorn via Sparse Lifted Optimal Transport

Ian Hsieh, Soumya Snigdha Kundu, Tom Vercauteren, Reuben Dorent

arXiv:2608.28262v1cs.LG

TL;DR

Dense EOT remains quadratic and uses an independent reference that can assign mass to higher-cost edges. SinkSLOT uses an expected lifted sliced OT plan to sparsify the Gibbs kernel, proving convergence and achieving substantial speedups while producing a divergence without debiasing.

  • Problem

    EOT remains limited by O(N^2) Sinkhorn iterations and an independent reference coupling that can bias mass toward higher-cost edges.

  • Method

    SinkSLOT replaces the dense independent reference with an expected lifted sliced OT plan and applies sparse Sinkhorn updates, with the SOT prior treated as fixed for differentiation.

  • Results

    SinkSLOT achieves substantial speedups over dense and sparse EOT methods while retaining accuracy comparable to dense and hardware-aware baselines.

  • Takeaways & Limitations

    The resulting divergence is a promising loss function for large-scale optimisation, as illustrated by a gradient flow experiment.

  • Takeaways & Limitations

    The slice count L is a trade-off: too few slices restrict the transport polytope, while too many reduce the sparsity advantage.

Abstract

from arXiv · show

Entropic optimal transport (EOT) has been shown to offer a computationally tractable approximation to exact optimal transport. However, the standard Sinkhorn-Knopp algorithm has two main limitations. First, given discrete measures with $N$ points, each iteration requires $O(N^2)$ operations, which restricts its use on large-scale datasets (e.g. $N\geq10^4$). Second, it uses the independent coupling as a reference measure for regularisation. This assigns mass to high-cost transport edges at moderate regularisation strengths. We propose SinkSLOT, which addresses both limitations by putting forth the expected sliced lifted transport plan as a natural way to sparsify the Gibbs kernel with a non-independent prior coupling. We prove that: 1) SinkSLOT converges; 2) with $L$ slices, each resulting sparse Sinkhorn iteration costs $O(LN)$; and 3) the resulting objective is a divergence requiring no debiasing. Experiments on synthetic benchmarks show that SinkSLOT delivers substantial speedups over state-of-the-art dense and sparse EOT methods. We also demonstrate the applicability of the proposed divergence in a gradient flow experiment. The code is publicly available at https://github.com/cai4cai/SinkSLOT.

Introduction

SinkSLOT targets the quadratic cost and cost-agnostic regularisation of dense EOT by using a sparse lifted sliced-transport prior. Its sparse Sinkhorn iterations reduce computation while retaining the improved coupling accuracy associated with SROT.

  • Motivation: Exact OT is computationally expensive for large discrete measures, while EOT reduces this cost but retains O(N^2) dense Sinkhorn iterations.Exact OT requires solving a costly optimisation problem, and dense EOT remains demanding for large datasets.
  • Related approaches: SOT is efficient through one-dimensional sorting but can underperform exact OT and does not directly provide an original-space transport plan.Lifting sliced matchings addresses the missing-plan issue, motivating sliced-regularised approaches.
  • Related approaches: SROT improves coupling accuracy over EOT but retains O(N^2) Sinkhorn iterations because smoothing destroys the sparsity of the lifted sliced plan.Its primary objective is coupling accuracy rather than preserving SOT’s computational efficiency.
  • Contribution: SinkSLOT restricts Sinkhorn iterations to the sparse support of the lifted SOT plan, reducing per-iteration complexity from O(N^2) to O(LN).The method replaces the dense independent reference with an expected lifted sliced OT plan.

Background and related work

This section reviews discrete OT, entropic regularisation, sparse-kernel alternatives, and sliced transport plans. It motivates SinkSLOT as a sparse, non-product-reference formulation that addresses dense computation and EOT’s biased discrepancy.

  • Optimal transport: Discrete OT seeks a cost-minimising coupling between measures with prescribed marginals, and its computational cost limits use in large-scale optimisation pipelines.The coupling is defined over support points with ground costs such as Cij = ∥xi − yj∥p.
  • Entropic optimal transport: EOT adds entropy regularisation to obtain a unique plan computable by Sinkhorn-Knopp, but dense iterations require O(N^2) time.Its reference coupling is the independent product a ⊗ b.
  • Entropic optimal transport: EOT regularisation toward a ⊗ b can produce dense, diffuse matchings with non-negligible mass on higher-cost edges at moderate ε.The standard EOT cost is also biased because OTε(µ, µ) is generally nonzero.
  • Sparsification techniques: Sparse-kernel methods reduce Sinkhorn iteration costs, but some require O(NM) preprocessing or impose sparsity independently of a transport-informed reference.These methods include adaptive truncation, locality-sensitive hashing, and importance sampling.
  • Sliced transport: SOT averages one-dimensional optimal transport over random projections, while lifted plans recover couplings on the original support with sparse support.Each one-dimensional plan contributes at most N + M − 1 nonzeros, and the aggregate support is much smaller than NM when L is small.
  • Sliced-regularised OT: SROT uses the lifted SOT plan as a more informative reference and improves coupling accuracy over EOT, but smoothing destroys sparsity and leaves O(NM) iterations.Its smoothed reference is γP_SOT + (1 − γ)a ⊗ b with γ > 0.

SinkSLOT: Sinkhorn via Sparse Lifted OT

SinkSLOT solves entropic transport with an unsmoothed lifted SOT reference, yielding a sparse Gibbs kernel and convergent sparse Sinkhorn iterations. Its complexity scales with the lifted support, while its objective is a divergence without debiasing and interpolates between sliced and support-restricted OT behaviour.

  • Problem formulation: SinkSLOT replaces the smoothed SROT reference with the unsmoothed expected sliced transport plan P_SOT.The resulting Gibbs kernel is Kε = P_SOT ⊙ exp(−C/ε), so finite-objective couplings are restricted to the SOT support.
  • Convergence: The sparse Sinkhorn iterations have a unique diagonal-scaling solution and converge to the unique coupling P⋆.Convergence is established for iterates initialized at u0 = v0 = 1.
  • Computational complexity: Each lifted slice contributes at most N + M − 1 nonzeros, so sparse matrix-vector products reduce each iteration to O(L(N + M)).The total complexity includes SOT construction, cost evaluation on the support, and T sparse Sinkhorn iterations.
  • Computational complexity: The total arithmetic complexity is O((N + M)(d + T) + N log N + M log M).Only the sparse kernel, scaling vectors, and support-restricted cost entries need to be stored.
  • Divergence property: SinkSLOT is a divergence without debiasing under p-norm costs and general-position support points.This avoids the self-transport correction required by standard EOT and smoothed SROT.
  • Convergence rate: When the support graph is connected, sparse Sinkhorn converges linearly, with the rate governed by the second singular value of the normalised plan matrix.As ε increases, convergence becomes faster; as ε approaches zero, the contraction degenerates toward one.
  • Limiting behaviour: As ε approaches zero, the solution tends to the KL-selected minimiser of support-restricted OT, while as ε approaches infinity it tends to P_SOT.Moderate ε is reported to offer a practical trade-off between OT proximity and convergence speed.

Gradient and Hessian of the functional

SinkSLOT differentiates through the transport cost while treating the sliced optimal transport plan as a fixed prior. This preserves sparse gradient and Hessian computations, with stop-gradient formally justified almost everywhere.

  • Gradient construction: SinkSLOT depends on source positions through both the cost matrix and the sliced reference plan, creating a direct term and a reference-plan term.The reference-plan term arises because sliced matchings depend on projection-based rank orders.
  • Gradient construction: Stop-gradient removes the reference-plan term; this is formally correct at differentiable points because sliced transport plans are locally constant almost everywhere.Rank orders are piecewise constant, so their derivative with respect to source positions is zero except at discontinuities.
  • Differentiation convention: The detached-prior convention gives SLOTε gradients and Hessians the same analytic form as EOT, replacing the dense plan with the sparse SinkSLOT plan.All computational gains therefore come from sparsity rather than a different differentiation formula.
  • Gradient cost: O(L(N+M)d) computes the sparse gradient matrix-vector product, compared with O(NMd) for dense EOT, without sorting or projection during differentiation.The sparse support contains |S| entries, with |S| scaling with the number of slices and points.
  • Second-order derivatives: O(|S|) per conjugate-gradient iteration computes Hessian-vector products through sparse linear solves, compared with O(NM) for dense EOT.The sensitivity matrix inherits the sparsity of the SinkSLOT transport plan.

Experiments

Experiments benchmark SinkSLOT against dense, smoothed-reference, and importance-sampled sparse Sinkhorn methods across geometric and dimensional settings. SinkSLOT reaches comparable accuracy with substantially lower runtime, while its slice count controls the accuracy–speed trade-off and its divergence supports faster gradient flows.

  • Benchmark setup: SinkSLOT is benchmarked against FlashSinkhorn, SROT, and Spar-Sink on five datasets spanning low-dimensional geometric structures and Gaussian data up to d=64.Experiments use an NVIDIA H100-80GB and identical early stopping based on marginal violation.
  • Runtime and accuracy: 2–45× speedups over Spar-Sink occur where Spar-Sink reaches the accuracy threshold, while Spar-Sink misses thresholds on several Gaussian and 8-Gaussians settings.Spar-Sink fails to reach any threshold on Gaussian d=3 and d=64 and cannot attain the ≤5% gap on 8-Gaussians.
  • Runtime and accuracy: 15–160× speedups over SROT occur across all datasets and transport-gap levels.SROT is reported as the slowest method overall in these experiments.
  • Runtime and accuracy: 3–36× speedups over FlashSinkhorn occur consistently, including 3–11× on the d=64 Gaussian setting.The largest gains occur on the low-dimensional benchmarks, while the experiments report comparable accuracy to dense methods.
  • Slice-count trade-off: Increasing L improves achievable accuracy but increases computation time, while reducing ε eventually reaches a support-restricted accuracy plateau.Very high precision at ∆≤1% requires large L and is associated with prohibitive computational cost.
  • Scalability: SinkSLOT remains faster than FlashSinkhorn as N and M increase and across dimensions up to d=1024, with 2–11× speedups at L=4096 and 4–36× at L=1024.The asymptotic comparison is O(L(N+M)) per iteration for SinkSLOT versus O(NM) for FlashSinkhorn at fixed L.
  • Gradient-flow experiment: In blob-to-crescent gradient flows, SLOTε reaches lower exact squared 2-Wasserstein distances earlier and rearranges the source into the target crescent shape by step 5.All four methods attain similar final W 2 values after 50 Euler steps, but SLOTε remains below SROT through step 30.

Conclusion

SinkSLOT replaces the dense independent reference coupling with an expected lifted sliced OT plan, yielding sparse Sinkhorn iterations and convergence. Experiments report reduced computation time, while the slice count creates a sparsity–transport-polytope trade-off.

  • SinkSLOT replaces the dense independent reference coupling with an expected lifted sliced OT plan, reducing per-iteration cost from O(NM) to O(L(N + M)).
  • The sparse Sinkhorn iterations converge, and the paper characterises limiting behaviour as ε varies and derives first- and second-order derivatives.
  • Experiments show accuracy comparable to dense and hardware-aware baselines while substantially reducing computation time.
  • Too few slices restrict the transport polytope, whereas too many reduce the sparsity advantage.
  • Substantial speedups are reported over state-of-the-art dense and sparse EOT methods, with the resulting divergence illustrated as a loss function in a gradient flow task.

A.2 Proof of Proposition 2 (convergence of sparse Sinkhorn)

The convergence proof establishes existence and uniqueness of the sparse KL-regularised plan, then identifies sparse Sinkhorn updates as alternating I-projections that converge to it. The argument also verifies the required support and scalability conditions.

  • The sparse plan P_SOT lies in Π(a, b) and shares the zero pattern of Kε, satisfying the support and scalability condition for Sinkhorn updates.
  • Each half-iteration is an I-projection onto row or column marginal constraints, preserving the diagonal-scaling structure and producing the stated u and v updates.
  • Csiszár’s alternating I-projection theorem implies convergence to the I-projection of Kε onto Π(a, b), namely P⋆ε.
  • The divergence is nonnegative, equals zero for identical measures, and is symmetric under exchanging the source and target measures.

A.4 Proof of Proposition 5 (linear convergence rate)

The linear-convergence proof links the Sinkhorn Jacobian spectrum to the sparse support graph. Connectivity makes the non-gauge spectral radius strictly below one, while each iteration costs O(|S|) = O(L(N+M)).

  • The Jacobian eigenvalues are the squared singular values of the normalised plan matrix A, with leading singular value σ1 = 1.
  • When the bipartite support graph GS is connected, Perron–Frobenius theory gives σ2 < 1; disconnected support instead yields σ2 = 1.
  • After quotienting out the gauge direction, the fixed-point Jacobian has spectral radius σ2 < 1, establishing linear convergence of the coupling.
  • The iteration count is given by (15), and the per-iteration cost is O(|S|) = O(L(N+M)).
  • As ε approaches zero, subsequential limits are cost minimisers over the support-restricted transport polytope.
  • The ε → 0 limit selects the unique minimum-KL cost minimiser, so the full family converges to that plan.

A.7 Proof of Proposition 8 (Hessian of SLOTε)

The Hessian analysis introduces a sensitivity matrix whose sparsity follows the SinkSLOT coupling. Under connected support, this matrix has a well-defined pseudoinverse, enabling sparse matrix–vector products in derivative computations.

  • The sensitivity matrix H⋆ is defined from the optimal coupling and has a simple zero eigenvalue under connected bipartite support.
  • Connected support makes H⋆ symmetric positive semidefinite, with null vector (1_N, −1_M) and a well-defined Moore–Penrose pseudoinverse.
  • Unlike dense EOT, SinkSLOT gives H⋆ sparse off-diagonal blocks with |S| nonzeros, reducing H⋆w matrix–vector products to O(|S| + N + M).
  • The auxiliary tensor B is built from point differences weighted by the optimal coupling, while E is block-diagonal across source points.
  • The Hessian proposition gives T = ∇^2 of SinkSLOT under the stop-gradient convention and squared Euclidean cost.

B.1 Synthetic Benchmark Experiments

The benchmark evaluates SinkSLOT and baselines under shared stopping, hardware, precision, feasibility, and hyperparameter protocols. Its implementation exploits sparse supports and fused kernels to reduce practical overhead.

  • Experimental protocol: All methods use the same problem instances, stopping rule, NVIDIA H100-80GB GPU, and float32 computations for runtime comparability.
  • Experimental protocol: 840 output plans satisfy both mass conservation and the marginal-violation feasibility criteria.The checks require mass error below 10^-3 and marginal violation below 10^-4.
  • Hyperparameters: The benchmark sweeps eight ε values, eight slice counts for SROT and SinkSLOT, and four sampling budgets for Spar-Sink.ε ranges are sampled relative to the ground-cost scale, while slice counts and sampling densities are specified separately by dataset.
  • Implementation: Triton kernels fuse sparse cost evaluation and segmented log-sum-exp reductions to avoid materialising dense intermediates.The implementation uses compressed sparse layouts, CSR and CSC storage, and streaming reductions over the support.
  • Gradient-flow setup: The gradient-flow experiment compares SOT, EOT, SROT, and SinkSLOT on a blob-to-crescent trajectory with N = M = 1000 and L = 100.SinkSLOT uses analytical gradients without backpropagation, while the baselines use their specified gradient procedures.

C.1 Impact of L

SinkSLOT generally offers a better accuracy–speed trade-off than FlashSinkhorn, but the useful slice count depends on the dataset and the advantage weakens in higher dimensions. Larger slice counts can improve low-error regimes, while ε controls a runtime–accuracy trade-off.

  • Low-dimensional datasets: SinkSLOT achieves a better accuracy–speed trade-off than FlashSinkhorn on the low-dimensional datasets.At a target cost gap it can be faster, and at a fixed runtime it can achieve a lower cost gap.
  • Low-dimensional datasets: For two-rings, L = 128 appears sufficient because its Pareto curve lies southwest of the curves for larger slice counts.
  • Low-dimensional datasets: For the other three low-dimensional datasets, SinkSLOT’s Pareto front combines points from multiple slice counts.
  • Gaussian d = 3: On Gaussian d = 3, decreasing ε reduces the cost gap but eventually increases runtime without improving the output plan at L = 128.Increasing L to 1024, and then to 4096 for cost gaps below 1%, is more effective.
  • High-dimensional dataset: In the high-dimensional setting, SinkSLOT retains the same qualitative advantage over FlashSinkhorn, but the advantage is less pronounced.Finding useful (L, ε) combinations may require substantial experimentation because lifted-plan approximation deteriorates with dimension.
  • Convergence: Except for SinkSLOT on Gaussian d = 3, all selected runs converge within 10^4 iterations, with SinkSLOT requiring iterations comparable to SROT.
  • Memory: At N = 5 × 10^4, SinkSLOT remains manageable with 4096 slices, whereas SROT runs out of memory because its dense matrix grows quadratically with N.

C.4 Stop-gradient experiments

The stop-gradient treatment of the sliced prior is supported by piecewise constancy and small jumps, while solver sensitivity remains negligible during the transport phase. Analytic gradients therefore track fully unrolled gradients until the objective reaches its O(ε) floor.

  • Gradient validation: Term (I) is validated as the gradient of SLOT_ε under a frozen-prior convention, while term (II) and solver sensitivity are assessed separately.
  • Gradient validation: With the prior frozen, finite differences match term (I) to relative error 1.1 × 10^-9 at the first flow step.Later reported error growth reflects the derivative decaying toward ∼7 × 10^-6 against a fixed noise floor.
  • Prior dependence: Rebuilding the sliced prior produces an O(h^-1) divided-difference growth near rank flips, indicating jumps rather than a missing gradient term.At h = 10^-7, when no rank flip occurs, the rebuilt difference returns to the analytic value.
  • Prior dependence: Rebuilding the prior changes roughly 0.3% of support entries and changes the value by at most 3.2 × 10^-4 relative at the last step.The evidence supports stop-gradient in practice through piecewise constancy and small jumps, not differentiability everywhere.
  • Solver sensitivity: The residual remains nearly constant while term (I) decays by a factor of 10^2, with their ratio growing from 0.029 to 1.12.The transport cost and exact W2 distance plateau at the same step, indicating that disagreement occurs after the flow has effectively arrived.
  • Solver sensitivity: During transport, the analytic gradient stays within a few percent of the fully unrolled gradient and diverges only after the objective reaches its O(ε) floor.The flow protocol rebuilds the reference plan from the current source positions and shares it across estimators.
Loading 2608.28262v1…