01 - Transformer Architecture: attention, encoder, decoder

Read the sentence "the bag would not fit in the suitcase because it was too small". To know that "it" refers to the suitcase and not to the bag, you have to connect two words that sit far apart, and you do it in one go. A recurrent network instead reads the words one after another and carries a single memory vector with it: the information about the first word has to cross all the others to reach the last one. The Transformer, proposed by Vaswani and colleagues in 2017, removes that forced passage: every position in the sequence looks at all the others directly, and decides how much weight to give each of them with a computation that fits in two lines of linear algebra. This article builds that computation, checks it against PyTorch, and then verifies one by one the four ideas that make it usable: the scale factor, multiple heads, the encoding of positions and the decoder mask.

Every listing was run twice with identical output, with numpy 2.4.6 and torch 2.12.0 on CPU, on random numbers with a fixed seed. There is no trained model and no text: attention can be studied on plain numbers, and it is the fastest way to understand it.

The problem: carrying information from far away

Table 1 of the original paper (Vaswani et al., Attention Is All You Need) compares layer types with three measures. With n the sequence length and d the dimension of the representation, the recurrent layer costs O(n·d²) operations, requires O(n) sequential steps and has a maximum path between two positions of O(n). Self-attention costs O(n²·d), but requires O(1) sequential steps and the maximum path is O(1): any two positions are connected in a single layer. The price is the n² term, which will be measured further on; the gain is that all positions are computed together, so an accelerator can use all of its parallelism, while recurrence is sequential by construction.

Attention as a soft lookup in a dictionary

A Python dictionary returns the value whose key matches the one you look up. Attention does the same, but with a graded match. Each position produces three vectors: a query (what am I looking for), a key (what am I about) and a value (what do I offer if you pick me). The dot product between the query of one position and the keys of all positions gives a score for each; the softmax turns the scores into positive weights that sum to one; the output is the average of the values weighted with those weights. In the formula of section 3.2.1 of the paper, Attention(Q, K, V) = softmax(QKᵀ/√d_k)·V, the division by √d_k is the only thing a dictionary does not have, and the next section explains what it is for.

The listing computes attention for five positions with d_k = 8 and d_v = 4. The softmax function subtracts the maximum of each row before exponentiating: it does not change the result (the softmax is invariant to a common shift), but it prevents exp from overflowing with large scores. The mask argument is used further on. The second, independent computation is PyTorch's scaled_dot_product_attention, the function that the library's API dedicates to this formula: if the two implementations coincide, the formula is written correctly.

import numpy as np
import torch
import torch.nn.functional as F

torch.set_num_threads(2)
rng = np.random.default_rng(0)


def softmax(z):
    z = z - z.max(axis=-1, keepdims=True)      # numerical stability: does not change the result
    e = np.exp(z)
    return e / e.sum(axis=-1, keepdims=True)


def attention(Q, K, V, mask=None):
    d_k = Q.shape[-1]
    scores = Q @ K.T / np.sqrt(d_k)            # (n_query, n_key)
    if mask is not None:
        scores = np.where(mask, scores, -np.inf)
    w = softmax(scores)
    return w @ V, w


n, d_k, d_v = 5, 8, 4
Q = rng.standard_normal((n, d_k))
K = rng.standard_normal((n, d_k))
V = rng.standard_normal((n, d_v))

out, w = attention(Q, K, V)
ref = F.scaled_dot_product_attention(
    torch.from_numpy(Q)[None], torch.from_numpy(K)[None], torch.from_numpy(V)[None]
)[0].numpy()

print("weights shape:", w.shape, "- sum of each row:", np.round(w.sum(axis=1), 12))
print("max difference from torch:", f"{np.abs(out - ref).max():.2e}")
print("weights of the first row:", np.round(w[0], 3))

# Output:
# weights shape: (5, 5) - sum of each row: [1. 1. 1. 1. 1.]
# max difference from torch: 3.33e-16
# weights of the first row: [0.367 0.063 0.113 0.115 0.342]

The difference is of the order of 10^-16, that is the rounding of double precision, and every row of weights sums to one. The weight is a distribution: the first position gives 0.367 to itself, 0.342 to the last one and little to the other three. The output of each position is therefore a mixture of values, and which mixture depends on the content, not on the distance: this is why the first and the last position can talk to each other without intermediaries.

Why we divide by √d_k

If the components of q and k are independent variables with mean 0 and variance 1, their dot product has mean 0 and variance d_k: section 3.2.2 of the paper says so, in a footnote. The standard deviation of the scores therefore grows like √d_k. The softmax of widely spread numbers is almost a "winner takes all": a single weight close to one and the others close to zero, and in that region the derivatives are tiny. The authors write that they suspect this as the cause of the worse result of unscaled attention for large d_k; here it is measured directly.

The listing draws 4000 query-key pairs for each d_k, with 16 keys per query, and measures three things. The standard deviation of the raw scores is compared with √d_k (a second calculation, in closed form). The mean maximum weight says how "switched on" the softmax is on a single key. The Frobenius norm of the Jacobian matrix of the softmax, diag(p) − ppᵀ, says how much the outputs react to a change in the scores: if it is small, the gradient that crosses the layer is small and learning slows down. The last two columns repeat the measurement after the division.

import numpy as np

rng = np.random.default_rng(1)
n_keys, trials = 16, 4000


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


def jac_norm(p):
    # Frobenius norm of the softmax Jacobian matrix: diag(p) - p p^T
    J = np.einsum("ti,ij->tij", p, np.eye(p.shape[1])) - np.einsum("ti,tj->tij", p, p)
    return np.sqrt((J ** 2).sum(axis=(1, 2)))


print(f"{'d_k':>5} {'std(q.k)':>9} {'sqrt(d_k)':>9} | {'max weight':>9} {'|J| without':>9} | {'max weight':>9} {'|J| with':>9}")
for d_k in (4, 16, 64, 256, 1024):
    q = rng.standard_normal((trials, 1, d_k))
    k = rng.standard_normal((trials, n_keys, d_k))
    s = (q * k).sum(axis=-1)                       # raw dot products, shape (trials, n_keys)
    p_raw = softmax(s)
    p_scl = softmax(s / np.sqrt(d_k))
    print(f"{d_k:>5} {s.std():>9.3f} {np.sqrt(d_k):>9.3f} | "
          f"{p_raw.max(axis=1).mean():>9.3f} {jac_norm(p_raw).mean():>9.4f} | "
          f"{p_scl.max(axis=1).mean():>9.3f} {jac_norm(p_scl).mean():>9.4f}")

# Output:
#   d_k  std(q.k) sqrt(d_k) | max weight |J| without | max weight  |J| with
#     4     1.986     2.000 |     0.426    0.3221 |     0.234    0.2967
#    16     3.991     4.000 |     0.681    0.2805 |     0.241    0.3024
#    64     7.973     8.000 |     0.846    0.1732 |     0.245    0.3042
#   256    15.978    16.000 |     0.923    0.0987 |     0.246    0.3049
#  1024    31.927    32.000 |     0.963    0.0504 |     0.248    0.3053

The standard deviation of the scores follows √d_k to within one decimal digit (31.9 against 32 for d_k = 1024). Without the scale, going from d_k = 4 to d_k = 1024 the mean maximum weight rises from 0.426 to 0.963 and the norm of the Jacobian falls from 0.32 to 0.050: with 1024 dimensions the softmax is almost a hard maximum and lets through about one sixth of the gradient that would pass with the scale (0.050 against 0.305). With the scale, the maximum weight stays around 0.24 and the Jacobian around 0.30 whatever d_k is, because the variance of the scores is one by construction. The correction therefore does not make attention "better": it keeps it in the same regime as the dimension changes. One caveat remains: the measurement holds for vectors with independent components and variance one, that is at initialization; in trained models the components are no longer like that, and this experiment does not say how they behave.

Multiple heads: several small attentions instead of one large one

A single attention produces a single mixture of values per position: in the opening sentence "it" must look at the suitcase, but the same layer should also look at the verb. Multiple heads (section 3.2.2) run h attentions in parallel, each on a different linear projection of dimension d_k = d_model/h, then concatenate the outputs and project them again with a matrix W^O. In the authors' base model d_model = 512 and h = 8, so d_k = 64.

The listing checks two things against torch.nn.MultiheadAttention, with d = 16 and h = 4. The first: my hand-written implementation, which reads the weights of the PyTorch module (in_proj_weight stacks the three projections Q, K, V) and gives each head a slice of dh columns, reproduces the output and the attention weights of every head. The second: the number of parameters does not depend on h. The heads divide the dimension, they do not multiply it, so with 1, 4 or 16 heads the parameters are always 4d² + 4d (three projections, the output one, the four biases).

import numpy as np
import torch

torch.set_num_threads(2)
torch.manual_seed(0)
d, h, n = 16, 4, 6
dh = d // h

mha = torch.nn.MultiheadAttention(d, h, batch_first=True, dtype=torch.float64)
x = torch.randn(1, n, d, dtype=torch.float64)
with torch.no_grad():
    ref, ref_w = mha(x, x, x, need_weights=True, average_attn_weights=False)

Wi = mha.in_proj_weight.detach().numpy()       # (3d, d): the three projections Q, K, V stacked
bi = mha.in_proj_bias.detach().numpy()
Wo = mha.out_proj.weight.detach().numpy()
bo = mha.out_proj.bias.detach().numpy()
X = x[0].numpy()


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


Q = X @ Wi[:d].T + bi[:d]
K = X @ Wi[d:2 * d].T + bi[d:2 * d]
V = X @ Wi[2 * d:].T + bi[2 * d:]
heads, weights = [], []
for i in range(h):                              # each head reads its own slice of dh columns
    sl = slice(i * dh, (i + 1) * dh)
    w = softmax(Q[:, sl] @ K[:, sl].T / np.sqrt(dh))
    weights.append(w)
    heads.append(w @ V[:, sl])
out = np.concatenate(heads, axis=1) @ Wo.T + bo  # concatenate and project with W^O

print("max difference on the output:", f"{np.abs(out - ref[0].detach().numpy()).max():.2e}")
print("max difference on the weights: ", f"{np.abs(np.stack(weights) - ref_w[0].detach().numpy()).max():.2e}")

for heads_n in (1, 4, 16):
    m = torch.nn.MultiheadAttention(d, heads_n)
    print(f"{heads_n:>2} heads -> parameters: {sum(p.numel() for p in m.parameters())}")
print("formula 4*d*d + 4*d =", 4 * d * d + 4 * d)

# Output:
# max difference on the output: 1.11e-16
# max difference on the weights:  1.11e-16
#  1 heads -> parameters: 1088
#  4 heads -> parameters: 1088
# 16 heads -> parameters: 1088
# formula 4*d*d + 4*d = 1088

The output and the weights coincide to 10^-16 for all four heads. The parameters are 1088 for any number of heads, equal to the formula. So the choice of h is not a memory cost of the weights: it is a choice about how to split the dimension. Table 3 of the original paper (rows A) shows that 8 and 16 heads give the same quality on their test bench (BLEU 25.8 on English-to-German translation), while a single head at full dimension loses 0.9 points and 32 heads drop to 25.4: too many heads mean keys that are too small. These are measurements of that task and that model, not a rule for yours.

Position: attention does not know the order

One consequence of attention is counterintuitive and can be shown in three lines: if you shuffle the order of the input tokens, the output is shuffled in the same way, and that is all. The layer does not see the order; it sees a set. For a language this is a problem, because "the dog bites the man" and "the man bites the dog" have the same tokens. The authors add to each embedding a positional encoding (section 3.5): for the position pos and the pair of dimensions 2i, 2i+1, PE = (sin(pos/10000^2i/d), cos(pos/10000^2i/d)). Each pair is a sinusoid with a different frequency, from the shortest wavelengths (2π) to the longest (10000·2π). The stated reason is that for any fixed offset k, PE_pos+k is a linear function of PE_pos.

The listing verifies both things. First, with MultiheadAttention without positions: the comparison of y_perm with y[:, perm] says whether the output of the permuted batch is the permutation of the output. Then, adding the encoding to the tokens and keeping the positions fixed, the same equality must fail. Finally, for k = 5 it builds by hand the 2×2 rotation matrix R_k with angle ω·k and checks that R_k·PE_pos coincides with PE_pos+k over 30 positions and 8 frequencies: it is the numerical demonstration of the authors' claim, and it shows that the operation is a rotation in each plane.

import numpy as np
import torch

torch.set_num_threads(2)
torch.manual_seed(0)
d, h, n = 16, 4, 7
mha = torch.nn.MultiheadAttention(d, h, batch_first=True, dtype=torch.float64)


def run(x):
    with torch.no_grad():
        return mha(x, x, x, need_weights=False)[0]


def pe(n, d):
    pos = np.arange(n)[:, None]
    i = np.arange(d // 2)[None, :]
    ang = pos / 10000 ** (2 * i / d)
    out = np.zeros((n, d))
    out[:, 0::2] = np.sin(ang)
    out[:, 1::2] = np.cos(ang)
    return out


x = torch.randn(1, n, d, dtype=torch.float64)
perm = torch.randperm(n, generator=torch.Generator().manual_seed(3))

y = run(x)
y_perm = run(x[:, perm])
print("without position: output(permuted) == permuted(output)?",
      torch.allclose(y_perm, y[:, perm]), f"(error {float((y_perm - y[:, perm]).abs().max()):.1e})")

P = torch.from_numpy(pe(n, d))[None]
z = run(x + P)
z_perm = run(x[:, perm] + P)                     # same tokens, different order, fixed positions
print("with positional encoding: same equality?",
      torch.allclose(z_perm, z[:, perm]), f"(error {float((z_perm - z[:, perm]).abs().max()):.1e})")

# PE(pos + k) = R_k PE(pos): each (sin, cos) pair rotates by an angle that depends only on k
k = 5
big = pe(40, d)
worst = 0.0
for pos in range(0, 30):
    for i in range(d // 2):
        w = 1 / 10000 ** (2 * i / d)
        R = np.array([[np.cos(w * k), np.sin(w * k)], [-np.sin(w * k), np.cos(w * k)]])
        pred = R @ big[pos, 2 * i:2 * i + 2]
        worst = max(worst, np.abs(pred - big[pos + k, 2 * i:2 * i + 2]).max())
print(f"PE(pos+{k}) = R_{k} PE(pos): max error over 30 positions and {d // 2} frequencies = {worst:.1e}")

# Output:
# without position: output(permuted) == permuted(output)? True (error 1.1e-16)
# with positional encoding: same equality? False (error 2.1e-01)
# PE(pos+5) = R_5 PE(pos): max error over 30 positions and 8 frequencies = 8.9e-16

Without positions the equality holds to 10^-16: the layer is equivariant to permutations. With the encoding it fails, with a difference of 0.21, because the same token in different positions now carries a different contribution. The last line confirms the reason for the sinusoidal form: a shift of 5 positions is the same rotation for every pos, to 10^-15. The original paper also tried learned positional embeddings and found nearly identical results (table 3, row E); it chose the sinusoidal version because it might extrapolate to sequences longer than those seen, a hypothesis that was not put to the test here. Other position schemes have been proposed since 2017, but this article does not compare them.

Encoder, decoder and the causal mask

The same attention layer appears in three roles, described in sections 3.1 and 3.2.3 of the paper:

• encoder self-attention: queries, keys and values come from the output of the previous layer, and every position looks at all the others;
• masked decoder self-attention: every position looks only at itself and the previous ones;
• encoder-decoder attention: the queries come from the decoder, keys and values from the encoder output, and every decoder position can look at the whole input.

Only where Q, K, V come from and which connections are allowed changes; the formula is always the same. The encoder of the paper is a stack of N = 6 identical layers, each with multi-head attention and a position-wise feed-forward network, with a residual connection around each sub-part followed by layer normalization: the output is LayerNorm(x + Sublayer(x)). The decoder has the same layers plus the third role in the list, and the output is shifted by one position because the prediction at i must depend only on the outputs known before i. The residual requires all sub-parts to produce the same dimension d_model (512 in the paper).

The mask is implemented in the formula itself: the illegal scores are set to −∞ before the softmax, and the weight becomes exactly zero. The listing verifies it in the most direct way: it modifies only positions 5, 6 and 7 of the input, and measures how much the output changes at each position, with and without a triangular mask. If the mask works, positions 0-4 must not notice the change; if there is none, the change must propagate to all of them.

import torch

torch.set_num_threads(2)
torch.manual_seed(0)
d, h, n = 16, 4, 8
mha = torch.nn.MultiheadAttention(d, h, batch_first=True, dtype=torch.float64)
causal = torch.triu(torch.ones(n, n, dtype=torch.bool), diagonal=1)   # True = not allowed to look


def run(x, mask):
    with torch.no_grad():
        return mha(x, x, x, attn_mask=mask, need_weights=False)[0]


x = torch.randn(1, n, d, dtype=torch.float64)
x2 = x.clone()
x2[:, 5:] = torch.randn(1, n - 5, d, dtype=torch.float64)     # change only positions 5, 6, 7

for nome, mask in (("no mask", None), ("causal mask", causal)):
    diff = (run(x, mask) - run(x2, mask)).abs().amax(dim=-1)[0]
    print(f"{nome:>22}: output change per position", [f"{v:.1e}" for v in diff.tolist()])

# Output:
#                no mask: output change per position ['2.3e-01', '3.4e-01', '3.3e-01', '2.9e-01', '2.5e-01', '3.4e-01', '2.8e-01', '2.8e-01']
#            causal mask: output change per position ['0.0e+00', '0.0e+00', '0.0e+00', '0.0e+00', '0.0e+00', '4.1e-01', '2.0e-01', '2.8e-01']

Without the mask the change alters the output of every position, with differences between 0.23 and 0.34. With the causal mask positions 0-4 stay at exact zero, and only positions 5-7, which "see" the modified token, change. The practical consequence is that in training the decoder can compute the prediction for all positions in a single pass without cheating, and in generation it must instead proceed one token at a time, because the following tokens do not exist yet.

What the n² term costs

The score of each pair of positions takes up one cell: an n × n matrix per head and per layer. The listing measures on the CPU, with d = 64 and float32, the size of that matrix and the mean time of five runs of a single attention computation, for n from 512 to 4096.

import time
import numpy as np

rng = np.random.default_rng(0)
d = 64
print(f"{'n':>5} {'n*n entries':>12} {'MiB (float32)':>14} {'mean time (s)':>16}")
for n in (512, 1024, 2048, 4096):
    Q = rng.standard_normal((n, d)).astype(np.float32)
    K = rng.standard_normal((n, d)).astype(np.float32)
    V = rng.standard_normal((n, d)).astype(np.float32)
    t = []
    for _ in range(5):
        t0 = time.perf_counter()
        S = Q @ K.T / np.sqrt(np.float32(d))
        S -= S.max(axis=1, keepdims=True)
        np.exp(S, out=S)
        S /= S.sum(axis=1, keepdims=True)
        out = S @ V
        t.append(time.perf_counter() - t0)
    print(f"{n:>5} {n * n:>12} {n * n * 4 / 2**20:>14.1f} {np.mean(t):>16.3f}")

# Output:
#     n  n*n entries  MiB (float32)    mean time (s)
#   512       262144            1.0            0.001
#  1024      1048576            4.0            0.004
#  2048      4194304           16.0            0.017
#  4096     16777216           64.0            0.084

Each doubling of n quadruples the cells (1, 4, 16, 64 MiB) and the time grows by about four times per doubling: from 512 to 4096 the cells grow by 64 times, and the time, in three runs, went from about one millisecond to 0.07-0.09 s. Times change from one machine and one run to another (and with the linear algebra library used), so the order of magnitude counts and not the value. With a single head and one layer, 4096 positions cost 64 MiB of scores; with 8 heads and 6 layers, if they were all kept for training, they would be 3 GiB for every example in the batch. This is why optimized layers compute the softmax in blocks without materializing the whole matrix, but the quadratic count of operations remains.

What you can do now

You can write dot-product attention with scaling in a few lines of NumPy and check it against PyTorch, you can measure why the factor 1/√d_k is needed, you can rebuild by hand a multi-head attention from a library module, you can show that positions must be added and that the sinusoidal encoding is a rotation, and you can check that a causal mask really prevents looking ahead. How much of all this holds at the scale of billions of parameters is not the subject of this article: the next one, 02 - LLM scaling laws and emergent capabilities, studies how the error changes with the size of the model and of the data. For those coming from classical machine learning, the article on the bias and variance decomposition is a good prerequisite, and the full index is on the series page.

Later in the series you will also find 08 - On-prem inference with Ollama, vLLM, llama.cpp, which takes up the cost of attention from the inference side, with the memory of the key-value cache.