GPU-Accelerated Bregman Douglas-Rachford Splitting for Discrete Optimal Transport

Yifan Xu; Shiqian Ma

math.OCOct 3, 2026

GPU-Accelerated Bregman Douglas-Rachford Splitting for Discrete Optimal Transport

Authors: Yifan Xu, Shiqian Ma

Organizations: Department of Computational Applied Mathematics and Operations Research, Rice University, Houston, TX, USA · Data Science and AI Institute, and Department of Applied Mathematics and Statistics, Johns Hopkins University, Baltimore, MD, USA

Abstract

We present GPU-accelerated Bregman Douglas--Rachford splitting algorithm (BDRS) for discrete optimal transport problem in three input formats: an explicit cost matrix, a point cloud with a ground cost between them, and a separable cost on a regular grid. For each input format, we propose hardware-aware designs of mathematically equivalent representations for the BDRS iterations to enhance numerical stability and empirical runtime. We benchmark the three proposed implementations against eight GPU baseline solvers from the literature on the same device. We demonstrate that our implementations of BDRS achieve state-of-the-art performance on their respective input formats. To the best of our knowledge, this is the first cross-solver study of GPU DOT solvers with a unified measure of optimality.

Figures & tables

Format	Solver	Memory	Flops per iteration
General	General (ours)	$O(mn)$	$O(mn)$
	DROT Mai et al. (2021)	$O(mn)$	$O(mn)$
	PDOT Lu and Yang (2024)	$O(mn)$	$O(mn)$
	cuPDLPx Lu et al. (2025)	$O(mn)$	$O(mn)$
	Sinkhorn Cuturi (2013) ; Flamary et al. (2021)	$O(mn)$	$O(mn)$
Point Cloud	Point Cloud (ours)	$O((m+n)d)$	$O(mnd)$

Table 1: Device memory requirement and per-iteration FLOPs, by format.

d

is the dimension of the point clouds, and typically

d\ll\min(m,n)

s

is a parameter dependent on data sparsity, with

s\leq mn

and typically

s\ll mn

Our solvers
Component	Version
Operating system	Linux-5.14.0-570.58.1.el9_6.x86_64-x86_64-with-glibc2.34
Python	3.11.16
PyTorch	2.13.0+cu130
CUDA (PyTorch runtime)	13.0
CUDA toolkit ( nvcc )	Build cuda_13.0.r13.0/compiler.36424714_0

Table 2: Software versions recorded on the machine that experiments in Section 5 are performed on.

Solver	Format	Largest $n_{g}$ solved	Why not larger
Grid BDRS (ours)	Grid	$1024$	full ladder
HOT	Grid	$512$	device OOM
Conv. Sinkhorn	Grid	$1024$	full ladder
Point Cloud BDRS (ours)	Point Cloud	$256$	timeout
HALO	Point Cloud	$512$	timeout
M3S	Point Cloud ∗	$512$	gate

Table 3: The list of GPU-accelerated DOT solvers considered and their associated format, as well as the resolution of the largest DOTmark they can solve and why they cannot solve larger ones. The rows of our three solvers are for the stopping tolerance

\mathrm{tol}=10^{-4}

. Full ladder : every size was solved. Timeout : the next size hit the

1800

seconds limit on all or most instances attempted. Device OOM and host OOM : the next size exceeds the device or the host memory allocated to a solve. Gate : the next size was declined by the solver. ∗ M3S requires its supports to be a grid but does not require a separable cost, so we classify it as a Point Cloud format solver for the purposes of our comparisons.

Figure 1: Comparison of all solvers across all three formats on the same

414

co-solved DOTmark instances at resolution

n_{g}=64

with

\ell_{2}^{2}

cost. We use a different color for each unique underlying algorithm, and use a different shape to indicate the format that the implementation consumes.

Figure 2: Our BDRS General solver against the baseline General solvers on lifted DOTmark. In particular, the

\ell_{1},\ell_{2},

and

\ell_{\infty}

norms were used at each size

n_{g}\in\{32,64,128,256\}

. Every point aggregates all

450

instances of its choice of size and cost on DOTmark, except at

n_{g}=128

with the

\ell_{1}

and

\ell_{\infty}

costs, where the co-solved sets hold the

447

and

442

instances on which PDOT finished inside the time limit. The cross on each panel marks cuPDLPx’s failed attempt at

n_{g}=128

due to a host-side OOM. Moreover, at

n_{g}=256

, our General solver solve all instances across all norms to

\mathrm{tol}=10^{-4}

, while all baseline solvers fail the solve due to memory constraints.

Figure 3: Our General BDRS solver against the baseline General solvers on random cost matrices whose entries are drawn independently from the uniform distribution (left) and from the half-normal distribution (right), with ten instances per cost family at each size

m=n=N\in\{2^{10},2^{12},2^{14},2^{16}\}

, with our solver at

\mathrm{tol}\in\{10^{-4},10^{-6}\}

. On the half-normal cost, the points at

N=2^{12}

and

2^{14}

aggregate only the

9

and

6

of the

10

instances on which PDOT finished inside the time limit. At

N=2^{16}

, the half-normal points aggregate the

8

of the

10

instances on which our solver did not exceed device memory, and on the uniform cost, our

\mathrm{tol}=10^{-6}

arm exceeds the time limit on all

10

Panel of Figure 4 — Figure 4: Our Point Cloud solver with $\mathrm{tol}\in\{10^{-4},10^{-6}\}$ against (a) HALO on ModelNet10 and (b) HALO and M3S on lifted DOTmark. (a) Our Point Cloud solver and HALO in its tree mode on ModelNet10 with the squared Euclidean cost, thirty pairs at each size from $N=2^{7}$ to $2^{16}$ points per marginal. The cross at the left edge represents the fact that HALO rejects the $2^{7}$ and $2^{8}$ instances. (b) Our Point Cloud solver, HALO, and M3S on lifted DOTmark with the $\ell_{2}^{2}$ cost, with resolutions $n_{g}\in\{32,64,128,256\}$ . The co-solved set holds $450$ instances at $n_{g}=32$ and $64$ , $439$ at $128$ and $370$ at $256$ . The excluded instances are those on which M3S aborted the solve due to “ too many absorptions ” at $n_{g}=128$ and $n_{g}=256$ . Our $\mathrm{tol}=10^{-6}$ arm reaches $N=2^{15}$ in (a) and $n_{g}=128$ in (b), the largest sizes at which it completes within the $1800$ s budget. One of the thirty instances at $N=2^{15}$ exceeds it and is excluded.

Figure 5: The Grid solvers on DOTmark at

n_{g}=32

1024

, with the

\ell_{2}^{2}

cost (left) and the

\ell_{1}

cost (right). HOT only accepts the

\ell_{2}^{2}

cost, and therefore appears only on the left panel. Our Grid solver is drawn at

\mathrm{tol}\in\{10^{-4},10^{-6},10^{-8}\}

, and the convolutional Sinkhorn at

\varepsilon\in\{10^{-2},10^{-3},10^{-4}\}

. Each point aggregates all

450

co-solved instances of its size. The

\varepsilon=10^{-4}

Conv. Sinkhorn points with the

\ell_{2}^{2}

cost aggregate

285

366

instances per size, with the rest being solve failures due to the returned transport plan being non-finite. With the

\ell_{1}

cost it return non-finite values on more than half of every size and therefore is not drawn. The cross marks HOT at

n_{g}=1024

, out of device memory on every instance. HOT’s certificate lines are hidden under its markers at this scale. The four remaining crosses mark our Grid solver at the first size each fails to complete within the

1800

s budget.

Figure 6: The median of various residuals against iteration count for the three Douglas–Rachford solvers (BDRS, DROT, HOT), on the same

50

DOTmark instances, with the

\ell_{2}^{2}

cost and

\mathrm{tol}=10^{-4}

, at

n_{g}=32

in panels (a), (b), and (c) (resp.

n_{g}=64

in panels (d), (e), and (f)). Each solver run for a fixed

6000

(resp.

12000

) iterations with no stopping rule and scored at

101

(resp.

163

) checks from iteration

1

. On all six panels, the blue, green, and orange curves are the medians of

r_{P}^{k}

r_{D}^{k}

, and

|g^{k}|

for the triple

(X^{k},\alpha^{k},\beta^{k})

, respectively. The magenta curve is the median KKT residual of the corrected triple

(X^{k},\alpha^{\prime},\beta^{\prime})

, whose duals are the c-transformed corrected duals . Note that the corrected duals are never fed back into the iterations and only serves as a termination check proxy.

Panel of Figure 7 — Figure 7: Sensitivity of our solvers to the (a) step size $\eta$ and (b) check period $T_{\mathrm{chk}}$ , on DOTmark at $\mathrm{tol}=10^{-4}$ . Each solve’s iteration count in (a), and its solve wall time in (b), is divided by the same quantity of the same instance at the deployed value of the swept constant, and the solid horizontal line marks a ratio of one. For both panels, one curve is plotted per choice of ground cost and resolution. (a) resolutions $n_{g}\in\{64,128\}$ was used, with $50$ instances each (five sampled pairs from all ten classes). The dashed vertical line marks the production value. (b) The production $\eta$ was used at $T_{\mathrm{chk}}$ ranging from $50$ to $5000$ . Resolutions $n_{g}\in\{32,64,128\}$ was used, with $30$ instances each (three sampled pairs from all ten classes). The rings mark the deployed period of each resolution, $100$ , $200$ and $500$ at $n_{g}=32$ , $64$ and $128$ , where each curve is one by construction.

Optimal Transport

Optimal transport compares and transforms probability distributions by finding low-cost ways to move mass between them.

See all 284 papers

Sep 27, 2026

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

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.

Samuel J. K. Chin, Maximilian Schiffer

MIT · HEC Paris

Primal-Dual Optimization Optimal Transport

Ask PDF

Apr 4, 2026

Fast Log-Domain Sinkhorn Optimal Transport with Warp-Level GPU Reductions

Entropic regularized optimal transport (OT) via the Sinkhorn algorithm has become a fundamental tool in machine learning, yet existing implementations either suffer from numerical instability for small regularization parameters or incur significant overhead from deep learning frameworks. We present FastSinkhorn, a lightweight, native CUDA implementation of the log-domain Sinkhorn algorithm that combines warp-level shuffle reductions with shared-memory tiling to achieve high GPU utilization without sacrificing numerical stability. Our solver operates entirely in the log-domain, enabling robust computation for regularization parameters as small as epsilon = 10^{-4} where standard-domain methods fail. On dense OT problems with n = m = 8192, our implementation achieves 12x speedup over the widely-used POT library and 5.9x speedup over GPU-accelerated PyTorch baselines, while consuming only 256 MB of GPU memory. We validate our solver on image color transfer, 3D point cloud matching, and convergence analysis, demonstrating that native CUDA kernels with careful numerical treatment provide a practical and efficient foundation for large-scale optimal transport computation.

Hao Xiao

ATLAS AI Lab

Optimal Transport GPU Acceleration GPU Kernel Optimization

Ask PDF

Oct 1, 2026

FlashSinkhorn 2: Block-Sparse Entropic Optimal Transport

Streaming GPU solvers for entropic optimal transport (EOT), such as FlashSinkhorn, avoid storing the dense kernel but still evaluate all

n\times m

point pairs in every Sinkhorn iteration. We present FlashSinkhorn~2 (FS2), a solver for squared-Euclidean cost on low-dimensional point clouds that solves large discrete EOT problems to a prescribed marginal residual on a single GPU by coupling two stages. A coarse stage solves on cell centroids, lifts the potentials to every point and, when a sampled marginal check rejects the lift, continues on the centroids, replacing most point-level updates. A block-sparse fine stage then removes the centroid error that coarse updates cannot. Its Morton-ordered blocks support screening and fused tensor-core execution, and a threshold set by the block masses bounds each omitted tile's contribution to every row and column. On synthetic benchmarks, FS2 reaches the target residual on all 32 problems and GeomLoss multiscale on 10. On one A100, FS2 solves discrete EOT between two

1.34\times10^8

-particle measures from a cosmological

N

-body simulation, at an entropic blur equal to the mean interparticle distance, to an all-particle marginal residual below 0.01 in under 2.5 hours. To our knowledge, it is the largest discrete EOT problem solved to this accuracy within hours. For reproducibility, we release an open-source implementation at https://github.com/ot-triton-lab/flash-sinkhorn

Felix X. -F. Ye, Yu Chin Fabian Lim, Naigang Wang +1

Department of Mathematics & Statistics, University at Albany, Albany, NY, USA · IBM T. J. Watson Research Center, Yorktown Heights, NY, USA

Optimal Transport GPU Kernel Optimization

Ask PDF

Abstract

Optimal Transport

Explore similar work

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

Fast Log-Domain Sinkhorn Optimal Transport with Warp-Level GPU Reductions

FlashSinkhorn 2: Block-Sparse Entropic Optimal Transport