The attention mechanism

The "明" in "明月" (bright moon) and the "明" in "明天" (tomorrow) mean different things. How to understand a character depends on the characters around it.

At each position, the model we'll train has the job of predicting the next character. To write "光" (light) after "床前明月" ("before my bed, the bright moon"), it must know that "明月" appeared earlier. So every position needs a way to gather information from the characters before it.

Attention is that way, and it's the core of the Transformer. This lesson builds it up step by step, starting from the simplest approach.

python attention.py

The simplest approach: take the average

Each character is first represented by a vector (Module 01, Lesson 5 covered embeddings). The simplest way to "gather information from before" is to average the vectors of the character itself and all the characters before it.

This can be done with a single matrix multiplication. Build a lower-triangular matrix where each row has values only in the first few positions, adding up to 1:

x = torch.randn(T, 4)  # 5 个词，每个词一个 4 维向量（先随便取）
weights = torch.ones(T, T).tril()  # 下三角：第 i 行只有前 i+1 个位置是 1
weights = weights / weights.sum(dim=1, keepdim=True)  # 每行除以个数，变成平均
out = weights @ x

== 1. 用一个矩阵乘法，让每个位置得到自己和前面所有位置的平均
tensor([[1.00, 0.00, 0.00, 0.00, 0.00],
        [0.50, 0.50, 0.00, 0.00, 0.00],
        [0.33, 0.33, 0.33, 0.00, 0.00],
        [0.25, 0.25, 0.25, 0.25, 0.00],
        [0.20, 0.20, 0.20, 0.20, 0.20]])
  第 3 个位置（'明'）的结果 = 前三个向量的平均？True

Row i says how much the i-th character takes from each position: the 1st character can only see itself, the 3rd takes an even share of the first three, the 5th takes an even share of all five. weights @ x computes the result for every position in one go.

Remember this matrix multiplication form; attention takes the same form later: a weight matrix times a set of vectors. The only difference is where the weights come from.

The problem with averaging is obvious: every character is treated the same. "光" wants to know whether "月" (moon) appeared earlier; "床" (bed) and "前" (before) matter less to it. The weights should be decided by content.

Queries and keys: let each character decide whom to look at

The approach is to give each character two vectors:

• Query (q): what kind of information I'm looking for.
• Key (k): what kind of information I have.

How much one character "attends" to another is scored by the dot product of its own q with the other's k. The closer the two vectors' directions, the larger the dot product.

First let's see the effect with hand-made vectors. Suppose the vectors have two dimensions: the first means "related to the sky" and the second "related to position". The query for "光" is [1, 0], meaning it's looking for something related to the sky:

k = torch.tensor([[0.1, 0.9],   # 床：一个地点
                  [0.0, 1.0],   # 前：一个方位
                  [0.9, 0.1],   # 明：和天空、光有关
                  [1.0, 0.0],   # 月：天空里的东西
                  [0.8, 0.2]])  # 光
q_guang = torch.tensor([1.0, 0.0])  # "光"想找的是：和天空有关的东西
scores = k @ q_guang

== 2. 点积打分：查询（q）和每个键（k）越像，分数越高
  '光' 对每个词的分数： {'床': 0.1, '前': 0.0, '明': 0.9, '月': 1.0, '光': 0.8}
  softmax 之后的权重：   {'床': 0.12, '前': 0.11, '明': 0.26, '月': 0.29, '光': 0.23}
  分数放大 5 倍再 softmax： {'床': 0.01, '前': 0.0, '明': 0.3, '月': 0.5, '光': 0.18}

The scores pass through softmax to become weights (Module 08, Lesson 5 covered softmax) that add up to 1. "月" has the highest score and gets the largest weight.

Note the last line: scale all the scores up 5 times, and the weights after softmax become much more concentrated, with "月" taking half on its own. The size of the scores determines whether attention is "spread evenly" or "fixed on one". This matters later when we cover scaling.

In a real model, q and k aren't hand-made; they're computed from each character's vector by two matrices, whose parameters are learned in training. The model learns for itself what kind of character should look for what kind of information.

Values: what's actually passed along

Once the scores decide how much to take from each position, what is it that's taken? A third vector: the value (v), meaning "if you attend to me, this is the information I give you".

Why not take the original vector directly? Because "information used for matching" and "information to be passed along" aren't necessarily the same. Separating them makes the model more flexible.

Put the three vectors together, and that's all of attention:

权重 = softmax(q 和每个 k 的点积)
输出 = 用这些权重，把每个位置的 v 加起来

The causal mask: no peeking ahead

During training, the model predicts the next character at every position. If position 3 could see the 4th character, it could simply "copy the answer" and learn nothing. So each position may see only itself and the positions before it.

This is done by setting the scores of later positions to negative infinity before softmax. The exponential of negative infinity is 0, so after softmax the weight is 0:

scores = q @ k.T / math.sqrt(8)
mask = torch.ones(T, T, dtype=torch.bool).tril()
scores = scores.masked_fill(~mask, float("-inf"))

== 3. 因果掩码：把后面位置的分数设成负无穷，softmax 之后权重就是 0
tensor([[1.00, 0.00, 0.00, 0.00, 0.00],
        [0.96, 0.04, 0.00, 0.00, 0.00],
        [0.24, 0.26, 0.50, 0.00, 0.00],
        [0.09, 0.39, 0.43, 0.09, 0.00],
        [0.02, 0.14, 0.02, 0.09, 0.73]])

Compare with the averaging matrix in the first section: the same shape, both lower-triangular, each row adding up to 1. The difference is that each row's weights are no longer even, but decided by q and k. The q and k here are random, so the weights look patternless; after training, they become meaningful.

Because of this mask, this kind of attention is called causal self-attention: self-attention means q, k and v all come from the same text, and causal means it can see only the past, not the future. Models like GPT that "write one character after another" all use it.

Why divide by the square root of the dimension

The code above has a / math.sqrt(8), where 8 is the dimension of q and k. This step is called scaling; let's see what happens without it:

== 4. 为什么要除以 √d：维度越大，点积的数值越大，softmax 会变得非常极端
  d=   16  点积的标准差   3.89（√d =  4.00）  每行最大权重的平均：不缩放 0.79，除以 √d 后 0.53
  d=   64  点积的标准差   8.17（√d =  8.00）  每行最大权重的平均：不缩放 0.86，除以 √d 后 0.47
  d=  256  点积的标准差  16.08（√d = 16.00）  每行最大权重的平均：不缩放 0.95，除以 √d 后 0.47
  d= 1024  点积的标准差  32.53（√d = 32.00）  每行最大权重的平均：不缩放 0.92，除以 √d 后 0.39

The dot product of two random vectors is a sum of d products. The more terms are added, the more the result fluctuates, with a standard deviation of about the square root of d: at d=256 the dot product's standard deviation is 16.

With scores fluctuating that much, softmax becomes extreme, just like the "scaled up 5 times" in the second section: without scaling, the largest weight in each row averages 0.8 to 0.95, with nearly all attention going to one position. That's bad at the start of training: where weights are close to 0 or 1, softmax's gradient is tiny, and the model struggles to learn.

Divide by the square root of d, and the scores' standard deviation returns to about 1, the largest weight in each row is between 0.4 and 0.5, attention is spread out, gradients are healthy, and the model can gradually learn whom to attend to.

A complete attention head

Putting it all together:

d_model, d_head = 16, 8
x = torch.randn(T, d_model)
W_q, W_k, W_v = (torch.randn(d_model, d_head) / math.sqrt(d_model) for _ in range(3))
q, k, v = x @ W_q, x @ W_k, x @ W_v
scores = (q @ k.T / math.sqrt(d_head)).masked_fill(~mask, float("-inf"))
out = F.softmax(scores, dim=-1) @ v

== 5. 一个完整的注意力头：q、k、v 都由同一个输入经过不同的矩阵得到
  输入 (5, 16) → q、k、v 各 (5, 8) → 输出 (5, 8)
  和 PyTorch 自带的 scaled_dot_product_attention 比，最大差别 2.4e-07

Five lines of code: three matrices turn the input into q, k and v; score, scale and mask; softmax; weighted sum. The result matches PyTorch's built-in scaled_dot_product_attention (the difference is in the 7th decimal place, floating-point error from computing in a different order).

W_q, W_k and W_v are the parameters this attention head learns. During training they're updated like any other parameters, through backpropagation and gradient descent.

This formula comes from the 2017 paper "Attention Is All You Need", and in mathematical form it's softmax(QKᵀ/√d)V. Now you know what every symbol in it does.

The cost of attention

Every position scores every position before it. With a sequence of length T, the score table is T×T. Double the length, and the computation and the memory the table takes quadruple.

This is one of the root reasons long contexts are expensive (Module 01, Lesson 4 covered context and cost). To support contexts of hundreds of thousands or millions of tokens, model makers have done a great deal of optimisation on how attention is computed, but the basic idea is still these few lines from this lesson.

Exercises

1. In the second section, give "光" a different query vector, [0, 1] (looking for something related to position). How do the weights change?
2. In the fifth section, remove the causal mask and compare against scaled_dot_product_attention(..., is_causal=False).
3. Change the fifth section's sequence length from 5 to 1,000, 2,000 and 4,000, time each with time.time(), and see whether it grows roughly with the square.

Self-check

1. What do queries, keys and values each do?

The query says "what I'm looking for", the key says "what I have". The dot product of one position's query with every position's key gives its attention score for each position, which softmax turns into weights. The value says "the information you get by attending to me", and the output is the sum of every position's value weighted by those weights.

2. How is the causal mask implemented, and why is it needed?

Before softmax, the scores for positions after each position are set to negative infinity, so after softmax those positions' weights are 0. It's needed because the model is learning to predict the next character; if it could see later characters, it could simply copy the answer and learn nothing.

3. Why are attention scores divided by the square root of the dimension?

The dot product of two d-dimensional vectors has a standard deviation of about the square root of d, so the higher the dimension, the more the scores fluctuate, and the weights after softmax concentrate extremely on one position, making gradients tiny and the model hard to train. Dividing by the square root of d brings the scores' standard deviation back to about 1.