Books › AI Foundations: from zero

The transformer, step by step

By Kapil Kumar Gupta

The problem attention solves

Read this: "The animal was tired, so it slept." To know what *it* means, you must connect it back to *animal*. Any design for understanding language has to be good at connecting words that may be far apart.

The older designs read a sentence one word at a time, from left to right, passing along a short memory note that summarises everything so far.

That design had two problems.

• Fading. Everything the model knew about the start of the sentence had to be squeezed through one small note, over and over. By the end of a long paragraph, early words had faded.
• Waiting. Each word had to wait for the word before it to be processed. That made it hard to use many chips at once, so training on huge libraries of text was slow.

Attention removes both problems. Instead of one note passed along a chain, each word reaches out to all the others and decides how much each one counts. Far-away words are one step away, and nothing has to wait, so the whole sentence can be handled at once.

One attention step with real numbers

Each word is first turned into three lists of numbers, using weights the model has learned:

• a question (query): what am I looking for?
• a label (key): what do I offer, so others can find me?
• a content (value): the information I carry.

Think of a library. Your question is what you are looking for. Every book has a label on its spine. You compare your question with every label, and then read mostly the books whose labels fit best.

We will use three words, with two numbers in each list. The question we follow belongs to the word it.

The question of *it* is (1, 0).

Step 1: compare the question with each label

Multiply matching numbers and add them up (the dot product).

• The: (1 × 0) + (0 × 1) = 0
• animal: (1 × 1) + (0 × 0) = 1
• it: (1 × 0.5) + (0 × 0.5) = 0.5

Then divide each by the square root of the list length. The lists have 2 numbers, and the square root of 2 is about 1.41. This stops the scores from growing too big when the lists are long. The scores become 0, 0.71 and 0.35.

Step 2: turn the scores into shares

We want shares that add up to 100%. The tool for this is softmax. For each score, calculate e raised to that score (e is about 2.718), and divide by the total.

• e⁰ = 1.000
• e⁰·⁷⁰⁷ ≈ 2.028
• e⁰·³⁵⁴ ≈ 1.425
• Total ≈ 4.453

(These use the scores with a few more digits, 0.707 and 0.354, before rounding to 0.71 and 0.35 for display.)

So the shares are 1.000 ÷ 4.453 ≈ 22.5% for *The*, 2.028 ÷ 4.453 ≈ 45.6% for *animal*, and 1.425 ÷ 4.453 ≈ 32.0% for *it*. The word *animal* gets the biggest share, because its label matched the question best.

Step 3: mix the content

Now collect the contents, weighted by those shares.

• First number: 0.225 × 1 + 0.456 × 0 + 0.320 × 1 = 0.545
• Second number: 0.225 × 0 + 0.456 × 2 + 0.320 × 1 = 1.232

The new list for *it* is about (0.54, 1.23). Notice the second number is large. That comes mostly from *animal*, whose content was (0, 2). The word *it* now carries a good deal of the meaning of *animal*. That is what "understanding that *it* means the animal" looks like in numbers.

Every word does this at the same time, each asking its own question.

Many heads at once

One attention step can only track one kind of link. So a transformer runs several side by side, called heads. Each has its own learned weights, so each can learn to look for something different: one for grammar, one for who did what, one for topic. Their results are joined together.

Where are the words?

Notice something odd: attention never mentions order. If we shuffled the words, the shares would stay the same, so "dog bites man" and "man bites dog" would look identical. To fix this, a code for each word's position is added to its list before attention starts. Different models use different kinds of position codes, but the idea is always the same: tell the model where each word sits.

Putting a block together

One transformer block has two main parts, plus two helpers.

1. Multi-head attention: words gather information from each other. 2. Feed-forward layer: a small network (Chapter 4) is applied to each word separately. It does the "thinking" about what attention gathered. 3. Add and normalise (twice): the input of each part is added back to its result (a shortcut, also called a residual connection), and the numbers are rescaled into a sensible range (normalisation). The shortcut means information is never lost, and it helps blame flow backwards during learning.

Stack many of these blocks. Early blocks tend to pick up simple patterns, such as grammar. Later blocks work with deeper meaning. Models range from a few blocks to many dozens. At the very top, a final layer produces a score for every possible next word, and softmax turns those scores into chances (Chapter 7).

Not peeking: the mask

When a model writes, it must only use the words so far, not the words that come next. Otherwise it could cheat during training by looking at the answer.

So for writing models, attention is masked. In our example, if the word is *animal* (the second word), its score for *it* (the third word) is replaced by a huge negative number. After softmax, that word gets a share of 0%. Each word can only look backwards.

Three flavours

Writer-only models can be trained on any text with a simple game, "predict the next word", and they scale well. That is why the chatbots you know are mostly of this kind.

The cost of looking at everyone

There is a price. If there are n words, every word compares itself with every other word: about n × n comparisons. Double the text and the work roughly quadruples. This is why models can read only a limited amount at once (the context window), and why long inputs are expensive. Engineers have invented many ways to reduce the cost, such as looking only at nearby words, and reorganising the calculation so less data moves in and out of memory.

Back to Asha

Asha's customer writes: "I ordered the lemon cake and the rolls. Can you make it eggless?" The word *it* has to point back to *lemon cake* or *rolls*. In a model, *it* asks its question, compares it with the label of every earlier word, and gives the biggest share to the best match. If two matches are close, the safest helper asks: "Do you mean the cake or the rolls?"

Words explained

• Query, key, value: the question a word asks, the label it offers, and the content it carries.
• Softmax: a recipe that turns a list of scores into shares that add up to 100%.
• Head: one attention process. A transformer runs several side by side.
• Position code: numbers added to a word to tell the model where it sits.
• Shortcut (residual connection): adding a part's input back to its result so nothing is lost.
• Normalisation: rescaling numbers into a sensible range.
• Mask: blocking a word from looking at words that come after it.

The limits of this chapter

• Real models use many more numbers per list, and the query, key and value lists are made by multiplying each word's embedding by learned weights. We skipped that step to keep the example short.
• Attention shares show where a model looked. They do not fully explain why it answered as it did.
• There are many variations of every part. This is the classic design.

For engineers

• Scaled dot-product attention: Attention(Q, K, V) = softmax(Q Kᵀ / √d_k) V. For a sequence of n tokens and model width d_model, with h heads, each head typically has d_k = d_model / h. Head outputs are concatenated and projected back to d_model.
• A block (in the common "pre-norm" arrangement) is: x = x + Attention(Norm(x)); x = x + FeedForward(Norm(x)). The feed-forward layer is usually two linear layers with a non-linearity (for example a gated variant of ReLU), expanding the width by a factor of about four in the classic design.
• Position information: fixed sinusoidal codes, learned embeddings, or rotary position embeddings applied inside attention. Grouped-query attention shares key and value heads across several query heads to cut KV-cache memory.
• Decoder-only models apply a causal mask: positions j > i get −∞ before the softmax. Encoder-only models (the BERT family) are bidirectional and are widely used for classification and embeddings. Encoder-decoder models add cross-attention from decoder to encoder outputs.
• Self-attention costs O(n²) time and memory in sequence length. Efficient kernels (such as FlashAttention) keep the exact result but reduce memory traffic. Sparse or windowed attention changes which pairs are compared.

Check yourself

• Why did attention replace the one-word-at-a-time designs? Answer: It lets every word look directly at every other word, so nothing fades, and the whole text can be processed at the same time, which makes training much faster.
• In the worked example, why did "animal" get the biggest share? Answer: Its label (1, 0) matched the question of "it" (1, 0) best, giving the highest score, 1.
• What does the mask do in a writing model? Answer: It stops each word from looking at words that come after it, so the model cannot peek at the answer.