You Need Better Attention Priors

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.

AI models process information using a mechanism called attention, which helps them decide which parts of a text or image are most relevant. Currently, these models start with a blank slate, assuming all pieces of information are equally important. This naive approach limits their ability to understand very long documents and requires massive amounts of computing memory. We propose a new method that gives the AI a better initial intuition about how data is organized. Instead of starting from scratch, our model uses a flexible, trainable guide that naturally grasps relationships like distance and order. It also creates a safe fallback for the model when it is unsure where to focus. Our upgraded attention mechanism allows AI to seamlessly understand text, images, and biological DNA sequences that are much longer than those it was originally trained on, all while using significantly less memory.