---
title: Generalized Optimal Transport Attention (GOAT)
url: https://www.emergentmind.com/topics/generalized-optimal-transport-attention-goat
type: topic
---

# Generalized Optimal Transport Attention (GOAT)

Generalized Optimal Transport Attention (GOAT) is a reformulation of the attention mechanism in neural networks, generalizing standard scaled-dot product attention by introducing a trainable, continuous prior under the Entropic Optimal Transport (EOT) framework. GOAT replaces the implicit uniform prior in classical attention with a learnable prior that enables more expressive structural inductive bias, providing enhanced stability, extrapolation, and efficiency while remaining fully compatible with fast attention kernels such as FlashAttention. This formulation yields a closed-form solution for the optimal transport problem and offers new insights into the behavior of attention sinks, incorporating positional and spatial priors directly into the attention computation [2601.15380].

## 1. Attention as Entropic Optimal Transport

Standard attention mechanisms compute a probability distribution over value vectors for each query via softmax normalization. GOAT casts this process as an instance of one-sided Entropic Optimal Transport. Let $L$ denote context length and $\bm{s}\in\mathbb{R}^L$ the vector of unnormalized dot-product scores, $s_j = \langle q_{c,i}, k_{c,j} \rangle / \sqrt{d_c}$. The EOT objective with Shannon entropy regularization seeks
\[
\bm{p}^\star = \arg\min_{\bm{p}\in\Delta^{L-1}} 
\big\{\langle \bm{p}, -\bm{s}\rangle - \tau\,H(\bm{p})\big\}, \quad H(\bm{p})=-\sum_j p_j\log p_j,~\tau>0.
\]
This is equivalent to minimizing
\[
\min_{p\in\Delta}~\langle p, -s \rangle + \tau KL(p \| \mathcal{U}),
\]
where $\mathcal{U}$ is the uniform prior. The solution recovers the conventional softmax:
\[
p_j^\star = \frac{\exp(s_j/\tau)}{\sum_{k} \exp(s_k/\tau)}.
\]
This establishes that standard attention is an instance of EOT regularized by a uniform prior.

## 2. Generalizing with a Trainable Prior

GOAT introduces a generalized EOT objective that replaces the uniform prior with a learnable continuous prior, $\bm{\pi} \in \Delta^{L-1}$:
\[
\bm{p}^\star = \arg\min_{p\in\Delta^{L-1}} \left\{ -\langle p, s \rangle + \tau KL(p \| \pi) \right\}.
\]
The analytic solution is
\[
p_j^\star = \frac{\pi_j \exp(s_j / \tau)}{\sum_k \pi_k \exp(s_k / \tau)},
\]
which can be written as $p^\star = \operatorname{softmax}(s/\tau + \log \pi)$. This design retains a closed-form solution in the one-sided scenario, removing the need for iterative Sinkhorn steps that arise in the general two-marginal EOT case. The trainable prior $\bm{\pi}$ enables the network to learn context- or position-dependent structural information, surpassing the limitations of a uniform prior.

## 3. EOT Perspective on Attention Sinks

In attention, sinks are locations that retain persistent attention weights regardless of context signal. From the EOT view, total logits are defined as $z_{ij} = s_{ij} + \mathcal{K}_{ij}$, with normalized prior $\pi_{ij} \propto e^{\mathcal{K}_{ij}}$. The dynamic-range signal $\omega_i = \max_k s_{ik} - \min_k s_{ik}$ quantifies the strength of the content-based scores.

Theorem 5.1 (Collapse to Prior) shows:
\[
\pi_{ij}\,e^{-\omega_i} \leq p_{ij} \leq \pi_{ij}\,e^{\omega_i}, \quad
\lim_{\omega_i \to 0} p_{ij} = \pi_{ij}.
\]
In the low-signal regime, the posterior collapses onto the prior. Sinks are characterized by margin (Definition 5.2): for query $i$, key $j^\star$ is a sink if $m_i(j^\star) = \min_{k \neq j^\star}(z_{ij^\star}-z_{ik}) > 0$, ensuring $p_{ij^\star}$ is lower-bounded independently of context length $L$.

Theorem 5.3 contrasts prior types for context sensitivity. With uniform prior, in the low-signal limit $\Psi=1-p_{ij^\star}\to (L-1)/L\to 1$ as $L$ grows. In contrast, a peaked prior with margin $\delta$ imposes
\[
\Psi \leq \frac{L-1}{e^\delta + L-1} \longrightarrow 0 ~\text{as}~ e^\delta \gg L,
\]
demonstrating that explicit learned sinks can exponentially suppress context noise.

## 4. FlashAttention Compatibility and Implementation

GOAT leverages a parameterization where the log prior $\mathcal{K}_{ij}$ is embedded directly within the attention mechanism, maintaining compatibility with highly optimized kernels such as FlashAttention. Each attention head is split into a content subspace ($d_c$ dimensions) and a positional subspace ($d_p$ dimensions), with $d_h = d_c + d_p$. The augmented queries and keys are:
\[
\bm{q}'_i = \big[~\bm{q}_{c,i} \sqrt{d_h/d_c} ~\big|~ \bm{q}_{\mathrm{rel},i} \sqrt{d_h} ~\big|~ \sqrt{d_h} \big]^\top,
\]
\[
\bm{k}'_j = \big[~\bm{k}_{c,j}~|~\bm{k}_{\mathrm{rel},j}~|~u(j)\big]^\top,
\]
so that a standard SDPA call $\mathrm{softmax}((Q'K'^\top)/\sqrt{d_h})$ automatically realizes
\[
\frac{\langle\bm{q}'_i, \bm{k}'_j\rangle}{\sqrt{d_h}} = \langle \bm{q}_{c,i}, \bm{k}_{c,j}\rangle/\sqrt{d_c} + \mathcal{K}_{ij},
\]
with no additional asymptotic or memory cost. GOAT thereby functions as a drop-in replacement in standard and acceleration-optimized attention layers.

## 5. Structural Priors and Length Extrapolation

GOAT’s prior decomposes into a sum of a truncated Fourier (spectral) series and a key-only sink:
\[
\mathcal{K}_{ij} = \sum_{r=1}^R [\alpha_r \cos(\omega_r(i-j)) + \beta_r \sin(\omega_r(i-j))] + u(j).
\]
The relative term, $\mathcal{K}^{\rm rel}_{ij}$, is constructed via Bochner’s theorem, with basis vectors
\[
\bm{k}^{(r)}_{\mathrm{rel},j} = [\cos(\omega_r j),\,\sin(\omega_r j)]^\top,~
\bm{q}^{(r)}_{\mathrm{rel},i} = [\alpha_r\cos(\omega_r i)+\beta_r\sin(\omega_r i),~ \alpha_r\sin(\omega_r i)-\beta_r\cos(\omega_r i)]^\top,
\]
such that the dot-product recovers the $r$-th series component. This design is translation-equivariant and generalizes across sequence lengths. The key-only sink $u(j)$, implemented as a rank-one “lane,” provides a minimal, content-disentangled query-independent bias. The learned spectral prior is stable under arbitrary sequence scaling.

Empirically, on C4 language modeling (trained to 2048 tokens), GOAT matches or surpasses RoPE and ALiBi in in-distribution perplexity, and extrapolates up to $16\times$ context length without degradation. On synthetic long-context retrieval tasks, GOAT maintains near-perfect accuracy well beyond the training window. The learned Fourier prior remains robust under sequence length extension.

## 6. Comparative Analysis with Other Attention Variants

A comparative summary of attention mechanisms:

| Mechanism   | Prior Type              | Structural Properties                              | Extrapolation/Inductive Bias                             |
|-------------|------------------------|----------------------------------------------------|----------------------------------------------------------|
| Softmax     | Uniform                | No structure, emergent content-norm sinks          | Poor extrapolation, generic                             |
| RoPE        | Multiplicative rotations| Content-structure entanglement                     | Catastrophic out-of-distribution degradation             |
| ALiBi       | Linear slope in $j$    | Fixed, underfits in-distribution, rigid            | Extrapolates, limited adaptation                        |
| GOAT        | Learned additive (Fourier+sink) | Fully expressive, translation-equivariant, disentangled | State-of-the-art extrapolation, stable, plug-and-play    |

Empirical results indicate that in vision (e.g., ViT-Small at $224^2$), GOAT achieves higher zero-shot accuracy at elevated resolutions than absolute embeddings. In genomics, GOAT matches RoPE in speed, reduces peak GPU memory by 36%, and decreases bits/base metric.

## 7. Theoretical and Empirical Properties

GOAT encompasses several formal and empirical properties:

- Provides a closed-form attention solution $p\propto \pi \exp(s/\tau)$, sidestepping iterative Sinkhorn procedures.
- The spectral plus sink parameterization is identified as the unique, finite-dimensional, SDPA-compatible, translation-equivariant, bounded prior (Theorem 7.1), maximizing entropy recency (Theorem 7.3) and furnishing minimal-rank sinks (Theorem 7.4).
- Implementation incurs no additional asymptotic computational or memory cost and interfaces directly with performance-optimized attention kernels.
- Empirical tests resolve key trade-offs among expressivity, stability, and efficiency, substantiating state-of-the-art length generalization, stable attention sinks, and improved outcomes in cross-modal domains [2601.15380].

Source: https://www.emergentmind.com/topics/generalized-optimal-transport-attention-goat