# DrainSinkhorn: Safe Elimination for Batched Entropic Optimal Transport

arXiv: https://arxiv.org/abs/2607.24741 · HTML: https://mlsanity.com/paper/2607.24741/

## 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 (most cited first, 13 of 13)

- [Sinkhorn Distances: Lightspeed Computation of Optimal Transportation Distances](https://mlsanity.com/paper/1306.0895/) (2013)
- [Multisample Flow Matching: Straightening Flows with Minibatch Couplings](https://mlsanity.com/paper/2304.14772/) (2023)
- [Near-linear time approximation algorithms for optimal transport via Sinkhorn iteration](https://mlsanity.com/paper/1705.09634/) (2017)
- [Optimal Transport Tools (OTT): A JAX Toolbox for all things Wasserstein](https://mlsanity.com/paper/2201.12324/) (2022)
- [Low-Rank Sinkhorn Factorization](https://mlsanity.com/paper/2103.04737/) (2021)
- [Linear-Time Graph Neural Networks for Scalable Recommendations](https://mlsanity.com/paper/2402.13973/) (2024)
- [Meta Optimal Transport](https://mlsanity.com/paper/2206.05262/) (2022)
- [Overrelaxed Sinkhorn-Knopp Algorithm for Regularized Optimal Transport](https://mlsanity.com/paper/1711.01851/) (2017)
- [Rethinking Initialization of the Sinkhorn Algorithm](https://mlsanity.com/paper/2206.07630/) (2022)
- [On Fitting Flow Models with Large Sinkhorn Couplings](https://mlsanity.com/paper/2506.05526/) (2025)
- [FlashSinkhorn: IO-Aware Entropic Optimal Transport on GPU](https://mlsanity.com/paper/2602.03067/) (2026)
- [Automatically Batching Control-Intensive Programs for Modern Accelerators](https://mlsanity.com/paper/1910.11141/) (2019)
- [A Truncated Newton Method for Optimal Transport](https://mlsanity.com/paper/2504.02067/) (2025)

## 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.