In §4.1 the dictionary was given to us. The keys were apartment sizes, the values were rents, and the query came from outside. Inside a Transformer there is no outside. A sentence of $T$ tokens arrives as $T$ vectors, and every token needs to look things up in the other tokens. In "The animal didn't cross the street because it was too tired", the token "it" needs to find "animal", and the token "tired" needs to find whatever "it" turned out to be.

So where do the queries, keys, and values come from? The obvious answer is the token vectors themselves, which is what you did in Exercise 3.6d. This section shows why that obvious answer fails, why the fix is three separate learned projections, and exactly what each of the three is for. By the end you will have built a single attention head by hand, with numbers small enough to check on paper, and you will know the shape of every matrix involved, which is the first thing an interviewer will ask you to write down.

The naive version, where each token is its own query, key, and value

Stack the $T$ token vectors as the rows of a matrix $\mathbf{X} \in \R^{T \times d}$, where $d$ is the model width (the embedding dimension of §1.2). Row $i$ is token $i$'s vector, which we write $\mathbf{x}_i$. (Throughout this chapter vectors belonging to tokens are stored as rows, exactly as they are in code.) The naive soft lookup of §3.6 uses $\mathbf{x}_i$ as the query, every $\mathbf{x}_j$ as a key, and every $\mathbf{x}_j$ as a value,

Three things go wrong.

Every token mostly finds itself. The score of a token with itself is $\lVert\mathbf{x}_i\rVert^2$. In high dimensions this is far larger than a typical score with another token. If the entries are roughly independent with unit variance, $\lVert\mathbf{x}_i\rVert^2 \approx d$ while $\mathbf{x}_i^\top\mathbf{x}_j$ is of order $\sqrt d$ (Fact 1 of §1.3). The lookup returns the token you started with, which is useless. Word2vec ran into the same problem and solved it the same way we are about to, by giving each word two different vectors (§1.5).

Relevance is forced to be symmetric. The score $\mathbf{x}_i^\top\mathbf{x}_j$ equals $\mathbf{x}_j^\top\mathbf{x}_i$. But relevance is not symmetric. "It" needs "animal" to resolve what it refers to, while "animal" has no particular need for "it". A verb looks for its subject, and the subject is not looking for the verb in the same way.

One vector is doing three jobs. What a token is looking for, what it advertises about itself, and what it should hand over when someone finds it are different pieces of information. A pronoun is looking for a noun phrase, advertises that it is a pronoun, and should hand over, once resolved, whatever it refers to. Forcing a single vector to play all three roles means none is done well.

Here is the first problem, measured on random vectors.

With raw vectors, every token's best match is itself and it puts 99.5% of its weight there. With two different projections, even random ones, the diagonal stops being special (averaged over many draws, a token's weight on itself is about $1/T$), and the scores are no longer symmetric. Now we just have to make the projections learned.

Three projections, three roles

Give the layer three weight matrices and compute, for every token at once,

with $\mathbf{W}_Q, \mathbf{W}_K \in \R^{d \times d_k}$ and $\mathbf{W}_V \in \R^{d \times d_v}$. Row $i$ of $\mathbf{Q}$ is token $i$'s query $\mathbf{q}_i \in \R^{d_k}$, row $j$ of $\mathbf{K}$ is token $j$'s key $\mathbf{k}_j \in \R^{d_k}$, and row $j$ of $\mathbf{V}$ is token $j$'s value $\mathbf{v}_j \in \R^{d_v}$. (Bold $\mathbf{V}$ is the value matrix. Italic $V$ remains the vocabulary size.) Then run the soft lookup of §4.1 with these,

In matrix form, with the softmax applied to each row separately, $\mathbf{S} = \mathbf{Q}\mathbf{K}^\top$, $\mathbf{A} = \operatorname{softmax}(\mathbf{S})$, and $\mathbf{O} = \mathbf{A}\mathbf{V}$. (The real formula divides $\mathbf{S}$ by $\sqrt{d_k}$. We leave that out until §4.3, where it gets a section of its own.)

In plain words, each token asks a question (its query), every token answers how well it matches (its key), and each token collects a blend of everyone's information (their values), weighted by the match.

The three roles are worth stating precisely, because interviewers ask.

• The query $\mathbf{q}_i$ encodes what position $i$ is looking for.
• The key $\mathbf{k}_j$ encodes what position $j$ can be found by. It is an advertisement, used only to compute scores.
• The value $\mathbf{v}_j$ encodes what position $j$ contributes once it is found. It never affects who attends to whom.

A library is a decent analogy. You walk in with a question (the query). Each book has a catalog entry with subject headings (the key). What you take home is the book itself (the value). The catalog entry and the book are different objects, and they should be. A good catalog entry is short and easy to match. A good book is full of information.

The separation between keys and values is the less obvious one, so here is a concrete case. A previous-token head, one of the most common heads in trained language models, makes every position attend to the position just before it. To find the previous position, the keys must encode position. But what the head should deliver is the previous token's identity. With $\mathbf{k}_j = \mathbf{v}_j$ the head would have to deliver the position code it used for matching. With separate projections, the key reads the position part of $\mathbf{x}_j$ and the value reads the token part. You will build exactly this head in Exercise 4.2b, and in §4.7 it becomes the first half of an induction head.

The similarity is learned, and it is low rank

Substitute the projections into the score,

with $\mathbf{W}_{QK} = \mathbf{W}_Q\mathbf{W}_K^\top \in \R^{d\times d}$.

In plain words, a query-key pair defines a learned bilinear similarity between token vectors, given by one $d \times d$ matrix.

That single observation explains most of the design.

• It generalizes the dot product. With $\mathbf{W}_{QK} = \mathbf{I}$ we recover the naive $\mathbf{x}_i^\top\mathbf{x}_j$. Learning $\mathbf{W}_{QK}$ lets the head decide which features matter for matching, and how features of the query should line up with features of the key.
• It can be asymmetric. $s_{ij} \ne s_{ji}$ unless $\mathbf{W}_{QK}$ is symmetric, so "it" can care about "animal" without the reverse.
• It is low rank. $\operatorname{rank}(\mathbf{W}_{QK}) \le d_k$, because it is a product through a $d_k$-dimensional space. The head can only compare tokens along $d_k$ directions. That is a restriction, but also a bargain, since the factored form costs $2dd_k$ parameters instead of $d^2$, and §4.5 will spend the savings on several heads that each look along different directions.
• Only the product matters. For any invertible $\mathbf{R} \in \R^{d_k\times d_k}$, replacing $\mathbf{W}_Q$ by $\mathbf{W}_Q\mathbf{R}$ and $\mathbf{W}_K$ by $\mathbf{W}_K\mathbf{R}^{-\top}$ leaves $\mathbf{W}_{QK}$ unchanged. The individual query and key vectors therefore have no meaning on their own, only their pairing does. This is why interpretability work studies $\mathbf{W}_{QK}$ directly, and it is a good answer to "what do the query vectors mean?".

The values have a matching story. In a full layer, the output of the head is multiplied by an output matrix $\mathbf{W}_O \in \R^{d_v\times d}$ (with several heads, this head's block of rows of a larger matrix) that maps it back to width $d$ (§4.5). What a head moves from token $j$ into token $i$ is then $\mathbf{x}_j^\top\mathbf{W}_V\mathbf{W}_O$, governed by the product $\mathbf{W}_{OV} = \mathbf{W}_V\mathbf{W}_O \in \R^{d\times d}$ of rank at most $d_v$. Elhage et al. (2021) call these the QK circuit, which decides where to look, and the OV circuit, which decides what to move. A head is fully described by these two $d\times d$ low-rank matrices, a point of view that pays off when we build an induction head in §4.7 and when circuit analysis reverse-engineers trained models.

Shapes, written out

Here is every shape in one head, for a batch of $B$ sequences. The queries may come from a sequence of length $T_q$ and the keys and values from a sequence of length $T_k$. In self-attention both are the same sequence and $T_q = T_k = T$. In cross-attention they differ (§4.6).

Three facts are worth memorizing. Queries and keys must have the same width $d_k$, because they meet in a dot product. Values can have any width $d_v$, because they are only averaged. And the output has one row per query, with the value width, so the key length $T_k$ disappears from the result. Notice also that the weights $\mathbf{W}_Q, \mathbf{W}_K, \mathbf{W}_V$ do not depend on $T$. The same three matrices serve a sequence of 5 tokens or 100,000, which is what lets one model handle inputs of any length.

def attention_shapes(B, T_q, T_k, d, d_k, d_v):
    raise NotImplementedError

A worked example you can check by hand

Let's build a head that does something recognizable, with numbers small enough to verify with a pencil. The sentence is "the cat sat". Each token is a 4-dimensional vector whose coordinates are hand-made features, (is a determiner, is a noun, is a verb, is an animal),

We want two behaviors from one head. A verb should look for a noun (its subject), and a noun should look for a determiner. And whatever is found should report whether it is an animal, plus whether it is a determiner. With $d_k = d_v = 2$ we can write the weights directly. Rows of each weight matrix correspond to the four input features.

Read them like this. In $\mathbf{W}_Q$, the verb feature produces the query $(3, 0)$, meaning "I am looking for a noun", and the noun feature produces $(0, 3)$, meaning "I am looking for a determiner". In $\mathbf{W}_K$, the noun feature produces the key $(1, 0)$, "I am a noun", and the determiner feature produces $(0, 1)$, "I am a determiner". In $\mathbf{W}_V$, the value reports (is an animal, is a determiner). The factor 3 in the queries sets how sharp the lookup is, which is the temperature of §4.1 in disguise.

Multiply out. Each row of $\mathbf{Q} = \mathbf{X}\mathbf{W}_Q$ picks out the rows of $\mathbf{W}_Q$ for the features that token has,

Row by row, "the" has an all-zero query, so all its scores are 0. "cat" scores 3 against "the" (noun looking for determiner). "sat" scores 3 against "cat" (verb looking for noun). Apply the softmax to each row. A row with one score of 3 and two of 0 gives $e^3/(e^3 + 2) = 0.909$ on the winner and $1/(e^3+2) = 0.045$ on each of the others, so

Read the outputs. "sat" now carries $(0.909, 0.045)$, meaning "my subject is an animal". The verb has learned something about its subject, which is exactly the information it needed. "cat" carries $(0.045, 0.909)$, meaning "I have a determiner". And "the", which had nothing to ask, gets the plain average of everyone's values. That last behavior is typical. A query that matches nothing spreads its weight uniformly, because softmax has to put its weight somewhere. Trained models often handle this by parking the weight on a harmless token, the attention sink we meet in §4.4.

Notice too that "cat" attends to "the", which comes before it, and "sat" attends to "cat", also before it. Nothing in this head stops a token from looking ahead. In a language model it would have to be stopped, which is the job of the causal mask in §4.4.

Now drag the vectors yourself. The widget below has two queries and four keys in a 2-D plane, and four values in a second plane. The heatmap is the $2\times 4$ attention matrix, and each output is drawn as the weighted average of the values.

import numpy as np

def previous_token_head(n_vocab, T):
    d = n_vocab + T
    W_Q = np.zeros((d, T)); W_K = np.zeros((d, T)); W_V = np.zeros((d, n_vocab))
    return W_Q, W_K, W_V      # placeholder: all-zero weights attend uniformly

The whole head, batched

Real code handles a batch of sequences, and usually computes the three projections with one matrix multiply. Concatenate the three weight matrices side by side into $\mathbf{W}_{QKV} = [\mathbf{W}_Q \mid \mathbf{W}_K \mid \mathbf{W}_V] \in \R^{d \times (2d_k + d_v)}$, multiply once, and split the result along the last axis. One large matmul runs faster on a GPU than three small ones, which is why GPT-2's code has a single c_attn layer producing all three. The batch dimension rides along untouched, because every operation either acts on the last axis or is a batched matrix product over the last two.

One head with GPT-2's sizes ($d = 768$, $d_k = d_v = 64$) has $3 \times 768 \times 64 = 147{,}456$ weights in its three projections. GPT-2 has 12 such heads per layer, and §4.5 counts everything, including the output projection.

import numpy as np

def fused_self_attention(X, W_qkv, scale):
    raise NotImplementedError

Biases, and the one that does nothing

Many implementations add a bias to each projection, $\mathbf{q}_i = \mathbf{W}_Q^\top\mathbf{x}_i + \mathbf{b}_Q$ and so on. GPT-2 and BERT do. LLaMA drops all of them, while Qwen2 keeps the query, key, and value biases and drops the rest. Here is a fact that surprises people in interviews. The key bias has no effect whatsoever on the output. Expand the score,

The second term depends on the query $i$ but not on the key $j$. It adds the same constant to every score in row $i$, and softmax is invariant to that (§4.1). So the weights, the output, and therefore the loss do not depend on $\mathbf{b}_K$, and its gradient is exactly zero. (One caveat matters for modern models. With rotary position embeddings, §5.5, each key is rotated by an angle that depends on its position after the bias is added, so the bias term differs from key to key and is no longer redundant.) The query bias is different, since $\mathbf{b}_Q^\top\mathbf{k}_j$ varies with $j$ and acts as a learned, query-independent preference for certain keys. The value bias passes straight through, because the weights sum to 1, so $\sum_j a_{ij}(\mathbf{v}_j + \mathbf{b}_V) = \mathbf{o}_i + \mathbf{b}_V$. It is equivalent to a bias added after attention, which is why it can be folded into the output projection's bias.

In PyTorch

A production single head is a fused linear layer and a call to the attention kernel. We write the attention by hand here, and switch to torch.nn.functional.scaled_dot_product_attention in §4.3.

Note that nn.Linear stores its weight as (out_features, in_features) and computes x @ W.T. So head.qkv.weight has shape (3 d_k, d), the transpose of our $\mathbf{W}_{QKV}$. Mixing up the two conventions is a frequent source of shape errors when porting weights between frameworks.

Read more

• Vaswani et al. (2017). Attention Is All You Need. Section 3.2 defines queries, keys, and values and the projections used throughout this chapter.
• Elhage et al. (2021). A Mathematical Framework for Transformer Circuits. The QK and OV circuits, the non-uniqueness of the factorization, and why heads are best studied through $\mathbf{W}_Q\mathbf{W}_K^\top$ and $\mathbf{W}_V\mathbf{W}_O$.
• Phuong & Hutter (2022). Formal Algorithms for Transformers. Precise pseudocode for every attention variant, with all shapes stated. The reference to check your notation against.
• Sukhbaatar et al. (2015). End-To-End Memory Networks. Separate input (key) and output (value) memory embeddings, an early statement of why keys and values should differ.
• Alammar. The Illustrated Transformer. Step-by-step pictures of the query, key, and value computation for a single token.
• Bloem (2019). Transformers from Scratch. Builds self-attention from the naive version up, the same route as this section.

Watch next

▶

▶