# A Unified Framework for Implicit Sinkhorn Differentiation

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

## Key facts

- Cited by 28 papers in the map (data through 14 Oct 2026).
- Builds on 30 of its 48 references that are in the map.
- First posted on arXiv 13 May 2022.

## Builds on (most cited first, 30 of 30)

- [Adam: A Method for Stochastic Optimization](https://mlsanity.com/paper/1412.6980/) (2014)
- [PyTorch: An Imperative Style, High-Performance Deep Learning Library](https://mlsanity.com/paper/1912.01703/) (2019)
- [PointNet: Deep Learning on Point Sets for 3D Classification and Segmentation](https://mlsanity.com/paper/1612.00593/) (2016)
- [Sinkhorn Distances: Lightspeed Computation of Optimal Transportation Distances](https://mlsanity.com/paper/1306.0895/) (2013)
- [SuperGlue: Learning Feature Matching with Graph Neural Networks](https://mlsanity.com/paper/1911.11763/) (2019)
- [OptNet: Differentiable Optimization as a Layer in Neural Networks](https://mlsanity.com/paper/1703.00443/) (2017)
- [Learning Generative Models with Sinkhorn Divergences](https://mlsanity.com/paper/1706.00292/) (2017)
- [Differentiation of Blackbox Combinatorial Solvers](https://mlsanity.com/paper/1912.02175/) (2019)
- [Fast Computation of Wasserstein Barycenters](https://mlsanity.com/paper/1310.4375/) (2013)
- [Efficient and Modular Implicit Differentiation](https://mlsanity.com/paper/2105.15183/) (2021)
- [Learning with a Wasserstein Loss](https://mlsanity.com/paper/1506.05439/) (2015)
- [RPM-Net: Robust Point Matching using Learned Features](https://mlsanity.com/paper/2003.13479/) (2020)
- [Learning Latent Permutations with Gumbel-Sinkhorn Networks](https://mlsanity.com/paper/1802.08665/) (2018)
- [Improving GANs Using Optimal Transport](https://mlsanity.com/paper/1803.05573/) (2018)
- [Differentiable Ranks and Sorting using Optimal Transport](https://mlsanity.com/paper/1905.11885/) (2019)
- [Deep Shells: Unsupervised Shape Correspondence with Optimal Transport](https://mlsanity.com/paper/2010.15261/) (2020)
- [Faster Wasserstein Distance Estimation with the Sinkhorn Divergence](https://mlsanity.com/paper/2006.08172/) (2020)
- [Differential Properties of Sinkhorn Approximation for Learning with Wasserstein Distance](https://mlsanity.com/paper/1805.11897/) (2018)
- [Solving the Blind Perspective-n-Point Problem End-To-End With Robust Differentiable Geometric Optimization](https://mlsanity.com/paper/2007.14628/) (2020)
- [Wasserstein Dictionary Learning: Optimal Transport-based unsupervised non-linear dictionary learning](https://mlsanity.com/paper/1708.01955/) (2017)
- [DeepPermNet: Visual Permutation Learning](https://mlsanity.com/paper/1704.02729/) (2017)
- [Sinkhorn Barycenters with Free Support via Frank-Wolfe Algorithm](https://mlsanity.com/paper/1905.13194/) (2019)
- [Learning Permutations with Sinkhorn Policy Gradient](https://mlsanity.com/paper/1805.07010/) (2018)
- [Inverse optimal transport](https://mlsanity.com/paper/1905.03950/) (2019)
- [Wasserstein Discriminant Analysis](https://mlsanity.com/paper/1608.08063/) (2016)
- [Super-efficiency of automatic differentiation for functions defined as a minimum](https://mlsanity.com/paper/2002.03722/) (2020)
- [Differentiable Top-k Operator with Optimal Transport](https://mlsanity.com/paper/2002.06504/) (2020)
- [Mapping in a cycle: Sinkhorn regularized unsupervised learning for point cloud shapes](https://mlsanity.com/paper/2007.09594/) (2020)
- [Learning 2D-3D Correspondences To Solve The Blind Perspective-n-Point Problem](https://mlsanity.com/paper/2003.06752/) (2020)
- [Supervised Quantile Normalization for Low-rank Matrix Approximation](https://mlsanity.com/paper/2002.03229/) (2020)

## Built on by (most cited first, 28 of 28)

- [Re-basin via implicit Sinkhorn differentiation](https://mlsanity.com/paper/2212.12042/) (2022)
- [Scalable Optimal Transport Methods in Machine Learning: A Contemporary Survey](https://mlsanity.com/paper/2305.05080/) (2023)
- [Image Clustering via the Principle of Rate Reduction in the Age of Pretrained Models](https://mlsanity.com/paper/2306.05272/) (2023)
- [Generalized Linear Mode Connectivity for Transformers](https://mlsanity.com/paper/2506.22712/) (2025)
- [Unsupervised Manifold Linearizing and Clustering](https://mlsanity.com/paper/2301.01805/) (2023)
- [Flow Factorized Representation Learning](https://mlsanity.com/paper/2309.13167/) (2023)
- [Robot Fleet Learning via Policy Merging](https://mlsanity.com/paper/2310.01362/) (2023)
- [FlashSinkhorn: IO-Aware Entropic Optimal Transport on GPU](https://mlsanity.com/paper/2602.03067/) (2026)
- [Differentiable Transportation Pruning](https://mlsanity.com/paper/2307.08483/) (2023)
- [The derivatives of Sinkhorn-Knopp converge](https://mlsanity.com/paper/2207.12717/) (2022)
- [Benchmarking PtO and PnO Methods in the Predictive Combinatorial Optimization Regime](https://mlsanity.com/paper/2311.07633/) (2023)
- [The quest for the GRAph Level autoEncoder (GRALE)](https://mlsanity.com/paper/2505.22109/) (2025)
- [Gaussian Processes on Distributions based on Regularized Optimal Transport](https://mlsanity.com/paper/2210.06574/) (2022)
- [Your contrastive learning problem is secretly a distribution alignment problem](https://mlsanity.com/paper/2502.20141/) (2025)
- [JPmHC Dynamical Isometry via Orthogonal Hyper-Connections](https://mlsanity.com/paper/2602.18308/) (2026)
- [SOTAlign: Semi-Supervised Alignment of Unimodal Vision and Language Models via Optimal Transport](https://mlsanity.com/paper/2602.23353/) (2026)
- [Unsupervised Deep Graph Matching Based on Cycle Consistency](https://mlsanity.com/paper/2307.08930/) (2023)
- [Scalable unsupervised alignment of general metric and non-metric structures](https://mlsanity.com/paper/2406.13507/) (2024)
- [Improving the matrix multiplication exponent with modern optimization and AlphaEvolve](https://mlsanity.com/paper/2608.16884/) (2026)
- [Learning low-dimensional representations of ensemble forecast fields using autoencoder-based methods](https://mlsanity.com/paper/2502.04409/) (2025)
- [Avoiding Premature Collapse: Adaptive Annealing for Entropy-Regularized Structural Inference](https://mlsanity.com/paper/2601.23039/) (2026)
- [How to Backpropagate through Hungarian in Your DETR?](https://mlsanity.com/paper/2211.14448/) (2022)
- [Optimal Transport under Group Fairness Constraints](https://mlsanity.com/paper/2601.07144/) (2026)
- [AutoSew: A Geometric Approach to Stitching Prediction with Graph Neural Networks](https://mlsanity.com/paper/2602.22052/) (2026)
- [JKO-RAG: Distributional Retrieval as Wasserstein Free-Energy Gradient Flow](https://mlsanity.com/paper/2607.24776/) (2026)
- [Support Topology and Gradient Mixing in Sinkhorn Layers](https://mlsanity.com/paper/2609.07954/) (2026)
- [Proximal Residual Value Functions for Consistent Planning and Real-Time Execution](https://mlsanity.com/paper/2609.23242/) (2026)
- [ColNanoVDR: Document-Free Query Distillation for Multi-Vector Visual Document Retrieval via Optimal Transport](https://mlsanity.com/paper/2609.34899/) (2026)

## Abstract (arXiv)

The Sinkhorn operator has recently experienced a surge of popularity in computer vision and related fields. One major reason is its ease of integration into deep learning frameworks. To allow for an efficient training of respective neural networks, we propose an algorithm that obtains analytical gradients of a Sinkhorn layer via implicit differentiation. In comparison to prior work, our framework is based on the most general formulation of the Sinkhorn operator. It allows for any type of loss function, while both the target capacities and cost matrices are differentiated jointly. We further construct error bounds of the resulting algorithm for approximate inputs. Finally, we demonstrate that for a number of applications, simply replacing automatic differentiation with our algorithm directly improves the stability and accuracy of the obtained gradients. Moreover, we show that it is computationally more efficient, particularly when resources like GPU memory are scarce.