Optional Foundation — Build a Small Language Model

Attention as learned Q, K and V

Data and parameters·Advanced·7 min read

You know what attention computes from the inference chapters. This chapter asks the training question: the Q, K and V projections are matrices that start random — so what, exactly, does the model learn when it learns to attend, and how is that learning shaped by the causal mask and multiple heads?

Attention was introduced on the inference side as a computation — score every query against every key, softmax, blend the values (see Transformer and attention intuition). Here we look at the learning side. The Q, K and V projections are just weight matrices, and like every other matrix they start random. So the real question is: when a model "learns to pay attention," which numbers are moving, and what are they settling into?

What you will understand by the end

  • That Q, K, V (and the output projection) are trainable matrices, not fixed operations.
  • What the model actually learns when it learns attention — and what that has to do with the KV cache and attention variants.
  • Why the causal mask is a training constraint, not a runtime detail.
  • Why multiple heads let one layer learn several relationships at once.

Q, K and V are three learned projections

For each token's hidden vector, the model produces three vectors by multiplying it through three learned matrices:

  • W_Q turns the hidden vector into a query — "what am I looking for?"
  • W_K turns it into a key — "what do I offer to others?"
  • W_V turns it into a value — "what do I contribute if attended to?"

A fourth matrix, W_O, projects the blended result back out. All four begin as random numbers. Nothing about "query," "key," or "value" is built in — those roles emerge because training rewards weights that make useful next-token predictions. Early in training the attention is meaningless; by the end, W_Q and W_K have settled into values that make relevant tokens score highly against each other.

Key idea

The model does not learn which tokens to attend to directly — it learns the projections W_Q, W_K, W_V that, once applied, cause the right tokens to line up. Attention weights are computed fresh every forward pass from these learned matrices; the matrices are the memory, the attention pattern is the momentary result.

What "learning to attend" buys

The point of attention is to let a token pull in information from the specific earlier tokens that matter — resolving what a pronoun refers to, connecting a verb to its subject, matching a bracket to its opener. None of that is programmed; the W_Q/W_K matrices are nudged until scoring one token's query against another's key produces high weight exactly when the connection is useful for predicting what comes next.

This is also where the inference story you already know comes from:

  • Because attention needs every previous token's key and value, and those don't change once computed, inference caches them — the KV cache.
  • Because the K and V projections are what get cached and streamed, shrinking or sharing them across heads is exactly what MQA/GQA/MLA do. Those variants are decisions about the shape of these learned matrices.

So the matrices this chapter trains are the very things the serving chapters spend their effort caching, sharing, and compressing.

The causal mask is a training constraint

During training, the model sees a whole sequence at once and is scored on every position in parallel (from the data chapter). That creates a danger: when predicting position i, the model must not be allowed to look at positions i+1 and beyond — those are the answers. The causal mask enforces this by blocking each position from attending to any later one.

Watch out

The causal mask is not a cosmetic runtime detail — it is what makes training honest. Without it, the model could "predict" the next token by simply looking at it, learning nothing. The mask is why a model trained on full sequences still generalises to generating one token at a time: at every position it only ever used the past.

Multiple heads learn several relationships at once

A single attention head can only learn one way of relating tokens. Real language has many relationships in play simultaneously — grammatical, referential, positional. So a layer runs several heads in parallel, each with its own learned W_Q/W_K/W_V, and concatenates their outputs. Each head is free to specialise; training decides what each one ends up tracking, and it is rarely something a human labelled in advance.

Key idea

Multi-head attention gives a layer several independent sets of Q/K/V matrices, so it can learn several relationships at once. This is why cutting heads has a quality cost (MQA/GQA): fewer key/value heads means fewer distinct "views" the model can keep — a direct link between a training choice and a serving cost.

Mental model

Attention is not a fixed operation the model performs; it is a set of matrices (W_Q, W_K, W_V, W_O) the model learns. Training moves those numbers until querying one token against another surfaces the tokens that help predict the next one. The causal mask keeps that learning honest, and multiple heads let one layer learn many relationships at once.

Common mistakes

  • Thinking the model learns an attention pattern. It learns the projections; the pattern is recomputed every pass from them and the current tokens.
  • Treating the causal mask as an inference trick. It is a training necessity — without it the model cheats and learns nothing useful.
  • Assuming heads are assigned roles. No one tells a head to "track syntax"; the roles emerge from training, and many are not human-interpretable.
  • Forgetting the tie to serving. The K/V matrices you train here are the ones cached, shared, and compressed at inference — they are the same objects.

Practical guidance

  • When reasoning about a model's attention, remember the learnable objects are W_Q, W_K, W_V, W_O per head — that is what a fine-tune or a LoRA adapter actually modifies.
  • Read an "attention variant" (MQA, GQA, MLA) as a decision about the shape of the K/V projections made at design/training time, then paid back at inference.
  • Trust the causal mask to be why full-sequence training transfers to one-token-at-a-time generation — the model only ever conditioned on the past.
  • Don't over-interpret individual heads; treat multi-head capacity as "room to learn several relationships," not a fixed assignment.

Summary

  • Q, K, V and O are learned matrices that start random; "learning to attend" means fitting them so relevant tokens score highly.
  • The model learns the projections, not the attention pattern, which is recomputed each forward pass.
  • The causal mask makes training honest and is why full-sequence training transfers to autoregressive generation.
  • Multiple heads learn several relationships at once — which is exactly why sharing/cutting K/V heads (MQA/GQA/MLA) trades quality for a smaller cache.

Knowledge check

"The model learned to attend to the subject of the sentence." What did it actually learn, mechanically?

It learned values for W_Q and W_K such that the query of the current token and the key of the subject token produce a high score after the dot-product and softmax. The model never stores "attend to the subject"; it stores projection matrices that, applied to the tokens, make that attention weight come out high. The pattern is a computed consequence of the learned matrices.

Why can't you train a next-token model on full sequences without a causal mask?

Because when scoring the prediction at position i, an unmasked model could attend to position i+1 — the token it is supposed to predict — and "succeed" by copying the answer, learning nothing generalisable. The causal mask blocks every position from seeing later ones, forcing predictions to use only the past, which is also the only information available during real one-token-at-a-time generation.

Related chapters