06 Module 2 · Building the Model
Chapter 6 of 17 Run it in your browser

Chapter 6: Self-Attention (Single Head)

Figure 6.1: Attention builds context by looking across the sequence.

Chapter 5 gave our tokens meaning through embeddings, but reading one word at a time is not enough. To understand "it" in "the trophy didn't fit in the suitcase because it was too big", you have to connect "it" back to "trophy". In this chapter you will:

• Learn how tokens "look" at each other to build context.
• Use the Query, Key, Value (QKV) system to find relevant information.
• See how attention is a weighted average of numbers.
• Apply a causal mask to hide the future.

Words to Know

• Self-Attention: The step where each token looks at the tokens before it and decides how much each one matters to its own meaning.
• Query (Q): A list of numbers that says what a token is looking for.
• Key (K): A list of numbers that says what a token has to offer.
• Value (V): A list of numbers holding the information a token passes on when it is picked.
• Causal Mask: A filter that hides the tokens that come later, so a token sees only itself and the past.

Theory

The Problem: Context Matters

If a model only considers one token at a time, it has no memory, so it cannot connect subjects to verbs or adjectives to nouns. Every token needs to "look at" the other tokens and decide which ones matter most to its own meaning.

The QKV System

Self-attention solves this using three vectors for every token:

• Query (Q): What am I looking for?
• Key (K): What do I offer?
• Value (V): What information do I actually contain?

Imagine you are at a networking event. You are looking for a marketing expert (your Query), and someone is wearing a badge that says "Marketing Director" (their Key). When your Query matches their Key, you start a conversation and absorb their advice (their Value).

In our model, every token creates its own Q, K, and V vectors by multiplying its embedding by learned weights. A token acts as a seeker (Query) and a source (Key and Value) at the same time.

Attention is a Weighted Average

Figure 6.2: Token 5 blends information from previous tokens.

The key idea in this chapter is that attention is a weighted average (Figure 6.2). When token 5 calculates its final value, it does not pick the single best token to look at. It mixes them all.

First, we multiply all Queries by all Keys (q @ k.transpose) to get an attention score for every pair. We scale these down by dividing by sqrt(head_size) (the size of each attention head, which is 32 here) so the numbers do not get too large. Then we apply softmax to turn these scores into percentages (weights) that sum to 1.0 (or 100%).

For example, when token 5 looks at the sequence, it might assign weights like this (Figure 6.3). These values are rounded for display, so they add up to about 100% (100.1%):

• Token 0: 19.9%
• Token 1: 20.4%
• Token 2: 12.2%
• Token 3: 12.1%
• Token 4: 21.4%
• Token 5: 14.1%

Figure 6.3: A plotted bar chart of the attention weights.

Token 5's final output is 19.9% of the value of token 0, plus 20.4% of the value of token 1, and so on for all six tokens. The "Numbers Machine" builds context by blending the numbers of the most relevant past tokens. On our map (Figure 6.1), gathering context in this way completes the "Attention" stage.

The Causal Mask

Figure 6.4: A lower triangular mask ensures tokens only see the past.

The catch is that the model processes all tokens at once. If we let token 5 look at token 6, it would be cheating, because during training it would copy the answer from the future instead of learning to predict it.

We apply a causal mask (a triangle of negative infinities, shown in Figure 6.4) to the scores before the softmax step. Softmax turns -infinity into exactly 0.0, so token 5 pays 0% attention to tokens 6, 7, 8, and 9 and can only see itself and the past.

In Business

When parsing a customer support chat, a simple keyword search treats every word independently. A context-aware system uses self-attention to connect the word "refund" in message 4 back to the "broken screen" mentioned in message 1. That way it understands the full conversation history before it generates a reply in the company's house style.

Code

Figure 6.5: The Q, K, V vectors are created, scored, masked, and then multiplied.

We use src/ch05_self_attention.py to build the SingleHeadAttention class (Figure 6.5).

        q = self.query(x)
        k = self.key(x)
        v = self.value(x)

        # Compute attention scores and scale them to keep variance stable
        scale = HEAD_SIZE ** -0.5
        scores = q @ k.transpose(-2, -1) * scale

        # Mask future tokens by setting their scores to negative infinity
        scores = scores.masked_fill(self.tril[:T, :T] == 0, float("-inf"))

        # Convert scores to probabilities and apply dropout
        weights = F.softmax(scores, dim=-1)
        weights = self.dropout(weights)

        # Compute final output by taking weighted sum of values
        out = weights @ v

Running the script prints the attention weights:

Terminal

$ python src/ch05_self_attention.py
--- Attention weights (what token 5 attends to) ---
Token 5 attends to tokens 0..5 (future tokens masked):
  token 0: 0.199  #####
  token 1: 0.204  ######
  token 2: 0.122  ###
  token 3: 0.121  ###
  token 4: 0.214  ######
  token 5: 0.141  ####
  token 6: 0.000
  token 7: 0.000
  token 8: 0.000
  token 9: 0.000

Self-attention done! Ready for Chapter 7.

What just happened:

• Lines 1 to 3 created independent Query, Key, and Value vectors.
• Line 7 computed the attention scores between all tokens.
• Line 10 masked the future tokens (notice tokens 6 to 9 have exactly 0.000 weight).
• Line 17 created the output (for example, token 5 created its output by taking a weighted average of tokens 0 through 5).

Shape Check: Table 6.1 lists the shapes of the attention variables.

Table 6.1: Tensor shapes during the self-attention calculation.

Variable	Shape	Meaning
q, k, v	[B, T, head_size]	Queries, Keys, and Values. head_size is 32.
scores	[B, T, T]	Attention scores between every pair of tokens.
weights	[B, T, T]	The softmax probabilities (summing to 1 per row).
out	[B, T, head_size]	The final context-aware output vectors.

Try It

Try It

Remove the scale division in the demo at the bottom of src/ch05_self_attention.py by changing scale = HEAD_SIZE ** -0.5 to scale = 1.0. Run it again. Notice how the weights become much more extreme: token 4 rises from 0.214 to 0.383, and tokens 2 and 3 fall to about 0.015. The scaling is what keeps the model flexible and learning smoothly.

Terminal

$ python src/ch05_self_attention.py
...
  token 0: 0.257  #######
  token 1: 0.293  ########
  token 2: 0.016
  token 3: 0.015
  token 4: 0.383  ###########
  token 5: 0.036  #
...

Watch Out

Be careful with the transpose step, because we only want to swap the last two dimensions (Time and head_size) to compute the dot product properly. If you use .T, it might flip the Batch dimension too, crashing your shape calculations, so always use .transpose(-2, -1).

Terminal

$ python src/ch05_self_attention.py
...
RuntimeError: The size of tensor a (2) must match the size of tensor b (32) at non-singleton dimension 0

Key Takeaways

• Self-attention builds context by letting tokens "look" at other tokens.
• The QKV system works like a search engine, where Queries match with Keys to retrieve Values.
• Attention is a weighted average of the Values.
• A causal mask sets future scores to negative infinity, preventing the model from cheating.

Check Your Understanding

1. What is the difference between a Query and a Key?
2. Why do we divide the attention scores by the square root of the head size?
3. What happens to a score of negative infinity when passed through softmax?
4. Why is the causal mask necessary for a language model?

Assignment

Assignment 6: New numbers, same mask

About 15 minutes. You change one number and write no new code. Work in this chapter's Colab notebook or on your own laptop.

1. Run it. Run src/ch05_self_attention.py as it is. Copy the ten weights it prints for token 5.
2. Change one number. Near the bottom of the file, find torch.manual_seed(42). The seed fixes the random starting numbers, so every run prints the same weights. Change 42 to any whole number you like and run the script again.
3. Compare. Which of the ten weights changed? Which ones did not? What do the weights for tokens 0 to 5 add up to now?
4. Explain it. In two or three sentences a manager could follow, say why four of the rows print the same value under any seed.

Hand in: both sets of ten weights, your answers to step 3, and your sentences.

Further Reading

The first attention. Translation models of the day read the whole source sentence, squeezed it into a single fixed vector, and wrote the translation from that. Long sentences did not survive the squeeze. The fix: let the model look back over every input word and decide, at each output word, which ones matter right now. That weighted look-back is attention. The Transformer three years later kept this idea and threw out everything around it.

The architecture this book builds. Reading a sequence one step at a time is slow, because step 500 cannot start until step 499 has finished, and distant words stay hard to connect. This paper removed the step-by-step reading entirely and kept only attention, plus a note of each token's position. Every token can then be processed at once, which is what made training on very large amounts of text practical. The model you build in Chapters 6 to 10 is this design, made small.

Bahdanau, D., Cho, K., & Bengio, Y. (2014). Neural machine translation by jointly learning to align and translate (arXiv:1409.0473). arXiv. https://doi.org/10.48550/arXiv.1409.0473

Vaswani, A., Shazeer, N., Parmar, N., Uszkoreit, J., Jones, L., Gomez, A. N., Kaiser, L., & Polosukhin, I. (2017). Attention is all you need (arXiv:1706.03762). arXiv. https://doi.org/10.48550/arXiv.1706.03762

Next up: Chapter 7: Multi-Head Attention