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

Sliced-Regularized Optimal Transport
Apr 27, 2026Khai NguyenGradient
The University of Texas at Austin
Dual-guided Hierarchical Edge Localization for Large-scale Optimal Transport Across Dimensions
Sep 14, 2026Wenzhou Xia, Qiaoqiao Ding, Jingwei Liang +1Differentiable Optimal Transport Optimal Transport Approach
School of Mathematical Sciences, Shanghai Jiao Tong University, Shanghai 200240, China · Institute of Natural Sciences, Shanghai Jiao Tong University, Shanghai 200240, China
Accelerated Algorithm for Sparse Regularized Partial Optimal Transport
Sep 30, 2026Khoa Nguyen, Dung T. Nguyen, Thong Huynh +4Differentiable Optimal Transport Regularization
Aalto University, Konemiehentie 2, 02150 Espoo, Finland · Linköping University, 581 83 Linköping, Sweden · Hanoi University of Science and Technology, Hanoi, Vietnam +4

Apr 27, 2026, stat.MLJ/K move · Enter open · S save

Sliced-Regularized Optimal Transport

Khai Nguyen

The University of Texas at Austin

Gradient

Open paper AskPDF arXiv

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.

Abstract

Explore similar work

Sliced-Regularized Optimal Transport

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

Accelerated Algorithm for Sparse Regularized Partial Optimal Transport