7 - Attention and the Transformer, from scratch

Before 2017, language models read text word by word with recurrent networks (RNNs), carrying a memory from one word to the next. The book tells how that approach hit a wall with long texts, and how the paper Attention Is All You Need replaced it with the Transformer, whose central piece is attention: every token can look directly at every other token and decide how much each one matters. In this notebook we build it piece by piece with NumPy, following Martha's notes (sections 7.7 and 7.8):

1. Import the libraries.
2. The wall of recurrent networks: memories that fade.
3. Scaled dot-product attention, step by step.
4. Why divide by $\sqrt{d_k}$.
5. Self-attention over a sentence, and the causal mask.
6. Multi-head attention.
7. Positional encoding with sines and cosines.
8. A complete Transformer block.

Step 1: Import the libraries

---

import numpy as np
import matplotlib.pyplot as plt

np.set_printoptions(precision=3, suppress=True)
rng = np.random.default_rng(0)

Step 2: The wall of recurrent networks

---

An RNN updates a hidden state with each new token, $h_t = \tanh(W h_{t-1} + U x_t)$, so the information about the first token has to survive every later step. During training, the gradient that tells the network how the first token affected the result is a product of one derivative per step, and a product of many numbers smaller than one shrinks exponentially: the vanishing gradient of chapter 6 again, now along the sequence. We measure how much the final state still depends on the first token for sequences of increasing length:

d_hidden = 16
W = rng.normal(0, 0.9 / np.sqrt(d_hidden), (d_hidden, d_hidden))     # recurrent weights
U = rng.normal(0, 1 / np.sqrt(d_hidden), (d_hidden, d_hidden))       # input weights
inputs = rng.normal(0, 1, (60, d_hidden))                            # a sequence of 60 token vectors

h, jacobian, influence = np.zeros(d_hidden), np.eye(d_hidden), []
for t in range(60):
    h = np.tanh(W @ h + U @ inputs[t])
    jacobian = (np.diag(1 - h ** 2) @ W) @ jacobian                   # d h_t / d h_0, by the chain rule
    influence.append(np.linalg.norm(jacobian))

for T in [1, 5, 10, 20, 40, 60]:
    print(f"After {T:>2} tokens, size of d h_T / d h_0: {influence[T - 1]:.2e}")

After  1 tokens, size of d h_T / d h_0: 2.63e+00
After  5 tokens, size of d h_T / d h_0: 2.29e-01
After 10 tokens, size of d h_T / d h_0: 1.44e-02
After 20 tokens, size of d h_T / d h_0: 2.20e-04
After 40 tokens, size of d h_T / d h_0: 7.85e-08
After 60 tokens, size of d h_T / d h_0: 1.62e-13

After one token the influence of the first token is of order 1; after 20 tokens it has fallen to $2.2 \times 10^{-4}$, and after 60 tokens to $1.6 \times 10^{-13}$, in practice zero. The network can hardly learn that a word at the end of a long text depends on one at the beginning. As the book explains, LSTMs and GRUs were developed to soften this problem, but they still read one token at a time, which also makes them slow to train.

Step 3: Scaled dot-product attention

---

Attention takes three matrices built from the tokens: the queries $Q$ (what each token is looking for), the keys $K$ (what each token offers to be found) and the values $V$ (the information each token passes on):

$$\text{Attention}(Q, K, V) = \text{softmax}\left(\frac{QK^\top}{\sqrt{d_k}}\right)V$$

Let's follow Martha's intuition with small vectors that we design by hand. Each token has 4 features: cardiac, body location, person and function word. For now the queries, keys and values are the vectors themselves, and we look at what the token "pain" attends to:

tokens = ['the', 'patient', 'has', 'chest', 'pain', 'arm']
#                  cardiac  location  person  function
X_toy = 2 * np.array([[0.0,    0.0,      0.0,    1.0],    # the
                      [0.1,    0.0,      1.0,    0.0],    # patient
                      [0.0,    0.0,      0.2,    0.8],    # has
                      [0.9,    0.8,      0.0,    0.0],    # chest
                      [1.0,    0.3,      0.0,    0.0],    # pain
                      [0.5,    0.9,      0.0,    0.0]])   # arm

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

def attention(Q, K, V, mask=None):
    d_k = K.shape[-1]
    scores = Q @ K.T / np.sqrt(d_k)                  # similarity of every query with every key
    if mask is not None:
        scores = np.where(mask, scores, -np.inf)      # forbidden positions get zero weight
    weights = softmax(scores, axis=-1)               # each row sums to 1
    return weights @ V, weights                       # weighted average of the values

q = X_toy[tokens.index('pain')]
scores = X_toy @ q / np.sqrt(4)
weights = softmax(scores)
for tok, s, w in zip(tokens, scores, weights):
    print(f"pain -> {tok:<8} score {s:5.2f}   weight {w:.3f}")
output = weights @ X_toy
print(f"\nNew vector for 'pain' (weighted average of the values): {output}")

pain -> the      score  0.00   weight 0.038
pain -> patient  score  0.20   weight 0.046
pain -> has      score  0.00   weight 0.038
pain -> chest    score  2.28   weight 0.369
pain -> pain     score  2.18   weight 0.334
pain -> arm      score  1.54   weight 0.176

New vector for 'pain' (weighted average of the values): [1.516 1.107 0.107 0.136]

"Pain" attends above all to "chest" (0.369), to itself (0.334) and to "arm" (0.176), and hardly at all to "the", "patient" or "has", the tokens without cardiac or body-location features. Its new vector is a mixture of those three, so after attention "pain" carries information about where it hurts. As Martha puts it, softmax decides the focus.

Step 4: Why divide by $\sqrt{d_k}$

---

Martha explains that, as the dimension $d_k$ grows, dot products become very large and push the softmax into regions where the gradients are tiny. We can measure it: with random queries and keys, we compute the attention weights of one query over 10 keys, with and without the scaling, for increasing dimensions.

for d_k in [4, 64, 512]:
    q, K = rng.normal(size=d_k), rng.normal(size=(10, d_k))
    raw = K @ q
    w_raw, w_scaled = softmax(raw), softmax(raw / np.sqrt(d_k))
    print(f"d_k = {d_k:>3}: spread (std) of the scores {raw.std():6.2f} | largest weight without scaling {w_raw.max():.3f}, "
          f"with scaling {w_scaled.max():.3f}")

d_k =   4: spread (std) of the scores   0.64 | largest weight without scaling 0.228, with scaling 0.160
d_k =  64: spread (std) of the scores   6.13 | largest weight without scaling 0.831, with scaling 0.222
d_k = 512: spread (std) of the scores  32.88 | largest weight without scaling 1.000, with scaling 0.326

With $d_k = 4$ the scores are small and the weights stay spread out with or without scaling. But the spread of the scores grows with the dimension (6.13 with 64 dimensions, 32.88 with 512), and without scaling the softmax saturates: with $d_k = 512$ a single key takes all the weight (1.000) and the rest get practically zero, so their gradients vanish. Dividing by $\sqrt{d_k}$ brings the scores back to a comparable spread in every dimension.

Step 5: Self-attention over a sentence, and the causal mask

---

In a real Transformer, $Q$, $K$ and $V$ are not the token vectors themselves: each one is obtained by multiplying the matrix of token vectors $X$ by a learned matrix, $Q = XW^Q$, $K = XW^K$, $V = XW^V$. With the whole sentence at once, attention produces a matrix of weights: row $i$ says how much token $i$ attends to each token $j$. Here the projections are random, as before training.

Language models like GPT generate text from left to right, so a token must not look at the tokens that come after it. This is done with a causal mask, which gives zero weight to every future position:

sentence = ['the', 'patient', 'has', 'chest', 'pain', 'radiating', 'to', 'the', 'arm']
n, d_model, d_k = len(sentence), 8, 8
X = rng.normal(size=(n, d_model))                                   # token embeddings (untrained)
W_Q, W_K, W_V = (rng.normal(0, 1 / np.sqrt(d_model), (d_model, d_k)) for _ in range(3))

out_full, w_full = attention(X @ W_Q, X @ W_K, X @ W_V)
causal = np.tril(np.ones((n, n), dtype=bool))                       # token i may only see tokens 0..i
out_causal, w_causal = attention(X @ W_Q, X @ W_K, X @ W_V, mask=causal)
print(f"Output shape: {out_full.shape} (one new vector per token)")
print(f"Each row of weights sums to 1: {np.allclose(w_full.sum(axis=1), 1)}")
print(f"With the mask, weight given to future tokens: {w_causal[~causal].sum():.1f}")

fig, axes = plt.subplots(1, 2, figsize=(13, 5))
for ax, w, title in [(axes[0], w_full, 'Self-attention (every token sees every token)'),
                     (axes[1], w_causal, 'Causal self-attention (no looking ahead)')]:
    ax.imshow(w, cmap='Blues', vmin=0, vmax=1)
    ax.set_xticks(range(n), sentence, rotation=45, ha='right')
    ax.set_yticks(range(n), sentence)
    ax.set_xlabel('attends to (keys)')
    ax.set_ylabel('token (queries)')
    ax.set_title(title)
plt.tight_layout()
plt.show()

Output shape: (9, 8) (one new vector per token)
Each row of weights sums to 1: True
With the mask, weight given to future tokens: 0.0

With random weights the pattern means nothing yet: training is what adjusts $W^Q$, $W^K$ and $W^V$ so that related tokens produce high scores. What matters here is the structure. Every token is connected to every other token in a single step, whatever the distance between them, which is exactly what the RNN was missing. And the causal mask leaves only the lower triangle.

Step 6: Multi-head attention

---

A single attention computes one notion of "relevance". Multi-head attention computes $h$ of them in parallel, each with its own projections, concatenates the results and mixes them with a final matrix $W^O$:

$$\text{head}_i = \text{Attention}(XW_i^Q, XW_i^K, XW_i^V) \qquad \text{MultiHead}(X) = \text{Concat}(\text{head}_1, \dots, \text{head}_h)\,W^O$$

With the usual choice $d_k = d_v = d_{\text{model}} / h$, the concatenation has $d_{\text{model}}$ columns again.

def init_multihead(d_model, n_heads, rng):
    d_k = d_model // n_heads
    heads = [tuple(rng.normal(0, 1 / np.sqrt(d_model), (d_model, d_k)) for _ in range(3)) for _ in range(n_heads)]
    W_O = rng.normal(0, 1 / np.sqrt(d_model), (d_model, d_model))
    return heads, W_O

def multihead(X, heads, W_O, mask=None):
    outputs = [attention(X @ Wq, X @ Wk, X @ Wv, mask)[0] for Wq, Wk, Wv in heads]
    return np.concatenate(outputs, axis=1) @ W_O

heads, W_O = init_multihead(d_model=8, n_heads=2, rng=rng)
mh = multihead(X, heads, W_O, mask=causal)
n_params = sum(w.size for h in heads for w in h) + W_O.size
print(f"2 heads of dimension {heads[0][0].shape[1]}; output shape {mh.shape}; parameters {n_params}")

2 heads of dimension 4; output shape (9, 8); parameters 256

Step 7: Positional encoding

---

Martha's notes point out a surprising property: attention on its own is blind to word order. If we shuffle the tokens, the outputs are shuffled in exactly the same way, but they do not change. Let's check it:

perm = rng.permutation(n)
out = multihead(X, heads, W_O)                       # without mask, so that order is not imposed by the mask
out_shuffled = multihead(X[perm], heads, W_O)
print(f"Shuffling the tokens only shuffles the outputs: {np.allclose(out[perm], out_shuffled)}")

Shuffling the tokens only shuffles the outputs: True

"The patient treated the doctor" and "the doctor treated the patient" would look the same. The Transformer solves this by adding to each token vector a positional encoding made of sines and cosines of different frequencies:

$$PE_{(pos, 2i)} = \sin\left(\frac{pos}{10000^{2i/d_{\text{model}}}}\right) \qquad PE_{(pos, 2i+1)} = \cos\left(\frac{pos}{10000^{2i/d_{\text{model}}}}\right)$$

Martha proves that, for any offset $k$, the encoding of position $pos + k$ is a rotation of the encoding of $pos$, with a matrix that depends only on $k$. We build the encoding, draw it, and check that property numerically:

def positional_encoding(n_positions, d_model):
    pos = np.arange(n_positions)[:, None]
    omega = 1 / 10000 ** (np.arange(0, d_model, 2) / d_model)       # one frequency per pair of dimensions
    pe = np.zeros((n_positions, d_model))
    pe[:, 0::2] = np.sin(pos * omega)
    pe[:, 1::2] = np.cos(pos * omega)
    return pe, omega

PE, omega = positional_encoding(100, 16)
plt.figure(figsize=(10, 4))
plt.imshow(PE.T, aspect='auto', cmap='RdBu')
plt.xlabel('position in the sequence')
plt.ylabel('dimension')
plt.title('Sinusoidal positional encoding: fast frequencies at the top, slow ones at the bottom')
plt.colorbar()
plt.show()

k = 3
R = np.zeros((16, 16))
for i, w in enumerate(omega):                        # block-diagonal matrix of 2x2 rotations
    R[2 * i:2 * i + 2, 2 * i:2 * i + 2] = [[np.cos(k * w), np.sin(k * w)], [-np.sin(k * w), np.cos(k * w)]]
error = np.max(np.abs(PE[k:] - PE[:-k] @ R.T))
print(f"Largest difference between PE(pos + {k}) and R_{k} PE(pos), over all positions: {error:.1e}")

PE8 = positional_encoding(n, 8)[0]                  # the encoding for our 8-dimensional tokens
out_pe = multihead(X + PE8, heads, W_O)
out_pe_shuffled = multihead(X[perm] + PE8, heads, W_O)
print(f"With positional encoding, shuffling only shuffles the outputs: {np.allclose(out_pe[perm], out_pe_shuffled)}")

Largest difference between PE(pos + 3) and R_3 PE(pos), over all positions: 2.4e-15
With positional encoding, shuffling only shuffles the outputs: False

The rotation property holds to within $2.4 \times 10^{-15}$, that is, exactly up to rounding errors: moving 3 positions forward is the same rotation wherever we are in the sequence. And once the positional encoding is added, shuffling the tokens really does change the outputs, so the model can now tell "the patient treated the doctor" from "the doctor treated the patient".

Step 8: A complete Transformer block

---

A Transformer block combines everything, as in Martha's notes: multi-head attention, which mixes information between tokens; a feed-forward network applied to each token separately, $\text{FFN}(x) = \text{ReLU}(xW_1 + b_1)W_2 + b_2$; and, around each of them, a residual connection plus layer normalization, $\text{LayerNorm}(x + \text{Sublayer}(x))$. The residual connection is a "highway" for the gradients, and the normalization gives each token vector mean 0 and variance 1. Large language models stack dozens of these blocks.

def layer_norm(x, eps=1e-5):
    return (x - x.mean(axis=-1, keepdims=True)) / np.sqrt(x.var(axis=-1, keepdims=True) + eps)

def transformer_block(x, heads, W_O, W1, b1, W2, b2, mask=None):
    x = layer_norm(x + multihead(x, heads, W_O, mask))            # attention + Add & Norm
    ffn = np.maximum(0, x @ W1 + b1) @ W2 + b2                    # feed-forward, token by token
    return layer_norm(x + ffn)                                     # FFN + Add & Norm

d_model, d_ff = 8, 32
W1, b1 = rng.normal(0, 1 / np.sqrt(d_model), (d_model, d_ff)), np.zeros(d_ff)
W2, b2 = rng.normal(0, 1 / np.sqrt(d_ff), (d_ff, d_model)), np.zeros(d_model)

x_in = X + PE8
y = transformer_block(x_in, heads, W_O, W1, b1, W2, b2, mask=causal)
n_block = n_params + W1.size + b1.size + W2.size + b2.size
print(f"Input {x_in.shape} -> output {y.shape}")
print(f"Mean of each output token: {y.mean(axis=1).round(3)}")
print(f"Variance of each output token: {y.var(axis=1).round(3)}")
print(f"Parameters of this small block: {n_block}")

Input (9, 8) -> output (9, 8)
Mean of each output token: [-0.  0. -0. -0.  0. -0.  0.  0. -0.]
Variance of each output token: [1. 1. 1. 1. 1. 1. 1. 1. 1.]
Parameters of this small block: 808

The block keeps the shape of the sequence: 9 tokens of 8 dimensions go in and 9 tokens of 8 dimensions come out, each one now mixed with its context. After the last layer normalization every token has mean 0 and variance 1. This small block has 808 parameters; in a large language model the same structure is repeated in dozens of layers with thousands of dimensions, which is how models reach billions of parameters.

Conclusion

---

• Recurrent networks pass information from token to token, and its influence fades exponentially with distance. Attention connects every token with every other one in a single step.
• Scaled dot-product attention compares queries with keys, turns the scores into weights with softmax and mixes the values with those weights. Dividing by $\sqrt{d_k}$ keeps the softmax from saturating.
• The causal mask keeps a token from looking at the future, which is what lets a model generate text from left to right. Multi-head attention computes several notions of relevance in parallel.
• Attention is blind to order, so the Transformer adds a positional encoding, whose sines and cosines turn a shift of $k$ positions into a rotation.
• A Transformer block adds a feed-forward network, residual connections and layer normalization. Everything here has random weights: training on huge amounts of text, as the book explains, is what turns this machinery into a language model.