Transformers & attention
The architecture behind modern language models: stacked layers in which every token can look back at earlier tokens and pull in what's relevant.
Read first: Tokens & tokenization.
D1 · StoryFor young children and total beginners
Read this: “The cat sat on the mat because it was tired.” You know “it” means the cat, not the mat.
Attention is how the AI works that out: for each word it looks back at the earlier words and decides which ones matter most.
Like: Highlighting the words that matter.
D2 · PictureFor curious beginners
A transformer is a stack of identical layers. Each layer has two parts: attention, which mixes in information from other positions, and a feed-forward network, which processes each position on its own.
In attention, each position asks a question (a query), compares it with a label on every earlier position (its key), and takes a weighted blend of what those positions carry (their values). Many attention heads do this in parallel, looking for different patterns.
In models that generate text, a causal mask stops any position from looking ahead at tokens that haven't been written yet.
Like: Highlighting the words that matter.
Where it breaks: You highlight a few words once. Attention gives every earlier token some weight, dozens of heads do it at once in every layer, and the result is spread across many layers rather than one clear highlight.
D3 · PlayFor motivated learners
Each word builds a query (what it is looking for) and a key (what it offers). The query is compared with the keys of earlier words using a dot product, and softmax turns those scores into attention weights that add up to 100%.
A transformer runs many heads side by side. Each head has its own queries and keys, so each looks for something different. Choose a word, then switch heads and watch where its attention goes.
Try the interactive version in the map.
Things to try
- In the “Who is it?” head, choose “it”. Which earlier word gets the most weight?
- Switch to “Any noun” with “it” still chosen. Why is the weight now split between “cat” and “mat”?
- Switch to “Previous word” and step through the sentence with the arrow keys. What does every word look at?
D4 · MechanismFor practitioners
Stack every token's query into a matrix Q, its key into K and its value into V. Attention is softmax(QKᵀ ÷ √d) · V, where d is the length of each query and key.
Row i of QKᵀ holds the dot product of token i's query with every key. Dividing by √d stops the scores growing just because the vectors are longer. The causal mask sets scores for later tokens to −∞ so softmax gives them zero weight. Each output row is a weighted blend of the value vectors.
Edit the vectors below; every table is computed live from the formula.
Try the interactive version in the map.
Things to try
- Look at the row for “it”. Why does its output come out close to the value of “cat”?
- Change the first number of the key for “mat” from 0.4 to 2.5. What happens to the weight “it” gives to “cat”?
- Turn the causal mask off. Which rows change, and which row stays the same?
def attention(Q, K, V, causal=True):
d = Q.shape[-1] # length of each query and key
scores = Q @ K.T / sqrt(d) # one score for every (query, key) pair
if causal:
scores[j > i] = -inf # hide tokens that come later
weights = softmax(scores, axis=-1) # each row adds up to 1
return weights @ V # blend the values
# multi-head: run several of these with different learned Q, K, V
# projections, then join the outputs and mix them with one more matrixTrade-offs
- Every token is compared with every other, so the score table grows with the square of the context length.
- The causal mask lets a model train on a whole sentence at once while still learning to predict each word without peeking.
- More heads let a layer look for more patterns at once, but each head works with shorter vectors.
- Attention weights show where a head looked, not why the model gave its answer; reading them as explanations can mislead.
D5 · FrontierFor experts
Open problems: attention's cost grows with the square of the context length, which drives work on efficient kernels, sparse and linear attention, and hybrids with state-space layers; and interpretability research into what individual heads and layers actually compute.