GOAT: entropic optimal transport, minus the transport

GOAT — Generalized Optimal transport Attention with Trainable priors — says it redefines attention through entropic optimal transport and "materializes a solution" for attention sinks. Both halves of that sentence are true. What surprised me reading the paper (Litman & Guo, Stanford, arXiv 2601.15380, accepted to ICML 2026) is which half of optimal transport it uses. If you already know Sinkhorn-based attention — this site covered a doubly-stochastic Sinkhorn loss in Marigold V2 — you'll expect row-and-column balancing, iterative projections, a real algorithmic cost. GOAT has none of that. It is one softmax call with a smarter bias, dressed in optimal-transport language because the derivation genuinely comes from there. That is not a knock — the derivation is the useful part — but the "we fixed attention sinks with OT" framing undersells how narrow the fix actually is.

I'm assuming you know what queries, keys, and softmax attention are — see how self-attention works if not — and I'm not re-deriving the mask/KV-cache taxonomy from the attention field guide. This piece picks up where that guide's one line on sinks leaves off, and goes through GOAT's actual math, its actual released code, and what its figures actually show.

Why a sink is a real problem, briefly

The attention field guide already covered the canonical sink story: softmax weights must sum to 1 even when a query has nothing worth attending to, so models learn to dump the leftover mass somewhere — usually the first token. Xiao et al.'s StreamingLLM (arXiv 2309.17453) showed evicting that token from a rolling KV cache collapses perplexity, and keeping just four sink tokens fixes it. Gu et al. (arXiv 2410.10781, ICLR 2025) went further empirically: sinks emerge reliably during pretraining, they act "more like key biases" than genuine content, and — the part that matters for production — that same behavior is exactly what makes quantization hard, because a key vector that has to dominate every query's softmax regardless of content ends up with a permanently huge norm, and huge-norm outliers are what break low-bit quantization schemes. Three real costs, then: wasted representational capacity, a KV cache you can't safely prune, and an activation outlier that survives every input. GOAT's whole pitch is aimed at the middle failure mode — the entanglement between "be a sink" and "be a normal content vector" — not at the KV-cache-eviction or quantization failure modes directly, and it doesn't test either of the latter two. More on that below.

The general argument: two-sided OT actually forbids a sink

Before GOAT, here's the argument that motivates using optimal transport for attention at all, and it's worth building from scratch because GOAT's own machinery is a deliberate departure from it.

Treat attention as a transport plan: mass from each query has to go somewhere among the keys. Softmax attention imposes exactly one constraint on that plan — each row (each query's outgoing weights) sums to 1. Nothing constrains a column. So if key 0 has a large enough dot product with every query (the classic "generic, high-norm key" StreamingLLM and Gu et al. describe), every row can dump its mass there simultaneously, and the column sum for key 0 grows without bound as more queries pile on. That's the sink, mechanically: an unconstrained column.

Entropic optimal transport, in its classic two-sided form — the one behind Sinkformers (Sander, Ablin, Blondel & Peyré, 2022) and the SinkLoss mechanism in Marigold V2 — adds the missing constraint. Solve for a transport plan that satisfies both a row marginal and a column marginal (each key gets a fixed capacity too), regularized by entropy, and the closed-form solution is no longer plain softmax: it's a matrix you reach by alternately rescaling rows and then columns of an exp(score / τ) kernel until both marginals hold — Sinkhorn–Knopp. That alternation is doing real work. No single key can absorb more than its column budget, no matter how large its score gets, because the column-normalize step at each iteration forcibly redistributes the excess.

Both grids above start from the same content scores. Slide the bias up and softmax's column 0 climbs toward its ceiling (every query dumping essentially all its mass on one key) while the OT plan's column 0 stays pinned at exactly 1.0 — the fixed budget makes a sink structurally impossible, not just unlikely. Notice the sharper thing this actually buys you, though: at high bias, the OT plan doesn't just cap key 0 — every column gets forced to exactly its 1.0 share, including keys that might genuinely deserve more attention across the batch on content grounds. Two-sided OT doesn't distinguish "an undeserved sink" from "a key the content scores say is broadly useful"; it flattens both. That bluntness is a real cost of the two-sided fix, and it's one honest reason production attention doesn't do this.

Getting there costs a pass over the score matrix per iteration, alternating which axis you normalize:

Step 1 above is exactly softmax — one row-normalize, nothing else — and you can watch column 0 sitting at several times its fair share right there. Each further half-step squeezes the other axis without fully undoing the one just fixed, and the two marginals converge together. That's the mechanical cost real two-sided OT attention pays that plain softmax doesn't: not one extra pass, but one extra pass per Sinkhorn iteration, and convergence to a tight doubly-stochastic matrix generally takes several.

What GOAT actually does instead

Here's the pivot. GOAT does not use any of the above. Its own Appendix A is titled, plainly, "Proof of Attention as One-Sided Entropic Optimal Transport" — and that one-sidedness is the paper's real starting point, first established in a solo preprint by the same first author (Litman, arXiv 2508.08369, 2025). Fix a single query as a unit mass, add one Lagrange multiplier for the one constraint (its own weights sum to 1), solve the stationarity condition, and you get back exactly softmax:

There is no column constraint anywhere in this derivation — it's one query at a time, against a fixed total-mass-one budget on itself. Standard attention, in other words, already is an entropic OT solution; it's just the entropy term equals a KL divergence against an implicit uniform prior over keys (Shannon entropy is negative KL against uniform, up to a constant). GOAT's actual move is generalizing that one fixed term, not adding a second constraint:

Swap the flat prior $\pi$ for a learned one and you get softmax over content scores shifted by a learned log-prior — additively, per key. That's the entire mechanism. Every positional encoding you already know — RoPE, ALiBi, sinusoidal — is reframed here as a hand-picked approximation to this one missing term, and GOAT's contribution is parameterizing it properly and training it.

The prior, precisely

GOAT builds the log-prior $\mathcal{K}_{ij}$ from two learned pieces, added together, both engineered to be additive logit biases — never multiplicative, unlike RoPE's rotation — so a sink can be created without inflating any content vector's norm:

A relative, shift-invariant term. A truncated Fourier series in the distance $i-j$, with learnable amplitudes $\alpha_r,\beta_r$ per frequency and fixed geometric frequencies $\omega_r$:

Trainable means exactly this: the frequencies are the fixed geometric ladder every Fourier-feature positional scheme uses, but how much each frequency contributes — and whether it attracts or actively suppresses (negative $\alpha_r,\beta_r$ are allowed) — is learned from data, not hand-set the way RoPE's rotation angles or ALiBi's linear slope are.

An absolute, key-only sink term. A small MLP over each key's normalized position, feeding a single learned bias $u(j)$ that every query sees identically:

Both pieces are linearized into extra query/key dimensions via angle-sum identities (the relative term) and a broadcast constant channel (the sink term), then concatenated onto the content query/key vectors before the usual dot product — Algorithm 1 in the paper. That's the trick that keeps this a drop-in: the composite vectors go straight into an ordinary scaled_dot_product_attention call, so the prior rides for free inside whatever fused kernel PyTorch picks (FlashAttention, mem-efficient, cuDNN). No extra passes, no materialized L×L bias matrix — the opposite of what the Sinkhorn machinery above requires.

expand

Sinks are optimal, not eliminated

This is the part the "materializes a solution" framing can mislead you on. GOAT does not make sinks go away. Its own Theorem 5.1 ("Collapse to Prior") proves the opposite: as a query's content signal weakens ($\omega_i \to 0$, the score spread across keys), its attention distribution provably converges to whatever the prior says — sink included. A sink is the correct response to an uninformative query under this objective, not a bug to be optimized out.

What GOAT changes is where the sink lives. Decompose the margin that makes key $j^\star$ a sink into a content part and a prior part: $z_{ij^\star}-z_{ik} = (s_{ij^\star}-s_{ik}) + (\mathcal{K}_{ij^\star} - \mathcal{K}_{ik})$. Standard attention has no prior term to spend, so the entire margin has to come from content — the model is forced to learn a generic, high-norm key vector, exactly the outlier Gu et al. found empirically. GOAT can spend the prior's budget instead, via the unconstrained bias $u(j^\star)$, leaving the content key free to represent whatever the token actually is.

The paper backs this with a stability bound (Theorem 5.4) that I found genuinely clarifying: define "context sensitivity" as the total probability mass not on the sink — how much a query's output can be perturbed by noise elsewhere in the context. Under a uniform prior (standard attention), that sensitivity provably converges to 1 as context length grows — the sink's relative share shrinks, and noise resistance decays away — unless the content margin keeps growing with it. Under a peaked, prior-driven sink, the same sensitivity stays bounded for any margin $\delta$, and $\delta$ only needs to grow logarithmically with length to hold the line. That's a plausible first-principles account of something StreamingLLM and Gu et al. only observed empirically: why a sink token's activation magnitude tends to grow with how much context it has to keep suppressing, and therefore why it's such a persistent outlier for quantization. I want to be precise about what's tested and what's argued, though — GOAT runs no quantization experiment and no KV-cache-eviction experiment. The stability theorem is a real, checkable piece of math; whether a disentangled sink actually behaves better under 4-bit quantization or mid-stream cache eviction is a plausible prediction from it, not a result in this paper.

Does it actually generalize length?

Three different evaluations, and they tell different stories.

Perplexity (125M params, C4, 4B tokens, trained at 2,048, tested to 32,768 — 16×). The in-distribution number is exact and real: GOAT 30.95 vs. ALiBi 32.49 vs. RoPE 33.24 — a genuine but modest 1.55-point win over the stronger baseline. RoPE's extrapolation is where "catastrophic" earns its word: perplexity climbs from about 30 at 2,048 to roughly 500 by 32,768. But reading the actual extrapolation curve (no exact numbers are given past the in-distribution point, so these are chart reads): ALiBi stays close to flat, in the low-to-mid 30s across the whole range, while GOAT's own line drifts up too — to somewhere around 55–65 by 32,768. GOAT is dramatically better than RoPE and still meaningfully worse than ALiBi at the longest length tested, despite winning in-distribution. "Combines the fidelity of a learned prior with the robustness of a fixed encoding" is the paper's framing; the more precise version is that GOAT trades a little of ALiBi's flatness for a real in-window improvement, not a strict improvement on both axes simultaneously.

expand

Passkey retrieval (trained at 1,024, tested to 16,384 — 16×, span-token accuracy). This is where GOAT looks strongest, and where I'd trust perplexity least as a proxy — perplexity averages over a whole sequence; passkey retrieval asks whether one specific fact survives. Reading the figure (again, no data table given): GOAT stays around 0.90–1.00 across the entire range. ALiBi degrades steadily from about 0.88 to 0.49 — more than half its accuracy gone — which is worth sitting with, because ALiBi looked essentially flat on perplexity over the same range. Flat perplexity does not mean flat retrieval; they're measuring different things, and a model can average well while still losing a specific fact. RoPE with position interpolation is the biggest surprise: it starts near 1.0 at the training length and falls to about 0.30 by 2× — collapsing faster than plain rotary RoPE, despite position interpolation existing specifically to help extrapolation. Sinusoidal absolute encoding is worst throughout, roughly 0.10–0.15 past training length.

Needle-in-a-haystack (same models, context 512–16,384 × depth 0.1–0.9, span-token accuracy heatmap.) GOAT's heatmap is close to uniformly yellow (roughly 0.85–1.0) across the whole grid. But all four baselines look uniformly mediocre — not "fine at short context, collapsing with length" the way passkey showed, but flatly teal-to-dark (roughly 0.2–0.5 by the color scale) even at 512 tokens, shorter than their own training length. I can't fully reconcile that with the passkey plot's baselines starting near 1.0 at comparable lengths, and the paper doesn't discuss it — worth flagging as something I could not independently verify rather than something I can explain away.

Vision (ViT-Small, ImageNet-1k, trained at 224×224, evaluated up to higher resolutions, zero-shot). The learned-PE baseline degrades sharply past its training resolution; GOAT holds up substantially better, and the 2D log-prior it learns is visibly local and shift-invariant despite starting from a uniform initialization. This is the one modality where the paper doesn't compare against a strong, already-extrapolation-robust baseline like ALiBi — the comparison is against ordinary learned absolute position embeddings, which is a lower bar to clear.

What it costs

The design is genuinely clever on cost: because the prior is injected as extra dimensions on the query and key vectors rather than an added L×L bias matrix, GOAT's forward pass is one call to torch.nn.functional.scaled_dot_product_attention — the same call a plain-softmax implementation makes, just with a few extra lanes in the head dimension. It is not Sinkhorn, so the "how many iterations, and what does that cost" question that matters for real two-sided OT attention (the machinery in the stepper above) simply doesn't apply here — there are zero iterations, by construction.

The one place the paper measures this is the DNA task, RoPE vs. GOAT, same training budget:

expand

That memory drop is a genuinely interesting side effect — plausibly because GOAT's prior replaces a separate learned/rotary positional pipeline with a few extra head-dim lanes computed inline, rather than because of anything OT-specific — and the throughput hit is close to noise. Both baselines here run through the same optimized SDPA call, so this is a fair fused-kernel-vs-fused-kernel comparison, not GOAT-optimized-versus-baseline-naive. But it's also the only throughput or memory number in the whole paper. The flagship C4 language-modeling experiment — the one carrying the length-generalization headline — reports zero wall-clock numbers. I'd want that filled in before trusting "minimal computational overhead" as a general claim rather than a DNA-task-specific one.

Scale, and what's not tested

Every experiment here is small: 125M-parameter decoder-only transformers for both the C4 and DNA tasks, ViT-Small (roughly 22M parameters) for ImageNet. Trained on 4B tokens (C4) — the DNA task's context length and token count aren't stated in the paper at all, which is its own small gap next to the detailed hyperparameter tables C4 and ImageNet get in the appendices. There is no scaling curve — one operating point per modality, no test at 350M, 1B, or beyond. That's a real limitation on how far any of this extrapolates, not a scandal; it puts GOAT in the same boat as most attention-replacement papers, which also rarely clear a billion parameters before publication. The paper is peer-reviewed (ICML 2026, per the released package's own citation), which at least means these specific numbers passed a review pass — it doesn't substitute for a scale test.

More importantly for the sink story specifically: no quantization experiment, no KV-cache-eviction experiment. Given that Xiao et al.'s streaming-eviction result and Gu et al.'s quantization link are the two production-relevant reasons sinks matter — cited in GOAT's own introduction — their absence from the evaluation is the biggest gap between what the framing promises and what's actually measured. The stability theorem gives a good reason to expect GOAT would help on both; expecting is not measuring.

The code

The package is real, installable, and does not require the paper's training setup to try:

pip install goat-attention

# what GOAT actually ships — one SDPA call, no Sinkhorn
# (trimmed from src/goat/attention.py, GoatAttention.forward)
q_total[..., :D_c] = q_content / math.sqrt(D_c)   # content lane, scaled as usual
k_total[..., :D_c] = k_content
q_total[..., D_c:] = q_pos                         # prior lane: spectral rotation + sink bias
k_total[..., D_c:] = k_pos                         # prior lane: raw Fourier features + u(j)
 
attn_output = F.scaled_dot_product_attention(
    q_total * math.sqrt(D_total), k_total, v_total,
    attn_mask=merged_mask, is_causal=use_internal_causal,
)

Next to plain softmax attention, for scale — this is the entire operation GOAT is replacing the cost model of, not the interface:

# standard scaled dot-product attention — one softmax, one constraint (rows sum to 1)
weights = torch.softmax(q @ k.transpose(-2, -1) / math.sqrt(d_k), dim=-1)
out = weights @ v

And the generic two-sided entropic-OT attention GOAT deliberately does not implement — this is the exact update the stepper above walks through, spelled out as code, the kind of thing Sinkformers-style attention actually runs per layer:

# generic two-sided entropic-OT attention (Sinkhorn–Knopp) — NOT what GOAT does.
# Illustrative background, not from any repo: what a real column-constrained
# OT attention layer pays per forward pass, on top of the same score matrix.
K = torch.exp(scores / tau)                 # scores: (queries, keys)
for _ in range(n_iters):                    # each iteration = one more full pass
    K = K / K.sum(dim=-1, keepdim=True)     # row-normalize  (row marginal = 1)
    K = K / K.sum(dim=-2, keepdim=True)     # col-normalize  (col marginal = 1)
plan = K                                    # doubly stochastic: no column can dominate

One more thing worth knowing if you're picking hyperparameters: the same knob has three different recommended values depending on where you look. The README suggests pos_rank=4 as "a good starting point" for GPT-style models; the package's own GoatAttention.for_gpt() factory defaults to pos_rank=0 (the spectral relative term disabled entirely, leaving only the sink bias and an ALiBi-seeded recency slope); the paper's own Appendix G reports the actual C4 experiment ran with pos_rank=2. None of these is wrong, exactly — research code drifts after a paper ships — but it's the kind of "does the shipped default match the paper's own config" mismatch worth checking before trusting a factory default. No pretrained weights are released either way; the repository ships the module, not a checkpoint that reproduces any number in this article.

The take

The derivation is the best part of this paper, and it's worth taking on its own terms: attention already was a one-sided entropic-OT solution with an implicit uniform prior, every positional encoding you use is already an ad-hoc stand-in for a prior term this framing makes explicit, and generalizing that one term is a clean, cheap, kernel-compatible move. The sink theory is real math with a real, if untested, implication. What I'd push back on is the packaging: "entropic optimal transport" reads, to anyone who knows Sinkformers or this site's own SinkLoss, like the two-sided machinery that structurally forbids a sink — and GOAT is careful, in its own proofs, to say the opposite: a sink is optimal, and what's being fixed is where it lives, not whether it exists. The perplexity win is real but modest and doesn't hold at the longest length tested; the retrieval wins are larger and more convincing; the cost story is genuinely good but measured on one task out of three; and the two experiments that would tell you whether any of this matters where sinks actually hurt — quantization, KV-cache eviction — aren't in the paper. Read the math, use the package if the prior sounds useful, and don't expect it to have quietly solved the production problems StreamingLLM and Gu et al. described.

Primary source: "You Need Better Attention Priors" (Litman & Guo; Stanford, 2026, arXiv 2601.15380), accepted to ICML 2026. The one-sided EOT framing originates in Litman's solo preprint (arXiv 2508.08369, 2025). Code: github.com/elonlit/goat / goat-attention on PyPI. Background: StreamingLLM (Xiao et al., 2023, arXiv 2309.17453); "When Attention Sink Emerges in Language Models" (Gu et al., ICLR 2025, arXiv 2410.10781); Sinkformers (Sander, Ablin, Blondel & Peyré, 2022, arXiv 2110.11773). Figures are reproduced from the paper's own assets and PDF; chart values not given exact numbers in the text are called out as such above. The interactive diagrams are original, illustrative implementations of the general Sinkhorn–Knopp algorithm, not GOAT's own code.