Attention

Read this sentence: "The animal didn't cross the street because it was too tired." What does "it" refer to? You knew instantly: the animal. Now swap one word: "…because it was too wide." Suddenly "it" is the street. Nothing about the word "it" changed. Its meaning came entirely from the words around it.

That is the problem attention solves. After the embedding layer, every token is a fixed vector: "it" is the same vector in both sentences, and "bank" is the same vector whether you are fishing or depositing a cheque. A language model cannot predict the next word well with vectors like that. It needs a way for each token's vector to be updated using the other tokens in the sentence, and it needs to decide, per token, per sentence, which other tokens matter.

Attention is that mechanism. It is not complicated. It is three matrix multiplications and a softmax. But it is the single idea that made the transformer work, and it is worth understanding so deeply that you could rebuild it from memory with a pencil. That is the goal here.

The problem: a token's meaning depends on context

Let's be precise about what we need. We have a sequence of $L$ token vectors $x_1, x_2, \dots, x_L$, each of width $d_{model}$. We want to produce a new sequence $y_1, \dots, y_L$ where each $y_i$ is a version of $x_i$ that has "absorbed" the relevant context. The mechanism must satisfy four requirements:

• Content-based. Which tokens matter to "it" depends on what the other tokens say, not on fixed positions. Sometimes the referent is two words back, sometimes twenty.
• Variable-length. The same weights must work for a 5-token prompt and a 5,000-token document.
• Parallel. We want to process every token at once on a GPU, not one after another.
• Differentiable. Every step must have a gradient so the whole thing can be learned from data.

Before 2017 the dominant answer was the recurrent network: read tokens one by one, squeezing everything seen so far into a single hidden state. That state is a bottleneck. By the time "it" arrives, "animal" has been compressed, overwritten, and half-forgotten. Bahdanau et al. (2014) patched this for translation by letting the decoder "look back" over all encoder states with learned weights, and called it attention. Vaswani et al. (2017) then asked the radical question: what if we drop the recurrence entirely and use only attention?

The library analogy: query, key, value

Here is the mental picture that makes attention obvious. Imagine walking into a library with a question in your head. Every book has a label on its spine. You compare your question against every label, pull out the books whose labels match best, and read them. Three things are in play:

• Your query $q$: what you are looking for.
• Each book's key $k_j$: the label on the spine, what the book advertises.
• Each book's value $v_j$: the content inside, what you actually take away when you read it.

Attention is exactly this, with two twists. First, it is soft: instead of picking one book, you read all of them, weighted by how well each label matched, so a book that matched 70% contributes 70% of what you take away. Second, every token is simultaneously a reader and a book. Token $i$ produces a query (to go looking), and also a key and a value (so that others can find it and read from it).

Where do $q$, $k$, $v$ come from? Each is a learned linear projection of the token's embedding. Three weight matrices, $W_Q$, $W_K$, $W_V$, are shared by every token and learned during training. That is all the "knowledge" attention has: it learns what kinds of questions tokens should ask and what kinds of labels they should wear.

Anatomy of attention: one token's point of view

Let's build attention for a single token, "sat", in the four-token sentence "The cat sat down". We will use vectors of width $d = 3$ so every number fits on a napkin. There are five steps.

Step 1 — Project: make a query, keys and values

Suppose the embedding of "sat" is $x_{\text{sat}} = [0,\ 1,\ 1]$, and the learned query matrix is

The query is the row vector times the matrix. Multiply and add:

So "sat" is now asking a question, encoded as $[1, 2, 0]$. Every token does the same thing with $W_K$ and $W_V$ to get its key and its value. To keep the arithmetic short, here are the results (you can invent any $W_K$, $W_V$ and embeddings that produce them; the mechanism does not care):

Step 2 — Score every key

How well does the question match each label? We use the dot product, because it is cheap and it is large when two vectors point the same way. The score of token $j$ for our query is $s_j = q \cdot k_j$:

"cat" scores highest. Its key $[1,2,0]$ happens to be exactly what the query was looking for. That is the whole idea: the query and key spaces are learned so that "verb looking for a subject" lands near "noun that can be a subject".

Step 3 — Scale

We divide every score by $\sqrt{d}$, here $\sqrt{3} \approx 1.732$. We will see exactly why in a moment; for now, think of it as keeping the numbers in a sensible range no matter how wide the vectors are.

Step 4 — Softmax: turn scores into weights

Scores are arbitrary real numbers. We want weights: non-negative, summing to one, so they can be used as "how much to read from each token". Softmax does this by exponentiating each score and normalising:

Plugging in: $e^{0}=1$, $e^{2.89}\approx 17.9$, $e^{0.58}\approx 1.78$, $e^{1.15}\approx 3.17$; their sum is $\approx 23.9$. So

"sat" will take 75% of what it reads from "cat", 13% from "down", and a little from itself and "The". Softmax also has a useful property: because of the exponential, gaps between scores turn into ratios between weights. A score lead of 2 becomes a weight ratio of $e^2 \approx 7.4$.

Step 5 — Weighted sum of values

Finally we read. The output for "sat" is the values, blended by the weights:

Look at what happened. The output is mostly the value of "cat" (the second coordinate is 0.88), with a bit of "down" mixed in. The token "sat" now carries, inside its vector, the information "my subject is a cat". Downstream layers can use that. And notice that nothing in the five steps depended on there being exactly four tokens: with forty tokens we would do the same thing, just with forty scores.

The matrix form

We just did attention for one query. But every token is asking a question at the same time, and a GPU is happiest doing everything at once. So we stack. Put the $L$ token embeddings as rows of a matrix $X$ of shape $(L, d_{model})$. Then all the queries, keys and values fall out of three matrix multiplications:

where $W_Q, W_K$ have shape $(d_{model}, d_k)$ and $W_V$ has shape $(d_{model}, d_v)$. Now $Q$, $K$ and $V$ each have one row per token. Every dot product between every query and every key is a single product, $QK^\top$, which has shape $(L, L)$: row $i$ holds the scores of query $i$ against all $L$ keys. Scale, softmax each row, multiply by $V$, and you have the famous formula:

Read it as a sentence: "compare every query with every key, turn each row of comparisons into weights, and use those weights to blend the values." The five steps from before are all in there; the matrix form just does them for every token in one shot.

Why divide by √d_k?

The scaling looks like a fussy detail. It is not; without it attention barely trains at wide $d_k$. Here is the argument. Suppose the entries of $q$ and $k$ are independent with mean 0 and variance 1 (roughly true at initialisation). Their dot product is a sum of $d_k$ terms, each a product of two such numbers:

Each term has mean 0 and variance $1 \times 1 = 1$, and the variance of a sum of independent terms is the sum of the variances. So the dot product has variance $d_k$, and standard deviation $\sqrt{d_k}$. With $d_k = 64$ (GPT-2's per-head width), typical raw scores are around $\pm 8$; with $d_k = 512$ they are $\pm 22$.

Now recall that softmax turns score gaps into weight ratios. A gap of 16 between the best and second-best key (easily produced at $d_k=64$) gives a ratio of $e^{16} \approx 9{,}000{,}000$. The softmax is effectively a hard argmax: one weight is 1.0, the rest are 0.0. That has two bad consequences. The model can no longer blend, and the gradient through softmax (which is proportional to $a_j(1-a_j)$) is essentially zero everywhere, so nothing learns.

Dividing by $\sqrt{d_k}$ brings the standard deviation back to 1 for every width. The formula is

and now a typical gap is around 1–2, giving weight ratios of $e^1$ to $e^2$, a gentle preference rather than a verdict. The model can still learn to be sharp when it needs to (by growing $W_Q$ and $W_K$), but it starts soft.

Causal masking: no peeking at the future

A language model is trained to predict the next token. Position 3 should predict token 4 using only tokens 1–3. But plain attention lets every token read every other token, including the ones to its right. If "sat" could read "down", the model would learn to simply copy the answer from the future, and at generation time, when the future does not exist yet, it would fall apart.

The fix is a causal mask (also called a look-ahead mask). Before the softmax, we set every score $S_{ij}$ with $j > i$ to $-\infty$. Since $e^{-\infty} = 0$, those entries get exactly zero weight, and the remaining weights in each row still sum to 1. Token $i$ can only read tokens $1, \dots, i$.

Two practical notes. First, the mask is a fixed pattern that depends only on $L$, so it costs nothing to learn; you build it once with torch.tril. Second, the diagonal is always allowed, which guarantees every row has at least one finite score. A row that is entirely $-\infty$ would make softmax divide zero by zero and produce not-a-number values, and that is a classic bug when people add padding masks carelessly.

Multi-head attention

Here is a limitation of what we have built. One attention layer produces one weight vector per token. "sat" can attend 75% to "cat", but it cannot at the same time attend mostly to "cat" for one purpose (who is the subject?) and mostly to "down" for another (what kind of sitting?). A single softmax has to compromise.

The fix is almost embarrassingly simple: run several attention operations in parallel, each with its own $W_Q, W_K, W_V$, and concatenate the results. Each parallel copy is a head. Head 1 can learn to look for subjects; head 2 can learn to look at the previous word; head 3 can learn to track punctuation. Since each head has its own projections, each can ask a different question of the same context.

Splitting the model width into heads

If we simply ran $h$ full-width heads, the cost would multiply by $h$. Instead the transformer keeps the total cost the same by making each head narrower: with $d_{model} = 768$ and $h = 12$ heads, each head works in $d_{head} = 768/12 = 64$ dimensions. Concretely:

Each head outputs $(L, d_{head})$. Concatenating $h$ of them side by side gives $(L, h\cdot d_{head}) = (L, d_{model})$ again. The final output projection $W_O$ mixes the heads' findings together, so the answer from head 3 can be combined with the answer from head 7 before being written back into the token's vector.

What heads learn

Nothing in the architecture tells head 3 to do anything in particular. Yet when you train a transformer on text and then look at its attention weights, you find heads with recognisable, stable jobs. This is one of the most satisfying findings in interpretability research, and it is worth knowing the main characters.

Previous-token heads and positional heads

Many heads in early layers attend almost entirely to the token immediately before the current one, or to a fixed offset like two back. Their attention matrix is a bright sub-diagonal. They are not "understanding" anything; they are giving each token a copy of its neighbour's information, which turns out to be a building block that later heads need (see induction heads below). Some heads attend mostly to the very first token, which usually means the head has nothing useful to do for that query and parks its weight on a harmless "sink" (Xiao et al., 2023, "Efficient Streaming Language Models with Attention Sinks").

Induction heads

Suppose the context contains "… Mr Dursley … Mr" and the model must predict what follows the second "Mr". A smart strategy: find the earlier "Mr", look at what came after it, and copy that. Olsson et al. (2022), "In-context Learning and Induction Heads", found that transformers reliably learn exactly this circuit, using two heads in sequence. A previous-token head in an earlier layer first stamps every token with "the token before me was X". Then an induction head in a later layer, sitting on the second "Mr", asks "who has a previous token equal to me?", attends to "Dursley", and copies its value forward. Their attention pattern looks like a shifted copy of the token-match pattern. Olsson et al. showed that these heads appear abruptly during training, at the same moment that the model's ability to do in-context learning jumps.

Attention patterns as pictures

The easiest way to build intuition is to look at attention matrices as images: rows are queries, columns are keys, brightness is weight. Once you have seen a few, you can recognise a previous-token head or an induction head at a glance.

The cost: attention is quadratic in sequence length

Now for the bad news. Look again at Step 2: $S = QK^\top$ is an $(L, L)$ matrix. Computing it takes $L^2 d_k$ multiply-adds, or about $2L^2 d_k$ FLOPs. Multiplying $A$ by $V$ costs the same again. So the core of attention costs roughly $4L^2 d$ FLOPs per layer, and it needs $L^2$ numbers of memory per head to store the scores (before FlashAttention-style tricks). Double the context and both quadruple.

Compare that to everything else in the layer. The four projections cost $8Ld^2$ and the feed-forward block (next chapter) costs about $16Ld^2$: linear in $L$. So attention's quadratic term is small for short sequences and dominant for long ones. The crossover, where $4L^2d = 24Ld^2$, is at $L = 6d$: about $4{,}600$ tokens for GPT-2 small ($d = 768$), about $25{,}000$ tokens for a $d = 4096$ model. Modern models with 128k-token contexts are far past the crossover.

The KV cache: a preview

At generation time the situation is a little different. We produce one token at a time, so there is only ever one new query, but it must attend to every key and value so far. Recomputing all previous keys and values at every step would be wasteful, since they never change (each depends only on its own token's vector, and the causal mask guarantees earlier tokens never see later ones). So we cache them: the KV cache. Generation then costs $O(L)$ per new token in compute, but memory grows with $L$: per token, per layer, we store $2 \times h \times d_{head}$ numbers. That memory is what forces the trick in the next section. The inference chapter goes deep on this.

Grouped-query and multi-query attention

Take Llama-2 7B: 32 layers, 32 heads, $d_{head} = 128$. The KV cache costs $2 \times 32 \times 128 = 8{,}192$ numbers per token per layer, times 32 layers, times 2 bytes in fp16: about 512 KB per token. A single 4,096-token conversation needs 2 GB of cache; serve 32 users at once and the cache alone is 64 GB, more than the weights. Worse, at generation time every step has to read the entire cache from GPU memory, and memory bandwidth, not compute, is what limits decoding speed.

The observation that saves the day: the queries are what make heads different, but the keys and values are just "labels" and "content" of the context. Do we really need 32 separate labellings of the same context? Multi-query attention (MQA; Shazeer, 2019, "Fast Transformer Decoding: One Write-Head is All You Need") says no: keep $h$ query heads but share a single key head and value head among all of them. The cache shrinks by a factor of $h$.

MQA is a little too aggressive: quality drops slightly and training can be less stable. Grouped-query attention (GQA; Ainslie et al., 2023, arXiv) is the compromise. Use $h_{kv}$ key/value heads, with $h_{kv}$ between 1 and $h$, and let each KV head serve a group of $h / h_{kv}$ query heads. Llama-2 70B and all Llama-3 models use $h_{kv} = 8$: with 64 query heads, each KV head serves 8 query heads, and the cache is 8× smaller than full multi-head attention at nearly identical quality.

The maths of a GQA layer is the same as before with one extra step: after projecting $K$ and $V$ with $h_{kv}$ heads, repeat each KV head $h/h_{kv}$ times so it lines up with its group of query heads (repeat_interleave in PyTorch), then run ordinary per-head attention. The parameter count of $W_K$ and $W_V$ drops from $d_{model}\times d_{model}$ to $d_{model} \times h_{kv} d_{head}$ each.

Implementation walkthrough in PyTorch

Let's build it, matching code/lumen/attention.py. First a single head, as a function.

import math
import torch
import torch.nn as nn

def scaled_dot_product_attention(q, k, v, mask=None):
    """q: (..., L_q, d_k)   k: (..., L_k, d_k)   v: (..., L_k, d_v)
    mask: bool tensor broadcastable to (..., L_q, L_k); True = may attend.
    Returns (output (..., L_q, d_v), weights (..., L_q, L_k))."""
    d_k = q.size(-1)
    scores = q @ k.transpose(-2, -1) / math.sqrt(d_k)          # step 2 + 3: (..., L_q, L_k)
    if mask is not None:
        scores = scores.masked_fill(~mask, float("-inf"))       # step 4: -inf where NOT allowed
    weights = torch.softmax(scores, dim=-1)                     # step 5: rows sum to 1
    return weights @ v, weights                                 # step 6: blend the values

def causal_mask(L, device=None):
    """(L, L) lower-triangular bool matrix: row i may see columns 0..i."""
    return torch.tril(torch.ones(L, L, dtype=torch.bool, device=device))

Line by line. k.transpose(-2, -1) swaps the last two axes so that the product is $QK^\top$; using negative axis indices means the same code works with or without batch and head dimensions in front. Dividing by math.sqrt(d_k) is Step 3. masked_fill(~mask, -inf) writes $-\infty$ where the mask is False (note the ~: our mask says where attention is allowed). softmax(dim=-1) normalises each row, i.e. over keys. Finally weights @ v is $AV$. We return the weights too, which is handy for plotting.

Now multi-head. The whole trick is in two view/transpose calls, so let's not use einsum and instead watch the shapes.

class MultiHeadAttention(nn.Module):
    def __init__(self, d_model, n_heads, bias=True):
        super().__init__()
        assert d_model % n_heads == 0, "d_model must divide evenly into heads"
        self.n_heads = n_heads
        self.d_head = d_model // n_heads
        self.qkv = nn.Linear(d_model, 3 * d_model, bias=bias)   # W_Q, W_K, W_V fused side by side
        self.out = nn.Linear(d_model, d_model, bias=bias)       # W_O

    def forward(self, x, mask=None):
        B, L, D = x.shape                                       # batch, seq len, d_model
        q, k, v = self.qkv(x).split(D, dim=-1)                  # three (B, L, D) tensors
        # (B, L, D) -> (B, L, H, d_head) -> (B, H, L, d_head): heads become a batch dim
        q = q.view(B, L, self.n_heads, self.d_head).transpose(1, 2)
        k = k.view(B, L, self.n_heads, self.d_head).transpose(1, 2)
        v = v.view(B, L, self.n_heads, self.d_head).transpose(1, 2)
        if mask is None:
            mask = causal_mask(L, x.device)                     # broadcasts over (B, H)
        y, _ = scaled_dot_product_attention(q, k, v, mask)     # (B, H, L, d_head)
        y = y.transpose(1, 2).contiguous().view(B, L, D)        # concat heads: (B, L, H*d_head)
        return self.out(y)                                      # mix heads with W_O

Walk through the shapes with GPT-2 small and a batch of 2 sequences of 1024 tokens. x is $(2, 1024, 768)$. The fused qkv linear produces $(2, 1024, 2304)$; split(768) cuts it into three $(2, 1024, 768)$ tensors. This fused layer is mathematically identical to three separate linears; it is just one bigger, faster matmul. view(B, L, 12, 64) reinterprets each 768-vector as 12 chunks of 64 with no data movement. transpose(1, 2) gives $(2, 12, 1024, 64)$, so that the attention function sees "batch × heads" as leading dimensions and does 24 independent attentions in one call. The mask $(1024, 1024)$ broadcasts against the scores $(2, 12, 1024, 1024)$. After attention, transpose(1,2) puts the heads back next to their $d_{head}$ columns, contiguous() makes the memory layout match (a transpose is only a view, and view needs contiguous memory), and view(B, L, D) is the concatenation: head 0's 64 numbers, then head 1's, and so on. self.out is $W_O$.

Grouped-query attention needs only a small change: separate projections with fewer KV heads, and a repeat_interleave before attention.

class GroupedQueryAttention(nn.Module):
    def __init__(self, d_model, n_heads, n_kv_heads):
        super().__init__()
        assert n_heads % n_kv_heads == 0
        self.n_heads, self.n_kv_heads = n_heads, n_kv_heads
        self.d_head = d_model // n_heads
        self.q_proj = nn.Linear(d_model, n_heads * self.d_head, bias=False)
        self.k_proj = nn.Linear(d_model, n_kv_heads * self.d_head, bias=False)   # smaller!
        self.v_proj = nn.Linear(d_model, n_kv_heads * self.d_head, bias=False)   # smaller!
        self.o_proj = nn.Linear(n_heads * self.d_head, d_model, bias=False)

    def forward(self, x, mask=None):
        B, L, _ = x.shape
        q = self.q_proj(x).view(B, L, self.n_heads, self.d_head).transpose(1, 2)      # (B, H, L, dh)
        k = self.k_proj(x).view(B, L, self.n_kv_heads, self.d_head).transpose(1, 2)   # (B, Hkv, L, dh)
        v = self.v_proj(x).view(B, L, self.n_kv_heads, self.d_head).transpose(1, 2)
        rep = self.n_heads // self.n_kv_heads
        k = k.repeat_interleave(rep, dim=1)   # (B, H, L, dh): each KV head serves `rep` query heads
        v = v.repeat_interleave(rep, dim=1)
        if mask is None:
            mask = causal_mask(L, x.device)
        y, _ = scaled_dot_product_attention(q, k, v, mask)
        y = y.transpose(1, 2).contiguous().view(B, L, -1)
        return self.o_proj(y)

With n_kv_heads == n_heads this is ordinary multi-head attention; with n_kv_heads == 1 it is multi-query. The repeat_interleave is only there for clarity; production kernels read the shared KV head directly without copying it.

Common confusions

Practice

Further reading

• Vaswani et al. (2017). Attention Is All You Need. The original; Section 3.2 is the formula and the $\sqrt{d_k}$ footnote.
• Bahdanau, Cho, Bengio (2014). Neural Machine Translation by Jointly Learning to Align and Translate. Where soft attention started.
• Olsson et al. (2022). In-context Learning and Induction Heads. The clearest story of what heads learn and when.
• Elhage et al. (2021). A Mathematical Framework for Transformer Circuits. Attention as reading and writing to the residual stream; heads as QK and OV circuits.
• Shazeer (2019). Fast Transformer Decoding: One Write-Head is All You Need. Multi-query attention.
• Ainslie et al. (2023). GQA: Training Generalized Multi-Query Transformer Models from Multi-Head Checkpoints. Grouped-query attention.
• Dao et al. (2022). FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness. Exact attention without materialising the $(L,L)$ matrix.
• Jain and Wallace (2019). Attention is not Explanation. Why attention maps are not the model's reasoning.
• Xiao et al. (2023). Efficient Streaming Language Models with Attention Sinks. Why so much weight lands on the first token.