How Attention Moves Information

The question attention answers

Most explanations of attention open with the formula \(\operatorname{softmax}(QK^\top/\sqrt{d})V\) and then work backwards, assigning a story to each symbol. That order is exactly wrong. The formula is the answer to a question, and the question is older and simpler than Transformers.

Here it is. You have a set of representations

and you are trying to produce a better representation of position \(i\). Better means contextual: the vector at position \(i\) should carry information that only exists elsewhere in the set. A convolution answers this by fixing a neighbourhood in advance. A recurrent net answers it by pushing everything through a bottleneck state. Both decide the routing before they see the data.

Attention makes a different bet. Let each position choose, at run time, which other positions to read from. Restrict the output to a weighted average of what the other positions offer:

That constraint is the whole design. It says the output lives in the convex hull of the values, so no single source can dominate by growing its norm, and the "how much" is separated cleanly from the "what". Everything else in this article, every variant with a three-letter acronym, is a different answer to one question:

The figure below is the mechanism with nothing else attached. Watch what happens when the weights are uniform, and why that failure forces the rest of the design.

Two things are worth naming before moving on. First, \(\alpha\) depends on the input, which is what people mean when they call attention a form of dynamic routing or fast weights. Second, the operation is permutation-equivariant: shuffle the inputs and the outputs shuffle with them. Order has to be injected separately, through position encodings or through the mask. Attention itself does not know what "next to" means.

Attention before Transformers

Attention was not invented to make Transformers work. It was invented to fix a bottleneck in sequence-to-sequence translation, and knowing the original shape of the problem makes cross-attention feel obvious later.

An encoder RNN once compressed a whole source sentence into one fixed vector, and a decoder generated the translation from it. The vector had to hold a forty-word sentence as well as a four-word one. The first attention mechanism removed that bottleneck: keep every encoder state \(h_1,\ldots,h_n\), and let the decoder build a fresh context vector at each output step,

Why the field chose dot products. Both scores solved the fixed-vector bottleneck. Dot products won because all pairwise scores can be computed as one matrix multiplication, which scales cleanly on accelerators. The representational flexibility did not disappear; Q and K projections moved it into learned vector spaces, which is the next step in the derivation.

Notice the shape of this. The queries came from the decoder; the keys and values came from the encoder. Attention was cross-attention first. Self-attention, where all three come from the same sequence, arrived later as the special case that turned out to be more useful. Keep that ordering in mind — it makes vision-language models look like a return to the original form rather than an exotic addition.

Why three projections

The standard gloss is that the query asks, the key advertises and the value carries. True, but it explains nothing about why we need three learned matrices. Let us derive the need.

Start with the cheapest possible scoring function: use the representations directly, \(s_{ij} = x_i^\top x_j\), so the whole layer is \(\operatorname{softmax}(XX^\top)X\). No parameters at all. Two failures follow immediately.

The score matrix is symmetric. \(x_i^\top x_j = x_j^\top x_i\), so if it attends to animal then animal must attend to it with the same pre-softmax score. But the relations we care about are directional. A pronoun needs its antecedent; the antecedent does not need the pronoun. A verb needs its subject; an adjective needs its noun. Symmetric scores cannot express a directed graph.

Self-similarity dominates. \(x_i^\top x_i = \lVert x_i \rVert^2\) is the largest entry in its row for any well-spread set of vectors, so every token mostly attends to itself and the layer approximates the identity.

Both problems disappear with two different linear maps:

The scoring form is now a general bilinear form \(x_i^\top M x_j\) with \(M = W_QW_K^\top\), and a general \(M\) is not symmetric. The model can learn a low-rank asymmetric relation. It can also make \(x^\top M x\) small, so a token need not attend to itself.

The third projection answers a separate question. Having decided who to read using a match in one space, why should the model be forced to copy the raw thing it matched on? Splitting off \(V = XW_V\) lets the retrieval key and the retrieved payload live in different spaces: match on "I am a noun phrase in subject position", transmit "I refer to a large mammal". In circuits language, \(W_QW_K^\top\) forms the QK circuit that decides where information moves, and \(W_VW_O\) forms the OV circuit that decides what is written when it does.

One practical consequence: \(W_Q\) and \(W_K\) only ever appear as the product \(W_QW_K^\top\), so their individual factorisation is not identifiable, and the rank of that product — not \(d_{\text{model}}\) — is what limits how expressive the scores can be. This is also why fused implementations are free to absorb one into the other, a trick that reappears in a big way when we get to latent attention.

Scaled dot-product attention

We have scores. Two questions remain: how do we turn scores into a distribution, and why is there a \(1/\sqrt{d_k}\) bolted onto the front?

The first is easy to motivate. We need weights that are non-negative and sum to one, from scores that are unbounded reals. Exponentiating and normalising is the canonical map with that signature, it is smooth everywhere, its gradient is cheap, and it has the useful property that only score differences matter: adding a constant to a row changes nothing. That last fact is not just aesthetic — it is what makes numerically-stable and streaming softmax possible, which is the hinge the entire FlashAttention section turns on.

The second question deserves an actual derivation. Take \(q\) and \(k\) with independent, zero-mean, unit-variance components. Then

so the logits have standard deviation \(\sqrt{d_k}\). At \(d_k = 128\) that is roughly \(\pm 11\) before the model has learned anything. Softmax of logits that far apart is nearly a one-hot vector, and the Jacobian of softmax at a one-hot output is nearly zero, so gradients through the attention weights vanish at initialisation. Dividing by \(\sqrt{d_k}\) puts the logit variance back at 1 regardless of head dimension, which is exactly what you want if \(d_k\) is a hyperparameter you intend to change.

Slide the head dimension in the figure below and watch the entropy collapse. The button toggles the scaling factor so you can see the difference at the same \(d_k\).

Two footnotes. The unit-variance assumption is about initialisation, not about a trained network — trained models happily learn large logits when they want a sharp lookup, and some run into the opposite problem, where a few enormous activations exist mainly to let heads produce a near-zero attention output. And the constant is a convention, not a law: some architectures fold a learned per-head temperature in, and \(\mu\)P-style parameterisations argue for \(1/d_k\) instead once you are scaling width seriously.

Assembling the pieces gives the operation in full:

That \(M\) is the additive mask, and it is the next design axis.

Masks: who can I see?

A mask is an additive matrix of zeros and \(-\infty\) applied before the softmax. It is the crudest possible mechanism and it does an enormous amount of work, because "which pairs are allowed to interact" is what distinguishes an encoder from a decoder, a packed batch from a padded one, and a linear-cost model from a quadratic one.

In practice nobody adds \(-\infty\) to a materialised matrix any more. Modern kernels take a predicate, evaluate it blockwise, and skip whole tiles whose predicate is false everywhere — which is why PyTorch's BlockMask and similar structures are the real interface. A mask that is expressible as a small number of rectangular blocks is nearly free; an arbitrary boolean mask of the same density is not.

Bidirectional. No mask at all. Every position sees every other. This is what BERT-style encoders and every vision Transformer use, and it is the right default whenever the whole input is available at once and you are producing representations rather than continuations.

Causal. \(j \le i\). This is the one structural thing that makes a Transformer a language model. It lets you compute the loss at all \(n\) positions from one forward pass while guaranteeing that position \(i\) never saw its own target — teacher forcing and parallel training in a single triangular matrix.

Padding. Sequences in a batch have different lengths, the tensor is rectangular, so the slack is masked out. Note that padding masks are per-sequence and column-shaped, while causal masks are per-position and triangular; they compose by elementwise AND, which is why real kernels take a mask function rather than a single fixed pattern.

Packed / block-diagonal. Padding wastes compute proportional to the length variance in your batch. Concatenating several documents into one sequence and masking across document boundaries wastes none. This is standard in serious pretraining pipelines, and it is also the mask most often gotten subtly wrong — forget it and documents silently attend to each other.

Prefix (non-causal prefix, causal suffix). The prompt gets to be bidirectional while the generated part stays autoregressive. Useful for encoder-decoder-flavoured behaviour in a decoder-only stack, and structurally identical to how many VLMs treat image tokens: the picture is not a sequence, so there is no reason for its tokens to be causally ordered among themselves.

Sliding window, with or without sinks. Now we are changing complexity rather than semantics, which is the topic of a later section — but note that it is the same mechanism. Nothing about the attention rule changed; we deleted entries from \(M\).

Self-attention and cross-attention

The next axis is the simplest of all: where do \(Q\), \(K\) and \(V\) come from?

Self-attention takes all three from the same sequence, \(Q,K,V \leftarrow X\), so the output has the same length as the input and the cost is \(O(n^2)\). Cross-attention takes \(Q \leftarrow X_{\text{tgt}}\) and \(K,V \leftarrow X_{\text{src}}\), so the output has the length of the target and the cost is \(O(n_{\text{tgt}} \cdot n_{\text{src}})\).

That asymmetry is more consequential than it looks. The output sequence length is set by the queries alone. If you can afford few queries and many keys, you get a cheap read over an enormous source, and several well-known architectures are exactly this observation applied hard.

A vision-language model has a choice here, and the choice is architectural rather than cosmetic. Early fusion projects image patches into the text embedding space and concatenates them, so every layer is self-attention over a mixed sequence — simple, no new modules, but the image tokens cost you full quadratic context. Cross-attention fusion, as in Flamingo, keeps the image tokens in a side buffer and inserts cross-attention layers that read from it, so the text stream stays short. Bottleneck resampling, as in BLIP-2's Q-Former or a Perceiver latent array, uses a small fixed set of learned query vectors to compress hundreds of patch tokens into a few dozen before the language model ever sees them.

The learned-query trick deserves a moment. In a Q-Former the queries are parameters, not functions of the input. They do not depend on the image at all; the image only enters through \(K\) and \(V\). Each query vector is effectively a learned question asked of every image, and the answers are the tokens the LLM receives. That is why the number of learned queries is such a load-bearing hyperparameter: it is the width of the channel between the two modalities.

Why more than one head

A single attention layer computes one distribution per position. One distribution means one notion of relevance, and one notion is not enough: at the same position a model may want to track the syntactic head, the coreferent, the position two tokens back, and the topic of the paragraph. Averaging four incompatible answers into one softmax gives you none of them.

So run \(h\) of them in parallel on slices of the representation:

The usual convention is \(d_h = d_{\text{model}}/h\), which makes multi-head attention cost the same as single-head attention with the same \(d_{\text{model}}\). The projections do not add parameters, they partition them. You are trading the ability to compute one attention pattern over \(d\) dimensions for \(h\) patterns over \(d/h\) dimensions each, and empirically that trade is very favourable.

Two honest caveats, because this is where popular explanations overclaim. Heads do not divide neatly into human-nameable jobs. Some are legible — previous-token heads, delimiter heads, induction heads — and some are not, and the ones that look interpretable in a diagram are a biased sample of the ones that were easy to visualise. Also, most heads are not doing much: Michel, Levy and Neubig found that large fractions of heads can be pruned after training at little cost, with a small number of heads carrying most of the work.

Which raises the question that drives the next section. If the heads are individually cheap in parameters and largely redundant in function, what exactly is expensive about them?

MHA, MQA, GQA

Nothing about multi-head attention is expensive during training. The parameters are fixed, the FLOPs are the same as single-head, and the memory for activations is dominated by other things. The problem appears at inference, and it appears for one specific reason: autoregressive decoding caches keys and values.

When you generate token \(t+1\) you need its query against every previous key. Recomputing those keys each step is quadratic waste, so you store them. The cache is

and it grows linearly in both batch size and sequence length, without bound, while the weights stay fixed. For a 7B-class model at 8K context in fp16 that is 4 GiB of cache for a single sequence. Serve sixty-four of them and the cache alone dwarfs the model.

Worse, this is a bandwidth problem, not just a capacity problem. Decoding one token requires reading the entire cache from HBM to compute one column of attention scores. Arithmetic intensity is close to one FLOP per byte, so the GPU idles while memory streams. The cache size is the token latency.

The single term in that product you can touch without changing the model's shape is \(H_{\mathrm{KV}}\).

Multi-query attention takes it to the floor: keep all \(H\) query heads, share one key head and one value head across all of them. The cache shrinks by a factor of \(H\). Every head still computes a different attention pattern, because the queries still differ, but they all match against the same key space and read from the same value space. Quality drops, and training gets less stable.

Grouped-query attention interpolates. Choose \(1 < H_{\mathrm{KV}} < H_Q\), partition the query heads into \(H_{\mathrm{KV}}\) groups, and give each group its own KV head. MHA and MQA are the endpoints \(H_{\mathrm{KV}} = H_Q\) and \(H_{\mathrm{KV}} = 1\). The GQA paper also showed you can uptrain an existing MHA checkpoint into GQA by mean-pooling the KV heads within each group and continuing training for a small fraction of the original budget, which is a large part of why adoption was so fast.

\(H_{\mathrm{KV}} = 8\) has become the default across open model families for a reason visible in the figure: the first factor-of-four is nearly free in quality and removes three quarters of the cache, while the last factor-of-eight down to MQA costs real accuracy for a much smaller absolute saving. There is also a systems argument. With eight KV heads and eight-way tensor parallelism, each rank owns exactly one KV head and no cache is replicated — go below eight and you either replicate the cache across ranks or complicate the sharding.

Note what this section did not change: the attention rule. Every query head still computes a softmax over scaled dot products. We only changed how many distinct key and value projections exist. That is a recurring pattern worth watching for.

Local and sliding-window attention

Back to the mask, now used to attack complexity rather than semantics. Dense attention costs \(O(n^2 d)\) FLOPs. Restricting each position to a window of \(w\) neighbours gives \(O(n w d)\), linear in sequence length.

The objection writes itself: surely a token sometimes needs a token far away. Two answers.

The first is empirical. Attention distributions in trained models are heavily concentrated on nearby positions, plus a small number of special positions — often the very first tokens, which act as attention sinks where heads dump probability mass when they have nothing to retrieve. Discarding the long tail of near-zero weights costs less than the quadratic term suggests.

The second is structural and more important. Local attention is not myopic once you stack it. A token attends to \(w\) neighbours in layer one; each of those attended to \(w\) neighbours of their own, so after \(L\) layers the receptive field is roughly \(L \cdot w\) wide. Mistral 7B's window of 4096 across 32 layers gives a theoretical reach of 131K tokens. Information travels; it just takes layers to do it, exactly like stacked convolutions.

The practical designs mix local with something global. Longformer pairs a sliding window with a handful of positions that attend everywhere and are attended by everyone. StreamingLLM keeps the first few tokens permanently in the window and gets stable generation far past the training length. Gemma 2 and several recent models alternate: local layers for most of the stack, full attention every few layers, so the cheap layers do the routine work and the expensive ones handle long-range routing. That alternation also has a pleasant consequence for serving, since only the global layers need a full-length KV cache.

Sparse attention

A sliding window is one sparsity pattern. The general move is to pick any structured subset \(\mathcal{S} \subset \{1..n\}^2\) of pairs and compute only those. Rather than cataloguing the two dozen published patterns, it is more useful to hold the design axes, because every pattern is a point in this space:

Fixed or content-dependent? A fixed pattern is decided before you see the data — windows, strides, blocks. It is kernel-friendly and predictable. A content-dependent pattern routes based on the input, as in clustering or routing attention, and recovers more of dense attention's flexibility at the cost of a data-dependent gather that hardware dislikes.

What guarantees connectivity? If two positions are never in the same pattern at any layer, no information can flow between them. Fixed patterns need an explicit story for this: strided attention interleaves a local pattern with a strided one so any pair is reachable in two hops; Big Bird adds random edges and proves the result is a expander-like graph that remains a universal approximator.

Is the sparsity block-structured? This is the axis that decides whether a paper's FLOP savings become wall-clock savings. GPUs multiply tiles, not elements. Skipping 90% of individual entries scattered at random buys nothing; skipping 90% of \(128\times128\) blocks buys almost the full 10×. Any sparse method that survives contact with a real kernel is block-sparse in practice.

Linear attention

Sparsity keeps the softmax and deletes pairs. Linear attention keeps every pair and deletes the softmax, which sounds worse and is actually a deeper change.

The quadratic cost has a specific cause: we compute \(QK^\top\) first. Matrix multiplication is associative, so with \(Q \in \mathbb{R}^{n\times d}\), \(K \in \mathbb{R}^{n \times d}\), \(V\in\mathbb{R}^{n\times d}\),

Same result, different order, and for \(n \gg d\) the right-hand side is dramatically cheaper. The intermediate is a \(d\times d\) matrix whose size does not depend on sequence length at all.

The reason we cannot simply use it is the softmax sitting between the two multiplications. \(\operatorname{softmax}(QK^\top)V\) requires the full score matrix to exist before any normalisation happens. So Katharopoulos et al. replace \(\exp(q^\top k)\) with a factorisable similarity \(\phi(q)^\top\phi(k)\) for some feature map \(\phi\) with non-negative outputs:

Both sums are shared across all queries, so you compute them once. Performer goes further and chooses \(\phi\) as a random feature map that makes \(\phi(q)^\top\phi(k)\) an unbiased estimator of the softmax kernel, so the approximation is to the original operation rather than a redefinition of it.

Now the part that matters most. In the causal case those sums become running sums, and the layer becomes a recurrence:

Generation is now \(O(1)\) per token with \(O(d^2)\) state, and there is no KV cache at all. That is the real prize, and it is also the real cost: a fixed-size state cannot store an arbitrarily long history. Softmax attention keeps every token verbatim and pays \(O(n)\) memory for perfect recall; linear attention keeps a lossy summary and pays \(O(1)\). Tasks that need exact retrieval from far back — the ones induction heads solve — are exactly where pure linear attention degrades.

This is the same trade-off that RetNet, the gated linear attention family and the modern SSMs are all negotiating, mostly by making the state update forget selectively instead of summing uniformly. It is also why nearly every production model in this family is a hybrid: mostly linear or recurrent layers for the streaming work, a few full-attention layers kept for exact recall. That comparison deserves its own article.

Multi-head latent attention

GQA shrinks the cache by reducing the number of key-value heads. There is another direction: keep the heads and reduce the rank.

The observation behind DeepSeek-V2's multi-head latent attention is that the concatenation of all per-head keys and values for a token is a very wide vector produced by a linear map of one hidden state. Wide, but not high-rank. So compress it, cache the compressed version, and reconstruct on the fly:

with \(d_c \ll H \cdot d_h\). Only \(c_t^{KV}\) is stored. The up-projections are weights, not cache, so they cost nothing per token.

Stated that way it sounds like it just moves work from memory to compute — you would decompress every cached token at every step. The trick that makes it genuinely cheap is absorption. The score is

and the bracketed product is a constant matrix that can be folded into the query projection at load time. The per-head keys are never materialised at all; attention scores are computed directly in latent space. The same trick absorbs \(W^{UV}\) into \(W_O\) on the output side.

RoPE breaks it. A rotary embedding applies a position-dependent rotation to \(k_j\), and a position-dependent matrix cannot be folded into a fixed weight, because it sits between the two things you wanted to multiply together. MLA's answer is surgical: split the key into a content part that lives in latent space and stays absorbable, and a small decoupled rotary part carried by its own narrow projection, shared across heads and cached separately. In DeepSeek-V2 that is 512 latent dimensions plus 64 rotary dimensions, 576 numbers per token per layer, against 32,768 for the equivalent full MHA.

The reported headline is a 93.3% KV cache reduction relative to their own dense-MHA 67B model, with quality that DeepSeek reports as better than MHA rather than merely close. Two caveats worth keeping: the comparison is first-party and against MHA rather than a tuned GQA baseline, and MLA is considerably more intricate to implement well — the absorption has to be done deliberately, and a naive implementation gives up most of the benefit.

Attention over images

Everything so far assumed a sequence. Images are grids, and the grid structure changes which patterns are natural.

A vision Transformer flattens patches into a sequence and runs ordinary bidirectional self-attention, which is the honest baseline: it makes no assumption about 2D locality and lets the model learn it. The cost is that patch count grows quadratically with resolution while attention grows quadratically with patch count, so resolution enters at the fourth power. At 224 pixels with 14-pixel patches you have 256 tokens; at 1024 pixels you have 5,329, and the score matrix has 28 million entries per head per layer.

Hence the structured alternatives, each of which is a mask expressed in grid coordinates rather than sequence coordinates.

Windowed and shifted windows. Swin partitions the grid into non-overlapping windows and attends within each, then shifts the partition by half a window in the next layer so information crosses boundaries. Cost becomes linear in the number of patches, and the shift is doing exactly the job the receptive-field argument from the sliding-window section describes.

Axial. Attend along rows, then along columns. Two layers give a full receptive field over the grid at \(O(N\sqrt{N})\) instead of \(O(N^2)\), and it decomposes the mask into two dense operations, which is friendlier to kernels than an irregular 2D neighbourhood.

Bottleneck resampling. The Perceiver and Q-Former pattern from earlier, which for vision is often the most useful of all: the expensive quadratic term is over patches, so if a fixed latent array does the reading, the cost becomes linear in patch count and the downstream stack never sees the full grid.

For vision-language models the practical question is usually not which of these to use inside the encoder, but how many tokens survive the connector. A 7B language model does not care whether your patches were attended axially; it cares that it has been handed 729 image tokens per tile and that they occupy its context and its KV cache like any other token.

FlashAttention is not a new attention rule

This distinction causes more confusion than anything else in the area, so it gets its own section.

MHA, GQA, MLA, sliding windows and linear attention change what function is computed. Give them the same inputs and they return different numbers, because the model is different. FlashAttention changes how the same function is executed. Given the same inputs it returns the same numbers, to floating-point reassociation. It is not an approximation, and it is not an architecture choice.

The problem it solves is arithmetic intensity. The textbook implementation materialises \(S = QK^\top\) in HBM, reads it back to softmax it, writes \(P\), reads it again to multiply by \(V\). Those \(n \times n\) round trips dominate: at \(n = 8192\) each intermediate is 128 MB in fp16, and the kernel spends its life moving bytes while the tensor cores idle.

FlashAttention tiles the computation so that a block of queries meets a block of keys and values inside SRAM, accumulating the output without ever writing the score matrix out. That requires normalising a softmax you have not finished seeing, which is possible because softmax only depends on score differences: carry a running maximum \(m\) and a running denominator \(\ell\), and rescale the accumulator by \(e^{m_{\text{old}} - m_{\text{new}}}\) whenever the maximum moves.

Memory traffic drops from \(O(n^2)\) to \(O(n^2 d / B)\) for on-chip block size \(B\), and activation memory for the backward pass drops from \(O(n^2)\) to \(O(n)\) because the scores are recomputed from the saved statistics instead of stored. FlashAttention-2 rebalanced the loops and parallelised over sequence length to fill more SMs; FlashAttention-3 overlapped the matmuls with the softmax and exploited Hopper's asynchrony and FP8 paths.

The same category error appears one level up the stack. Paged attention is a memory-management scheme: it stores the KV cache in fixed-size blocks with an indirection table, so a serving system can avoid fragmentation and share prefixes between requests. Continuous batching is a scheduling scheme. Neither changes the mathematics. When you read that a system got a large speedup, the first question to ask is which of the three layers it touched — the model, the kernel, or the scheduler — because they compose, and confusing them makes it impossible to reason about what to try next.

A taxonomy

A chronological list of attention papers is nearly useless, because the papers are answers to different questions and only look comparable in a table sorted by date. Sorted by the question instead:

The last row is the one to internalise. Seven of these eight rows are model changes that need to be trained and evaluated. One of them is free.

When a new attention paper appears, the fastest way to place it is to ask which row it belongs to, and then to ask what it gave up — because every row above the last one gives up something, and a paper that claims otherwise has usually chosen a baseline carefully.

References

Grouped by topic. Every entry links to its primary source.