Accelerated Algorithm for Sparse Regularized Partial Optimal Transport
Authors: Khoa Nguyen, Dung T. Nguyen, Thong Huynh, Hoang-Hiep Nguyen-Mau, Anh Nguyen, Minh Ngoc Dinh, Juho Kannala
Organizations: Aalto University, Konemiehentie 2, 02150 Espoo, Finland · Linköping University, 581 83 Linköping, Sweden · Hanoi University of Science and Technology, Hanoi, Vietnam · High School for the Gifted, Ho Chi Minh City, Vietnam · VinUniversity, Hanoi, Vietnam · Rochambeau French International School 9600 Forest Rd, 20814 Bethesda, Maryland, USA · School of Computing Technologies, RMIT University, Melbourne, Australia
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.
Figures & tables
Notation
Description
X∗
Solution of original POT: X∗∈argX∈Umin⟨C,X⟩
Xη
Solution of RPOT: Xη∈argX∈Uminfη(X)
Xη,α
Solution of P-RPOT: Xη,α∈argX∈SminFη(X,α)
Xˉ
Final output of ROUND-POT algorithm
Xk
Output at iteration k of PNAG-POT solver
Table 1: Notation and Definitions of Main Variables
Notation
Description
ε
Objective gap between the rounded and original POT solution: ⟨C,Xˉ⟩−⟨C,X∗⟩≤ε
ε′
Difference between P-RPOT and RPOT objectives: 0≤fη(Xη)−Fη(Xη,α,α)≤ε′
Algorithm 3 Proximal Nesterov’s Accelerated Gradient for RPOT (PNAG-POT)
Task
Method
Final Cost
Sparsity
Iterations
Time (s)
Color Transfer
CVXPY
0.028
0.9848
50,000
6007.00
APDAGD
0.044
0.9805
–
–
PNAG-POT
0.027
0.9958
1,800
673.50
Point Cloud Registration
CVXPY
0.015
0.4351
11,650
820.14
APDAGD
0.012
0.8424
–
–
PNAG-POT
0.0084
0.9229
2,940
198.12
Table 3: Comparison of the methods under quadratic regularization.
Figure 1: Original source and target images and the source image with colors transferred using the compared POT methods.
Figure 2: Heatmaps of transport plans for color transfer obtained by (a) CVXPY, (b) APDAGD, and (c) PNAG-POT.
Figure 3: (a) Original point clouds, (b) transported point clouds using CVXPY, and (c) transported point clouds using PNAG-POT. The point clouds are visualized in 2D instead of 3D.
Figure 4: Heatmaps of transport plans for point cloud registration obtained by (a) CVXPY, (b) APDAGD, and (c) PNAG-POT.
Figure 5: Decision boundaries of the SVM on the transported data obtained by (a) CVXPY, (b) APDAGD, and (c) PNAG-POT.
Figure 6: Heatmaps of transport plans for domain adaptation obtained by (a) CVXPY, (b) APDAGD, and (c) PNAG-POT.
Task
Method
Final Cost
Sparsity
Iterations
Time (s)
Color Transfer
CVXPY
0.028
0.9847
50,000
7607.85
PNAG-POT
0.027
0.9959
1,800
620.22
Point Cloud Registration
CVXPY
0.015
0.4218
11,675
489.58
PNAG-POT
0.0084
0.9281
2,940
204.49
Domain Adaptation
CVXPY
0.036
0.4644
8,025
70.64
PNAG-POT
0.036
0.9181
1,420
8.46
Table 4: Comparison between CVXPY and PNAG-POT under Elastic-Net regularization.
Figure 7: Domain adaptation results under Elastic-Net regularization obtained by (a) CVXPY, (b) APDAGD, and (c) PNAG-POT.
Appendix figures & tables6 assets
Supplementary material from the paper’s appendix.
Appendix
Figure 8: Color-transfer results under Elastic-Net regularization. The figure shows the source and target images together with the transferred results obtained by the compared POT methods.
Figure 9: Transport-plan heatmaps for color transfer under Elastic-Net regularization.
Figure 10: Point-cloud registration results under Elastic-Net regularization.
Figure 11: Transport-plan heatmaps for point cloud registration under Elastic-Net regularization.
Figure 12: Domain-adaptation results under Elastic-Net regularization.
Figure 13: Transport-plan heatmaps for domain adaptation under Elastic-Net regularization.
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.
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×1024 images, BDRS attains a lower repaired transport cost than MDOT-TNT with a 9× speed up, reaching a relative duality gap of 1.59% in 24 minutes. We further demonstrate a color transfer with 4238×2365 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.
Optimal transport (OT) has emerged as a fundamental tool in modern machine learning, yet its computational cost remains a significant bottleneck for large-scale applications. While harnessing the massive parallelism of modern GPU hardware is critical for efficiency, the de facto standard Sinkhorn algorithm, despite its ease of parallelization, often suffers from slow convergence in challenging problems. More recently, the sparse-plus-low-rank quasi-Newton method offers a balance between convergence rate and per-iteration complexity; however, its efficiency on GPUs is severely hindered by the serial nature of sparse matrix symbolic analysis and irregular memory access patterns. To bridge this gap, we present cuRegOT, a high-performance GPU solver tailored for entropic-regularized OT. We introduce a suite of algorithmic and architectural optimizations, including an amortized symbolic analysis strategy to mitigate CPU bottlenecks, an asynchronous Sinkhorn iterates generation mechanism, and a fused kernel for bandwidth-efficient gradient evaluation. These strategies are backed by rigorous theoretical guarantees ensuring algorithmic convergence. Extensive numerical experiments demonstrate that cuRegOT achieves significant speedups over state-of-the-art GPU-based solvers across a variety of benchmark tasks.
Yixuan Qiu
School of Statistics and Data Science & Institute of Big Data Research Shanghai University of Finance and Economics