Attention: Learning Where to Look

Why context changes meaning

Read these two sentences:

• The animal didn't cross the street because it was too tired.
• The animal didn't cross the street because it was too wide.

In the first, "it" is the animal. In the second, it is the street. One word changed, five words after "it", and the meaning of "it" flipped.

An embedding from Level 9 gives "it" the same vector in both sentences. To understand language, each token's vector has to be updated with information from the other tokens that matter to it. Attention does exactly that: every token looks at every other token, decides how relevant each one is, and takes in a weighted mix of their information.

An RNN passes information along one step at a time, so distant words fade. Attention connects any two tokens directly, however far apart they are.

Queries, keys and values

Attention borrows its vocabulary from searching a library:

Each token's embedding x is multiplied by three learned weight matrices to make three new vectors:

The same three matrices are used for every token. Training adjusts them so that a token's query lines up with the keys of the tokens it needs: a pronoun's query might come to match the keys of nouns, and a verb's query might match its subject.

Scaled dot-product attention, step by step

For one token, attention takes four steps:

1. Score: take the dot product of its query with every key. Vectors pointing the same way give large scores.
2. Scale: divide every score by √d_k, the square root of the length of the keys.
3. Softmax: turn the scores into weights between 0 and 1 that add up to 1.
4. Blend: add up the value vectors, each multiplied by its weight.

Here is the query of "it", q = [2, 1], against four keys. The keys have 2 numbers, so the scores are divided by √2 ≈ 1.41:

Each weight is its e^x divided by the column total, 46.71, so the weights add up to 1. If the values are animal [4, 0], cross [2, 2], street [0, 4] and it [1, 1], the output is

0.735 × [4, 0] + 0.043 × [2, 2] + 0.043 × [0, 4] + 0.179 × [1, 1] = [3.20, 0.44]

The new vector for "it" is now mostly made of "animal".

Every token at once: the attention matrix

In self-attention every token does this at the same time. Stack all the queries as the rows of a matrix Q, the keys as K and the values as V, and the whole calculation is two matrix multiplications and a softmax:

Here K^T is K transposed, its rows turned into columns, so that Q × K^T holds the dot product of every query with every key. After the softmax this is the attention matrix: one row per token, showing how much it attends to every token. Every row adds up to 1.

Because it is all matrix multiplication, a GPU computes attention for every token in parallel. That is the speed advantage over an RNN, which must read one token after another.

Try it yourself: drag the query and the keys, and watch the scores, the softmax and the blended value change.

Masking: no peeking at the future

A model that generates text predicts each token from the tokens before it. In training it is shown whole sentences at once, so something must stop each token attending to the tokens after it, or it could simply copy the answer it is meant to predict.

The fix is a causal mask. Before the softmax, every score where the key comes after the query is set to minus infinity. Since e to the power of minus infinity is 0, those weights become exactly 0, and the remaining weights in each row still add up to 1.

Transformers use attention in three ways:

Cross-attention is the 2014 idea from Level 9: while writing each word, the decoder looks back at the whole source sentence.

Multi-head attention

One attention calculation gives each token one weighted mix. But a word has many relationships at once: "it" needs the noun it refers to, a verb needs its subject, and every word benefits from its neighbours. Multi-head attention runs several attention calculations, called heads, side by side, each with its own W_Q, W_K and W_V.

1. Each of the h heads projects the tokens down to d_model ÷ h numbers and computes attention on its own.
2. The h outputs are joined end to end, concatenated, back to d_model numbers per token.
3. A final matrix, W_O, mixes what the heads found.

The original transformer used 8 heads of 64 numbers, so d_model = 512. Because each head works on fewer numbers, 8 heads cost about the same as one big one.

Researchers who look inside trained models find heads with recognisable jobs: heads that attend to the previous token, heads that link a pronoun to its noun, and induction heads, which spot that a pattern appeared earlier in the text and predict that it will repeat.

Hands-on: attention from scratch

One query, four keys

This is the worked example from the table above, in code. Check every number against the table.

import math

def softmax(xs):
    m = max(xs)                      # subtract the largest for numerical safety
    exps = [math.exp(x - m) for x in xs]
    total = sum(exps)
    return [e / total for e in exps]

words  = ["animal", "cross", "street", "it"]
keys   = [[2, 1], [0, 1], [1, -1], [1, 1]]
values = [[4, 0], [2, 2], [0, 4], [1, 1]]
query  = [2, 1]                      # the query of "it"

d_k = len(query)
scores = [sum(q * k for q, k in zip(query, key)) / math.sqrt(d_k) for key in keys]
weights = softmax(scores)
output = [sum(w * v[i] for w, v in zip(weights, values)) for i in range(2)]

for word, s, w in zip(words, scores, weights):
    print(f"{word:7} score {s:5.2f}  weight {w:.3f}  " + "#" * round(w * 40))
print("weights add up to", round(sum(weights), 6))
print("new vector for 'it':", [round(x, 2) for x in output])

Self-attention with matrices, with and without a mask

Every token gets a query, a key and a value from the same three matrices, and the whole attention matrix comes out at once. The weights here are random, so the pattern means nothing yet: training is what gives it meaning.

import math, random

def matmul(A, B):
    return [[sum(a * b for a, b in zip(row, col)) for col in zip(*B)] for row in A]

def transpose(A):
    return [list(col) for col in zip(*A)]

def softmax(xs):
    m = max(xs)
    exps = [math.exp(x - m) for x in xs]
    return [e / sum(exps) for e in exps]

def self_attention(X, W_Q, W_K, W_V, causal=False):
    Q, K, V = matmul(X, W_Q), matmul(X, W_K), matmul(X, W_V)
    d_k = len(Q[0])
    scores = matmul(Q, transpose(K))
    for i in range(len(X)):
        for j in range(len(X)):
            scores[i][j] /= math.sqrt(d_k)
            if causal and j > i:
                scores[i][j] = float("-inf")    # e to the power of -inf is 0
    A = [softmax(row) for row in scores]
    return A, matmul(A, V)

random.seed(4)
tokens = ["the", "cat", "sat", "down"]
d_model, d_k = 4, 3
X = [[random.uniform(-1, 1) for _ in range(d_model)] for _ in tokens]
W_Q, W_K, W_V = ([[random.uniform(-1, 1) for _ in range(d_k)] for _ in range(d_model)]
                 for _ in range(3))

for causal in (False, True):
    A, out = self_attention(X, W_Q, W_K, W_V, causal)
    print("causal mask:", causal)
    print("        " + "".join(f"{t:>7}" for t in tokens))
    for t, row in zip(tokens, A):
        print(f"{t:>7} " + "".join(f"{w:7.2f}" for w in row) + f"   sum {sum(row):.2f}")
    print("output:", len(out), "tokens x", len(out[0]), "numbers")
    print()

With the mask, everything above the diagonal is 0, and "the" can only attend to itself. The last row is the same in both runs, because the last token may see everything either way.

Several heads

import math, random

def matmul(A, B):
    return [[sum(a * b for a, b in zip(row, col)) for col in zip(*B)] for row in A]

def softmax(xs):
    m = max(xs)
    exps = [math.exp(x - m) for x in xs]
    return [e / sum(exps) for e in exps]

def head(X, d_head, seed):
    rnd = random.Random(seed)          # every head has its own weights
    W = [[[rnd.uniform(-1.5, 1.5) for _ in range(d_head)] for _ in X[0]] for _ in range(3)]
    Q, K, V = (matmul(X, w) for w in W)
    A = [softmax([sum(q * k for q, k in zip(Q[i], K[j])) / math.sqrt(d_head)
                  for j in range(len(X))]) for i in range(len(X))]
    return A, matmul(A, V)

random.seed(0)
tokens = ["she", "saw", "the", "bird"]
d_model, h = 6, 3
X = [[random.uniform(-1, 1) for _ in range(d_model)] for _ in tokens]
outputs = []
for n in range(h):
    A, out = head(X, d_model // h, seed=10 + n)
    outputs.append(out)
    focus = [tokens[max(range(len(tokens)), key=lambda j: A[i][j])] for i in range(len(tokens))]
    print(f"head {n + 1}: each token attends most to", dict(zip(tokens, focus)))
concat = [sum((o[i] for o in outputs), []) for i in range(len(tokens))]
print("each head gives", d_model // h, "numbers; joined:", len(concat[0]), "numbers per token")

Try this

1. In the first program, change the query to [6, 3]. The scores triple. What happens to the weight on "animal", and why does that show the reason for scaling?
2. Change the key of "street" to [2, 1] so it matches "animal". How do the weights and the output change?
3. In the second program, count the zeros in the masked matrix. How many would there be for 10 tokens?

Level check

• Explain why a word's meaning depends on its context, and how attention uses context
• Describe the roles of queries, keys and values, and how each is made
• Calculate attention weights from scores using softmax, and the blended output
• Explain why scores are divided by √d_k
• Explain what a causal mask does and why generation needs it
• Distinguish self-attention, masked self-attention and cross-attention
• Describe how multi-head attention splits, computes and recombines

Key terms you must be able to define

Attention

A way for each token to build its new vector from a weighted mix of the information in other tokens.

Query

The vector a token uses to search the other tokens: what it is looking for.

Key

The vector each token is matched against queries with: what it offers.

Value

The information a token passes on when other tokens attend to it.

Scaled dot-product attention

softmax(Q × K^T ÷ √d_k) × V: dot-product scores, scaled, turned into weights, used to blend the values.

Attention weights

The softmax outputs: for each token, numbers between 0 and 1 that add up to 1.

Self-attention

Attention in which the queries, keys and values all come from the same sequence.

Cross-attention

Attention in which the queries come from one sequence and the keys and values from another.

Causal mask

Setting the scores for later positions to minus infinity, so no token can attend to the future.

Multi-head attention

Several attention heads with their own weights, run side by side and combined.

Head

One attention calculation inside multi-head attention, with its own W_Q, W_K and W_V.

Flashcards

Go further: notes for future AI engineers

Attention as a soft dictionary

A Python dictionary does hard lookup: a key matches exactly or not at all, and you get one value back. Attention is a soft lookup: every key matches a little, and you get back a blend of all the values, weighted by how well each key matched. Because the blend is smooth, gradient descent can adjust the queries and keys, so the network learns what to look up.

The cost of attention

Every token attends to every token, so n tokens need an n × n attention matrix: double the context and the work quadruples. That quadratic cost long limited context windows to a few thousand tokens. FlashAttention (2022) computes exactly the same result in tiles that fit the GPU's fastest memory, and sliding-window and sparse attention let each token attend to only some of the others. Together with similar engineering, they are how modern models handle contexts of hundreds of thousands of tokens.

Misconceptions to avoid