You Need Better Attention Priors
Not on the map. Why?
Abstract (arXiv)
We generalize the attention mechanism by viewing it through the lens of Entropic Optimal Transport, revealing that standard attention corresponds to a transport problem regularized by an implicit uniform prior. We introduce Generalized Optimal transport Attention with Trainable priors (GOAT), a new attention mechanism that replaces this naive assumption with a learnable, continuous prior. This prior maintains full compatibility with optimized kernels such as FlashAttention. GOAT also provides an EOT-based explanation of attention sinks and materializes a solution for them, avoiding the representational trade-offs of standard attention. Finally, by absorbing spatial information into the core attention computation, GOAT learns an extrapolatable prior that combines the flexibility of learned positional embeddings with the length generalization of fixed encodings.
Key facts. Cited by 2 papers in the map (data through 7 Oct 2026). Builds on 26 of its 34 references that are in the map. First posted on arXiv 21 Jan 2026.
Builds on
26 of 34 references.
References · 26
- 2017Attention Is All You Needcited by 59,932
- 2020Language Models are Few-Shot Learnerscited by 32,199
- 2020An Image is Worth 16x16 Words: Transformers for Image Recognition at Scalecited by 29,915
- 2020Denoising Diffusion Probabilistic Modelscited by 20,851
- 2019Exploring the Limits of Transfer Learning with a Unified Text-to-Text Transformercited by 14,491
- 2021RoFormer: Enhanced Transformer with Rotary Position Embeddingcited by 4,581
- 2014Neural Machine Translation by Jointly Learning to Align and Translatecited by 4,530
- 2023Lost in the Middle: How Language Models Use Long Contextscited by 3,383
- 2020Longformer: The Long-Document Transformercited by 3,362
- 2023FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioningcited by 2,551
- 2013Sinkhorn Distances: Lightspeed Computation of Optimal Transportation Distancescited by 2,277
- 2023Efficient Streaming Language Models with Attention Sinkscited by 2,115
14 more
- 2020Transformers are RNNs: Fast Autoregressive Transformers with Linear Attentioncited by 1,979
- 2019Transformer-XL: Attentive Language Models Beyond a Fixed-Length Contextcited by 1,830
- 2020Big Bird: Transformers for Longer Sequencescited by 1,673
- 2020Reformer: The Efficient Transformercited by 1,658
- 2020Rethinking Attention with Performerscited by 1,584
- 2021Decision Transformer: Reinforcement Learning via Sequence Modelingcited by 1,528
- 2018Self-Attention with Relative Position Representationscited by 1,112
- 2021Train Short, Test Long: Attention with Linear Biases Enables Input Length Extrapolationcited by 1,085
- 2015Effective Approaches to Attention-based Neural Machine Translationcited by 1,024
- 2023Extending Context Window of Large Language Models via Positional Interpolationcited by 721
- 2024An Image is Worth 1/2 Tokens After Layer 2: Plug-and-Play Inference Acceleration for Large Vision-Language Modelscited by 629
- 2024FlashAttention-3: Fast and Accurate Attention with Asynchrony and Low-precisioncited by 507
- 2023Learning a Fourier Transform for Linear Relative Positional Encodings in Transformerscited by 10
- 2025Scaled-Dot-Product Attention as One-Sided Entropic Optimal Transportcited by 7
Built on by
Citing papers per year: 2026: 2.
2 papers.