AIM-03-004-P · Primer — Self-Attention: Q, K, V

The idea in one sentence

Self-attention passes each token's embedding through three separate learned linear layers to produce a query, a key, and a value — the same input wearing three different hats.

Why it exists

If a token used its raw embedding both to decide who's relevant and to supply what it actually contributes, those two jobs would be stuck sharing one representation — the network would have no room to shape "how I get found" independently from "what I hand over once found." Think of a conference name badge system with three separate slots: one line says what you're looking for right now (query — "seeking: ML infra people"), a second line says how you present yourself to be matched (key — "10 years, distributed systems"), and what you actually say once someone approaches you (value — your real pitch, which might emphasize something totally different, like a side project) is a third thing entirely. Self-attention gives every token exactly that: its own learned query, key, and value, each shaped by backprop to be useful for its own job — searching, advertising, and contributing content — rather than forcing one vector to do all three.

How it works

Given a sequence of token embeddings stacked as rows of a matrix \(X\) (shape \(n \times d_{model}\)), self-attention learns three weight matrices \(W_Q \in \mathbb{R}^{d_{model}\times d_k}\), \(W_K \in \mathbb{R}^{d_{model}\times d_k}\), and \(W_V \in \mathbb{R}^{d_{model}\times d_v}\), and computes \[Q = XW_Q,\qquad K = XW_K,\qquad V = XW_V.\] Every row of \(X\) (one token) now has three derived vectors: its query (what it's searching for), its key (how it advertises itself to be found), and its value (what it actually hands over once matched). \(Q\) and \(K\) must share dimension \(d_k\) so their dot product is defined, but \(V\)'s dimension \(d_v\) is independent — in the original Transformer both happen to be 64, but conceptually they don't have to match. Attention then proceeds exactly as scaled dot-product attention: \[\text{Attention}(Q,K,V) = \text{softmax}\!\left(\frac{QK^T}{\sqrt{d_k}}\right)V.\]

Concretely, take two tokens with 4-dim embeddings \(x_1 = [1,0,1,0]\), \(x_2 = [0,1,0,1]\), and \(d_k = d_v = 2\). Suppose (for illustration) \(W_Q = W_K\), each mapping rows \([1,0],[0,1],[1,0],[0,1]\), and \(W_V\) maps rows \([1,1],[1,-1],[0,2],[2,0]\). Projecting: \(q_1 = x_1 W_Q = [2,0]\), \(q_2 = x_2 W_Q = [0,2]\); \(k_1 = [2,0]\), \(k_2 = [0,2]\) (same matrix here); \(v_1 = x_1 W_V = [1,3]\), \(v_2 = x_2 W_V = [3,-1]\). To get token 1's output, score its query against both keys: \(q_1\!\cdot\!k_1 = 4\), \(q_1\!\cdot\!k_2 = 0\). Scale by \(\sqrt{d_k}=\sqrt{2}\approx 1.4142\): scores \(\approx 2.8284, 0\). Softmax: \(e^{2.8284}\approx 16.918\), \(e^{0}=1\), sum \(\approx 17.918\), giving weights \(\approx 0.9442\) and \(0.0558\). The output is the weighted blend of the values (not the keys!): \(0.9442\cdot[1,3] + 0.0558\cdot[3,-1] \approx [1.112, 2.777]\) — token 1 attends almost entirely to itself because its query happened to align tightly with its own key, but the number it actually mixes in is \(v_1\), a third, independently-learned vector, not \(k_1\) or \(x_1\) itself.

In practice \(W_Q, W_K, W_V\) are three nn.Linear layers (no bias, or with bias) trained jointly by backprop — there's no separate "matching loss," the whole network just learns projections that make useful tokens score high against each other. Multi-head attention repeats this whole recipe \(h\) times in parallel with smaller \(d_k = d_{model}/h\) per head, each head getting its own \(W_Q^{(i)}, W_K^{(i)}, W_V^{(i)}\), then concatenates the \(h\) outputs and applies one more linear layer \(W_O\) to mix them back to \(d_{model}\).

See it — One token's embedding, split into three learned projections, then recombined via attention

flowchart TD
    X["Token embedding x"] --> WQ["Project with W_Q"]
    X --> WK["Project with W_K"]
    X --> WV["Project with W_V"]
    WQ --> Q["Query q"]
    WK --> K["Key k"]
    WV --> V["Value v"]
    Q --> DOT["Dot q with<br/>every key"]
    K --> DOT
    DOT --> SC["Scale by<br/>sqrt(d_k)"]
    SC --> SM["Softmax<br/>over scores"]
    SM --> BLEND["Weighted blend<br/>of values"]
    V --> BLEND
    BLEND --> OUT["Output vector"]

The one gotcha

The misconception: "self-attention" does not mean \(Q\), \(K\), and \(V\) are the same vector — it means they all come from the same sequence of embeddings \(X\), each pushed through its own learned matrix (\(W_Q \ne W_K \ne W_V\) in general). If you set \(W_Q = W_K = W_V\), every token would compare itself to others using the exact same representation it hands back as content, collapsing the "am I relevant" question and the "what do I contribute" question into one — the model loses the ability to separate matching from content. Also don't confuse it with cross-attention: in encoder-decoder attention (e.g. the original Transformer's decoder, or T5), \(Q\) comes from the decoder's sequence while \(K\)/\(V\) come from the encoder's — that's still "the same three projections" trick, just applied across two different sequences instead of one.

The one-line takeaway

Same token, three learned projections: query, key, value.

Where you'll meet it

PyTorch: nn.MultiheadAttention's in_proj_weight (one packed matrix that splits into \(W_Q, W_K, W_V\)), or three explicit nn.Linear layers in a custom implementation; Hugging Face transformers: BertSelfAttention / GPT2Attention expose literal .query, .key, .value Linear submodules you can inspect with a debugger or hook; the mechanism underlies every attention layer inside models served through the OpenAI and Anthropic APIs, even though those APIs only expose the input/output text, not the internal \(Q\)/\(K\)/\(V\) tensors.

Named concepts: qkv