EasyCompilers

Vocabulary for this course

The terms the chapters assume, in plain words. You can open this from any chapter.

attention
A weighted blend of value rows, with weights decided by how well each query matches each key.
attention map
The table of scores showing how much each position attended to every other position.
causal mask
A rule that blocks each position from seeing later positions, so predictions can only use the past.
dimension
One axis of a list or grid of numbers; its length is how many values fit along that axis.
distribution
The pattern of how often each possible value appears, or how likely each one is.
dot product
Multiply matching numbers in two lists and add the results; a large total means the lists point the same way.
embedding
A learned list of numbers that stands for a token, with similar tokens placed near each other.
GPU
A processor built for many small calculations at once, which is why matrix work runs well on it.
gradient
A list of slopes, one per number, saying how much each nudge changes the result.
hard pointer
A choice that sends all the weight to one position, instead of spreading it across several.
head
One parallel slice of attention, with its own projections; a layer runs many so it can track different relations at once.
key
What a position advertises it can offer, compared against every query to decide who receives attention.
key axis
The direction in the key vectors whose width decides how far the scores are divided down.
layer
One stage of a model that reads the running vector, transforms it and passes it on.
layer norm
A per-position rescale that keeps the residual stream at a steady size before the next block.
logits
The raw scores a model produces before softmax turns them into probabilities.
matrix
A grid of numbers arranged in rows and columns, used to transform a list into another list.
negative infinity
A value smaller than every real number, used to mark a score that must never be chosen.
one-hot
A vector with a single one and zeros elsewhere, meaning all the weight sits on one choice.
output projection
The learned matrix applied after attention, which turns the mixed values back into the model's width.
parameter
A number inside a model that training adjusts, rather than an input it reads.
positional encoding
Extra numbers added to token embeddings so the model can tell which position each one came from.
projection
A learned matrix that turns a list of numbers into a different list, often shorter.
query
What a position is looking for, compared against every key to decide where to take information from.
residual stream
The running vector each block reads from and adds its output back into.
retrieval layer
A layer whose job is to pull specific stored information into the current position.
saturating
A score so large that the softmax gives almost all the weight to one position.
scale
The step that divides the scores by the square root of the key width, so they stay in a workable range.
score matrix
The grid holding every query's score against every key, one row per position.
soft alignment
Matching by degree rather than by a single choice, so several positions can contribute a little.
softmax
A rule that turns a list of scores into positive weights that sum to one.
token
A chunk of text — a word, part of a word, or a character — that the model treats as one unit.
value
The content a position passes on when attention selects it.
variance
The average squared distance of values from their average, so a spread-out set scores higher.
vector
An ordered list of numbers, where each position is one measured property.
weighted sum
Adding values together with different shares, where the shares come from the attention scores.
weights
The numbers a model multiplies inputs by before adding the results together.

Transformers from the ground up

Multi-head attention

One attention pattern per layer is not enough. Several heads, each with its own projections, look for different relations and are concatenated back into one vector.

Part 3 of 10

Advanced17 min readBack to Courses

Why one head is a bottleneck

A single softmax produces one distribution per token. A token may need to track several things at once. It might track the previous word, the subject of the sentence and a pronoun’s referent. Those relations compete for the same weights. The result is an average that serves none of them well.

Multi-head attention gives the layer several independent ways to look. Each head has its own query, key and value projections. Each can learn a different similarity function. The heads run in parallel and their outputs are concatenated, joined end to end. A later projection then mixes them. Until that point the layer keeps several relations separate.

The number of heads is a hyperparameter, a setting you choose. It is not a consequence of the data. More heads mean more independent patterns and more score matrices to store. Fewer heads mean each one has to average more relations. The total width is usually held fixed. The choice is therefore about how to divide the same budget. A head that is too narrow cannot represent a rich similarity. A head that is too wide leaves little room for the others.

Split, attend, concatenate

The model dimension is divided by the number of heads. With a width of 512 and 8 heads, each head works in 64 dimensions. The total compute stays close to that of one full-width head. Each head attends in a smaller subspace. What changes is the number of independent patterns the layer can represent.

After attention, the head outputs are concatenated back to the full width. They then pass through an output projection. That projection is what lets information from different heads interact. Without it, the heads would stay separate channels until the next layer. The block would be a sum of independent attention results.

The parameter count is easy to estimate. Each head has three projections of model width by head width. The output projection is model width by model width. With a fixed total width, the attention parameters are roughly the same either way. Few wide heads and many narrow heads cost about the same. The difference is in what the softmax can separate.

The built-in layer

import torch
from torch import nn

layer = nn.MultiheadAttention(embed_dim=512, num_heads=8, batch_first=True)
x = torch.randn(4, 32, 512)  # batch, tokens, d_model

out, weights = layer(x, x, x, need_weights=True)
print(out.shape)      # (4, 32, 512)
print(weights.shape)  # (4, 32, 32)

Two heads from one vector

Run the split, then give each head its own small projection matrix instead of slicing the input, and compare the two outputs.

Output
Click "Run" to execute your code

Heads are not modules

A head is often described as the one that tracks syntax or the one that resolves pronouns. Some heads do behave that way in some models. But heads are not designed for a job, and they are not guaranteed to keep one. Ablating a single head, which means removing it, often changes very little. Other heads carry redundant information.

The useful exceptions are worth knowing. Induction heads copy a token that followed the same prefix earlier. They appear in many trained transformers and seem to matter for in-context learning, the ability to use examples in the prompt. That is a discovered regularity, not a slot in the architecture. Treat head labels as hypotheses to test, not as documentation.

Head pruning is the practical version of this question. Remove a head and check the loss. If the loss barely moves, the head was redundant for that test set. If the loss rises, the head was carrying something. The result depends on the task. That is another reason a head’s name should not travel from one model to another.

When you change the model width, decide whether to keep the head dimension fixed. Doubling the width and doubling the head count keeps each head at the same size; doubling only the width makes every head wider and changes what the softmax can represent.

The cost of many heads

The score matrix is n × n per head. Attention memory therefore grows with the number of heads, even when the total width is fixed. For a long sequence that memory is often the limit, not the parameter count. Flash attention reduces the materialised matrix. It does not change the asymptotic cost.

Grouped-query attention is the common compromise. Several query heads share one key and value head. That shrinks the cache used during generation while keeping the query side expressive. Multi-query attention is the extreme case: one key and value head for all queries. The choice trades quality against the memory bandwidth of inference.

Two heads, two patterns

Run both heads on the same three tokens and compare their weight rows. Then swap the two query vectors and confirm the patterns swap with them.

Output
Click "Run" to execute your code

Check what you learned

One question at a time. A wrong answer explains itself — take another go.

Part 1 of 5

With a model width of 512 and 8 heads, how wide is each head?