Source-linked AI summary
Stabilized Sparse Scaling Algorithms for Entropy Regularized Transport Problems
Bernhard Schmitzer
TL;DR
The paper addresses numerical instability, slow convergence, and dense-kernel costs in entropy-regularized transport scaling algorithms. It combines log-domain stabilization, ε-scaling, adaptive kernel truncation, and a coarse-to-fine scheme, and develops a new convergence analysis. Numerical examples show robust execution across ε, reduced iterations with ε-scaling, and lower time and memory from sparse truncation, while the theory exposes limits of the convergence and ε-scaling arguments.
Problem
Scaling algorithms face impractical slow convergence and numerical limitations at small ε, while dense kernels make large problems costly in runtime and memory.
Method
The paper combines log-domain stabilization, ε-scaling, adaptive kernel truncation, a multi-scale scheme, and an auction-inspired convergence analysis.
Results
Numerical examples show robust solution across tested ε values, substantially fewer iterations with ε-scaling, and time-memory savings from adaptive sparse kernels.
Takeaways & Limitations
The combined algorithm solves larger transport problems with smaller regularization and reduced runtime and memory, while maintaining negligible truncation error.
Takeaways & Limitations
The convergence argument for ε-scaling requires additional assumptions because the analysis assumes β1 is known exactly while the algorithm computes it approximately.
Abstract
from arXiv · showhide
Scaling algorithms for entropic transport-type problems have become a very popular numerical method, encompassing Wasserstein barycenters, multi-marginal problems, gradient flows and unbalanced transport. However, a standard implementation of the scaling algorithm has several numerical limitations: the scaling factors diverge and convergence becomes impractically slow as the entropy regularization approaches zero. Moreover, handling the dense kernel matrix becomes unfeasible for large problems. To address this, we combine several modifications: A log-domain stabilized formulation, the well-known epsilon-scaling heuristic, an adaptive truncation of the kernel and a coarse-to-fine scheme. This permits the solution of larger problems with smaller regularization and negligible truncation error. A new convergence analysis of the Sinkhorn algorithm is developed, working towards a better understanding of epsilon-scaling. Numerical examples illustrate efficiency and versatility of the modified algorithm.
1. Introduction.
Optimal transport supports broad analytical and computational applications, but computational cost remains a bottleneck for large or dense problems. The paper addresses instability, slow convergence, and kernel-size limitations through four combined modifications and a new convergence analysis.
- Applications: Optimal transport lifts metrics to probability measures and supports applications in PDE analysis, image processing, computer vision, and machine learning.Wasserstein distances provide the analytical basis, while increased computational resources have expanded OT's numerical use.
- Computational bottleneck: Large and multi-marginal optimal transport problems remain challenging because of runtime and memory demands.Combinatorial solvers do not scale well for large dense problems, motivating alternative numerical methods.
- Scaling methods: Entropy regularization extends scaling methods to barycenters, multi-marginal problems, gradient flows, and unbalanced transport.The regularized formulation enables a family of Sinkhorn-like diagonal scaling algorithms.
- Limitations: Naive entropic transport implementations become increasingly limited as regularization decreases, including blur, diverging scaling factors, and slow convergence.The paper targets these limitations while seeking smaller regularization values.
- Contributions: The paper combines log-domain stabilization, ε-scaling, adaptive kernel truncation, and a multi-scale scheme, alongside a new auction-inspired convergence analysis.The modifications target numerical instability, iteration counts, memory demand, and large kernel matrices.
2. Entropy Regularized Transport-Type Problems and Diagonal Scaling Algorithms.
The paper formulates balanced, unbalanced, and broader transport-type problems through marginal functionals, then entropy-regularizes them to obtain diagonal scaling iterations. These iterations alternate KL-proximal updates using a cost-dependent kernel and reconstruct the primal solution from scaling factors.
- Generic transport formulation: Transport-type problems optimize a nonnegative coupling with a linear cost and functions acting on its marginals.The generic primal energy is FX(PXπ) + FY(PYπ) + ⟨c,π⟩ + ι+(π).
- Balanced and unbalanced cases: Standard optimal transport fixes both marginals, whereas KL-fidelity unbalanced transport replaces those hard constraints with weighted KL penalties and permits differing total masses.The unbalanced formulation uses λ · KL(PXπ|µ) and λ · KL(PYπ|ν), with λ > 0.
- Entropy regularization: Entropy regularization replaces nonnegativity enforcement with ε KL(π|ρ), and the cost is incorporated into a kernel K(x,y) = exp(−c(x,y)/ε) · ρ(x,y).The regularized generic energy is FX(PXπ) + FY(PYπ) + ε KL(π|K).
- KL-proximal updates: Alternating optimization updates each dual variable through a KL-proximal step, with the corresponding primal and dual optimizers linked by exponentiated scaling variables.For fixed β, σ† = proxεFX(K exp(β/ε)) and α† = ε log(σ† ⊘ (K exp(β/ε))).
- Diagonal scaling algorithm: The effective scaling iterations are u(ℓ+1) := proxdivεFX(Kv(ℓ)) and v(ℓ+1) := proxdivεFY(K⊤u(ℓ+1)), followed by a stopping test.The primal optimizer is reconstructed as π† = diag(u†)K diag(v†), while stopping may use a primal-dual gap, marginal error, or iteration limit.
- Scope and implementation: The framework supports efficiently computable KL proximal steps, especially when marginal functionals are separable and decompose into pointwise one-dimensional updates.The paper identifies entropy-regularized optimal transport as a special case and introduces four adaptations to address naive scaling limitations.
3. Stabilized Sparse Multi-Scale Algorithm.
The paper combines stabilization, ε-scaling, adaptive kernel truncation, and coarse-to-fine refinement to address numerical instability, slow convergence, and dense-kernel costs in entropy-regularized transport. It also analyzes Sinkhorn complexity and bounds truncation error, while noting a convergence caveat for truncated stabilized iterations.
- 3.1. Log-Domain Stabilization: Small ε can make kernel entries and scaling factors numerically extreme, while Sinkhorn convergence becomes slow as ε approaches zero.The paper addresses these issues with stabilized iterations and ε-scaling.
- 3.1. Log-Domain Stabilization: Redundant scaling-factor parametrization keeps auxiliary factors bounded and periodically absorbs them into dual variables.Stabilized iterations update the auxiliary factors; absorption iterations transfer their values into the dual variables and recompute the kernel.
- 3.1. Log-Domain Stabilization: The stabilized and absorption iterations are mathematically equivalent to the original scaling algorithm while preserving much of its matrix-multiplication structure.The combined implementation has less overhead than a direct log-domain formulation.
- 3.3. Kernel Truncation: Adaptive truncation re-estimates the active kernel support from current dual iterates and combines it with stabilized kernel evaluation.The truncation threshold defines a sparse support, and the resulting additional duality-gap error is bounded by ||ũ||∞·||ṽ||∞·θ·ρ(X × Y).
- 3.3. Kernel Truncation: The truncation bound is practically kept below the remaining primal-dual gap, but truncated stabilized iterations might not converge without sufficiently regular absorption steps.Regular absorption before the auxiliary factors become too large makes the resulting oscillations numerically negligible.
- 3.4. Multi-Scale Scheme: A hierarchical multi-scale scheme combines coarse-to-fine refinement with ε-scaling, reducing variables during early stages while initializing finer levels from coarser dual solutions.The paper also develops a new Sinkhorn complexity analysis, with fixed-ε iteration counts scaling as O(1/ε).
4. Analogy between Sinkhorn and Auction Algorithm.
The paper develops an analogy between auction and asymmetric Sinkhorn algorithms, using monotonicity and complementary slackness to analyze convergence and ε-scaling. It proves stability of dual solutions under changing regularization and derives iteration bounds, while identifying assumptions needed for rigorous ε-scaling guarantees.
- Auction algorithm: The auction algorithm maintains monotone dual variables, ε-complementary slackness, sub-feasible primal iterates, and terminates after at most N · (C/ε + 1) iterations.Its primal-dual gap is at most N · ε, and for integer costs with ε < 1/N the final coupling is optimal.
- Asymmetric Sinkhorn algorithm: The asymmetric Sinkhorn algorithm mirrors this structure: u and α increase, v and β decrease, q increases, and the primal iterate remains sub-feasible.A component of v remains unchanged throughout the iterations.
- Asymmetric Sinkhorn algorithm: For fixed ε, the asymmetric Sinkhorn iteration count scales as O(1/ε), and its iterates converge to a scaling solution with q(ℓ) → 1.The bound for reaching q(n) ≥ qtarget is n ≤ 2 + C/[ε·(1−qtarget)].
- Convergence measures: The q-criterion measures marginal error but can terminate prematurely, while Hilbert-metric convergence becomes impractical for small ε because its contraction factor is roughly 1 − 4 exp(−C/ε).The q-criterion is considered reasonable when 1 − qtarget ≪ δ.
- Stability and ε-scaling: Theorem 20 bounds changes in optimal dual variables under ε1 > ε2 independently of the cost function, supporting ε-scaling iteration estimates.Both dual-variable ranges are bounded by ε1 · N · (4 log N + 24 log M).
- Stability and ε-scaling: With exact initialization, subsequent ε-scaling steps require at most 2 + A/[λ(1−qtarget)] iterations, yielding a total bound of O(log(C/ε̂)).The argument has a gap because the theorem assumes an exact β1, whereas Algorithm 7 supplies only an approximation; additional assumptions are required for rigor.
5. Numerical Examples.
Numerical experiments show that the combined stabilized, epsilon-scaled, sparsified, and coarse-to-fine algorithm substantially improves efficiency while retaining versatility across transport-type applications. It supports larger problems at smaller regularization, with runtime, sparsity, and accuracy tunable through algorithmic parameters.
- Experimental setup: At ε = 0.1 · h2, most problems use very little blur while retaining a good approximation of the original transport problem.For the squared Euclidean cost, ε is measured in units of h2, and ε = h2 corresponds to blur on the scale of one pixel.
- Efficiency: The combined modifications achieve an average total speed-up of more than two orders of magnitude on the tested problem type.Log-domain stabilization enables robust small-ε computation; ε-scaling reduces iterations, sparse kernels reduce time and memory, and multi-scale processing reduces variables during early stages.
- Efficiency: Runtime scales approximately linearly with |X|, and the algorithm becomes faster than the adaptive sparse linear programming solver for large problems.The final sparse-kernel variable count is comparable to the solver’s, while larger ε during scaling requires more memory.
- Limitations and scope: For larger images, only the multi-scale Algorithm 5 can be run practically because of time and memory constraints.The runtime gap between the tested variants is expected to widen with image size, but only the multi-scale method was feasible at larger sizes.
- Efficiency: Choosing the stopping threshold and final ε tunes the trade-off between required precision and available runtime.Small ε reduces blur but requires more ε-scaling stages; larger ε can reduce or increase runtime depending on truncation efficiency.
- Versatility: The enhanced algorithm remains applicable to unbalanced transport, Wasserstein-Fisher-Rao distances, barycenters, and Wasserstein gradient flows.Experiments include WFR geodesics, Wasserstein barycenters, and porous-media gradient flows, including large problems with small regularization.
6. Conclusion.
The paper enhances scaling algorithms for entropy-regularized transport by combining stabilization, scaling, truncation, and coarse-to-fine techniques. Numerical examples support the method’s efficiency and broad applicability, while the analysis motivates further study of ε-scaling and extensions.
- Log-domain stabilization and ε-scaling address diverging scaling factors and slow convergence for small regularization.
- Adaptive kernel truncation and a coarse-to-fine scheme significantly reduce runtime and memory requirements.
- A new convergence analysis for the Sinkhorn algorithm was developed and its predicted scaling was confirmed numerically.
- Numerical examples demonstrated efficiency and sharp results across a wide range of transport-type problems.
- Future research includes deeper study of ε-scaling, log-domain stabilization, and applications to multi-marginal problems.
Appendix A. Additional Proofs.
The appendix begins by establishing existence of minimizers for the analyzed functional under a positive regularization parameter and a finite-dimensional coefficient setting.
- The proof first establishes existence of minimizers for some ε > 0.
- The coefficient condition is stated for d ∈ R^R×R.
- The result forms the starting point for the appendix’s additional proofs.
A.1. Proof of Lemma 23.
The proof bounds differences between maximizers by selecting a maximal component, recursively constructing child indices through reduced problems, and exploiting the resulting tree structure.
- The functional is invariant under constant shifts, so fixing β(1) = 0 yields a unique minimizer after strict convexity and coercivity are established.
- The proof selects an index where ∆β is maximal and derives another index as its child using optimality conditions.
- A reduced problem replaces the variables in I1 with one effective variable while retaining separate variables for indices in I2.
- The child-selection procedure is repeated until I1 contains all indices, producing a connected tree rooted at i1.
A.2. Proof of Theorem 20.
The proof of Theorem 20 uses assignment-graph components, representative points, and path estimates to control differences between dual maximizers. It derives explicit bounds in terms of ε1, N, and M.
- Strongly connected assignment-graph components provide partition cells with matched masses µ(Xi) = ν(Yi).
- The proof introduces reduced effective diagonal functionals using the partition cells and selected representatives.
- Assignment-graph edge estimates bound alternating-path changes in the dual variables by L = 2ε1 log M.
- Representative-based path arguments yield maxdiam(w) ≤ 16ε1N log M + 2ε1R log N.
- The reduced maximizers satisfy max ∆β̂ − min ∆β̂ ≤ ε1N (4 log N + 16 log M).
- Combining the reduced and representative bounds gives max ∆β − min ∆β ≤ ε1N (4 log N + 24 log M), with an analogous bound for ∆α.