Definition
An attention mechanism builds each output position out of a weighted mixture of every input position, with the weights computed from the data at runtime rather than fixed by the architecture. Concretely: every token emits a query, a key and a value; the relevance of B to A is the dot product of A's query with B's key; a softmax normalises those scores into weights that sum to 1; and A's output is the weighted sum of all the value vectors.
That is the entire mechanism — multi-head attention, causal masking and FlashAttention are variations on those steps, not additions. One consequence first: every token is scored against every other, so an n-token sequence needs n² scores, and doubling the context quadruples the work. Bahdanau et al. introduced the idea in 2014 for machine translation; "Attention Is All You Need" made it an architecture in 2017 by removing recurrence and keeping only attention.
How It Works
The bottleneck it removed
Before attention, a sequence-to-sequence model was two recurrent networks: the encoder compressed the whole input into one fixed-size hidden vector, and the decoder generated the output from that vector alone. Every word of a forty-word sentence had to survive inside the same few hundred numbers, and the first token's information had to be relayed through thirty-nine sequential update steps to reach the end. Attention deleted that bottleneck: the encoder keeps a vector per input position and the decoder reads from all of them at every step, so the path from any input to any output is one step, not n.
Query, key, value
Each token's embedding is multiplied by three learned weight matrices to produce a query, a key and a value. The surviving intuition is a lookup: the query is what a token is looking for, the key is what it offers, the value is what it hands over when selected. The match is soft — every key matches to some degree, and the degrees sum to 1. Note what is absent: nothing refers to position, so attention is permutation-equivariant, and word order reaches the model only through positional encodings the transformer adds around it.
The score between A and B is the dot product of A's query with B's key: large and positive when the
vectors point the same way, near zero when unrelated. For n tokens that is an n × n matrix of raw
scores. Each row is divided by √d (d is the head dimension, 64 in the base Transformer), softmaxed
into a probability distribution, and used to average the value vectors: softmax(QKᵀ / √d) · V.
Why the √d
If query and key components behave like independent numbers with mean 0 and variance 1, their dot product over d dimensions has variance d and a standard deviation of √d — at d = 64 that is 8, so raw scores land eight units apart rather than one, and softmax does not tolerate that. On two scores one unit apart it returns 0.731 and 0.269, a local gradient p(1−p) of 0.197. On two scores eight apart it returns 0.99966 and 0.00034, and that gradient collapses to 0.00034 — nearly 600 times smaller. The softmax has saturated: it has committed to one token and almost nothing flows back to correct the choice. Dividing by √d is a variance fix, not a normalisation convention.
Multi-head attention, and what it costs
One set of attention weights expresses one relationship, and language needs several at once — which noun a pronoun refers to, which verb governs a subject. Multi-head attention runs several attentions in parallel on lower-dimensional projections and concatenates the results. The arithmetic is the point: the base Transformer has a model dimension of 512 and 8 heads, so each head works in 512 / 8 = 64 dimensions, each head's scores cost n² × 64 multiply-adds, and eight of them cost n² × 512 — exactly what one 512-dimensional head would have cost. The split buys several distinct relations for the same compute budget, at the price of a narrower slice per head.
The quadratic cost, worked
The score matrix has n² entries. A 1,000-token prompt is 1,000,000 pairs; 2,000 tokens is 4,000,000 — double the input, quadruple the work; a 100,000-token document is 10,000,000,000 pairs, ten thousand times the work of the 1,000-token prompt for a hundred times the text.
The term is quadratic in sequence length but only linear in model width, so widening a model is cheap next to lengthening its context window.
Memory scales the same way and bites sooner. At 4,096 tokens one head's score matrix in 16-bit precision is 4,096² × 2 bytes = 33.5 MB, so a 32-head layer holds 1.07 GB of scores that exist only to be softmaxed and discarded. That number, not the arithmetic, is what FlashAttention attacks: it computes the identical result in tiles without ever writing the full matrix to high-bandwidth memory. The same pressure produced the KV cache and the whole inference optimization literature.
Types
Two distinctions matter, and they are independent of each other.
Self-attention versus cross-attention is about where the three vectors come from. In self-attention they are all projected from the same sequence, so each token reads its own neighbours — what every layer of a large language model does. In cross-attention the queries come from one sequence and the keys and values from another: a translation decoder queries the encoded source sentence, a vision-language model queries image patches.
Causal versus bidirectional masking is about which positions a token may look at. Bidirectional attention lets every token see the whole sequence. Causal attention sets every future score to −∞ before the softmax, so a token reads only what precedes it — mandatory for generation, since a model that could see the next token would copy it rather than predict it.
Real-World Applications
The useful examples are decisions attention changes, not products it appears in.
- Long-context pricing. Providers charge more per token at long context, or cap it, because the n² term is real: doubling a prompt more than doubles the cost of serving it.
- Machine translation. The original application — attention weights align target to source words without ever being told what an alignment is.
- Vision transformers. An image split into 16×16 patches becomes a sequence: a 224×224 image gives 196 patches and a tractable 196² matrix.
- Reranking. Cross-attention between query and document scores relevance more accurately than comparing two embeddings, and far more slowly.
Challenges
The n² term is the challenge everything else descends from, and it resists clean fixes. Sparse and linear variants cut the asymptotic cost but drop information: on a task needing one exact long-range lookup, an approximation that skips the relevant pair simply fails. Exact methods like FlashAttention are the safer trade, changing the memory access pattern and not the result.
Inference has a second problem the training-time picture hides. Generating token n+1 needs the keys and values of all n previous tokens, so they are cached rather than recomputed — and that cache grows with context and with concurrent users until it, not the weights, fills the GPU.
The third is interpretability. Attention weights show which positions were mixed together, not why or to what effect — a heatmap is evidence to check, not a reason.
Future Trends
Work on attention has split in two. One direction keeps the operation exact and attacks the hardware: successive FlashAttention kernels are co-designed with each GPU generation, and the gains come from arranging memory traffic, not changing the maths. That is what has actually shipped.
The other replaces it. Grouped-query and multi-query attention share keys and values across query heads to shrink the inference cache; latent-attention schemes compress them before caching; sliding-window layers interleaved with full-attention layers keep global reach at a fraction of the pairs. State-space models abandon the n² matrix entirely — which is why nearly every long-context design is a hybrid rather than a replacement.
Code Example
Scaled dot-product attention in NumPy with causal masking — four tokens, small enough to read the weight matrix directly.
import numpy as np
def softmax(x):
x = x - x.max(axis=-1, keepdims=True) # subtract the row max, for numerical stability
e = np.exp(x)
return e / e.sum(axis=-1, keepdims=True)
def attention(Q, K, V, mask=None):
scores = Q @ K.T / np.sqrt(Q.shape[-1]) # (n, n): every query against every key, over sqrt(d)
if mask is not None:
scores = np.where(mask, scores, -np.inf) # masking happens before the softmax
weights = softmax(scores) # each row now sums to 1
return weights @ V, weights
rng = np.random.default_rng(0)
n, d_k, d_v = 4, 8, 3
Q, K, V = rng.normal(size=(n, d_k)), rng.normal(size=(n, d_k)), rng.normal(size=(n, d_v))
causal = np.tril(np.ones((n, n), dtype=bool)) # token i may only look at tokens 0..i
out, w = attention(Q, K, V, mask=causal)
print(np.round(w, 3))
print("row sums:", w.sum(axis=1))
Output:
[[1. 0. 0. 0. ]
[0.827 0.173 0. 0. ]
[0.34 0.413 0.247 0. ]
[0.102 0.031 0.684 0.183]]
row sums: [1. 1. 1. 1.]
Row 0 is the first token, which can see only itself, so it attends to itself with weight 1 — the softmax has no choice. Row 3 spreads 0.102, 0.031, 0.684 and 0.183 across four positions, and the upper triangle is exactly zero because the mask set those scores to −∞. Every row sums to 1, which is what makes the output a weighted average. Scale n to 100,000 and this is the ten-billion-entry matrix the efficient-attention literature exists to avoid building.