Day 94: Self-attention mechanics: Q/K/V by hand
The heart of the transformer, worked by hand
Self-attention is three learned projections of each token — Query, Key, Value — and a simple recipe. Each token's query is compared (dot product) with every token's key to get attention scores; those scores are softmaxed into weights; then each token's output is the weighted sum of all values. A Stage 3 exit criterion is whiteboarding Q/K/V unprompted, so work this until it's second nature.
The recipe, step by step
- Project each token embedding into a Query, a Key, and a Value vector (learned linear layers).
- Compute scores:
Q · Kᵀ— how much each token's query aligns with every token's key (dot product, Stage 0 Day 7). - Scale by
√d(keeps scores from growing with dimension) and apply softmax so weights per token sum to 1. - Output = weighted sum of the Value vectors, using those attention weights.
import torch
import torch.nn.functional as F
def self_attention(x, Wq, Wk, Wv):
Q, K, V = x @ Wq, x @ Wk, x @ Wv # project to query/key/value
d = Q.shape[-1]
scores = (Q @ K.transpose(-2, -1)) / d ** 0.5 # scaled dot-product
weights = F.softmax(scores, dim=-1) # attention weights sum to 1
return weights @ V # weighted sum of valuesThe library analogy
Your query is what you're searching for. Each book's key is its topic label. You match your query against all keys to see which books are relevant (scores), then read a blend of the most-relevant books' content (values), weighted by relevance. Attention is soft, differentiable library lookup — and once this clicks, the whole transformer opens up.
Key terms
- Query / Key / Value
- Three learned projections of each token: query (what it seeks), key (what it offers for matching), value (what it contributes).
- Scaled dot-product attention
- Scores = softmax(Q·Kᵀ / √d), used to take a weighted sum of values.
- Softmax
- Normalizes scores into weights that sum to 1, turning raw alignments into a probability-like distribution.
In self-attention, how is a token's output vector computed?
