ASAP: Amortized Doubly-Stochastic Attention via Sliced Dual Projection
Not on the map. Why?
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.
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
36 of 38 references.
References · 36
- 2017Attention Is All You Needcited by 60,361
- 2020An Image is Worth 16x16 Words: Transformers for Image Recognition at Scalecited by 30,186
- 2022FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awarenesscited by 3,673
- 2015Character-level Convolutional Networks for Text Classificationcited by 2,330
- 2013Sinkhorn Distances: Lightspeed Computation of Optimal Transportation Distancescited by 2,316
- 2023Efficient Streaming Language Models with Attention Sinkscited by 2,186
- 2020Reformer: The Efficient Transformercited by 1,668
- 2020Rethinking Attention with Performerscited by 1,601
- 2020Linformer: Self-Attention with Linear Complexitycited by 1,474
- 2021Nystr"omformer: A Nystr"om-Based Algorithm for Approximating Self-Attentioncited by 471
- 2021Attention is Not All You Need: Pure Attention Loses Rank Doubly Exponentially with Depthcited by 419
- 2020Sparse Sinkhorn Attentioncited by 251
24 more
- 2021Sinkformers: Transformers with Doubly Stochastic Attentioncited by 138
- 2023Energy-Based Sliced Wasserstein Distancecited by 39
- 2022Spherical Sliced-Wassersteincited by 31
- 2022Meta Optimal Transportcited by 28
- 2024Sliced-Wasserstein Distances and Flows on Cartan-Hadamard Manifoldscited by 28
- 2022Revisiting Sliced Wasserstein on Images: From Vectorization to Convolutioncited by 25
- 2023Quasi-Monte Carlo for 3D Sliced Wassersteincited by 22
- 2024Expected Sliced Transport Planscited by 19
- 2022Rethinking Initialization of the Sinkhorn Algorithmcited by 14
- 2024Wasserstein Wormhole: Scalable Optimal Transport Distance with Transformerscited by 14
- 2024Stereographic Spherical Sliced Wasserstein Distancescited by 12
- 2026FlashSinkhorn: IO-Aware Entropic Optimal Transport on GPUcited by 11
- 2023Slicing Unbalanced Optimal Transportcited by 10
- 2025ESPFormer: Doubly-Stochastic Attention with Expected Sliced Transport Planscited by 10
- 2025Differentiable Generalized Sliced Wasserstein Planscited by 7
- 2025Quantum Doubly Stochastic Transformerscited by 6
- 2025Fast Estimation of Wasserstein Distances via Regression on Sliced Wasserstein Distancescited by 5
- 2025LOTFormer: Doubly-Stochastic Linear Attention via Low-Rank Optimal Transportcited by 4
- 2025Efficient Transferable Optimal Transport via Min-Sliced Transport Planscited by 4
- 2025Sliced Transport Planscited by 3
- 2025An Introduction to Sliced Optimal Transportcited by 2
- 2026Sinkhorn doubly stochastic attention rank decay analysiscited by 2
- 2026Amortized Optimal Transport from Sliced Potentialscited by 2
- 2026Sliced-Regularized Optimal Transportcited by 2
Built on by
Not cited yet.