# Sinkformers: Transformers with Doubly Stochastic Attention

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

## Key facts

- Cited by 138 papers in the map (data through 14 Oct 2026).
- Builds on 40 of its 64 references that are in the map.
- First posted on arXiv 22 Oct 2021.

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

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [Deep Residual Learning for Image Recognition](https://mlsanity.com/paper/1512.03385/) (2015)
- [Adam: A Method for Stochastic Optimization](https://mlsanity.com/paper/1412.6980/) (2014)
- [Language Models are Few-Shot Learners](https://mlsanity.com/paper/2005.14165/) (2020)
- [An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale](https://mlsanity.com/paper/2010.11929/) (2020)
- [Neural Machine Translation by Jointly Learning to Align and Translate](https://mlsanity.com/paper/1409.0473/) (2014)
- [Neural Ordinary Differential Equations](https://mlsanity.com/paper/1806.07366/) (2018)
- [Sinkhorn Distances: Lightspeed Computation of Optimal Transportation Distances](https://mlsanity.com/paper/1306.0895/) (2013)
- [fairseq: A Fast, Extensible Toolkit for Sequence Modeling](https://mlsanity.com/paper/1904.01038/) (2019)
- [Point Transformer](https://mlsanity.com/paper/2011.00931/) (2020)
- [An overview of gradient descent optimization algorithms](https://mlsanity.com/paper/1609.04747/) (2016)
- [Scaling Vision Transformers](https://mlsanity.com/paper/2106.04560/) (2021)
- [PCT: Point cloud transformer](https://mlsanity.com/paper/2012.09688/) (2020)
- [Unpaired Image-to-Image Translation using Cycle-Consistent Adversarial Networks](https://mlsanity.com/paper/1703.10593/) (2017)
- [Deep Equilibrium Models](https://mlsanity.com/paper/1909.01377/) (2019)
- [A Mean Field View of the Landscape of Two-Layers Neural Networks](https://mlsanity.com/paper/1804.06561/) (2018)
- [Differentiable Convex Optimization Layers](https://mlsanity.com/paper/1910.12430/) (2019)
- [Attention is Not All You Need: Pure Attention Loses Rank Doubly Exponentially with Depth](https://mlsanity.com/paper/2103.03404/) (2021)
- [Are Transformers universal approximators of sequence-to-sequence functions?](https://mlsanity.com/paper/1912.10077/) (2019)
- [Augmented Neural ODEs](https://mlsanity.com/paper/1904.01681/) (2019)
- [Invertible Residual Networks](https://mlsanity.com/paper/1811.00995/) (2018)
- [Sparse Sinkhorn Attention](https://mlsanity.com/paper/2002.11296/) (2020)
- [Large Kernel Matters -- Improve Semantic Segmentation by Global Convolutional Network](https://mlsanity.com/paper/1703.02719/) (2017)
- [Beyond Finite Layer Neural Networks: Bridging Deep Architectures and Numerical Differential Equations](https://mlsanity.com/paper/1710.10121/) (2017)
- [Efficient and Modular Implicit Differentiation](https://mlsanity.com/paper/2105.15183/) (2021)
- [Deep Neural Networks Motivated by Partial Differential Equations](https://mlsanity.com/paper/1804.04272/) (2018)
- [The Lipschitz Constant of Self-Attention](https://mlsanity.com/paper/2006.04710/) (2020)
- [On Differentiating Parameterized Argmin and Argmax Problems with Application to Bi-level Optimization](https://mlsanity.com/paper/1607.05447/) (2016)
- [Kernel Operations on the GPU, with Autodiff, without Memory Overflows](https://mlsanity.com/paper/2004.11127/) (2020)
- [A Mean-Field Optimal Control Formulation of Deep Learning](https://mlsanity.com/paper/1807.01083/) (2018)

## Built on by (most cited first, 30 of 138)

- [A mathematical perspective on Transformers](https://mlsanity.com/paper/2312.10794/) (2023)
- [The emergence of clusters in self-attention dynamics](https://mlsanity.com/paper/2305.05465/) (2023)
- [A Survey on LoRA of Large Language Models](https://mlsanity.com/paper/2407.11046/) (2024)
- [How Smooth Is Attention?](https://mlsanity.com/paper/2312.14820/) (2023)
- [Transformers are Universal In-context Learners](https://mlsanity.com/paper/2408.01367/) (2024)
- [Sparsity-Constrained Optimal Transport](https://mlsanity.com/paper/2209.15466/) (2022)
- [Clustering in Causal Attention Masking](https://mlsanity.com/paper/2411.04990/) (2024)
- [The Mean-Field Dynamics of Transformers](https://mlsanity.com/paper/2512.01868/) (2025)
- [Measure-to-measure interpolation using Transformers](https://mlsanity.com/paper/2411.04551/) (2024)
- [Do Residual Neural Networks discretize Neural Ordinary Differential Equations?](https://mlsanity.com/paper/2205.14612/) (2022)
- [Learning Gaussian Mixtures Using the Wasserstein-Fisher-Rao Gradient Flow](https://mlsanity.com/paper/2301.01766/) (2023)
- [Scalable Optimal Transport Methods in Machine Learning: A Contemporary Survey](https://mlsanity.com/paper/2305.05080/) (2023)
- [Quantitative Clustering in Mean-Field Transformer Models](https://mlsanity.com/paper/2504.14697/) (2025)
- [A multiscale analysis of mean-field transformers in the moderate interaction regime](https://mlsanity.com/paper/2509.25040/) (2025)
- [Adversarial Robustness in Graph Neural Networks: A Hamiltonian Approach](https://mlsanity.com/paper/2310.06396/) (2023)
- [Image Clustering via the Principle of Rate Reduction in the Age of Pretrained Models](https://mlsanity.com/paper/2306.05272/) (2023)
- [Setting the Record Straight on Transformer Oversmoothing](https://mlsanity.com/paper/2401.04301/) (2024)
- [Implicit regularization of deep residual networks towards neural ODEs](https://mlsanity.com/paper/2309.01213/) (2023)
- [How do Transformers perform In-Context Autoregressive Learning?](https://mlsanity.com/paper/2402.05787/) (2024)
- [Synchronization on circles and spheres with nonlinear interactions](https://mlsanity.com/paper/2405.18273/) (2024)
- [The Intrinsic Dimension of Prompts in Internal Representations of Large Language Models](https://mlsanity.com/paper/2501.10573/) (2025)
- [Unlocking Slot Attention by Changing Optimal Transport Costs](https://mlsanity.com/paper/2301.13197/) (2023)
- [Clustering in pure-attention hardmax transformers and its role in sentiment analysis](https://mlsanity.com/paper/2407.01602/) (2024)
- [Rethinking Initialization of the Sinkhorn Algorithm](https://mlsanity.com/paper/2206.07630/) (2022)
- [Neural ODE Transformers: Analyzing Internal Dynamics and Adaptive Fine-tuning](https://mlsanity.com/paper/2503.01329/) (2025)
- [Unsupervised Manifold Linearizing and Clustering](https://mlsanity.com/paper/2301.01805/) (2023)
- [SNEkhorn: Dimension Reduction with Symmetric Entropic Affinities](https://mlsanity.com/paper/2305.13797/) (2023)
- [Synchronization of mean-field models on the circle](https://mlsanity.com/paper/2507.22857/) (2025)
- [Designing Robust Transformers using Robust Kernel Density Estimation](https://mlsanity.com/paper/2210.05794/) (2022)
- [FlashSinkhorn: IO-Aware Entropic Optimal Transport on GPU](https://mlsanity.com/paper/2602.03067/) (2026)

## Abstract (arXiv)

Attention based models such as Transformers involve pairwise interactions between data points, modeled with a learnable attention matrix. Importantly, this attention matrix is normalized with the SoftMax operator, which makes it row-wise stochastic. In this paper, we propose instead to use Sinkhorn's algorithm to make attention matrices doubly stochastic. We call the resulting model a Sinkformer. We show that the row-wise stochastic attention matrices in classical Transformers get close to doubly stochastic matrices as the number of epochs increases, justifying the use of Sinkhorn normalization as an informative prior. On the theoretical side, we show that, unlike the SoftMax operation, this normalization makes it possible to understand the iterations of self-attention modules as a discretized gradient-flow for the Wasserstein metric. We also show in the infinite number of samples limit that, when rescaling both attention matrices and depth, Sinkformers operate a heat diffusion. On the experimental side, we show that Sinkformers enhance model accuracy in vision and natural language processing tasks. In particular, on 3D shapes classification, Sinkformers lead to a significant improvement.