Multi-Head Self-Attention
One linear projection, three views, H parallel heads, one mask. The attention block as the model actually uses it. Implement a batched Query/Key/Value projection as a single linear layer split into H heads. Compute scaled dot-product attention with the correct normalization and dtype handling. Apply a causal mask that prevents a position from attending to future positions. Inspect per-head attention weights for a fixed input and reason about what each head looks at. Train a small attention block on a toy task and watch the loss fall as the heads specialize. Attention is the function that lets a token's representation pull information from other tokens in the same sequence. Self-attention means queries, keys, and values are all derived from the same input. Multi-head means the projection is split into H parallel attention problems whose outputs are concatenated and projected back. The efficient implementation pattern is one linear layer that projects from D to 3 D and gets sliced into three views, then reshaped into H heads of size D // H each. The matmul, softmax, and weighted sum happen as batched tensor operations so the heads run in parallel on the accelerator. This lesson builds that block. It also adds the causal mask so the same code works as the attention layer in a decoder-only language model. The next lesson…
Multi-Head Self-Attention: One linear projection, three views, H parallel heads, one mask. The attention block as the model actually uses it.
This free lesson is part of the AI Engineering from Scratch curriculum. Read the full explanation, run the lesson code, and verify the result in the interactive reader or from the repository source.
Browse the complete course catalog or open this lesson on GitHub.