95 - Attention, Derived in Full: Query, Key, Value and the Transformer Block Written From Scratch

In the sentence "the technician handed the form to the customer because it was incomplete" the pronoun refers to the form, and the reader understands it effortlessly: they do not reread from the beginning and do not rely on a summary, they go back six words, compare "it was incomplete" with the two candidates and take the information from there. Formulating a question, comparing it with everything available, weighing the content by how well each piece answers: this is attention.

The mechanism was born in 2014 with Bahdanau, Cho and Bengio as a crutch for a recurrent network that translated long sentences badly (Neural Machine Translation by Jointly Learning to Align and Translate, arXiv:1409.0473); three years later Vaswani and colleagues remove the crutch and throw away the leg (Attention Is All You Need, arXiv:1706.03762). Here you derive it in full, count the parameters of a block, measure the quadratic cost, and train a decoder to copy and sort sequences.

How the examples are written. The code is pure NumPy: attention, multiple heads, block, hand-written backward pass and decoder training. Environment: Python 3.14.4, NumPy 2.3.5, fixed seeds. Every listing was run twice with identical output; the exceptions are the times, reported as an interval. The listings are to be read in order, because each one uses the names defined by the previous ones. PyTorch 2.12.0 appears in one place only, a numerical comparison of the formula with torch.nn.functional.scaled_dot_product_attention, run on CPU with the interpreter that contains it and without downloading anything.

The analogy: the question put to the filing cabinet

Imagine a filing cabinet. Each card has a heading on its edge - "invoices 2024", "expired contracts" - and the content inside. You arrive with a question: how much did we spend on transport last year. You do not open all the cards and you do not open just one: you run through the headings, give each a similarity score, and read the content of those with the highest score, weighing each by how much it resembled the question.

The dictionary is this: the question is the query, the headings are the keys, the contents are the values; similarity is a dot product and the passage to weights that sum to one is the softmax. Heading and content of the same card are two different things, obtained from the same sheet with different criteria: whoever compares the input vectors directly can only say "you are equal to me".

The wall that recurrent networks hit

Module 06 took recurrent networks apart piece by piece: the lesson on the vanishing and exploding gradient in time showed the product of Jacobians that contracts, the one on the LSTM cell with gates built the way out, the experiment over several seeds measured how much it is worth. Two limits remain that do not depend on optimization but on the shape of the computation.

The first limit is the number of sequential steps. The state hₜ depends on hₜ₋₁: for T tokens you need T operations in order. An attention computes all the scores with two matrix multiplications: sequential steps O(1). The second limit is the path between distant tokens. In a recurrent network the information from token 1 crosses 99 transformations to reach token 100 and the gradient retraces the road: path O(T). In attention it is 1. It is table 1 of arXiv:1706.03762, and it is the real argument of the paper.

The price is paid elsewhere. A recurrent network costs O(T · d²) and keeps a state of fixed size; attention costs O(T² · d) and keeps a T × T matrix in memory. We will measure it with a stopwatch.

Scaled dot-product attention, derived

We start from X of shape (T, d): T tokens of d numbers. Three parameter matrices project each token into three roles, and the similarity scores are a single matrix product:

Q = X W_Q      (T, d) · (d, d_k) → (T, d_k)      the queries
K = X W_K      (T, d) · (d, d_k) → (T, d_k)      the headings
V = X W_V      (T, d) · (d, d_v) → (T, d_v)      the contents
S = Q Kᵀ       (T, d_k) · (d_k, T) → (T, T)      sᵢⱼ = qᵢ · kⱼ = Σ_{m=1..d_k} qᵢₘ · kⱼₘ

Row i of Q is the question of token i, row j of K the heading of token j, row j of V its content; when all three come from the same X we speak of self-attention. The matrix S is square in the length and is not symmetric, because query and heading pass through different matrices.

The scores become weights that sum to one, row by row, with a softmax over the keys axis; the output is the average of the values with those weights, and together they give equation 1 of arXiv:1706.03762. I write the index explicitly because it is the axis that gets wrong: if you normalize by columns the code runs, the shapes match, and no exception warns you.

pᵢⱼ = exp(sᵢⱼ / √d_k) / Σ_{j′=1..T} exp(sᵢⱼ′ / √d_k)     the sum runs over j′, that is over the KEYS, with i fixed
Y = P V        (T, T) · (T, d_v) → (T, d_v)              yᵢ = Σ_{j=1..T} pᵢⱼ · vⱼ

Attention(Q, K, V) = softmax( Q Kᵀ / √d_k ) V

The listing translates the formula line by line, with one extra dimension for the batch B. The softmax subtracts the maximum before exponentiating: it does not change the result, because the softmax is invariant to adding a constant to all the scores, but it keeps exp from overflowing. The weight matrices are Gaussians divided by √d, so that the projections keep a variance close to one. The last lines are the checks: the shapes of the three matrices and the sum of each row of the weights, which must be one.

import numpy as np

def softmax(s, asse=-1):
    """Numerically stable softmax: the maximum is subtracted before exponentiating."""
    s = s - s.max(axis=asse, keepdims=True)
    e = np.exp(s)
    return e / e.sum(axis=asse, keepdims=True)

def attenzione(X, W_Q, W_K, W_V, scala=True):
    Q = X @ W_Q                                   # (B, T, d_k)
    K = X @ W_K                                   # (B, T, d_k)
    V = X @ W_V                                   # (B, T, d_v)
    d_k = Q.shape[-1]
    punteggi = Q @ K.transpose(0, 2, 1)           # (B, T, T)
    if scala:
        punteggi = punteggi / np.sqrt(d_k)
    pesi = softmax(punteggi, asse=-1)             # softmax OVER THE KEYS
    return pesi @ V, pesi, punteggi               # (B, T, d_v), (B, T, T), (B, T, T)

rng = np.random.default_rng(7)
B, T, d, d_k = 2, 5, 16, 8
X   = rng.normal(size=(B, T, d))
W_Q = rng.normal(size=(d, d_k)) / np.sqrt(d)
W_K = rng.normal(size=(d, d_k)) / np.sqrt(d)
W_V = rng.normal(size=(d, d_k)) / np.sqrt(d)

Y, P, S = attenzione(X, W_Q, W_K, W_V)
print("scores QK^T        ", S.shape)
print("weights (softmax)  ", P.shape)
print("output             ", Y.shape)
print("sum of each row of the weights:", np.round(P.sum(axis=-1), 12).ravel()[:5])

scores QK^T         (2, 5, 5)
weights (softmax)   (2, 5, 5)
output              (2, 5, 8)
sum of each row of the weights: [1. 1. 1. 1. 1.]

Before trusting a hand-written formula it is worth comparing it with an independent implementation. PyTorch has the function scaled_dot_product_attention, which computes the same expression; the listing applies it to the same Q, K, V as the NumPy version, with and without a causal mask, and prints the maximum deviation. The variables softmax and np come from the previous listing.

import torch
import torch.nn.functional as F

torch.manual_seed(0)
torch.set_num_threads(2)

rng = np.random.default_rng(7)
B, T, d_k = 2, 5, 8
Q, K, V = (rng.normal(size=(B, T, d_k)) for _ in range(3))
S = Q @ K.transpose(0, 2, 1) / np.sqrt(d_k)
mio = softmax(S) @ V
mio_causale = softmax(np.where(np.tril(np.ones((T, T), dtype=bool)), S, -np.inf)) @ V

Qt, Kt, Vt = (torch.from_numpy(a) for a in (Q, K, V))
libreria = F.scaled_dot_product_attention(Qt, Kt, Vt).numpy()
libreria_causale = F.scaled_dot_product_attention(Qt, Kt, Vt, is_causal=True).numpy()
print("torch", torch.__version__)
print("maximum deviation, no mask:", f"{np.abs(mio - libreria).max():.2e}")
print("maximum deviation, causal :", f"{np.abs(mio_causale - libreria_causale).max():.2e}")

torch 2.12.0+cu130
maximum deviation, no mask: 2.22e-16
maximum deviation, causal : 2.22e-16

What the numbers say. The deviation is 2.22 · 10⁻¹⁶, that is a rounding error in double precision: the formula above and the library's one are the same calculation. It does not prove that attention is the right choice for your problem; it proves that your code implements the one in the paper, and that the is_causal=True option is equivalent to the mask with minus infinity before the softmax.

Why you divide by the square root of d_k, measured

The square root is the only piece of arXiv:1706.03762 relegated to a footnote despite being decisive. The score is sᵢⱼ = Σ_{m=1..d_k} qᵢₘ · kⱼₘ; if the components of q and k are independent with mean 0 and variance 1, each term has mean 0 and variance 1, and the sum of d_k independent terms has variance d_k:

Var(qᵢ · kⱼ) = d_k          sd(qᵢ · kⱼ) = √d_k          sd( (qᵢ · kⱼ) / √d_k ) = 1

In words: the larger the dimension of the comparison, the wider the scores become. Without the division, at d_k = 512 the scores swing over tens and the softmax concentrates almost all the weight on a single key. The derivative of the softmax is ∂pᵢ/∂sⱼ = pᵢ(δᵢⱼ − pⱼ) and its trace Σⱼ pⱼ(1 − pⱼ) goes to zero when one weight is close to one: the gradient that teaches where to look does not arrive.

def misura(d_k, scala, seme, T=64, ripetizioni=400):
    """Statistics of the scores and of the weights, with and without the division by sqrt(d_k)."""
    rng = np.random.default_rng(seme)
    Q = rng.normal(size=(ripetizioni, T, d_k))        # independent components, mean 0, variance 1
    K = rng.normal(size=(ripetizioni, T, d_k))
    grezzi = Q @ K.transpose(0, 2, 1)                 # (r, T, T): the dot products
    punteggi = grezzi / np.sqrt(d_k) if scala else grezzi
    P = softmax(punteggi, asse=-1)
    sensibilita = (P * (1.0 - P)).sum(axis=-1)        # trace of the softmax Jacobian
    entropia = -(P * np.log(P + 1e-300)).sum(axis=-1)
    return (grezzi.var(), punteggi.std(), P.max(axis=-1).mean(),
            entropia.mean(), sensibilita.mean())

print("d_k  divided  Var(q.k)  sd scores  max weight  entropy  sensitivity")
for d_k in (8, 64, 512):
    for scala in (False, True):
        v, sd, pm, h, s = misura(d_k, scala, 100 + d_k)
        print(f"{d_k:4d}  {'yes    ' if scala else 'no     '} {v:8.2f} {sd:11.3f}"
              f" {pm:11.4f} {h:8.4f} {s:12.4f}")
print("maximum entropy over 64 keys = ln(64) =", round(float(np.log(64)), 4),
      "| maximum sensitivity = 1 - 1/64 =", round(1 - 1 / 64, 4))

d_k  divided  Var(q.k)  sd scores  max weight  entropy  sensitivity
   8  no          7.98       2.825      0.4240   2.0727       0.7203
   8  yes         7.98       0.999      0.1063   3.6936       0.9596
  64  no         63.99       8.000      0.7913   0.5835       0.2908
  64  yes        63.99       1.000      0.1072   3.6861       0.9610
 512  no        512.59      22.641      0.9303   0.1760       0.0998
 512  yes       512.59       1.001      0.1075   3.6848       0.9611
maximum entropy over 64 keys = ln(64) = 4.1589 | maximum sensitivity = 1 - 1/64 = 0.9844

What the numbers say. The Var(q.k) column confirms the derivation - 7.98, 63.99, 512.59 for d_k 8, 64, 512 - and the deviation of the scores goes from 2.825 to 22.641 without the division, while with the division it stays at 1.00. At d_k = 512 without scaling the maximum weight of a row is on average 0.9303 out of 64 keys and the sensitivity drops to 0.0998 against 0.9611: a factor of 9.6 of gradient lost before starting.

The causal mask: minus infinity before the softmax

A decoder generates one token at a time and cannot look at those it has not yet produced, otherwise it learns to copy the answer from the future. The delicate point is where the prohibition is applied: the forbidden scores go to −∞ before the softmax, you do not zero the weights after. With −∞ before, that term does not enter the denominator; zeroing afterwards, the denominator has already included the forbidden keys and the row is no longer an average.

np.set_printoptions(precision=3, suppress=True)
T = 5
rng = np.random.default_rng(3)
punteggi = rng.normal(size=(T, T))
causale = np.tril(np.ones((T, T), dtype=bool))            # True = I may look

giusto = softmax(np.where(causale, punteggi, -np.inf), asse=-1)
sbagliato = softmax(punteggi, asse=-1) * causale          # zeroing AFTER the softmax
print("causal mask (1 = visible):"); print(causale.astype(int))
print("BEFORE the softmax - row sums:", np.round(giusto.sum(-1), 12)); print(giusto)
print("AFTER the softmax  - row sums:", np.round(sbagliato.sum(-1), 4))
print("maximum deviation between the two matrices:", round(float(np.abs(giusto - sbagliato).max()), 4))

lunghezze = np.array([5, 3])                              # sequence 2 has 3 real tokens
valido = np.arange(T)[None, :] < lunghezze[:, None]       # (B, T) real keys
combinata = causale[None, :, :] & valido[:, None, :]      # (B, T, T)
print("valid keys in the batch:"); print(valido.astype(int))
print("combined mask of sequence 2 (causal AND padding):")
print(combinata[1].astype(int))

causal mask (1 = visible):
[[1 0 0 0 0]
 [1 1 0 0 0]
 [1 1 1 0 0]
 [1 1 1 1 0]
 [1 1 1 1 1]]
BEFORE the softmax - row sums: [1. 1. 1. 1. 1.]
[[1.    0.    0.    0.    0.   ]
 [0.859 0.141 0.    0.    0.   ]
 [0.462 0.259 0.278 0.    0.   ]
 [0.119 0.285 0.138 0.458 0.   ]
 [0.115 0.529 0.194 0.068 0.094]]
AFTER the softmax  - row sums: [0.733 0.031 0.759 0.874 1.   ]
maximum deviation between the two matrices: 0.8317
valid keys in the batch:
[[1 1 1 1 1]
 [1 1 1 0 0]]
combined mask of sequence 2 (causal AND padding):
[[1 0 0 0 0]
 [1 1 0 0 0]
 [1 1 1 0 0]
 [1 1 1 0 0]
 [1 1 1 0 0]]

What the code does. np.tril builds the lower triangular: row i sees itself and the past. The second row of the wrong matrix sums to 0.031 instead of 1 - that token would have given 97% of the weight to the future, thrown away without redistributing it - and the maximum deviation between the two matrices is 0.8317. Only the last row coincides, because for it there is no future.

Causal mask and padding mask are two different things, and confusing them is a silent error. The causal one depends only on positions and is the same for the whole batch: (T, T). The padding one depends on the content - real tokens against filler - and changes from sequence to sequence: (B, T) on the keys. They are combined with AND, and if you forget the second nothing breaks: the results get worse in a way that depends on how you ordered the batch.

Multiple heads: why h heads of d/h and not one of d

A single attention produces, for each token, a single convex combination of the values: a token that has to look at the subject of the sentence and at the last number mentioned has one set of weights for two tasks. Multi-head attention (§3.2.2 of arXiv:1706.03762) runs h attentions in parallel on distinct subspaces and concatenates the results.

The surprising choice is the dimension: each head works in d/h dimensions, not in d. With one full head you need W_Q, W_K, W_V, W_O of shape (d, d): 4d². With h heads each has matrices (d, d/h), the h of a family stack into a single (d, d), and W_O remains (d, d) too. Total 4d², identical: multiple heads are free in parameters.

def multi_testa(X, W_Q, W_K, W_V, W_O, h, maschera=None):
    B, T, d = X.shape
    d_h = d // h
    def dividi(M):                       # (B, T, d) -> (B, h, T, d/h)
        return M.reshape(B, T, h, d_h).transpose(0, 2, 1, 3)
    Q, K, V = dividi(X @ W_Q), dividi(X @ W_K), dividi(X @ W_V)
    punteggi = Q @ K.transpose(0, 1, 3, 2) / np.sqrt(d_h)
    if maschera is not None:
        punteggi = np.where(maschera, punteggi, -np.inf)
    pesi = softmax(punteggi, asse=-1)                       # softmax OVER THE KEYS
    ctx = pesi @ V                                          # (B, h, T, d/h)
    ctx = ctx.transpose(0, 2, 1, 3).reshape(B, T, d)        # stitches the heads back together
    return ctx @ W_O, pesi

rng = np.random.default_rng(11)
B, T, d = 1, 6, 64
X = rng.normal(size=(B, T, d))
W = [rng.normal(size=(d, d)) / np.sqrt(d) for _ in range(4)]
for h in (1, 2, 4, 8):
    Y, pesi = multi_testa(X, *W, h=h)
    print(f"h={h:2d}  d_head={d//h:3d}  weights {str(pesi.shape):16s} output {Y.shape}"
          f"  4*d^2 = {4*d*d}")

perm = rng.permutation(T)                                   # no mask, no position
Y, _ = multi_testa(X, *W, h=4)
Yp, _ = multi_testa(X[:, perm, :], *W, h=4)
print("max |attn(PX) - P attn(X)| =", f"{np.abs(Yp - Y[:, perm, :]).max():.3e}")
print("max |attn(PX) -   attn(X)| =", f"{np.abs(Yp - Y).max():.3e}")

h= 1  d_head= 64  weights (1, 1, 6, 6)     output (1, 6, 64)  4*d^2 = 16384
h= 2  d_head= 32  weights (1, 2, 6, 6)     output (1, 6, 64)  4*d^2 = 16384
h= 4  d_head= 16  weights (1, 4, 6, 6)     output (1, 6, 64)  4*d^2 = 16384
h= 8  d_head=  8  weights (1, 8, 6, 6)     output (1, 6, 64)  4*d^2 = 16384
max |attn(PX) - P attn(X)| = 4.441e-16
max |attn(PX) -   attn(X)| = 1.849e+00

What the code does. reshape(B, T, h, d_h) splits the last axis without moving a number in memory; transpose(0, 2, 1, 3) brings the head in front of time. On the way back you walk the inverse path in the same order: if you swap the two steps you get a tensor of the right shape with the numbers mixed between different heads, a bug that raises no exception.

Where position lives

There is a big hole in the formula, and it shows in two lines. Let Π be a permutation of the rows: then Q → ΠQ, K → ΠK, V → ΠV, hence S → Π S Πᵀ and P → Π P Πᵀ, because the softmax acts row by row; the output is P V → Π P Πᵀ Π V = Π P V, because Πᵀ Π = I. Permuting the input permutes the output and nothing else: attention is equivariant, and for it a sentence is a bag of words. ∎

Sinusoidal position. The choice of the original paper (§3.5): to the embedding you add a fixed vector of sines and cosines at geometrically scaled frequencies, PE(pos, 2i) = sin(pos / 10000^{2i/d}) and PE(pos, 2i+1) = cos(pos / 10000^{2i/d}), where the slow waves distinguish distant positions. Zero parameters, defined behavior beyond the length seen. Learned position. A table (T_max, d) trained like any embedding, which the same paper gives as "nearly identical": it is the one in the experiment, it costs T_max · d parameters, and at position T_max + 1 you have nothing.

RoPE, rotary position. Instead of adding, it rotates: the pairs of components of q and k turn by an angle proportional to the position, before the dot product (Su and colleagues, RoFormer: Enhanced Transformer with Rotary Position Embedding, arXiv:2104.09864). The dot product between a query at m and a key at n then depends only on m − n, that is on the relative distance: it is the de facto choice of recent models.

The complete block, and the count of its parameters

Attention alone is not enough: position by position it is a linear transformation of the values. You need a nonlinearity that works inside each position, you need residual connections so the stack can be deep, and you need layer normalization (Ba, Kiros and Hinton, Layer Normalization, arXiv:1607.06450). Together they are the Transformer block.

The feed-forward network is a two-layer network applied independently to each position: it expands from d to 4d, applies a nonlinearity, returns to d; the factor four is the paper's choice (d = 512, d_ff = 2048). The residual connections make the output of each sub-layer x + f(x): the gradient always finds a path that does not go through f, and a stack of 24 blocks remains trainable.

Here the choice that distinguishes the two families opens up. In the post-LN of the original paper you compute x_new = LayerNorm(x + f(x)), with the normalization after the sum; in the pre-LN you compute x_new = x + f(LayerNorm(x)) and the residual path stays clean. Today pre-LN is used, and the reason is documented: Xiong and colleagues (arXiv:2002.04745) show that in post-LN the gradients near the layers close to the output grow with depth, and you need a small step and a warm-up.

Let us count the parameters in closed form: four (d, d) matrices with biases give 4d² + 4d, the feed-forward 8d² + 5d, the two layer norms 4d. Total 12d² + 13d, and the number of heads does not appear. To check it against the inventory of tensors you need the real model, which is also needed by the experiment further on. I write it here in full, in four pieces: parameters and layer norm, attention with its backward pass, block and loss, optimizer.

Parameters, layer norm, heads

inizializza builds a dictionary name → array, with the names that the count below will use (b0.Wq, b0.g1, and so on). The matrices are Gaussians divided by the square root of the number of rows, the biases start at zero, γ at one; embedding and positions have scale 0.1, and the position is a learned table (t_max, d), the one in the experiment. ln_indietro uses the closed form of the layer norm gradient, dx = inv · (dx̂ − mean(dx̂) − x̂ · mean(dx̂ · x̂)): I have not derived it in the text, and the proof that it is right is the numerical check further down. teste and unisci are the two steps of the reshape and transpose seen before, one the inverse of the other.

def inizializza(vocab, d, h, n_strati, d_ff, t_max, rng):
    """All the decoder's tensors in a dictionary: name -> array."""
    n = lambda *forma: rng.normal(size=forma) / np.sqrt(forma[0])
    p = {"E": 0.1 * rng.normal(size=(vocab, d)),      # token embedding
         "P": 0.1 * rng.normal(size=(t_max, d))}      # learned position
    for l in range(n_strati):
        q = f"b{l}."
        for nome in ("Wq", "Wk", "Wv", "Wo"):
            p[q + nome] = n(d, d)
            p[q + "b" + nome[1].lower()] = np.zeros(d)
        p[q + "g1"], p[q + "be1"] = np.ones(d), np.zeros(d)
        p[q + "g2"], p[q + "be2"] = np.ones(d), np.zeros(d)
        p[q + "W1"], p[q + "b1"] = n(d, d_ff), np.zeros(d_ff)
        p[q + "W2"], p[q + "b2"] = n(d_ff, d), np.zeros(d)
    p["gf"], p["bef"] = np.ones(d), np.zeros(d)
    p["Wout"], p["bout"] = n(d, vocab), np.zeros(vocab)
    return p

def ln_avanti(x, g, b, eps=1e-5):
    inv = 1.0 / np.sqrt(x.var(-1, keepdims=True) + eps)
    xh = (x - x.mean(-1, keepdims=True)) * inv
    return g * xh + b, (xh, inv, g)

def ln_indietro(dy, c):
    xh, inv, g = c
    dxh = dy * g
    dx = inv * (dxh - dxh.mean(-1, keepdims=True) - xh * (dxh * xh).mean(-1, keepdims=True))
    return dx, (dy * xh).sum((0, 1)), dy.sum((0, 1))

def teste(M, h):                      # (B, T, d) -> (B, h, T, d/h)
    B, T, d = M.shape
    return M.reshape(B, T, h, d // h).transpose(0, 2, 1, 3)

def unisci(M):                        # (B, h, T, d/h) -> (B, T, d)
    B, h, T, dh = M.shape
    return M.transpose(0, 2, 1, 3).reshape(B, T, h * dh)

Attention, forward and backward

The forward pass keeps what the backward pass needs: the input, Q, K, V, the weights P and the concatenated context. The mask is applied with −∞ before the softmax, as in the dedicated section. Backwards the line that matters is dS = P · (dP − Σ dP · P): it is the softmax derivative row by row, the same pᵢ(δᵢⱼ − pⱼ) of the section on scale, contracted with the gradient dP that comes from the product with V. Where the mask put P = 0 also dS is zero, with no extra work: a position that did not influence the result receives no gradient.

def att_avanti(x, p, q, h, scala):
    Q, K, V = (teste(x @ p[q + w] + p[q + "b" + w[1].lower()], h) for w in ("Wq", "Wk", "Wv"))
    T = x.shape[1]
    S = Q @ K.transpose(0, 1, 3, 2)
    if scala:
        S = S / np.sqrt(Q.shape[-1])
    S = np.where(np.tril(np.ones((T, T), dtype=bool)), S, -np.inf)   # -inf BEFORE the softmax
    P = softmax(S, -1)
    C = unisci(P @ V)
    return C @ p[q + "Wo"] + p[q + "bo"], (x, Q, K, V, P, C)

def att_indietro(dy, c, p, q, h, g, scala):
    x, Q, K, V, P, C = c
    flat = lambda a: a.reshape(-1, a.shape[-1])
    g[q + "Wo"] += flat(C).T @ flat(dy); g[q + "bo"] += dy.sum((0, 1))
    dC = teste(dy @ p[q + "Wo"].T, h)
    dP = dC @ V.transpose(0, 1, 3, 2)
    dV = P.transpose(0, 1, 3, 2) @ dC
    dS = P * (dP - (dP * P).sum(-1, keepdims=True))
    if scala:
        dS = dS / np.sqrt(Q.shape[-1])
    dQ, dK = dS @ K, dS.transpose(0, 1, 3, 2) @ Q
    dx = 0
    for w, dM in (("Wq", dQ), ("Wk", dK), ("Wv", dV)):
        dM = unisci(dM)
        g[q + w] += flat(x).T @ flat(dM); g[q + "b" + w[1].lower()] += dM.sum((0, 1))
        dx = dx + dM @ p[q + w].T
    return dx

Block, loss, gradient

avanti builds the decoder, in the two variants: post_ln false gives the pre-LN, true the post-LN, and in both there is a final layer norm before the projection onto the vocabulary, so the number of parameters stays the same. perdita_e_gradienti is the cross-entropy weighted by pesa, divided by the sum of the weights: the initial gradient on the logits is (P − one-hot) · weight / sum of weights, and then the block is traversed in reverse order. For the embedding np.add.at is used and not an assignment, because the same token can appear several times in the batch and the contributions add up.

def avanti(p, tok, n_strati, h, scala=True, post_ln=False):
    B, T = tok.shape
    x = p["E"][tok] + p["P"][:T]
    cache = []
    for l in range(n_strati):
        q = f"b{l}."
        if post_ln:
            a, ca = att_avanti(x, p, q, h, scala)
            x1, cn1 = ln_avanti(x + a, p[q + "g1"], p[q + "be1"])
            z = x1 @ p[q + "W1"] + p[q + "b1"]; f = np.maximum(z, 0) @ p[q + "W2"] + p[q + "b2"]
            x2, cn2 = ln_avanti(x1 + f, p[q + "g2"], p[q + "be2"])
            cache.append((ca, cn1, x1, z, cn2))
        else:
            n1, cn1 = ln_avanti(x, p[q + "g1"], p[q + "be1"])
            a, ca = att_avanti(n1, p, q, h, scala)
            x1 = x + a
            n2, cn2 = ln_avanti(x1, p[q + "g2"], p[q + "be2"])
            z = n2 @ p[q + "W1"] + p[q + "b1"]; f = np.maximum(z, 0) @ p[q + "W2"] + p[q + "b2"]
            x2 = x1 + f
            cache.append((ca, cn1, n2, z, cn2))
        x = x2
    xf, cf = ln_avanti(x, p["gf"], p["bef"])
    return xf @ p["Wout"] + p["bout"], (tok, cache, xf, cf)

def perdita_e_gradienti(p, tok, ber, pesa, n_strati, h, scala=True, post_ln=False):
    logit, (tok, cache, xf, cf) = avanti(p, tok, n_strati, h, scala, post_ln)
    P = softmax(logit, -1)
    B, T, V = P.shape
    pr = np.take_along_axis(P, ber[..., None], -1)[..., 0]
    perdita = -(pesa * np.log(pr + 1e-12)).sum() / pesa.sum()
    dl = P.copy(); np.put_along_axis(dl, ber[..., None], np.take_along_axis(dl, ber[..., None], -1) - 1, -1)
    dl *= (pesa / pesa.sum())[..., None]
    g = {k: np.zeros_like(v) for k, v in p.items()}
    flat = lambda a: a.reshape(-1, a.shape[-1])
    g["Wout"] += flat(xf).T @ flat(dl); g["bout"] += dl.sum((0, 1))
    dx, g["gf"], g["bef"] = ln_indietro(dl @ p["Wout"].T, cf)
    for l in reversed(range(n_strati)):
        q = f"b{l}."
        ca, cn1, nn, z, cn2 = cache[l]
        if post_ln:
            dx2, g[q + "g2"], g[q + "be2"] = ln_indietro(dx, cn2)       # x2 = LN(x1 + f)
            x1 = nn
            g[q + "W2"] += flat(np.maximum(z, 0)).T @ flat(dx2); g[q + "b2"] += dx2.sum((0, 1))
            dz = (dx2 @ p[q + "W2"].T) * (z > 0)
            g[q + "W1"] += flat(x1).T @ flat(dz); g[q + "b1"] += dz.sum((0, 1))
            dx1 = dx2 + dz @ p[q + "W1"].T
            dsum, g[q + "g1"], g[q + "be1"] = ln_indietro(dx1, cn1)     # x1 = LN(x + a)
            dx = dsum + att_indietro(dsum, ca, p, q, h, g, scala)
        else:
            g[q + "W2"] += flat(np.maximum(z, 0)).T @ flat(dx); g[q + "b2"] += dx.sum((0, 1))
            dz = (dx @ p[q + "W2"].T) * (z > 0)
            g[q + "W1"] += flat(nn).T @ flat(dz); g[q + "b1"] += dz.sum((0, 1))
            dn2, g[q + "g2"], g[q + "be2"] = ln_indietro(dz @ p[q + "W1"].T, cn2)
            dx1 = dx + dn2
            dn1 = att_indietro(dx1, ca, p, q, h, g, scala)
            dn1, g[q + "g1"], g[q + "be1"] = ln_indietro(dn1, cn1)
            dx = dx1 + dn1
    np.add.at(g["E"], tok, dx)
    g["P"][:T] += dx.sum(0)
    return perdita, g

Optimizer

Adam with the usual values from the optimizer's original paper (β₁ = 0.9, β₂ = 0.999, ε = 10⁻⁸): one estimate of the mean and one of the square of the gradient, both corrected for the initial bias. It is the only piece that is not specific to the Transformer.

def stato_adam(p):
    return {k: [np.zeros_like(v), np.zeros_like(v)] for k, v in p.items()}

def adam(p, g, stato, lr, t, b1=0.9, b2=0.999, eps=1e-8):
    for k in p:
        m, v = stato[k]
        m *= b1; m += (1 - b1) * g[k]
        v *= b2; v += (1 - b2) * g[k] ** 2
        p[k] -= lr * (m / (1 - b1 ** t)) / (np.sqrt(v / (1 - b2 ** t)) + eps)

With the model in hand we can do the count that the formula promised. The listing creates a block for three widths (64, 256, 512), sums the sizes of every tensor whose name starts with b0. and compares them with 12d² + 13d; then, for d = 512, it breaks the total down entry by entry, to see where the parameters are.

def parametri_blocco(p, l):
    """Sums the actual sizes of all the tensors of block l."""
    q = f"b{l}."
    return sum(v.size for k, v in p.items() if k.startswith(q))

rng = np.random.default_rng(1)
print(" d    d_ff   formula 12d^2+13d   actual count")
for d in (64, 256, 512):
    p = inizializza(vocab=32, d=d, h=4, n_strati=1, d_ff=4 * d, t_max=16, rng=rng)
    print(f"{d:4d}  {4*d:6d}   {12*d*d + 13*d:17d}   {parametri_blocco(p, 0):12d}")

d = 512
voci = [("W_Q, W_K, W_V, W_O", 4 * d * d), ("biases of the four projections", 4 * d),
        ("W_1 (d -> 4d)", 4 * d * d), ("b_1", 4 * d),
        ("W_2 (4d -> d)", 4 * d * d), ("b_2", d),
        ("two layer norms (gamma and beta)", 4 * d)]
print("breakdown for d = 512:")
for nome, n in voci:
    print(f"  {nome:34s} {n:10d}")
print(f"  {'TOTAL':34s} {sum(n for _, n in voci):10d}")

 d    d_ff   formula 12d^2+13d   actual count
  64     256               49984          49984
 256    1024              789760         789760
 512    2048             3152384        3152384
breakdown for d = 512:
  W_Q, W_K, W_V, W_O                    1048576
  biases of the four projections           2048
  W_1 (d -> 4d)                         1048576
  b_1                                      2048
  W_2 (4d -> d)                         1048576
  b_2                                       512
  two layer norms (gamma and beta)         2048
  TOTAL                                 3152384

What the numbers say. Closed-form count and inventory coincide to the digit: 49,984, 789,760, 3,152,384. The breakdown for d = 512 shows three blocks of 1,048,576 parameters - attention, W_1, W_2 - plus 6,656 parameters of biases and scales, 0.21% of the total. This is why in back-of-the-envelope estimates you write 12d² and ignore the rest.

You have already derived the chain rule in the lesson on backpropagation; what matters here is how you verify a hand-written gradient. You perturb a parameter by ±ε and compare (L(θ+ε) − L(θ−ε)) / 2ε with the analytic gradient. The listing does it on a tiny model (d = 16, 2 layers, 3 sequences), on the first six coordinates of each tensor where the gradient exceeds 10⁻⁶, and reports the worst relative deviation.

def scarto_massimo(post_ln, eps, seme=0):
    """Compares the analytic gradient with the centred finite difference, tensor by tensor."""
    rng = np.random.default_rng(seme)
    p = inizializza(21, 16, 4, 2, 64, 13, rng)
    tok = rng.integers(0, 21, size=(3, 12))
    ber = rng.integers(0, 21, size=(3, 12))
    pesa = np.zeros((3, 12)); pesa[:, 6:] = 1.0
    _, g = perdita_e_gradienti(p, tok, ber, pesa, 2, 4, True, post_ln)
    peggiore = 0.0
    for nome in p:
        for i in np.argwhere(np.abs(g[nome]) > 1e-6)[:6]:       # only where the gradient is measurable
            i = tuple(i); v0 = p[nome][i]
            p[nome][i] = v0 + eps; piu = perdita_e_gradienti(p, tok, ber, pesa, 2, 4, True, post_ln)[0]
            p[nome][i] = v0 - eps; meno = perdita_e_gradienti(p, tok, ber, pesa, 2, 4, True, post_ln)[0]
            p[nome][i] = v0
            numerico = (piu - meno) / (2 * eps)
            peggiore = max(peggiore, abs(numerico - g[nome][i]) / max(abs(numerico), abs(g[nome][i])))
    return peggiore

for post_ln in (False, True):
    for eps in (1e-6, 1e-4):
        print(f"{'post-LN' if post_ln else 'pre-LN '}  eps={eps:.0e}  maximum relative deviation {scarto_massimo(post_ln, eps):.2e}")

pre-LN   eps=1e-06  maximum relative deviation 2.56e-06
pre-LN   eps=1e-04  maximum relative deviation 9.94e-02
post-LN  eps=1e-06  maximum relative deviation 4.90e-05
post-LN  eps=1e-04  maximum relative deviation 6.60e-03

What the numbers say. With ε = 10⁻⁶ the maximum relative deviation is 2.56 · 10⁻⁶ for pre-LN and 4.90 · 10⁻⁵ for post-LN: the hand-written gradient is right in both variants, within the limits of a sample of coordinates. With ε = 10⁻⁴ the deviation rises to 9.94 · 10⁻² and 6.60 · 10⁻³. If the check gets worse when you increase ε, the problem is ε: the finite difference is an approximation of a derivative, and if the perturbation crosses a corner of the ReLU it stops measuring it. This last one is the explanation consistent with the behavior; I have not verified it coordinate by coordinate.

The count that decides everything: O(T²)

The score matrix has T² elements, and this sentence explains most of the engineering of recent language models. The time for Q Kᵀ and P V is O(T² · d), and the memory to keep P in the forward pass - it is needed in the backward one - is O(T²) per head and per batch element.

import numpy as np, time

def attenzione_causale(Q, K, V):
    T, d_k = Q.shape
    S = (Q @ K.T) / np.sqrt(d_k)
    S = np.where(np.tril(np.ones((T, T), dtype=bool)), S, -np.inf)
    return softmax(S, asse=-1) @ V

def tempo_minimo(Q, K, V, prove):
    """The MINIMUM over several trials: machine load can lengthen a time,
    never shorten it, so the minimum is the least polluted estimate."""
    attenzione_causale(Q, K, V)                        # warm-up round
    misure = []
    for _ in range(prove):
        t0 = time.perf_counter()
        attenzione_causale(Q, K, V)
        misure.append((time.perf_counter() - t0) * 1000)
    return min(misure)

rng = np.random.default_rng(0)
d_k = 64
print(" T     trials   time (ms)   TxT matrix (MiB)   Q,K,V together (MiB)")
tempi, lunghezze, prove = [], (128, 512, 2048), (400, 60, 12)
for T, n in zip(lunghezze, prove):
    Q, K, V = (rng.normal(size=(T, d_k)) for _ in range(3))
    ms = tempo_minimo(Q, K, V, n)
    tempi.append(ms)
    print(f"{T:5d} {n:8d} {ms:11.2f} {T*T*8/2**20:17.2f} {3*T*d_k*8/2**20:22.2f}")

pendenza = np.polyfit(np.log(np.array(lunghezze, float)), np.log(np.array(tempi)), 1)[0]
print(f"slope of log(time) against log(T): {pendenza:.2f}  (a T^2 law gives 2)")
for i in range(1, 3):
    print(f"  from T={lunghezze[i-1]} to T={lunghezze[i]}: length x4, "
          f"time x{tempi[i]/tempi[i-1]:.1f}, matrix memory x16")

 T     trials   time (ms)   TxT matrix (MiB)   Q,K,V together (MiB)
  128      400        1.28              0.12                   0.19
  512       60       25.92              2.00                   0.75
 2048       12      367.22             32.00                   3.00
slope of log(time) against log(T): 2.04  (a T^2 law gives 2)
  from T=128 to T=512: length x4, time x20.2, matrix memory x16
  from T=512 to T=2048: length x4, time x14.2, matrix memory x16

What the numbers say. The output shown is one of three consecutive runs, on a machine shared with other processes, and that is why the times are intervals: 1.28-1.38 ms at T = 128, 24.7-25.9 ms at 512, 340-367 ms at 2048; the slope in the logarithmic plane is between 1.99 and 2.04, against the exact 2 of a quadratic law. Quadrupling the length multiplies the time by 17.8-20.2 in the first jump and by 13.2-14.2 in the second. That the second stays below 16 may depend on the processor cache, but I have not verified it. The last column is the memory of Q, K and V together, 3.00 MiB: the score matrix weighs more than ten times the tensors that generated it.

The numbers are in float64; in float16 they are divided by four, but the dependence stays quadratic - at T = 32,768 a single head takes 2 GiB for the matrix. This is why there are implementations that never materialize P but compute it in blocks, keeping only the partial sums of the softmax: it is what torch.nn.functional.scaled_dot_product_attention does with its fused kernel, as the PyTorch documentation describes. The formula is the same (the numerical comparison is above); what changes is where the intermediate numbers live.

The key-value cache: what is kept, and how much it weighs

In generation the model produces one token at a time and redoes the forward pass over everything it has written, but keys and values of the positions already seen do not change: the k₇ and v₇ of the token at position 7 are the same when the model is generating token 300. The only new thing at step t is the query qₜ.

From here the formula for the memory of the key-value cache, derived from what is really kept: two matrices, hence the factor 2; for each layer, because each layer has its own projections; for each key-value head; d_head components for each of the T positions, for the batch and for the bytes of the type. Note what is not there: not the queries, not the attention weights, not the feed-forward activations.

The listing builds a four-layer toy model, attention and residual only, with random weights: it serves to measure time, not quality. The two functions produce the same result, and the maximum deviation printed proves it. rinormalizza appears only because, feeding the output back in place of a token, the residual sums would grow at every step: it is not part of the cache mechanism. The prompt is 32 rows of noise, and 128 vectors are generated.

import time
import numpy as np

rng = np.random.default_rng(5)
d, h, n_strati = 256, 8, 4
d_h = d // h
Ws = [tuple(rng.normal(size=(d, d)) / np.sqrt(d) for _ in range(4)) for _ in range(n_strati)]

def dividi(M):                                      # (T, d) -> (h, T, d_h)
    return M.reshape(M.shape[0], h, d_h).transpose(1, 0, 2)

def prefisso_intero(X):
    """No cache: at every token K and V of the whole prefix are recomputed."""
    for Wq, Wk, Wv, Wo in Ws:
        T = X.shape[0]
        Q, K, V = dividi(X @ Wq), dividi(X @ Wk), dividi(X @ Wv)
        S = (Q @ K.transpose(0, 2, 1)) / np.sqrt(d_h)
        S = np.where(np.tril(np.ones((T, T), dtype=bool)), S, -np.inf)
        C = (softmax(S, -1) @ V).transpose(1, 0, 2).reshape(T, d)
        X = X + C @ Wo
    return X

def un_token(x, cache):
    """With cache: q, k, v are computed for the new token; k and v are appended."""
    for l, (Wq, Wk, Wv, Wo) in enumerate(Ws):
        q = (x @ Wq).reshape(h, 1, d_h)
        k = (x @ Wk).reshape(h, 1, d_h)
        v = (x @ Wv).reshape(h, 1, d_h)
        K = k if cache[l][0] is None else np.concatenate([cache[l][0], k], axis=1)
        V = v if cache[l][1] is None else np.concatenate([cache[l][1], v], axis=1)
        cache[l] = [K, V]
        S = (q @ K.transpose(0, 2, 1)) / np.sqrt(d_h)   # (h, 1, t): no mask needed
        c = (softmax(S, -1) @ V).reshape(1, d)
        x = x + c @ Wo
    return x

def memoria_cache(n_strati, n_teste_kv, d_testa, T, batch, byte):
    return 2 * n_strati * n_teste_kv * d_testa * T * batch * byte

def rinormalizza(x):
    """Keeps the toy in a stable numerical scale: without it, the residual sums grow."""
    return x / np.linalg.norm(x, axis=-1, keepdims=True) * np.sqrt(d)

def genera_senza_cache(prompt, n_nuovi):
    X = prompt.copy()
    for _ in range(n_nuovi):
        X = np.vstack([X, rinormalizza(prefisso_intero(X)[-1])])
    return X

def genera_con_cache(prompt, n_nuovi):
    cache = [[None, None] for _ in Ws]
    for riga in prompt[:-1]:                         # fills the cache with the prompt
        un_token(riga, cache)
    x, righe = prompt[-1], list(prompt)
    for _ in range(n_nuovi):
        x = rinormalizza(un_token(x, cache)[0])
        righe.append(x)
    return np.array(righe)

prompt = rng.normal(size=(32, d))
t0 = time.perf_counter(); senza = genera_senza_cache(prompt, 128); t_senza = time.perf_counter() - t0
t0 = time.perf_counter(); con = genera_con_cache(prompt, 128); t_con = time.perf_counter() - t0
print(f"{n_strati} layers, {h} heads, d = {d}, prompt of 32 tokens, 128 tokens generated")
print(f"  maximum deviation between the two generations: {np.abs(senza - con).max():.1e}")
print(f"  without cache: {t_senza:6.2f} s")
print(f"  with cache:    {t_con:6.2f} s")
print(f"  ratio:         {t_senza / t_con:6.1f}x")
print("memory = 2 * n_strati * n_teste_kv * d_testa * T * batch * byte")
print(f"  this model, T = 160, batch 1, float64: {memoria_cache(4, 8, 32, 160, 1, 8) / 2**20:.2f} MiB")
for kv in (32, 1):
    print(f"  32 layers, d_head 128, T = 4096, batch 1, float16, {kv:2d} K,V pairs: "
          f"{memoria_cache(32, kv, 128, 4096, 1, 2) / 2**30:.3f} GiB")

4 layers, 8 heads, d = 256, prompt of 32 tokens, 128 tokens generated
  maximum deviation between the two generations: 6.9e-15
  without cache:   8.26 s
  with cache:      0.15 s
  ratio:           56.6x
memory = 2 * n_strati * n_teste_kv * d_testa * T * batch * byte
  this model, T = 160, batch 1, float64: 2.50 MiB
  32 layers, d_head 128, T = 4096, batch 1, float16, 32 K,V pairs: 2.000 GiB
  32 layers, d_head 128, T = 4096, batch 1, float16,  1 K,V pairs: 0.062 GiB

What the numbers say. The maximum deviation between the two generations is 6.9 · 10⁻¹⁵: the cache does not change the result, it changes the cost. On the toy model, generating 128 vectors after a prompt of 32 costs 8.3-13.0 s without cache and 0.15-0.41 s with the cache, in three runs on a shared machine: the ratio goes from about 30 to about 65 (the line above shows one, 56.6). It grows with the generated length, because without the cache at every step keys and values of the whole prefix are recomputed. On a realistic configuration (32 layers, 32 key-value pairs, d_head = 128, T = 4096, batch 1, float16) the formula gives 2.000 GiB; with a single shared K, V pair, the multi-query attention of Shazeer (arXiv:1911.02150), 0.062 GiB. They are formulas, not measurements on a real model.

The experiment: copying and sorting sequences

To see attention at work you need a well-chosen task, and the choice is methodological. No text: the perplexity falls and you cannot say whether the model has learned the structure or the frequency of the words. No images: there is no order to reconstruct. The right task is synthetic and verifiable at a glance: given an alphabet of 20 symbols and a separator, copy the sequence after the separator or sort it. It requires exactly what attention can do: take a value at a precise position.

The symbols are already integers, that is already tokens: where they come from was shown by the lesson on tokenisation. The sequence given to the model is [s₁ … s_L, SEP, answer₁ … answer_L], of length 2L + 1, and the model predicts the next token at each position. The cost function is the cross-entropy of the lesson on cross-entropy and KL divergence, computed only on the L positions of the answer.

import numpy as np
import numpy as np

ALFABETO = 20
SEP = ALFABETO            # the 21st symbol separates question and answer
VOCAB = ALFABETO + 1

def lotto(rng, B, L, compito):
    s = rng.integers(0, ALFABETO, size=(B, L))
    risposta = np.sort(s, axis=1) if compito == "sort" else s
    seq = np.concatenate([s, np.full((B, 1), SEP), risposta], axis=1)   # (B, 2L+1)
    tok = seq[:, :-1]                      # input:  positions 0 .. 2L-1
    ber = seq[:, 1:]                       # target: positions 1 .. 2L
    pesa = np.zeros_like(tok, dtype=float)
    pesa[:, L:] = 1.0                      # the loss counts only the answer, L tokens
    return tok, ber, pesa

def addestra(compito, L, passi=1500, B=32, d=32, h=4, n_strati=2, lr=3e-3, seme=0,
             scala=True, post_ln=False, posizione=True):
    rng = np.random.default_rng(seme)
    p = inizializza(VOCAB, d, h, n_strati, 4 * d, 2 * L + 1, rng)
    if not posizione:
        p["P"][:] = 0.0                    # table set to zero...
    stato = stato_adam(p)
    for t in range(1, passi + 1):
        tok, ber, pesa = lotto(rng, B, L, compito)
        perdita, g = perdita_e_gradienti(p, tok, ber, pesa, n_strati, h, scala, post_ln)
        if not posizione:
            g["P"][:] = 0.0                # ...and never updated
        adam(p, g, stato, lr, t)
    return p, perdita

def valuta(p, compito, L, n=500, h=4, n_strati=2, seme=999, scala=True, post_ln=False):
    """True generation: the answer is built one token at a time, without targets."""
    rng = np.random.default_rng(seme)
    s = rng.integers(0, ALFABETO, size=(n, L))
    atteso = np.sort(s, axis=1) if compito == "sort" else s
    corrente = np.concatenate([s, np.full((n, 1), SEP)], axis=1)
    for _ in range(L):
        logit, _ = avanti(p, corrente, n_strati, h, scala, post_ln)
        corrente = np.concatenate([corrente, logit[:, -1, :].argmax(-1)[:, None]], axis=1)
    prodotto = corrente[:, L + 1:]
    return (prodotto == atteso).mean(), (prodotto == atteso).all(axis=1).mean()

The task is declared, the model is the one written above; what remains is the bookkeeping work. The second listing launches the 23 trainings (two tasks, three lengths, three seeds, plus five ablations) in separate processes, one thread each, and prints the tables of this and of the next two sections. I verified on one case that the result of a training does not depend on how it is launched: the final loss of copy with L = 6 and seed 0 is identical, at the digits printed (0.0002, accuracy 1.000), launched alone and inside the pool. Alone, a training took 35 s in the run I timed.

import os
os.environ["OMP_NUM_THREADS"] = "1"              # one process per training run, one thread each
os.environ["OPENBLAS_NUM_THREADS"] = "1"
from multiprocessing import Pool

SEQUENZA_MAPPA = np.array([[17, 3, 0, 12, 7, 9, SEP, 17, 3, 0, 12, 7]])

def prova(spec):
    nome, compito, L, seme, opzioni = spec
    p, perdita = addestra(compito, L, seme=seme, **opzioni)
    ok = {k: opzioni[k] for k in ("scala", "post_ln") if k in opzioni}
    per_token, per_sequenza = valuta(p, compito, L, **ok)
    esito = {"perdita": perdita, "token": per_token, "sequenza": per_sequenza}
    if nome in ("full", "no position") and compito == "copy" and L == 6:
        rng = np.random.default_rng(999)         # the same 500 sequences as valuta()
        s = rng.integers(0, ALFABETO, size=(500, L))
        corrente = np.concatenate([s, np.full((500, 1), SEP)], axis=1)
        for _ in range(L):
            logit, _ = avanti(p, corrente, 2, 4)
            corrente = np.concatenate([corrente, logit[:, -1, :].argmax(-1)[:, None]], axis=1)
        giusta = (corrente[:, L + 1:] == s).all(axis=1)
        distinti = np.array([len(set(r)) == L for r in s])
        esito["distinti"] = (int(distinti.sum()), giusta[distinti].mean(), giusta[~distinti].mean())
    if nome == "full" and compito == "copy" and L == 6 and seme == 0:
        _, (_, cache, _, _) = avanti(p, SEQUENZA_MAPPA, 2, 4)
        esito["attenzione"] = [c[0][4][0] for c in cache]          # per layer: (heads, T, T)
    return (nome, compito, L, seme), esito

if __name__ == "__main__":
    lavori = [("full", c, L, seme, {}) for c in ("copy", "sort") for L in (4, 6, 8) for seme in (0, 1, 2)]
    lavori += [("no scaling", "copy", 6, 0, {"scala": False}),
               ("no position", "copy", 6, 0, {"posizione": False}),
               ("post-LN", "copy", 6, 0, {"post_ln": True}),
               ("pre-LN, high lr", "copy", 6, 0, {"lr": 1e-2}),
               ("post-LN, high lr", "copy", 6, 0, {"lr": 1e-2, "post_ln": True})]
    with Pool(16) as pool:
        r = dict(pool.map(prova, lavori, chunksize=1))

    n_par = sum(v.size for v in inizializza(VOCAB, 32, 4, 2, 128, 13, np.random.default_rng(0)).values())
    print("Configuration: d=32, 4 heads (d_head=8), 2 layers, d_ff=128, batch 32,")
    print("Adam lr=3e-3, 1500 steps, alphabet of 20 symbols + separator (vocabulary 21).")
    print(f"Model parameters: {n_par}\n")
    print("task    L   per-token accuracy         per-sequence accuracy")
    print("            (mean over 3 seeds)        (mean over 3 seeds)")
    for compito in ("copy", "sort"):
        for L in (4, 6, 8):
            e = [r[("full", compito, L, s)] for s in range(3)]
            t, q = [x["token"] for x in e], [x["sequenza"] for x in e]
            print(f"{compito:<8} {L:>2}   {np.mean(t):.3f}  (min {min(t):.3f})        "
                  f"{np.mean(q):.3f}  (min {min(q):.3f})")

    print("\nAblations on the copy task, L = 6, seed 0, same hyperparameters unless stated:")
    for nome in ("full", "no scaling", "no position", "post-LN", "pre-LN, high lr", "post-LN, high lr"):
        e = r[(nome, "copy", 6, 0)]
        print(f"  {nome:<18} final loss {e['perdita']:.4f}   token {e['token']:.3f}   sequence {e['sequenza']:.3f}")

    print("\nCopy without position, 500 test sequences:")
    n_dist, a_dist, a_rip = r[("no position", "copy", 6, 0)]["distinti"]
    print(f"  on sequences whose symbols are all distinct:   {a_dist:.3f}   ({n_dist} cases)")
    print(f"  on sequences with at least one repeated symbol: {a_rip:.3f}   ({500 - n_dist} cases)")
    print(f"  share of sequences with distinct symbols: measured {n_dist / 500:.3f}, "
          f"expected 20*19*18*17*16*15/20^6 = {20*19*18*17*16*15 / 20**6:.3f}")

    print("\nsequence: [17 3 0 12 7 9 SEP 17 3 0 12 7]")
    print("the 6 output positions are 6..11 and must emit the symbols found at keys 0..5\n")
    print("layer  head   most-weighted key for outputs 6..11      deviation          mean weight")
    for l, pesi in enumerate(r[("full", "copy", 6, 0)]["attenzione"], start=1):
        for t in range(4):
            righe = pesi[t][6:]
            scelta = righe.argmax(-1)
            scarto = scelta - np.arange(6, 12)
            fisso = f"{scarto[0]:+d} constant" if (scarto == scarto[0]).all() else "variable"
            print(f"{l:>4} {t:>6}     {str(scelta):<36} {fisso:<18} {righe.max(-1).mean():.3f}")

Configuration: d=32, 4 heads (d_head=8), 2 layers, d_ff=128, batch 32,
Adam lr=3e-3, 1500 steps, alphabet of 20 symbols + separator (vocabulary 21).
Model parameters: 27253

task    L   per-token accuracy         per-sequence accuracy
            (mean over 3 seeds)        (mean over 3 seeds)
copy      4   1.000  (min 1.000)        1.000  (min 1.000)
copy      6   1.000  (min 1.000)        1.000  (min 1.000)
copy      8   1.000  (min 1.000)        1.000  (min 1.000)
sort      4   0.995  (min 0.989)        0.981  (min 0.958)
sort      6   0.993  (min 0.989)        0.970  (min 0.958)
sort      8   0.987  (min 0.979)        0.951  (min 0.918)

What the code does. The tensor pesa is the mask of the loss, not an attention mask. Training generates a new batch at every step, so there is no set to overfit: the space of sequences is 20⁶ ≈ 6.4 · 10⁷ and the model sees 48,000 of them, 0.075%. The evaluation is true generation, with the model receiving its own previous output.

What the numbers say. On copy the model is perfect at all three lengths, 1.000 per token and per sequence over three seeds: 27,253 parameters solve the task exactly. On sorting the per-token accuracy stays between 0.987 and 0.995, but the per-sequence one falls from 0.981 to 0.951, with the minimum at 0.918: sorting is a comparison between all pairs, and an error on one token out of six is enough to get the sequence wrong. 0.987 and 0.951 are the same model at L = 8: it is also a lesson about metrics.

The scale is small, and I declare it. Training runs in NumPy on CPU with the hand-written backward pass: it is not a model, it is a test bench. d = 32, 4 heads of 8, 2 blocks, d_ff = 128, batches of 32, 1500 steps. Three seeds are few: the minimum shown next to the mean serves to remember it. Better a small experiment that was run than a large one that was described.

Three ablations, and two that show nothing

An experiment serves above all to remove pieces. Relaunching the same training on copy with L = 6, seed 0 and everything else unchanged, one piece at a time:

The output of the ablations is in the same listing as the previous section, and I report it here. Each row is a training with seed 0; "no position" zeroes the position table and never updates it.

Ablations on the copy task, L = 6, seed 0, same hyperparameters unless stated:
  full               final loss 0.0002   token 1.000   sequence 1.000
  no scaling         final loss 0.0002   token 1.000   sequence 1.000
  no position        final loss 0.7342   token 0.656   sequence 0.204
  post-LN            final loss 0.0002   token 1.000   sequence 1.000
  pre-LN, high lr    final loss 0.0000   token 1.000   sequence 1.000
  post-LN, high lr   final loss 0.0000   token 1.000   sequence 1.000

Copy without position, 500 test sequences:
  on sequences whose symbols are all distinct:   0.246   (224 cases)
  on sequences with at least one repeated symbol: 0.170   (276 cases)
  share of sequences with distinct symbols: measured 0.448, expected 20*19*18*17*16*15/20^6 = 0.436

Removing position breaks the task. The loss goes from 0.0002 to 0.7342 and the per-sequence accuracy from 1.000 to 0.204. Yet 0.656 per token remains, well above 1/20 = 0.05: the model is not completely blind to position. A plausible explanation is that the causal mask lets it leak through, because row t sees exactly t + 1 keys, and for decoders without positional encoding there is literature in this sense (Haviv and colleagues, arXiv:2203.16634). I have not tested this explanation on my bench.

An indication, not a result. On sequences whose symbols are all distinct the model without position guesses 0.246 against 0.170 with a repetition (224 and 276 cases). The direction is the expected one, because a repeated symbol is a point where content alone is not enough to decide the place, but it is a single seed and two subsets of the same sample: I do not use it to support anything. The measured share of sequences with distinct symbols, 0.448, is the expected one up to sampling (0.436).

The other two ablations show nothing, and that is a result worth reporting. Removing the division by √d_k leaves the accuracy at 1.000: it does not contradict the section on scale, it contradicts the idea that the problem shows at any dimension. Here d_head = 8 and the deviation of the raw scores is √8 ≈ 2.83: at d_k = 8 the sensitivity drops from 0.96 to 0.72, not by an order of magnitude. The regime in which scaling decides starts at d_k = 64, as in the table above (sensitivity 0.29 against 0.96), and my bench does not reach it.

The same goes for pre-LN against post-LN: indistinguishable at lr = 3 · 10⁻³ and at lr = 10⁻² without warm-up, because two blocks are not a deep stack and the one in arXiv:2002.04745 is a depth problem. So I report the source and not a number of my own: a measurement on two blocks does not support a recommendation on twenty-four.

The attention map, and why it is not an explanation

On the copy task we can look directly at the weight matrix: if the model has learned the right mechanism, the position that emits the k-th symbol of the answer should weigh the position that contains the k-th symbol of the question, at distance L back.

The table is the last part of the output of the experiment listing: for each of the eight heads, the key with the highest weight in each of the six output positions, the deviation from the expected position (outputs 6..11 must look at keys 0..5, that is six back) and the mean maximum weight.

sequence: [17 3 0 12 7 9 SEP 17 3 0 12 7]
the 6 output positions are 6..11 and must emit the symbols found at keys 0..5

layer  head   most-weighted key for outputs 6..11      deviation          mean weight
   1      0     [0 1 2 3 4 5]                        -6 constant        0.835
   1      1     [0 3 6 6 6 5]                        variable           0.416
   1      2     [0 1 2 3 4 5]                        -6 constant        0.996
   1      3     [0 1 2 3 4 5]                        -6 constant        0.767
   2      0     [ 0  1  0  4 10 11]                  variable           0.355
   2      1     [ 0  7  8  5  0 11]                  variable           0.369
   2      2     [ 0  7  8  4 10 11]                  variable           0.365
   2      3     [ 0  1  0  4 10 11]                  variable           0.391

What the numbers say. Three of the four heads of the first layer (0, 2 and 3) have learned the right thing: for each of the six output positions the maximum weight falls exactly six positions back, deviation −6 constant, with mean weight 0.835, 0.996 and 0.767. Head 2, the one drawn, is an almost pure pointer; head 1 is not (0.416, variable deviation). The heads of the second layer are diffuse, with mean maximum between 0.355 and 0.391: the task, in this training, is solved by the first layer. Another seed might distribute the work in another way, and that is one more reason not to generalize from the figure.

And here comes the honest warning. A readable map on an artificial task does not authorize reading maps as explanations on a real task. Jain and Wallace (Attention is not Explanation, arXiv:1902.10186) build distributions very different from the learned one that produce the same prediction: if two contradictory maps lead to the same result, neither of the two is the reason for the result. Wiegreffe and Pinter (Attention is not not Explanation, arXiv:1908.04626) reply by distinguishing plausible and faithful explanation.

The rule to take away. Attention weights are an internal variable of the computation, not a justification. Here the map is readable because the task has a single mechanical solution; on a real task the same image can be a convincing and wrong story. To support why a model decided you need interventions on the input, not a heat map.

Trade-offs

• Short path against quadratic cost. Any two tokens at distance 1 and a parallelized pass, but T² in time and memory: slope about 2 measured (1.99-2.04). Mitigation: local windows, block attention, kernels that do not materialize the matrix - all give up some pairs.
• Long context against service memory. Doubling T quadruples the matrix and doubles the cache. In exercise 2 the cache weighs 10.6 times the weights: it is the context that decides how many simultaneous requests you can serve.
• Many heads against rich heads. At equal parameters, a large h gives many patterns in small subspaces and a small h few patterns in rich subspaces; the empirical balance is d/h between 32 and 128.
• Multi-query against capacity. Sharing a single K, V pair divides the cache by the number of heads - above, from 2.000 to 0.062 GiB - but the heads lose the distinct subspaces. It is adopted when the bottleneck is memory.
• Learned position against computed. The learned table is one line of code, but beyond T_max it does not exist; sinusoidal and RoPE cost no parameters and stay defined beyond the length seen.
• Pre-LN against post-LN. Pre-LN tolerates larger steps without warm-up, post-LN normalizes the output of each block. On my two-block bench there is no measurable difference: it is the limit of my measurement, not the absence of the phenomenon.

The six compromises have only two directions: those on cost (path against T², context against memory, multi-query) you measure with a stopwatch or with a formula, and they must be decided on your workload; those on quality (number of heads, type of position, position of the normalization) depend on the data and on the depth, and a small bench like mine does not decide them. Whoever chooses an architecture should know in which of the two groups each doubt falls.

The traps you actually fall into

1. Forgetting the division by √d_k. Symptom: the loss falls very slowly and the weights are almost all 0 or 1. Cause: the softmax saturates and loses the derivative with respect to the scores (sensitivity from 0.96 to 0.10 at d_k = 512). Remedy: divide, and print the mean maximum weight as a diagnostic.
2. Applying the mask after the softmax. Symptom: none. Cause: zeroing the weights does not renormalize and the rows sum to less than one - here a row sums to 0.031. Remedy: −∞ before the softmax, plus an assert on the row sums.
3. Getting the axis order wrong when stitching the heads back together. Symptom: no exception, because the shapes match. Cause: reshape before transpose mixes components of different heads. Remedy: inverse path in the same order, and with h = 1 the identity.
4. Including the padding in the softmax. Symptom: results that change if you reorder the batch. Cause: the padding mask is missing and the vectors that represent "nothing" take weight. Remedy: mask on the keys, in AND with the causal one.
5. Confusing causal mask and padding mask. Symptom: it works with sequences of the same length and breaks in production. Cause: the first is (T, T) and depends on positions, the second is (B, T) and depends on content. Remedy: two different names, combined in one place only.
6. Normalizing on the query axis. Symptom: training proceeds but never converges well. Cause: axis=-2 instead of axis=-1, that is sum to one by columns. Remedy: the check of the row sums.
7. Forgetting position. Symptom: much lower per-sequence accuracy, 0.204 against 1.000, with a bag-of-words behavior. Cause: attention is permutation-equivariant. Remedy: positional embedding, sinusoid or RoPE; as a test, the output must no longer permute.
8. Using an ε that is too large in the gradient check. Symptom: the check gets worse when you increase ε. Cause: the finite difference crosses a corner of the ReLU. Remedy: ε = 10⁻⁶ and only coordinates with a measurable gradient.
9. Presenting attention weights as the explanation of a decision. Symptom: a heat map in a presentation, and nobody objecting. Cause: plausibility confused with faithfulness (arXiv:1902.10186, arXiv:1908.04626). Remedy: intervene on the input and measure the output.

The nine traps have a common property: none raises an exception. The shapes match, the loss falls, the model produces a text; the defect shows only in a measurement that someone decided to make. This is why, next to each, there is a check that can be written in one line: the row sums, the identity with h = 1, the permutation test, the gradient by finite differences.

Four exercises to complete (variance of the dot product and softmax saturation, causal mask before and after normalization, parameters of a block in closed form, memory of the key-value cache) are in esercizi.py, Python 3.12+ and NumPy. The commented solutions are at the bottom of the file. I did not run it for this revision: the listings of the article yes, the downloadable file no.

What to take away

Attention is a question compared with headings to weigh contents, and the whole formula softmax(Q Kᵀ / √d_k) V is that sentence in algebra. The pieces that make it usable are not accessories: the division by √d_k keeps the softmax derivative alive, the causal mask goes before the normalization or the rows are no longer averages, position has to be injected from outside, and h heads cost as much as one. The block is 12d² + 13d parameters; the cost is T², slope measured between 1.99 and 2.04; the cache is worth a factor between about 30 and 65 in generation on my toy.

Two things to keep beyond the formula. The first is that an artificial task teaches a lot precisely because it is artificial: a decoder of 27,253 parameters copies sequences perfectly and sorts them entirely in 95-98% of cases, and from those two figures you understand the difference between positional retrieval and a comparison between pairs. The second is that two of my three ablations showed nothing, and the right answer was not to insist until finding a favorable number. You understand a mechanism when you also know under which conditions its effect does not show. The next lesson uses these vectors to search by meaning: the lesson on embeddings and semantic search.

The same series also includes Running a Model Locally: Quantisation, Memory and Latency, Measured.