Attention from first principles

Attention is how a transformer lets each token look up information from other tokens: a soft, differentiable lookup table built fresh for every input. This lesson derives the mechanism, explains its quadratic cost and its memory footprint at inference, and surveys how 2025–2026 models changed it to reach million-token contexts.

The problem attention solved

Early neural translation systems compressed an entire source sentence into one fixed-length vector before decoding. Bahdanau, Cho and Bengio saw that this vector was a bottleneck, and that it hurt most on long sentences. Their fix let the decoder, at every output step, compute a weighted average over all encoder states. The weights came from a learned score of how relevant each source position was.^[1] Vaswani et al. then made that operation the whole architecture, with no recurrence and no convolution, in "Attention Is All You Need".^[2]

Scaled dot-product attention

Each token's vector $x_i$ is projected three ways: a query $q_i = W_Q x_i$ (what am I looking for?), a key $k_j = W_K x_j$ (what do I contain?) and a value $v_j = W_V x_j$ (what do I hand over if selected?). Stack them into matrices $Q, K, V$. Then:

Row $i$ of the softmax is a probability distribution over positions $j$, and the output for token $i$ is the average of the values under that distribution. The match score is a dot product, so keys pointing in the query's direction get the most weight.

Why divide by $\sqrt{d_k}$? If the components of $q$ and $k$ are independent with mean 0 and variance 1, then $q\cdot k=\sum_{m=1}^{d_k} q_m k_m$ is a sum of $d_k$ terms each with variance 1, so its variance is $d_k$. Without scaling, logits for $d_k=128$ would have a standard deviation of about 11. The softmax would then saturate to nearly one-hot, and its gradients would vanish. Dividing by $\sqrt{d_k}$ brings the variance back to 1.^[2]

Worked example. Let $d_k=2$, query $q=(1,0)$, keys $k_1=(1,0)$, $k_2=(0,1)$, $k_3=(-1,0)$ and scalar values $v=(10,20,30)$. The scaled scores are $q\cdot k_j/\sqrt2 = (0.707,\,0,\,-0.707)$. Exponentiating gives $(2.028,\,1,\,0.493)$, which sum to $3.521$, so the weights are $(0.576,\,0.284,\,0.140)$. The output is

The value whose key matches best dominates, but every position contributes. That softness is what makes the lookup differentiable.

Heads and masks

Multi-head attention runs $h$ of these in parallel, each with its own smaller projections ($d_k = d_{\text{model}}/h$). It concatenates the results and projects them back, so different heads can track different relations: syntax, coreference, copying.^[2] For language modelling, a causal mask sets the score to $-\infty$ wherever $j>i$, so no token can attend to the future. That is what makes next-token training parallel.

Position has to be added

Attention as written is permutation-equivariant: shuffle the tokens and the outputs shuffle the same way. Word order has to be injected. Most current models use rotary position embeddings (RoPE). RoPE rotates each pair of query and key coordinates by an angle proportional to the token's position, so the dot product between positions $m$ and $n$ depends only on the offset $m-n$.^[3] Llama 3 uses RoPE with base frequency $\theta = 500{,}000$, a setting chosen to support long contexts.^[4]

What attention costs

For a sequence of length $n$, the score matrix $QK^\top$ has $n^2$ entries per head per layer. At $n=128{,}000$ that is about $1.6\times10^{10}$ scores. Compute grows as $O(n^2 d)$. A naive implementation also writes the $n\times n$ matrix to GPU memory, and on modern accelerators that traffic, not the arithmetic, is the bottleneck. FlashAttention computes exact attention in tiles that stay in on-chip SRAM, using an online softmax, so the full matrix never reaches high-bandwidth memory. The original paper proved the IO savings and reported large wall-clock speedups.^[5] FlashAttention-3 reached up to 740 TFLOP/s in FP16 on an H100, about 75% of peak.^[6] These kernels change constants and memory traffic. They do not change the $O(n^2)$ arithmetic.

The KV cache, and sharing heads

During generation, each new token needs the keys and values of every earlier token. Recomputing them would be wasteful, so they are stored in a KV cache. Its size per token is

where the factor 2 counts K and V, $L$ is the number of layers, $h_{kv}$ the number of key/value heads and $b$ the bytes per number.

Worked example: Llama 3 70B. It has 80 layers, 64 query heads, 8 key/value heads and $d_{\text{head}} = 8192/64 = 128$.^[4] In 16-bit precision:

A single 128k-token context therefore needs about 43 GB of cache. If every one of the 64 heads kept its own K and V, as in the original design, it would need 8× that, about 343 GB.

The saving comes from grouped-query attention (GQA). Shazeer first proposed multi-query attention, in which all heads share one K/V head, because incremental decoding was limited by the memory bandwidth spent reloading keys and values.^[7] GQA generalizes this: groups of query heads share a K/V head. Its authors reported quality close to full multi-head attention at speed close to multi-query.^[8] DeepSeek-V2 went further with multi-head latent attention (MLA), which caches a compressed latent vector and reconstructs K and V from it. It reported a 93.3% smaller KV cache than DeepSeek 67B.^[9] Inference economics explains why bytes in this cache translate directly into cost.

How attention changed in 2025–2026

As of November 2026, long-context frontier models rarely run plain dense attention in every layer. Four trends stand out.

• Sparse attention. DeepSeek-V3.2 added DeepSeek Sparse Attention: a small, cheap "lightning indexer" scores earlier tokens, and each query attends only to a selected subset. The report says this reduces long-context cost while preserving quality.^[10] DeepSeek-V4-Pro (2026) combines two compressed attention variants. Its model card reports that at one million tokens it needs 27% of the per-token inference FLOPs and 10% of the KV cache of V3.2.^[11]
• Linear-attention hybrids. These replace most attention layers with recurrent layers that keep a fixed-size state. Kimi Linear interleaves three "Kimi Delta Attention" layers per full-attention (MLA) layer. It reports up to 75% less KV cache and up to 6.3× faster decoding at 1M context than full MLA.^[12] Qwen3-Next uses a similar 3:1 mix of Gated DeltaNet and gated attention.^[13]
• Gating the output. Qwen researchers showed that a sigmoid gate after the attention output improves training stability and long-context extrapolation. It also largely removes the "attention sink", in which heads park weight on the first token. The share of attention on the first token fell from 46.7% to 4.8% in their models.^[14] The paper won a NeurIPS 2025 best paper award.^[15]
• Advertised vs usable context. The RULER benchmark found that most models claiming 32k-token or longer windows degraded well before their nominal limit on tasks harder than simple retrieval.^[16] A context window is a capacity, not a guarantee of use.

Key takeaways

• Attention is a softmax-weighted average of values, with weights from query–key dot products scaled by $\sqrt{d_k}$ to keep logits at unit variance.
• It has no built-in notion of order. RoPE supplies relative position through rotations.
• Cost is quadratic in sequence length. FlashAttention makes it IO-efficient but does not change the asymptotics.
• At inference the KV cache dominates memory. GQA, MQA and MLA shrink it by sharing or compressing keys and values.
• Frontier designs as of 2026 mix full attention with sparse, compressed or linear-recurrent layers to make million-token contexts affordable.

Sources