You Need Better Attention Priors

AI summary

Overview

Research area: Machine learning — Transformer architecture design, specifically the mathematical foundations of the self-attention mechanism (optimal transport theory, positional encoding, long-context generalization).

Technical level: Advanced. The paper builds on Entropic Optimal Transport (EOT), Kullback–Leibler divergence, and the internals of scaled dot-product attention. It is accessible to readers comfortable with the attention equation and softmax math, but the derivations and kernel-compatibility arguments assume real familiarity with Transformers.

Scope: The paper reframes standard attention as a transport problem with an implicit uniform prior, derives a generalized mechanism with a learnable log-prior, and validates it on language modeling, long-context retrieval, DNA, and vision tasks.

What This Paper Is About

Standard attention applies a softmax to dot-product scores as if the default were to spread probability evenly across all positions. The authors show that this uniform assumption is a hidden, restrictive design choice baked into the mechanism, and that many known failure modes — poor length generalization, attention sinks, the need for heuristic positional encodings — trace back to it. Their goal is to replace that naive assumption with a learnable prior that the model can shape, while keeping the computation cheap and compatible with optimized attention kernels.

Key Contributions

1. A generalized EOT formulation of attention. The authors prove that if the Shannon entropy regularizer in the optimal-transport objective is replaced by a KL divergence against an arbitrary prior π, the optimal attention distribution is simply a softmax over content scores shifted by the log-prior: p* = softmax(s/τ + log π). Standard attention emerges as the special case where π is uniform.

2. Goat: a kernel-compatible implementation. They introduce Generalized Optimal transport Attention with Trainable priors, which absorbs the log-prior directly into the query and key vectors so the whole prior is realized inside a single unmodified scaled dot-product attention call. No dense L×L bias matrix is materialized, and FlashAttention remains fully usable.

3. A first-principles account of attention sinks. Using the EOT view, sinks are shown to be the natural optimum when query signal is weak: the posterior collapses to the prior. Goat models sinks with a dedicated key-only bias term rather than by inflating content key norms, decoupling the structural default from the semantic representation.

4. Empirical validation across modalities and scales. Improvements are shown on a small synthetic task, 125M-parameter C4 models, 1.23B-parameter Llama-style models on FineWeb, long-context retrieval benchmarks, human genome modeling, and ImageNet-1k vision transformers.

Main Findings

• In-distribution fidelity plus length extrapolation. On 125M-parameter models trained on C4 with a 2048-token window, Goat lowers in-distribution perplexity by 1.55 points versus ALiBi while still generalizing to roughly 16× the training length, where RoPE degrades catastrophically.

• Billion-parameter scaling holds up. At 1.23B parameters on FineWeb (2048-token training window), RoPE perplexity jumps from 20.82 at 2k to 28.53 at 4k, 89.45 at 8k, and 417.43 at 16k. Goat matches RoPE in-distribution (20.75 at 2k) and degrades smoothly to only 24.77 at 16k.

• Sinks are learned automatically and explained theoretically. Theorem 5.1 shows the posterior converges pointwise to the prior as content signal vanishes. The learned prior bias u(j) spontaneously develops a sharp spike at position 0 plus a local recency rise around position ~2000 — exactly the structural defaults the theory predicts.

• Superior stability bounds. Theorem 5.4 proves that with a uniform prior, output sensitivity to context noise converges to 1 as sequence length grows, whereas a peaked prior bounds sensitivity by (L−1)/(exp(δ)+L−1). Stability therefore requires only logarithmic growth of the prior margin, which the unconstrained sink bias supplies for free.

• Retrieval at extreme lengths. On Passkey Retrieval and Needle-in-a-Haystack, Goat keeps near-perfect accuracy well beyond training length, while rotary, position-interpolated rotary, and sinusoidal absolute encodings degrade sharply. ALiBi is noted as especially brittle on NIAH because its fixed slopes impose large negative logits on distant needles.

• Efficiency gains in practice. In DNA modeling, peak CUDA memory drops from 2.86 GB to 1.83 GB (36% reduction) with comparable training throughput, alongside lower validation NLL and better agreement with ground-truth GC% statistics (Pearson r = 0.466 vs 0.320 for RoPE).

• Multimodal generality. Applied to Vision Transformers on ImageNet-1k, Goat learns a 2D shift-invariant prior despite uniform initialization and supports zero-shot extrapolation to higher input resolutions, where absolute positional embeddings degrade.

• Disentanglement beats structural entanglement. Because Goat adds the log-prior rather than injecting position multiplicatively through rotations (as RoPE does), the positional bias is not coupled to the magnitudes ‖q‖‖k‖ of the content vectors, so the model can express strong structural defaults without distorting semantics.

Methodology in Plain English

The authors start from an existing reinterpretation of attention as a transport problem: each query is a unit of mass that must be distributed across keys, paying a cost equal to negative similarity, with an entropy bonus for staying spread out. They note that the entropy bonus is mathematically the same as a penalty for deviating from a uniform distribution — a flat prior.

They then swap the uniform distribution for a general prior π and solve the resulting optimization. The answer is clean: take the ordinary softmax and add log π to the scores.

The engineering problem is making that addition cheap. Their trick is to split each head's dimensions into a "content" subspace and a "positional" subspace, then construct composite query and key vectors with carefully chosen scaling factors. When the standard dot product is taken, the content part contributes the usual scaled similarity and the positional part contributes the log-prior, unscaled. Because everything reduces to vector arithmetic, an off-the-shelf FlashAttention kernel computes the whole thing.

The prior itself has two pieces. A relative component uses a truncated Fourier series in the position difference i − j, with learnable coefficients on fixed geometric frequencies, capturing locality, periodicity, and even suppression at chosen distances. An absolute component is a learned per-key bias u(j) — a recency term plus a small MLP plus an optional first-token bump — which supplies the model's query-independent default, i.e. its attention sink. Both components are initialized so that the model begins at standard attention or a maximum-entropy recency baseline, learning structure only when it helps.

Why This Matters

Impact on research: The paper supplies a unifying explanation for several previously separate phenomena — the necessity of positional encodings, the emergence of attention sinks, and the length-generalization gap between fixed and learned positional schemes. It reframes positional encodings as heuristic approximations to a quantity that can be derived and learned, which is a meaningful reframing for architecture research.

Real-world applications:

• Long-document and long-context LLMs — retrieval over very long inputs (contracts, codebases, multi-hour transcripts) where fixed positional encodings currently break down.
• Retrieval-augmented generation — needle-in-a-haystack robustness is the core failure mode of RAG systems with large evidence pools.
• Genomics and computational biology — DNA sequences are long, structurally repetitive, and expensive to process; the demonstrated memory reduction and NLL gains matter directly.
• Vision systems at variable resolution — medical imaging, satellite imagery, and document understanding all involve inputs whose resolution differs from training.

Industry relevance: The mechanism is a drop-in replacement for standard multi-head attention with no extra kernel work, no dense bias matrices, and lower peak memory. Combined with the 1.23B-parameter results, that makes it plausible to adopt in production training pipelines rather than only in research prototypes. It also offers a more controllable lever for sink behavior, which practitioners have long struggled to manage manually in long-context serving.

Future Directions

• Learnable frequencies. The paper fixes the Fourier frequencies to a geometric ladder based on base 10,000 and reports ablations on this choice in the appendix. Making frequencies themselves trainable is a natural extension left partially explored.

• Scaling well beyond 1.23B parameters. Whether the extrapolation advantage persists at frontier model scales, and whether the prior stabilizes or drifts during very long training runs, remains open.

• Theory-driven margin scheduling. Theorem 5.4 shows stability requires only logarithmic growth of the sink margin in sequence length. Whether that growth can be scheduled explicitly rather than learned, or used to design length-extrapolating priors by construction, is an unanswered question.

• Extending the prior family beyond the trigonometric restriction. The paper proves in the appendix that SDPA-compatibility, translation equivariance, and stability jointly restrict admissible priors to a finite trigonometric family. Relaxing any of those constraints may open a wider design space.

• More modalities and structures. The 2D image extension suggests graph, video, or hierarchical 3D structures with their own natural shift-invariance groups could be handled by the same recipe.

Target Audience

This paper is best suited for machine learning researchers and engineers working on Transformer architecture, long-context modeling, or efficient attention kernels. It will also appeal to readers interested in the theoretical foundations of attention through the lens of optimal transport, and to practitioners in genomics or vision who need length or resolution generalization. Readers without a working grasp of softmax attention, KL divergence, and positional encoding schemes such as RoPE and ALiBi will find the derivations in Sections 2 through 5 demanding.

Authors’ abstract

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.