Inside a transformer block

A transformer block is the unit a language model is built from, and a model is the same block stacked many times. It holds masked attention with several heads, then a small network for each token, with shortcuts and rescaling around both.

Once you know one block, you know the shape of the whole model.

A language model learns to guess the next token. If a token could look at the tokens after it, it would see the answer.

So a causal mask lets each token look only at itself and the tokens before it. A later token gets a weight of 0. In practice its score is set to minus infinity before softmax, because e to minus infinity is 0.

import math

# scores[i][j]: how well token i's query matches token j's key.
scores = [[1, 2, 0], [2, 1, 1], [0, 1, 3]]

for i in range(3):
    powers = []
    for j in range(3):
        if j <= i:
            powers.append(math.exp(scores[i][j]))
        else:
            # A later token: masked, so it gets no weight at all.
            powers.append(0)
    total = sum(powers)
    weights = []
    for p in powers:
        weights.append(round(p / total, 3))
    print("token", i, "weights", weights)
# Change j <= i to True, so nothing is masked, and press Run.
# How does the first token's row change, and why is that cheating?

docker exec jupyter python3 -c "import numpy as np; S = np.array([[1, 2, 0], [2, 1, 1], [0, 1, 3]], dtype=float); S[np.triu(np.ones((3, 3), dtype=bool), 1)] = -np.inf; W = np.exp(S) / np.exp(S).sum(axis=1, keepdims=True); print(W.round(3).tolist())"

One attention step can only mix in one way. So a block runs several heads side by side, each with its own query, key and value tables.

One head might link a word to the word before it. Another might link a word to the thing it describes. Their outputs are joined end to end, then one more learnt table mixes them back to the usual length.

After attention, each token goes on its own through a small feed-forward step: two layers with ReLU between. ReLU turns every number below 0 into 0.

A residual connection adds each step's input to its output, so nothing already known is lost. Normalisation rescales each vector so its numbers average 0 and keep a steady spread. Both keep a deep stack trainable.

docker exec jupyter python3 -c "import numpy as np; x = np.array([2.0, 4.0, 6.0]); print(x.mean(), ((x - x.mean()) / x.std()).round(2).tolist())"