Annealed Sinkhorn with Momentum: Certified Unregularized Optimal Transport in Linear Memory

Samuel J. K. Chin; Maximilian Schiffer

math.OCSep 27, 2026

Annealed Sinkhorn with Momentum: Certified Unregularized Optimal Transport in Linear Memory

Authors: Samuel J. K. Chin, Maximilian Schiffer

Organizations: MIT · HEC Paris

Abstract

We characterize Bregman Douglas-Rachford splitting (BDRS) for unregularized discrete optimal transport and develop an anytime primal-dual certificate in linear memory. We first establish that BDRS coincides with warm-started Inexact Proximal point method for exact Optimal Transport (IPOT) using a single inner Sinkhorn iteration. By eliminating the primal transport plan from the updates, we derive an equivalent dual formulation that reveals BDRS as annealed Sinkhorn under an implicit inverse-linear temperature schedule, with an additional log-scaling momentum term and a cooler kernel. While this explains the role of the temperature parameter in BDRS as an initial temperature, it also reduces the solver's memory requirement from quadratic to linear. Utilizing this annealing perspective, we introduce overrelaxed BDRS, which combines annealing and overrelaxed scaling within a single recursion. We derive a primal-dual certificate for both methods that can be evaluated in linear memory without transport plan construction, thus providing a computable stopping rule. On the DOTmark benchmark, combining momentum with the cooler kernel produces substantially smaller optimality gaps than annealed Sinkhorn under the same schedule. For pixel-level color transfer between $1024\times1024$ images, BDRS attains a lower repaired transport cost than MDOT-TNT with a 9 $\times$ speed up, reaching a relative duality gap of $1.59\%$ in 24 minutes. We further demonstrate a color transfer with $4238\times2365$ images, yielding 10 million pixels per image and approximately one hundred trillion implicit transport entries, reaching a best relative duality gap of $2.41\%$ and $2.80\%$ within 35 hours in each direction on a single NVIDIA L40S GPU.

Figures & tables

Figure 1: Momentum and kernel ablation on all 450 DOTmark pairs at resolution $64\times 64$ , with $\eta=1$ . The optimality gap $f-f^{\star}$ is shown against iteration count (left) and median solver time, excluding feasibility rounding (right). Here $f$ is the cost after rounding and $f^{\star}$ is the reference value computed using POT’s network simplex solver. Lines show medians across pairs; shaded bands indicate the interquartile range.

Figure 2: Anytime certification on 450 DOTmark pairs at resolution $64\times 64$ , comparing Dual BDRS ( $\lambda=1$ ) and Dual O-BDRS ( $\lambda=1.99$ ). Left: aggregated certified gaps $G^{\mathrm{cert}}=U^{\mathrm{best}}-L^{\mathrm{best}}$ (solid) and reference gaps $E^{\mathrm{oracle}}=U^{\mathrm{best}}-p^{\star}$ (dashed). Right: empirical cumulative distributions of the certification delay $D_{i}(\tau)=k_{i}^{\mathrm{cert}}(\tau)/k_{i}^{\mathrm{oracle}}(\tau)$ for $\tau\in\{10^{-2},10^{-3},10^{-4}\}$ . Smaller delays indicate closer agreement with each method’s own oracle stopping rule.

Panel of Figure 3 — Figure 3: Large-scale pixel-level color transfer between two $4238\times 2365$ images using Dual BDRS. Top: original images A (left) and B (right). Bottom: transfers from A to B (left) and B to A (right). Each direction runs for 1000 iterations in float32 on a single NVIDIA L40S GPU and takes approximately 35 hours. The best relative duality gaps $(U-L)/U$ are $2.41\%$ and $2.80\%$ respectively. Transferred images are obtained by barycentric projection.

Appendix figures & tables11 assets

Supplementary material from the paper’s appendix.

Appendix

Row-update kernel	No momentum: $b^{k-1}$	Momentum: $\widetilde{b}^{k-1}$
Old: $K^{\odot k}$	Annealed Sinkhorn	Annealed Sinkhorn with momentum
Current: $K^{\odot(k+1)}$	Current-kernel Annealed Sinkhorn	Dual BDRS

Appendix

Table 1: Controlled ablation of kernel timing and log-scaling momentum. All variants share the same column update, annealing schedule, initialization, and transport plan reconstruction formula.

Figure 4: Classwise certification-delay distributions for Dual BDRS ( $\lambda=1$ ) and Dual O-BDRS ( $\lambda=1.99$ ) on DOTmark. Each panel contains 45 image pairs at resolution $64\times 64$ . Solid and dashed curves correspond to $\tau=10^{-2}$ and $\tau=10^{-3}$ , respectively; dotted red curves additionally show O-BDRS at $\tau=10^{-4}$ , with BDRS omitted as certification was not obtained within the iteration budget. The horizontal axis is $D_{i}(\tau)=k_{i}^{\mathrm{cert}}(\tau)/k_{i}^{\mathrm{oracle}}(\tau)$ and the vertical axis is its empirical cumulative distribution. The reference $D_{i}=1$ denotes agreement with the oracle stopping iteration. Curves further to the left indicate less additional iteration cost relative to each method’s own oracle.

Panel of Figure 5 — Figure 5: Visual comparison on the same $1024\times 1024$ color-transfer problem. All panels show barycentric reconstructions using all source and target pixels, float32 , and the same normalized squared Euclidean cost in RGB space. (a) BDRS after 10 iterations, requiring $13.55$ s of solver time ( $15.71$ s including setup and certificate evaluation). (b) MDOT-TNT at final inverse temperature $\gamma=2^{5}$ , requiring $1685.76$ s of solver time. (c) BDRS after 1,000 iterations, used as a reference solution with a reported relative duality gap of $1.59\%$ . BDRS uses $\eta=0.1$ and $\lambda=1$ . Solver times exclude image reconstruction. Observe how BDRS with 10 iterations is closer to the reference, where the brown parts of the image are more pronounced, whereas MDOT-TNT is still generally green, an indication of a diffuse transport plan. On the 10 iterations comparison, BDRS provides a 107 $\times$ speedup over MDOT-TNT . Best viewed digitally.

Panel of Figure 6 — Figure 6: Bidirectional color transfer between two $1024\times 1024$ images using BDRS. (a, b) Original images. (c) Image 1 transferred to the color distribution of image 2. (d) Image 2 transferred to the color distribution of image 1. Both runs use all source and target pixels and perform 1,000 iterations in float32 , with $\eta=0.1$ and $\lambda=1$ . The reported relative duality gaps are $1.59\%$ and $1.58\%$ for the forward and backward directions, respectively. Each direction ran in approximately 24 minutes with 156MB of memory. Transferred images are obtained by barycentric projection of the unrepaired transport plans.

Statistic	Forward ( $1\to 2$ )	Backward ( $2\to 1$ )
Primal upper bound $U$	$0.0496885$	$0.0496843$
Dual lower bound $L$	$0.0488973$	$0.0488973$
Relative duality gap $(U-L)/U$	$1.59\%$	$1.58\%$
Marginal $\ell_{1}$ error	$(2.94\times 10^{-5},\,6.35\times 10^{-5})$	$(2.54\times 10^{-5},\,7.34\times 10^{-5})$
Maximum absolute marginal error	$(2.06\times 10^{-10},\,7.87\times 10^{-10})$	$(2.06\times 10^{-10},\,7.87\times 10^{-10})$
Maximum relative marginal error	$(0.0216\%,\,0.0826\%)$	$(0.0216\%,\,0.0826\%)$

Appendix

Table 3: Accuracy, runtime, and memory of BDRS after 1,000 iterations for both

1024\times 1024

color transfer directions in float32 . Marginal errors refer to the unrepaired plan; paired entries are ordered as (row, column). Relative marginal errors are measured against the prescribed pixel mass

2^{-20}

. Runtime includes setup and certificate evaluation but excludes image reconstruction. Memory is the peak GPU allocation tracked by PyTorch, including reconstruction and validation, and excludes allocations outside PyTorch.

Statistic	Forward ( $1\to 2$ )	Backward ( $2\to 1$ )
Primal upper bound $U$	$0.0101706$	$0.0102118$
Dual lower bound $L$	$0.00992558$	$0.00992557$
Relative duality gap $(U-L)/U$	$2.41\%$	$2.80\%$
Marginal $\ell_{1}$ error	$(1.16\times 10^{-4},\,3.29\times 10^{-4})$	$(1.17\times 10^{-4},\,4.28\times 10^{-4})$
Maximum absolute marginal error	$(4.68\times 10^{-11},\,1.50\times 10^{-9})$	$(4.68\times 10^{-11},\,1.34\times 10^{-9})$
Maximum relative marginal error	$(0.0469\%,\,1.5023\%)$	$(0.0469\%,\,1.3414\%)$

Appendix

Table 4: Accuracy, runtime, and memory of BDRS for both

4238\times 2365

color transfer directions, with

\eta=0.01

\lambda=1

, and float32 pairwise computations. Images 1 and 2 are the canyon and outdoor images, respectively. Each run completes 1,000 iterations; reported accuracy corresponds to the selected iterate at iteration 840. Marginal errors refer to the unrepaired intermediate plan; paired entries are ordered as (row, column). Relative marginal errors are measured against the prescribed pixel mass

1/10{,}022{,}870

. Runtime covers the full run, including setup, certificate evaluation, and final verification, but excludes image reconstruction. Memory is the peak GPU allocation tracked by PyTorch, including reconstruction and validation, and excludes allocations outside PyTorch.

Panel of Figure 7 — Figure 7: Momentum and kernel ablations on the Cauchy Density (top) and Classic images (bottom) classes of DOTmark. Each class contains 45 image pairs at resolution $64\times 64$ , with $\eta=1$ . The left panels show the transport cost difference $f-f^{\star}$ against iteration count, and the right panels show the same quantity against median solver time, excluding rounding. Here $f$ is the evaluated transport cost and $f^{\star}$ is the reference value computed using POT’s network simplex solver. Lines show medians across image pairs; shaded bands indicate the 25th–75th percentiles.

Panel of Figure 8 — Figure 8: Momentum and kernel ablations on the GRF Moderate (top) and GRF Rough (bottom) classes of DOTmark. Each class contains 45 image pairs at resolution $64\times 64$ , with $\eta=1$ . The left panels show the transport cost difference $f-f^{\star}$ against iteration count, and the right panels show the same quantity against median solver time, excluding rounding. Here $f$ is the evaluated transport cost and $f^{\star}$ is the reference value computed using POT’s network simplex solver. Lines show medians across image pairs; shaded bands indicate the 25th–75th percentiles.

Panel of Figure 9 — Figure 9: Momentum and kernel ablations on the GRF Smooth (top) and Log-GRF (bottom) classes of DOTmark. Each class contains 45 image pairs at resolution $64\times 64$ , with $\eta=1$ . The left panels show the transport cost difference $f-f^{\star}$ against iteration count, and the right panels show the same quantity against median solver time, excluding rounding. Here $f$ is the evaluated transport cost and $f^{\star}$ is the reference value computed using POT’s network simplex solver. Lines show medians across image pairs; shaded bands indicate the 25th–75th percentiles.

Panel of Figure 10 — Figure 10: Momentum and kernel ablations on the Logit-GRF (top) and Microscopy Images (bottom) classes of DOTmark. Each class contains 45 image pairs at resolution $64\times 64$ , with $\eta=1$ . The left panels show the transport cost difference $f-f^{\star}$ against iteration count, and the right panels show the same quantity against median solver time, excluding rounding. Here $f$ is the evaluated transport cost and $f^{\star}$ is the reference value computed using POT’s network simplex solver. Lines show medians across image pairs; shaded bands indicate the 25th–75th percentiles. Here, we observe that O- BDRS has a performance degradation on Microscopy Images as the number of iterations increase. We suspect that this late-iteration degradation reflects numerical sensitivity at low effective temperatures, potentially exacerbated by the $10^{-15}$ mass assigned to zero-mass pixels, which are common in Microscopy Images and Shapes, but not the other classes.

Panel of Figure 11 — Figure 11: Momentum and kernel ablations on the Shapes (top) and White Noise (bottom) classes of DOTmark. Each class contains 45 image pairs at resolution $64\times 64$ , with $\eta=1$ . The left panels show the transport cost difference $f-f^{\star}$ against iteration count, and the right panels show the same quantity against median solver time, excluding rounding. Here $f$ is the evaluated transport cost and $f^{\star}$ is the reference value computed using POT’s network simplex solver. Lines show medians across image pairs; shaded bands indicate the 25th–75th percentiles. Here, we observe that O- BDRS has a performance degradation on Shapes as the number of iterations increase. We suspect that this late-iteration degradation reflects numerical sensitivity at low effective temperatures, potentially exacerbated by the $10^{-15}$ mass assigned to zero-mass pixels, which are common in Microscopy Images and Shapes, but not the other classes.

Explore similar work

Apr 27, 2026stat.ML

Sliced-Regularized Optimal Transport

We propose a new regularized optimal transport (OT) formulation, termed sliced-regularized optimal transport (SROT). Unlike entropic OT (EOT), which regularizes the transport plan toward an independent coupling, SROT regularizes it toward a smoothened sliced OT (SOT) plan. To the best of our knowledge, SROT is the first approach to leverage a version of SOT plan as a reference to improve classical OT. We provide a formal definition of SROT, derive its dual formulation, and provide a post-Bayesian interpretation of SROT. We then develop a Sinkhorn-style algorithm for efficient computation, retaining the same scalability advantages as EOT. By incorporating a scalable SOT plan as a prior, SROT yields more accurate approximations of the exact OT plan than EOT under the same level of regularization. Moreover, the resulting transport plan improves upon the reference SOT plan itself. We further introduce the corresponding OT divergence induced by SROT, named SROT divergence, and analyze its topological and computational properties. Finally, we validate our approach through experiments on synthetic datasets and color transfer tasks, demonstrating that SROT is better than both EOT and SOT in approximating exact OT. Additional experiments on gradient flows further highlight the advantages of SROT divergence.

Khai Nguyen

The University of Texas at Austin

Gradient

Ask PDF

Sep 14, 2026cs.LG

Dual-guided Hierarchical Edge Localization for Large-scale Optimal Transport Across Dimensions

Optimal transport (OT) compares distributions and aligns datasets in machine learning, yet unregularized discrete OT requires a linear program with quadratically many transport variables. We propose HELLO, a hierarchical solver that casts large-scale discrete OT as edge localization and uses dual potentials to guide both coarse-to-fine initialization and within-level refinement. Initialization propagates coarse dual potentials across a recursive subsampling hierarchy to assign candidate edges. Refinement then iteratively inserts the largest dual violators in each row and column until the relative KKT residual meets a prescribed tolerance, while budgeted pruning ensures linear memory complexity. For exact-arithmetic refinement, we prove finite termination at a global optimum under a symbolic lexicographic rule. At the million-point scale, HELLO attains lower transport objectives with order-of-magnitude runtime improvements over strong baselines across feature dimensions from single digits to thousands. It further scales to 1.28 million samples per marginal in 8192 dimensions on a single H100, using 41.6 GiB peak GPU memory while satisfying a full relative KKT residual below

10^{-6}

. Beyond standard discrete OT, the framework supports general pairwise costs and serves as a scalable balanced-OT oracle for semi-discrete OT, Gromov--Wasserstein, unbalanced OT, and OT-based Flow Matching.

Wenzhou Xia, Qiaoqiao Ding, Jingwei Liang +1

School of Mathematical Sciences, Shanghai Jiao Tong University, Shanghai 200240, China · Institute of Natural Sciences, Shanghai Jiao Tong University, Shanghai 200240, China

Differentiable Optimal Transport Optimal Transport Approach Primal-Dual Methods Neural Solvers Coarse-To-Fine Object Localization

Ask PDF

Sep 30, 2026cs.LG

Accelerated Algorithm for Sparse Regularized Partial Optimal Transport

Partial Optimal Transport (POT) extends the classical optimal transport problem by relaxing the strict mass conservation constraint, enabling its use in a wide range of real-world applications. In many of these settings, sparse transport plans are preferred for their interpretability and computational benefits. While smooth and strongly convex regularizers - such as quadratic or elastic net - have been vastly used in various machine learning applications to induce sparsity and accelerate computation, they have received less algorithmic attention compared to entropic approaches for computational POT. In this paper, we propose a new optimization framework that leverages these regularizers through a penalty-based reformulation, enabling efficient gradient-based updates while preserving the structure of the original problem. Our method accommodates a broad class of regularizers that promote structured and sparse transport plans. Building on this formulation, we design an accelerated first-order algorithm that alternates between smooth updates and simple projection steps. Through empirical benchmarks on color transfer, domain adaptation, and point cloud registration, our approach consistently outperforms established baselines - achieving lower transport cost, higher sparsity, and faster convergence - making it a practical and scalable solution for modern transport problems.

Khoa Nguyen, Dung T. Nguyen, Thong Huynh +4

Aalto University, Konemiehentie 2, 02150 Espoo, Finland · Linköping University, 581 83 Linköping, Sweden · Hanoi University of Science and Technology, Hanoi, Vietnam +4

Differentiable Optimal Transport Regularization

Ask PDF