Research map / Architectures and efficiency / Attention and transformers
ESPFormer: Doubly-Stochastic Attention with Expected Sliced Transport Plans
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Proposes ESPFormer, a doubly-stochastic attention mechanism built from Expected Sliced Transport Plans with soft sorting along axis-aligned slices, avoiding iterative Sinkhorn normalization.
Problem
Standard softmax self-attention can over-concentrate on a few tokens; doubly-stochastic attention fixes this but existing methods (Sinkformer) need iterative Sinkhorn normalization, which is computationally costly and sequential.
Approach
Treat attention as a transport plan between queries and keys, built as an expectation over axis-aligned 1D slices of lifted optimal transport plans obtained by differentiable soft sorting, with an inverse-temperature weighting over slices; this yields doubly-stochastic attention in parallel without Sinkhorn iterations.
Key results
- On Cats vs. Dogs with 1% data, ESPFormer (hard sort) reaches 58.52% vs. Sinkformer 55.07% and Transformer 49.71%; at 100% data, 81.23% vs. 79.12% for Sinkformer.
- On ModelNet40 with Set Transformer, ESPFormer achieves best 89.6% vs. Sinkformer 89.1% (reported from prior work).
- On IWSLT14 De-En fine-tune boost, ESPFormer reaches BLEU 34.64 (Transformer) and 34.83 (DiffTransformer), slightly above Sinkformer (34.61 and 34.81).
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| replace standard attention with ESPFormer hard sorting | ViT · Cats vs Dogs 1% | accuracy: 55.07 $\pm$ 3.34 -> 58.52 $\pm$ 3.73 | Helped |
| replace standard attention with ESPFormer hard sorting | ViT · Cats vs Dogs 10% | accuracy: 69.56 $\pm$ 0.32 -> 72.71 $\pm$ 0.36 | Helped |
| replace standard attention with ESPFormer hard sorting | ViT · Cats vs Dogs 25% | accuracy: 74.86 $\pm$ 0.17 -> 75.92 $\pm$ 0.28 | Helped |
| replace standard attention with ESPFormer hard sorting | ViT · Cats vs Dogs 100% | accuracy: 79.12 $\pm$ 0.17 -> 81.23 $\pm$ 0.11 | Helped |
| replace Sinkformer attention with ESPFormer | Set Transformer · ModelNet40 | mean accuracy: 88.3 -> 89.4 | Helped |
| replace DiffTransformer attention with ESPFormer | Point Cloud Transformer · ModelNet40 | mean accuracy: 92.7 -> 92.7 | No clear effect |
| replace DiffTransformer attention with ESPFormer | Transformer encoder · IMDb | mean accuracy: 85.45 -> 85.47 | No clear effect |
| replace DiffTransformer attention with ESPFormer | Transformer encoder · TweetEval | mean accuracy: 72.31 -> 72.36 | No clear effect |
| replace Transformer attention with ESPFormer and fine-tune | Transformer · IWSLT14 German-to-English | BLEU: 34.61 -> 34.64 | Helped |
| replace DiffTransformer attention with ESPFormer and fine-tune | DiffTransformer · IWSLT14 German-to-English | BLEU: 34.78 -> 34.83 | Helped |
| replace Transformer attention with ESPFormer plug-and-play | Transformer · IWSLT14 German-to-English | BLEU: 33.40 -> 33.38 | Hurt |
| use a learnable slicer instead of axis-aligned slicing (ablation) | ESPFormer ViT · Cats vs Dogs, $\tau=0.1$, $L=64$ | accuracy: 79.47 \pm 0.12 -> 78.40 \pm 0.26 | Hurt |
Surprising findings
- Gains over Sinkformer on NMT and sentiment analysis are very small (often within 0.1 points), suggesting limited practical difference in those settings.
- Learnable slicers help only when the number of slices L is small; axis-aligned slicing performs comparably without extra parameters.
Limitations
- Incompatible with strictly causal or autoregressive Transformers, since lower-triangularity plus doubly-stochasticity collapses to identity.
- Training memory scales linearly with the number of slices since each slice's permutation matrix is retained.
Experimental setup
Datasets: Cats and Dogs ModelNet40 IMDb TweetEval IWSLT14 German-to-English MNIST
Benchmarks: Cats and Dogs ModelNet40 IMDb TweetEval IWSLT14 German-to-English MNIST
Metrics: classification accuracy BLEU runtime
Baselines: Transformer Sinkformer DiffTransformer Set Transformer Point Cloud Transformer
Builds on: Sinkformer Expected Sliced Transport Plans SoftSort Differential Transformer FSPool
Closest papers by meaning
- Delta Attention: Fast and Accurate Sparse Attention Inference by Delta Correction (NeurIPS 2025 · similarity 0.84)
- AdaSplash: Adaptive Sparse Flash Attention (ICML 2025 · similarity 0.84)
- Token Statistics Transformer: Linear-Time Attention via Variational Rate Reduction (ICLR 2025 · similarity 0.84)
- Theory, Analysis, and Best Practices for Sigmoid Self-Attention (ICLR 2025 · similarity 0.84)
- Expected Sliced Transport Plans (ICLR 2025 · similarity 0.84)
- Sparse Attention as Compact Kernel Regression (NeurIPS 2026 · similarity 0.83)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.