Auto-attention

Consider "The animal did not cross the street because it was too tired." What does "it" refer to? The animal or the street? To resolve this, the model must connect a token to another token many positions earlier. Recurrent networks passed a hidden state hand to hand and slowly forgot; self-attention lets every position consult every other position directly, in one step.

Each token produces three vectors by multiplying its embedding with three learned matrices: a Query ("what am I looking for?"), a Key ("what do I offer?"), and a Value ("what will I actually hand over if chosen?").

• Q·Kᵀ produces an n×n matrix: how much every token cares about every other token.
• softmax converts each row into weights that sum to 1 — a distribution of attention.
• Multiplying by V blends the values using those weights: the new, context-aware representation.
• The result is a per-token embedding that now contains information from wherever it looked.

Dot products between high-dimensional vectors grow with dimension. With d = 512, raw scores can reach into the hundreds. Feed large numbers into softmax and it becomes a hard argmax: one weight becomes 1.0, the rest 0.0, and the gradient through that softmax is effectively zero. Dividing by √d keeps the scores in a friendly range where softmax stays soft and gradients flow.

The second detail is masking. A decoder-only LLM must never see the future, otherwise predicting the next token is cheating. So before the softmax, every position above the diagonal is set to −∞, which becomes exactly 0 after softmax. This is the causal mask, and it is the only thing that turns a bidirectional reader into an autoregressive generator.

Instead of one attention computation, the model runs several heads in parallel, each with its own Q, K and V projection matrices of reduced size. Their outputs are concatenated and mixed by an output projection. The total parameter count stays roughly constant — you are not adding cost, you are splitting the perspective.

• Different heads empirically specialise: some track adjacent tokens, some track verbs to their objects, some track rare-token identity.
• Head count and head dimension trade off: 12 heads × 64 dims = 768 model width.
• Multi-Query / Grouped-Query Attention shares one K and V across heads, shrinking the KV-cache dramatically — the trick that makes long-context inference affordable (Stage 11).

1. 1Input embeddings are combined with positional information, since attention itself is order-agnostic.
2. 2Multi-head self-attention lets every token gather context from everywhere it may look.
3. 3Residual connection adds the original vector back, so gradients survive deep stacks.
4. 4Layer normalisation re-centres and re-scales activations for stable training.
5. 5Feed-forward network (two linear layers with a non-linearity, typically 4× wider) processes each token position independently — this is where most parameters and most factual recall live.
6. 6Stack this block 12 to 100+ times and you have a transformer.

def block(x, mask):
    # 1. attention with a residual highway
    x = x + attention(norm(x), mask=mask)
    # 2. position-wise feed-forward, again with a highway
    x = x + mlp(norm(x))
    return x

# pre-norm ordering (above) trains more stably than post-norm
# at depth; most modern LLMs are pre-norm.