# ASAP: Amortized Doubly-Stochastic Attention via Sliced Dual Projection

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

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Builds on 36 of its 38 references that are in the map.
- First posted on arXiv 13 May 2026.

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

- [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)
- [FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness](https://mlsanity.com/paper/2205.14135/) (2022)
- [Character-level Convolutional Networks for Text Classification](https://mlsanity.com/paper/1509.01626/) (2015)
- [Sinkhorn Distances: Lightspeed Computation of Optimal Transportation Distances](https://mlsanity.com/paper/1306.0895/) (2013)
- [Efficient Streaming Language Models with Attention Sinks](https://mlsanity.com/paper/2309.17453/) (2023)
- [Reformer: The Efficient Transformer](https://mlsanity.com/paper/2001.04451/) (2020)
- [Rethinking Attention with Performers](https://mlsanity.com/paper/2009.14794/) (2020)
- [Linformer: Self-Attention with Linear Complexity](https://mlsanity.com/paper/2006.04768/) (2020)
- [Nystr"omformer: A Nystr"om-Based Algorithm for Approximating Self-Attention](https://mlsanity.com/paper/2102.03902/) (2021)
- [Attention is Not All You Need: Pure Attention Loses Rank Doubly Exponentially with Depth](https://mlsanity.com/paper/2103.03404/) (2021)
- [Sparse Sinkhorn Attention](https://mlsanity.com/paper/2002.11296/) (2020)
- [Sinkformers: Transformers with Doubly Stochastic Attention](https://mlsanity.com/paper/2110.11773/) (2021)
- [Energy-Based Sliced Wasserstein Distance](https://mlsanity.com/paper/2304.13586/) (2023)
- [Spherical Sliced-Wasserstein](https://mlsanity.com/paper/2206.08780/) (2022)
- [Meta Optimal Transport](https://mlsanity.com/paper/2206.05262/) (2022)
- [Sliced-Wasserstein Distances and Flows on Cartan-Hadamard Manifolds](https://mlsanity.com/paper/2403.06560/) (2024)
- [Revisiting Sliced Wasserstein on Images: From Vectorization to Convolution](https://mlsanity.com/paper/2204.01188/) (2022)
- [Quasi-Monte Carlo for 3D Sliced Wasserstein](https://mlsanity.com/paper/2309.11713/) (2023)
- [Expected Sliced Transport Plans](https://mlsanity.com/paper/2410.12176/) (2024)
- [Rethinking Initialization of the Sinkhorn Algorithm](https://mlsanity.com/paper/2206.07630/) (2022)
- [Wasserstein Wormhole: Scalable Optimal Transport Distance with Transformers](https://mlsanity.com/paper/2404.09411/) (2024)
- [Stereographic Spherical Sliced Wasserstein Distances](https://mlsanity.com/paper/2402.02345/) (2024)
- [FlashSinkhorn: IO-Aware Entropic Optimal Transport on GPU](https://mlsanity.com/paper/2602.03067/) (2026)
- [Slicing Unbalanced Optimal Transport](https://mlsanity.com/paper/2306.07176/) (2023)
- [ESPFormer: Doubly-Stochastic Attention with Expected Sliced Transport Plans](https://mlsanity.com/paper/2502.07962/) (2025)
- [Differentiable Generalized Sliced Wasserstein Plans](https://mlsanity.com/paper/2505.22049/) (2025)
- [Quantum Doubly Stochastic Transformers](https://mlsanity.com/paper/2504.16275/) (2025)
- [Fast Estimation of Wasserstein Distances via Regression on Sliced Wasserstein Distances](https://mlsanity.com/paper/2509.20508/) (2025)
- [LOTFormer: Doubly-Stochastic Linear Attention via Low-Rank Optimal Transport](https://mlsanity.com/paper/2509.23436/) (2025)

## Abstract (arXiv)

Doubly-stochastic attention has emerged as a transport-based alternative to row-softmax attention, with recent Transformer variants using it to reduce attention sinks and rank collapse while improving performance. In this family, the standard approach is Sinkhorn scaling, which trains more efficiently but still repeats matrix scaling in every inference forward pass. Sliced-transport attention removes the online iteration, but its soft sorting approximation materializes dense tensors for each slice, requiring substantially more training resources than Sinkhorn attention. We introduce ASAP: Amortized Doubly-Stochastic Attention via Sliced Dual Projection, a train-then-compile method that trains the doubly-stochastic layer with Sinkhorn, then replaces the iterative scaling loop at inference with a fixed sliced-dual operator. It learns a lightweight parametric map from exact one-dimensional Kantorovich potentials to the Sinkhorn query-side dual, then reconstructs the attention plan with a two-sided entropic c-transform. Across language and vision benchmarks, ASAP keeps the cheaper training setup and remains highly competitive with recent baselines. In the main frozen-layer benchmark, ASAP is 5.3 faster than the trained Sinkhorn teacher while matching its accuracy; in downstream replacements, ASAP recovers most of the teacher performance without any retraining.
