If one attention head mixes information from several words, why add another? Start with two heads that see the same tokens but make different mixtures. The difference between those mixtures is what the layer gains the opportunity to use.
Before you begin: You can explain self-attention, a dot product, and a matrix shape.
That is a capacity the model can learn to use. It does not mean head one is always a grammar expert and head two always understands meaning. Heads can overlap, change their behavior across inputs, or contribute little.
Follow one token through a head
Start with a table of token representations, called X. Three learned linear maps turn that table into queries, keys, and values. A query is used to score possible connections. A key participates in that score. A value supplies the information that will be mixed.
scores = Q @ K.T / sqrt(key_width)
weights = softmax(scores + mask)
output = weights @ V
Softmax is applied across the available key positions for each query. It makes their weights nonnegative and sum to one. Dividing by the square root of the key width helps control score magnitudes under the usual variance assumptions; it is not a universal guarantee against unstable training.
Suppose one query gives two scores, 2 and 0. Their softmax weights are approximately 0.881 and 0.119. If the values are 10 and 20, the output is about 11.19. The head blends values; it does not copy the most highly scored token alone.
# Runnable: Python 3, standard library.
import math
scores = [2.0, 0.0]
values = [10.0, 20.0]
peak = max(scores)
exp_scores = [math.exp(s - peak) for s in scores]
weights = [s / sum(exp_scores) for s in exp_scores]
output = sum(w * v for w, v in zip(weights, values))
assert math.isclose(sum(weights), 1.0)
assert math.isclose(output, 11.192029220221174)
print(round(output, 2))
Subtracting the largest score keeps exponentials manageable without changing softmax. This example isolates the weighted average; it does not train queries or keys.
Can two mixtures say something one average loses?
Imagine two value positions, 10 and 20. One head assigns weights 0.9 and 0.1, producing 11. Another assigns 0.1 and 0.9, producing 19. Keeping [11, 19] gives the output projection two signals to combine. Averaging them immediately to 15 erases which head preferred which position.
Predict what happens if both heads learn the first pattern. Does a second output automatically add a new relationship?
Inspect the redundant head
Both outputs are 11 in this toy setup, so the second supplies no new numerical distinction. Real heads have learned value projections too, but simply counting heads still cannot establish independent information. You would need an experiment, such as removing a head and measuring the effect on held-out tasks.
The numbers above are a thought experiment with fixed values, not measurements from a trained model. Now use the dimension calculation below to see what you pay for keeping several mixtures separate.
Keep the dimensions honest
For a conventional layer with model width 512 and eight heads, each head commonly uses width 64. Each produces a 64-number representation per token. Concatenating eight outputs gives 512 numbers, and an output projection mixes them.
This equal-width arrangement is a design convention. The algebra does not require every possible attention design to satisfy that exact arrangement. Increasing the head count while holding total width fixed narrows each head. It does not automatically increase quality or total projection parameters.
Grouped-query attention makes another tradeoff: several query heads share key and value heads. That can reduce the key/value cache during generation. It changes how information is shared, so it belongs in the model configuration and evaluation, not in an assumption about every Llama model.
Prevent a quiet training bug
In causal attention, a position must not read later target tokens. Add a mask before softmax so forbidden positions receive zero probability. If you mask every key in a row with negative infinity, ordinary softmax is undefined. Padding and causal masks therefore need joint testing.
Attention weights also are not a complete explanation of an answer. Values, output projections, residual paths, and later layers all affect the result.
Try it and check your understanding
Change the second score in the example from 0 to 2. Predict the output before running it. Then describe a test that would catch a missing causal mask.
Compare your reasoning
Equal scores give equal weights and an output of 15. A causal test changes only tokens after a chosen position and checks that the output at that position stays unchanged, with dropout disabled. That checks a behavioral property instead of merely checking tensor sizes.
You can now trace how attention changes a representation. Next, follow the path that lets a deep model keep that representation stable across many layers.
Sources
The attention definitions and multi-head construction come from Attention Is All You Need. Shared key/value heads are developed in GQA. The numerical example here is original and deliberately small.
Continue: Layer Normalization and Residual Connections.