Chapter 9: Attention Is All You Need

The problem with fixed windows

Recurrent neural networks process a sentence one token at a time. They keep a hidden state that theoretically remembers everything, but in practice gradients vanish and the model forgets distant words. Convolutional networks see only a fixed local window.

Attention solves this by letting every token look directly at every other token. There is no fixed window and no sequential bottleneck. The model learns how much each token should care about every other token, and it can do so in parallel across the whole sequence.

Query, Key, Value

For each token, we create three vectors from its embedding: a query, a key, and a value. The query represents what the token is looking for. The key represents what the token offers. The value is the actual information to be passed on.

X is the input sequence matrix where each row is a token embedding. The weight matrices W_Q, W_K, W_V are learned. Their dimensions are chosen so that every token ends up with a query, key, and value vector. The queries and keys have dimension d_k; the values have dimension d_v.

Attention as a soft dictionary

A useful intuition is to think of attention as a differentiable lookup. A normal dictionary lookup returns one value for one exact key. Attention instead returns a weighted mixture of all values, where the weights are determined by the similarity between the query and each key.

The query is like a question, the keys are like index cards, and the values are like the answers on those cards. If the query matches several keys, the output blends their values. If it matches one key strongly, that value dominates. The entire operation is differentiable, so the model can learn what each token should ask for and what each token should offer.

Scaled dot-product attention

To decide how much token i attends to token j, we compare the query of i with the key of j. The dot product measures similarity. We scale by the square root of the key dimension to keep gradients stable, then apply softmax.

The softmax produces a probability distribution over positions for each query. We use these weights to take a weighted sum of the value vectors. The result is a new representation for token i that blends information from the whole sequence. Because every token queries every other token, the operation can be written as one matrix multiplication and is highly parallel.

Multi-head attention

One attention function can only capture one kind of relationship. Multi-head attention runs h attention functions in parallel, each with its own learned projections. The outputs are concatenated and projected again. The total computation is roughly the same as a single large attention because each head uses a smaller dimension.

Different heads can specialize: one may track pronoun references, another may track verb-subject agreement, another may track nearby punctuation. The model learns these specializations from data. The final projection W_O mixes the heads back into the model dimension.

Causal masking for generation

When generating text, a token should not look at future tokens. A causal or autoregressive mask sets attention scores for future positions to negative infinity before the softmax, forcing the model to use only past tokens.

The mask M is added to the attention logits before softmax. Positions where i < j, the future, are set to a very negative number so that softmax gives them zero weight. Positions where i >= j, the present and past, are left alone.

import numpy as np

def softmax(x):
    e = np.exp(x - np.max(x, axis=-1, keepdims=True))
    return e / e.sum(axis=-1, keepdims=True)

def causal_attention(Q, K, V):
    scores = Q @ K.T / np.sqrt(Q.shape[-1])
    # mask future positions: set them to a large negative value
    mask = np.triu(np.ones_like(scores), k=1) * -1e9
    scores = scores + mask
    weights = softmax(scores)
    return weights @ V

A causal mask walkthrough

Imagine a three-token sentence: "A B C". When predicting token B, the model may look at A. When predicting C, it may look at A and B. The attention matrix is lower-triangular: each row is a probability distribution over positions up to and including the current one.

A     B     C
A    [ 1.0   0.0   0.0  ]  <- A only sees itself
B    [ 0.6   0.4   0.0  ]  <- B sees A and itself
C    [ 0.2   0.5   0.3  ]  <- C sees A, B, and itself

This triangular structure is what makes autoregressive generation possible. At training time the whole matrix is computed in one go because the correct target for every position is known. At inference time the model generates one token at a time, but each new token still only attends to previous tokens.

Position and RoPE

Attention by itself is permutation-invariant: it does not know whether "cat sat" or "sat cat" came first. Positional encoding, introduced in the embeddings chapter, gives each position a unique fingerprint. Modern models such as LLaMA and GPT use Rotary Position Embedding, or RoPE, which rotates query and key vectors by a position-dependent angle before their dot product.

RoPE bakes relative position directly into the dot product. The attention score between two tokens depends on how far apart they are, not just on their absolute positions. This helps models generalize to longer sequences than they saw during training.