# ESPFormer: Doubly-Stochastic Attention with Expected Sliced Transport Plans

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

## Key facts

- Cited by 10 papers in the map (data through 14 Oct 2026).
- Builds on 37 of its 63 references that are in the map.
- First posted on arXiv 11 Feb 2025.

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

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale](https://mlsanity.com/paper/2010.11929/) (2020)
- [Training data-efficient image transformers & distillation through attention](https://mlsanity.com/paper/2012.12877/) (2020)
- [Sinkhorn Distances: Lightspeed Computation of Optimal Transportation Distances](https://mlsanity.com/paper/1306.0895/) (2013)
- [Transformers are RNNs: Fast Autoregressive Transformers with Linear Attention](https://mlsanity.com/paper/2006.16236/) (2020)
- [Rethinking Attention with Performers](https://mlsanity.com/paper/2009.14794/) (2020)
- [fairseq: A Fast, Extensible Toolkit for Sequence Modeling](https://mlsanity.com/paper/1904.01038/) (2019)
- [Point Transformer](https://mlsanity.com/paper/2011.00931/) (2020)
- [Transformers in Vision: A Survey](https://mlsanity.com/paper/2101.01169/) (2021)
- [Going deeper with Image Transformers](https://mlsanity.com/paper/2103.17239/) (2021)
- [PCT: Point cloud transformer](https://mlsanity.com/paper/2012.09688/) (2020)
- [Recipe for a General, Powerful, Scalable Graph Transformer](https://mlsanity.com/paper/2205.12454/) (2022)
- [Graph Transformer Networks](https://mlsanity.com/paper/1911.06455/) (2019)
- [DeiT III: Revenge of the ViT](https://mlsanity.com/paper/2204.07118/) (2022)
- [TweetEval: Unified Benchmark and Comparative Evaluation for Tweet Classification](https://mlsanity.com/paper/2010.12421/) (2020)
- [Axial Attention in Multidimensional Transformers](https://mlsanity.com/paper/1912.12180/) (2019)
- [Graph Neural Networks with Learnable Structural and Positional Representations](https://mlsanity.com/paper/2110.07875/) (2021)
- [Cross Attention Network for Few-shot Classification](https://mlsanity.com/paper/1910.07677/) (2019)
- [Generalized Sliced Wasserstein Distances](https://mlsanity.com/paper/1902.00434/) (2019)
- [Geometric Transformer for Fast and Robust Point Cloud Registration](https://mlsanity.com/paper/2202.06688/) (2022)
- [Learning Latent Permutations with Gumbel-Sinkhorn Networks](https://mlsanity.com/paper/1802.08665/) (2018)
- [FLatten Transformer: Vision Transformer using Focused Linear Attention](https://mlsanity.com/paper/2308.00442/) (2023)
- [Exphormer: Sparse Transformers for Graphs](https://mlsanity.com/paper/2303.06147/) (2023)
- [Max-Sliced Wasserstein Distance and its use for GANs](https://mlsanity.com/paper/1904.05877/) (2019)
- [Differentiable Ranks and Sorting using Optimal Transport](https://mlsanity.com/paper/1905.11885/) (2019)
- [Stochastic Optimization of Sorting Networks via Continuous Relaxations](https://mlsanity.com/paper/1903.08850/) (2019)
- [Multimodal Optimal Transport-based Co-Attention Transformer with Global Structure Consistency for Survival Prediction](https://mlsanity.com/paper/2306.08330/) (2023)
- [SoftSort: A Continuous Relaxation for the argsort Operator](https://mlsanity.com/paper/2006.16038/) (2020)
- [Sliced Wasserstein Kernels for Probability Distributions](https://mlsanity.com/paper/1511.03198/) (2015)
- [FSPool: Learning Set Representations with Featurewise Sort Pooling](https://mlsanity.com/paper/1906.02795/) (2019)

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

- [FlashSinkhorn: IO-Aware Entropic Optimal Transport on GPU](https://mlsanity.com/paper/2602.03067/) (2026)
- [Quantum Doubly Stochastic Transformers](https://mlsanity.com/paper/2504.16275/) (2025)
- [LOTFormer: Doubly-Stochastic Linear Attention via Low-Rank Optimal Transport](https://mlsanity.com/paper/2509.23436/) (2025)
- [Constrained Sliced Wasserstein Embedding](https://mlsanity.com/paper/2506.02203/) (2025)
- [Sliced ReLU attention: Quasi-linear contextual expressivity via sorting](https://mlsanity.com/paper/2512.11411/) (2025)
- [Sinkhorn doubly stochastic attention rank decay analysis](https://mlsanity.com/paper/2604.07925/) (2026)
- [Sliced Inner Product Gromov-Wasserstein Distances](https://mlsanity.com/paper/2605.08546/) (2026)
- [Streaming Sliced Optimal Transport](https://mlsanity.com/paper/2505.06835/) (2025)
- [LUNA: Linear Universal Neural Attention with Generalization Guarantees](https://mlsanity.com/paper/2512.08061/) (2025)
- [ASAP: Amortized Doubly-Stochastic Attention via Sliced Dual Projection](https://mlsanity.com/paper/2605.12879/) (2026)

## Abstract (arXiv)

While self-attention has been instrumental in the success of Transformers, it can lead to over-concentration on a few tokens during training, resulting in suboptimal information flow. Enforcing doubly-stochastic constraints in attention matrices has been shown to improve structure and balance in attention distributions. However, existing methods rely on iterative Sinkhorn normalization, which is computationally costly. In this paper, we introduce a novel, fully parallelizable doubly-stochastic attention mechanism based on sliced optimal transport, leveraging Expected Sliced Transport Plans (ESP). Unlike prior approaches, our method enforces doubly stochasticity without iterative Sinkhorn normalization, significantly enhancing efficiency. To ensure differentiability, we incorporate a temperature-based soft sorting technique, enabling seamless integration into deep learning models. Experiments across multiple benchmark datasets, including image classification, point cloud classification, sentiment analysis, and neural machine translation, demonstrate that our enhanced attention regularization consistently improves performance across diverse applications. Our implementation code can be found at https://github.com/dariansal/ESPFormer.