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

• URL: http://arxiv.org/abs/2605.00837v1
• Date: Sat, 04 Apr 2026 16:06:27 GMT
• Title: Fast Log-Domain Sinkhorn Optimal Transport with Warp-Level GPU Reductions
• Abstract summary: We present FastSinkhorn, a lightweight, native implementation of the log-domain Sinkhorn algorithm.<n>Our solver operates entirely in the log-domain, enabling computation for regularization parameters as small as epsilon = 10-4 where standard-domain methods fail.<n>We validate our solver on image color transfer, 3D point cloud matching, and convergence analysis, demonstrating that native kernels with careful numerical treatment provide a practical and efficient foundation for large-scale optimal transport.
• Score: 1.6679662639178268
• License: http://creativecommons.org/licenses/by-nc-nd/4.0/
• Abstract: 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.

Related papers

• GraphGP: Scalable Gaussian Processes with Vecchia's Approximation [0.0]
  Vecchia's approximation is a sparse precision matrix approximation for stationary, decaying kernels.<n>We present GraphGP, an algorithm for Vecchia's approximation that scales to nearly a billion parameters with linear time and memory requirements.
  arXiv  Detail & Related papers  (2026-06-09T19:50:27Z)
• ChunkFT: Byte-Streamed Optimization for Memory-Efficient Full Fine-Tuning [58.54940026861599]
  textscChunkFT is a memory-efficient fine-tuning framework.<n>textscChunkFT enables gradient computation for arbitrary sub-tensors without modifying the network architecture.<n>textscChunkFT consistently outperforms existing memory-efficient baselines.
  arXiv  Detail & Related papers  (2026-05-20T13:44:44Z)
• cuRegOT: A GPU-Accelerated Solver for Entropic-Regularized Optimal Transport [8.189956090077677]
  cuRegOT is a high-performance GPU solver tailored for entropic-regularized OT.<n>We introduce a suite of algorithmic and architectural optimizations, including an amortized symbolic analysis strategy to mitigate CPU bottlenecks.<n>CuRegOT achieves significant speedups over state-of-the-art GPU-based solvers across a variety of benchmark tasks.
  arXiv  Detail & Related papers  (2026-05-09T08:27:39Z)
• Libra: Synergizing CUDA and Tensor Cores for High-Performance Sparse Matrix Multiplication [6.557224606759151]
  Modern accelerators are commonly equipped with cores and cores to accelerate sparse operators.<n>We show that utilizing one resource alone leads to inferior performance for sparse matrix multiplication, due to their respective limitations.<n>We propose a 2D-aware workload computation strategy find out the sweet point of task mapping operators, leveraging both the high performance of 2.9 cores and the low redundancy on cores.
  arXiv  Detail & Related papers  (2025-06-28T01:50:13Z)
• INR-Arch: A Dataflow Architecture and Compiler for Arbitrary-Order Gradient Computations in Implicit Neural Representation Processing [66.00729477511219]
  Given a function represented as a computation graph, traditional architectures face challenges in efficiently computing its nth-order gradient. We introduce INR-Arch, a framework that transforms the computation graph of an nth-order gradient into a hardware-optimized dataflow architecture. We present results that demonstrate 1.8-4.8x and 1.5-3.6x speedup compared to CPU and GPU baselines respectively.
  arXiv  Detail & Related papers  (2023-08-11T04:24:39Z)
• Efficient and Accurate Optimal Transport with Mirror Descent and Conjugate Gradients [13.848861021326755]
  We propose Mirror Descent Optimal Transport (MDOT) as a novel method for solving discrete optimal transport (OT) problems with high precision.<n>We solve each problem efficiently using a GPU-parallel nonlinear conjugate algorithm (PNCG) that outperforms traditional Sinkhorn iterations under weak regularization.
  arXiv  Detail & Related papers  (2023-07-17T14:09:43Z)
• Efficient Dataset Distillation Using Random Feature Approximation [109.07737733329019]
  We propose a novel algorithm that uses a random feature approximation (RFA) of the Neural Network Gaussian Process (NNGP) kernel. Our algorithm provides at least a 100-fold speedup over KIP and can run on a single GPU. Our new method, termed an RFA Distillation (RFAD), performs competitively with KIP and other dataset condensation algorithms in accuracy over a range of large-scale datasets.
  arXiv  Detail & Related papers  (2022-10-21T15:56:13Z)
• Communication-Efficient Adam-Type Algorithms for Distributed Data Mining [93.50424502011626]
  We propose a class of novel distributed Adam-type algorithms (emphi.e., SketchedAMSGrad) utilizing sketching. Our new algorithm achieves a fast convergence rate of $O(frac1sqrtnT + frac1(k/d)2 T)$ with the communication cost of $O(k log(d))$ at each iteration.
  arXiv  Detail & Related papers  (2022-10-14T01:42:05Z)
• Distributed Out-of-Memory NMF on CPU/GPU Architectures [1.0051474951635875]
  We propose an efficient out-of-memory implementation of the Non-negative Matrix Factorization (NMF) algorithm for HPC systems. Benchmark results show significant improvement of 32X to 76x speedup with the new implementation using GPU over the CPU-based NMFk.
  arXiv  Detail & Related papers  (2022-02-19T03:49:21Z)
• Instant Neural Graphics Primitives with a Multiresolution Hash Encoding [67.33850633281803]
  We present a versatile new input encoding that permits the use of a smaller network without sacrificing quality. A small neural network is augmented by a multiresolution hash table of trainable feature vectors whose values are optimized through a gradient descent. We achieve a combined speed of several orders of magnitude, enabling training of high-quality neural graphics primitives in a matter of seconds.
  arXiv  Detail & Related papers  (2022-01-16T07:22:47Z)
• Scalable Optimal Transport in High Dimensions for Graph Distances, Embedding Alignment, and More [7.484063729015126]
  We propose two effective log-linear time approximations of the cost matrix for optimal transport. These approximations enable general log-linear time algorithms for entropy-regularized OT that perform well even for the complex, high-dimensional spaces. For graph distance regression we propose the graph transport network (GTN), which combines graph neural networks (GNNs) with enhanced Sinkhorn.
  arXiv  Detail & Related papers  (2021-07-14T17:40:08Z)
• Efficient and Generic 1D Dilated Convolution Layer for Deep Learning [52.899995651639436]
  We introduce our efficient implementation of a generic 1D convolution layer covering a wide range of parameters. It is optimized for x86 CPU architectures, in particular, for architectures containing Intel AVX-512 and AVX-512 BFloat16 instructions. We demonstrate the performance of our optimized 1D convolution layer by utilizing it in the end-to-end neural network training with real genomics datasets.
  arXiv  Detail & Related papers  (2021-04-16T09:54:30Z)
• GPU-Accelerated Primal Learning for Extremely Fast Large-Scale Classification [10.66048003460524]
  One of the most efficient methods to solve L2-regularized primal problems, such as logistic regression and linear support vector machine (SVM) classification, is the widely used trust region Newton algorithm, TRON. We show that using judicious GPU-optimization principles, TRON training time for different losses and feature representations may be drastically reduced.
  arXiv  Detail & Related papers  (2020-08-08T03:40:27Z)
• Kernel methods through the roof: handling billions of points efficiently [94.31450736250918]
  Kernel methods provide an elegant and principled approach to nonparametric learning, but so far could hardly be used in large scale problems. Recent advances have shown the benefits of a number of algorithmic ideas, for example combining optimization, numerical linear algebra and random projections. Here, we push these efforts further to develop and test a solver that takes full advantage of GPU hardware.
  arXiv  Detail & Related papers  (2020-06-18T08:16:25Z)

This list is automatically generated from the titles and abstracts of the papers in this site.

This site does not guarantee the quality of this site (including all information) and is not responsible for any consequences.