mlsanityData to 14 Oct 2026

DrainSinkhorn: Safe Elimination for Batched Entropic Optimal Transport

arXiv 2607.24741 ↗First posted 27 Jul 2026cs.DC cs.LGcited by 0

Not on the map. Why?

Abstract (arXiv)

Fast entropic optimal transport backends reduce the cost of each Sinkhorn update, but static batches still run at full width until the slowest problem finishes. We introduce DrainSinkhorn, a verifier-gated active-packing layer for batches of independent Sinkhorn problems. It combines candidate-axis packing, a Sinkhorn-specific one-sided screen, verifier-gated retirement under the backend's configured two-sided residual check, and physical compaction of all candidate-indexed state. The EOT objective, per-instance Sinkhorn map, and stopping rule are unchanged; later kernels run only on unfinished problems. We characterize the removable work exactly. If completion depths differ within a packed window, active execution removes the padding between the static batch rectangle and the observed survival curve. A quotient nonlinear Perron-Frobenius analysis gives a local explanation for these finite-tolerance depth differences: convergence depends on the full modal spectrum and proposal alignment, not only on the slowest mode. DrainSinkhorn achieves state-of-the-art execution performance on the tested heterogeneous batched-EOT workloads within matched backend families. The complete Flash-backed OT path is 4.110x faster on MetroPT-3, 3.798x faster on ImageNet-32 feature couplings, and 1.250-1.270x faster across a five-tolerance Packer19 sweep. Independent implementations reach 2.600x on ImageNet-32 with OTT-JAX, 3.174x on A2D2 LiDAR with PyKeOps, and 1.415x on large ImageNet-32 PyKeOps couplings. End-to-end speedups remain 4.074x on MetroPT-3 and 2.786x on ImageNet-32 feature-space OT flow matching, with all reported residual, consumer-output, and training-quality checks passing.

Key facts. No citing papers in the map yet (data through 14 Oct 2026). Builds on 13 of its 18 references that are in the map. First posted on arXiv 27 Jul 2026.

Builds on

13 of 18 references.

References · 13

  1. 2013Sinkhorn Distances: Lightspeed Computation of Optimal Transportation Distancescited by 2,316
  2. 2023Multisample Flow Matching: Straightening Flows with Minibatch Couplingscited by 278
  3. 2017Near-linear time approximation algorithms for optimal transport via Sinkhorn iterationcited by 223
  4. 2022Optimal Transport Tools (OTT): A JAX Toolbox for all things Wassersteincited by 101
  5. 2021Low-Rank Sinkhorn Factorizationcited by 67
  6. 2024Linear-Time Graph Neural Networks for Scalable Recommendationscited by 31
  7. 2022Meta Optimal Transportcited by 28
  8. 2017Overrelaxed Sinkhorn-Knopp Algorithm for Regularized Optimal Transportcited by 18
  9. 2022Rethinking Initialization of the Sinkhorn Algorithmcited by 14
  10. 2025On Fitting Flow Models with Large Sinkhorn Couplingscited by 11
  11. 2026FlashSinkhorn: IO-Aware Entropic Optimal Transport on GPUcited by 11
  12. 2019Automatically Batching Control-Intensive Programs for Modern Acceleratorscited by 6
1 more
  1. 2025A Truncated Newton Method for Optimal Transportcited by 6

Built on by

Not cited yet.

Similar, not cited

Close in wording, no citation either way

  1. 2026In-Context Batch Black-box Optimization via Flow Matching
  2. 2025Masked Diffusion Models as Energy Minimization
  3. 2025The Curse of Conditions: Analyzing and Improving Optimal Transport for Conditional Flow-Based Generation
  4. 2026Faster Block-Diffusion Serving with Distribution-Free Risk Guarantees
  5. 2026Distribution-Aware Algorithm Design with LLM Agents
  6. 2021Score-based Generative Neural Networks for Large-Scale Optimal Transport
  7. 2026Expected Batch Optimal Transport Plans and Consequences for Flow Matching
  8. 2026OTCache: Optimal Transport for Geometry-Aware Caching in Diffusion Models