You Need Better Attention Priors

TL;DR

The paper reframes self-attention as a KL-regularized transport problem from Entropic Optimal Transport, highlighting that standard attention uses an implicit uniform prior. It introduces GOAT, a trainable-prior attention mechanism that preserves SDPA compatibility by decomposing queries/keys into content and positional subspaces and incorporating a learned log-prior consisting of a spectral relative-position component and an explicit key-only sink. Theoretical results show that priors determine sinks and stability, with finite-trigonometric priors arising under SDPA-compatibility and translation equivariance, and that sinks are optimal transport defaults in low-signal regimes. Empirically, GOAT improves language modeling perplexity and long-context retrieval, reduces memory, and extends zero-shot extrapolation in vision tasks, demonstrating broad applicability across domains including biology and computer vision.

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.

Figures (5)

• Figure 1: Learned Goat log-prior decomposition on a toy copy-mixture task.Task: train a small causal LM on synthetic sequences where each token is a mixture of copying the first token (global), copying the previous token (local), or random noise. (a) Key-only sink component $K_{\mathrm{sink}}(i,j)=u(j)$. (b) Translation-equivariant relative component $K_{\mathrm{rel}}(i,j)=\kappa(i-j)$. (c) Row-centered total log-prior $K$ (row shifts do not affect softmax). (d) Induced causal prior probabilities after masking to $j \le i$ and row-normalizing with softmax (axes shown: query position $i$, key position $j$).
• Figure 2: Goat combines the flexibility of learnable priors with robust length extrapolation. Comparison of 125M parameter models trained on 4B tokens of C4 with context length $L_{\text{train}} = 2048$. (a) Length Extrapolation. Perplexity evaluated on extended sequences. RoPE degrades catastrophically past $L_{\text{train}}$. While ALiBi maintains flat extrapolation, it underfits the training window; the inset reveals Goat improves in-distribution perplexity by 1.55 points over ALiBi. Goat provides the best trade-off: superior fidelity within the window and robust generalization to $16\times$ length. (b) Learned Prior Bias $u(j)$. The model spontaneously discovers a sharp spike at $j = 0$ (an explicit attention sink) and a rise at $j \approx 2000$ (local recency), validating that these structural needs can be decoupled from content. (c) Signal Strength vs. Sink Mass (Theorem \ref{['thm:collapse_prior']}). As signal $\omega$ increases, Goat smoothly sheds sink mass from $\approx 1$ (prior-dominated) to $\approx 0$ (content-driven), while ALiBi remains over-defaulted (fixed distance bias adds an $\mathcal{O}(md)$ margin) and RoPE only partially disengages.
• Figure 3: Comparison of Goat to other methods on long-context synthetic tasks.(a) Span-token accuracy on the passkey retrieval task as a function of context length. Each model is trained as a GPT-style language model on sequences up to the training context length (e.g., $L_{\text{train}}=1024$) and evaluated on an 8-digit passkey that appears once in the context and must be reproduced at the end of the sequence. The learned prior (Goat) maintains substantially higher accuracy than rotary su2021roformer, position-interpolated rotary chen2023positionalinterpolation, and sinusoidal absolute position encodings as the context length grows far beyond the training window. (b) Visualization of the learned log-prior for a Goat attention head on a $512$-token sequence. The left panel shows the unmasked prior, which allocates larger probability mass to later key positions; the right panel shows the same prior after applying the causal mask and row-renormalization, yielding a strong recency bias along the causal diagonal. This illustrates how the prior implements an explicit, interpretable inductive bias without modifying content scores. (c) Needle-in-a-haystack (NIAH) results: span-token accuracy heatmaps over context length and needle depth (fractional position of the needle) for the same attention variants. The learned prior maintains near-perfect retrieval across depths and lengths, while the rotary, interpolated rotary, and sinusoidal baselines degrade sharply as sequences become longer and the needle moves deeper into the context.
• Figure 4: Goat outperforms RoPE on DNA modeling.(a) Validation Negative Log-Likelihood (NLL) on the Human Reference Genome (125M parameter decoder-only LM). Goat achieves lower Validation NLL (solid bars, left axis) and bits/base (hatched bars, right axis) than RoPE under identical training budgets. (b) Computational Efficiency. While training throughput (solid bars, left axis) remains comparable, Goat significantly reduces Peak CUDA memory allocation (hatched bars, right axis), dropping from $2.86$ GB to $1.83$ GB ($36\%$ reduction). (c) Generative Quality. Top: Representative completion example and metrics table. The visualization shows a single sample, while reported metrics are averaged over $N=3$ continuations. Generated nucleotides are colored by alignment to the ground truth (green = match, orange = mismatch). Bottom: Sliding-window GC% trajectory (window=25) for the generated continuation. Goat (blue) tracks the ground-truth statistical profile (black) more accurately than RoPE (purple), evidenced by a higher Pearson correlation ($r=0.466$ vs. $0.320$).
• Figure 5: Goat enables robust zero-shot resolution extrapolation on ImageNet-1kdeng2009imagenet. Models are trained at $224{\times}224$ and evaluated at higher resolutions without fine-tuning. (a) Top-1 accuracy vs. input resolution: the baseline ViT with absolute positional embeddings dosovitskiy2021imageworth16x16words degrades as resolution increases, while Goat maintains substantially higher accuracy. (b) Learned log-prior $\mathcal{K}_{ij}$ (shown relative to the center patch) exhibits a local, shift-invariant inductive bias despite uniform initialization.

Theorems & Definitions (26)

• Definition 2.1: The EOT Objective
• Proposition 2.2
• Proposition 3.1: Attention with Priors
• proof
• Theorem 5.1: Collapse to Prior
• Definition 5.2: Sink by Margin
• Definition 5.3: Total Context Sensitivity
• Theorem 5.4: Context Sensitivity Bounds
• proof
• Theorem 1: Bochner