The transformer block

Two ways into the same idea

Before You Start

The Core Idea

In this lesson you build the block piece by piece: layer norm, then multiple heads, then the feed-forward, then the residual connection that ties them together.

1Layer norm

Layer norm takes each token's vector on its own and gives it mean zero and spread one. That way numbers neither explode nor vanish as they pass through dozens of blocks.

import numpy as np

np.set_printoptions(precision=2, suppress=True)
x = np.array([[2.0, 4.0, 6.0, 8.0],
              [100.0, 101.0, 99.0, 100.0]])   # two tokens on very different scales

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

y = layer_norm(x)
print(y)
print("mean per token:", y.mean(axis=-1) + 0.0)
print("std per token: ", y.std(axis=-1))

Two tokens, one with values from 2 to 8 and one near 100. After normalising both have mean 0 and spread 1, and only each one's shape is left: the first rises, the second has a peak and a dip. nn.LayerNorm also has two learned parameters that rescale and shift the result if the model needs to.

2Multiple heads

One head gives one blend. When we split the width between several heads, each learns a different kind of relation, at about the same cost.

Multiple heads are not separate matrices, just a reshape of the same one.

import numpy as np

T, d_model, heads = 5, 8, 2
d_head = d_model // heads
q = np.arange(T * d_model, dtype=float).reshape(T, d_model)

split = q.reshape(T, heads, d_head).transpose(1, 0, 2)
print("one sequence:", q.shape)
print("split into heads:", split.shape, "-> (heads, tokens, d_head)")
print("head 0 sees columns 0..3 of token 0:", split[0, 0])
print("head 1 sees columns 4..7 of token 0:", split[1, 0])
back = split.transpose(1, 0, 2).reshape(T, d_model)
print("joined back equals the original:", bool((back == q).all()))

reshape then transpose turn (tokens, d_model) into (heads, tokens, d_head), so each head computes attention over its own quarter of the columns, then we reverse it to join the results. This exact line is in every transformer implementation, and ordering mistakes in it are a classic cause of models that "train" but do not learn.

Look at the picture: each head sees only half of every token's columns and computes its attention on them, then the results are joined back to the full width.

3The feed-forward

After attention, each word passes alone through a small network: a layer widening d to 4d, an activation, and a layer back to d. Attention is the only place tokens exchange information; the feed-forward processes each token with what it gathered.

Look at the picture: attention holds 4 × d² parameters and the feed-forward 8 × d². Together that is 12 × d², the rule you use to estimate any transformer's size.

4The residual: x + f(x)

The block does not replace its input; it adds a change to it. So the gradient crosses dozens of layers without vanishing, and the shape stays the same, so we can stack any number of blocks.

All the pieces together: norm, four-head masked attention, add; norm, feed-forward, add.

import numpy as np

rng = np.random.default_rng(1)
T, d, heads, d_ff = 6, 16, 4, 64
dh = d // heads

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

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

def init(*shape):
    return rng.normal(0, 0.02, size=shape)

Wqkv, Wo = init(d, 3 * d), init(d, d)
W1, W2 = init(d, d_ff), init(d_ff, d)
mask = np.triu(np.ones((T, T), dtype=bool), 1)

def attention(x):
    q, k, v = np.split(x @ Wqkv, 3, axis=-1)
    q, k, v = (m.reshape(T, heads, dh).transpose(1, 0, 2) for m in (q, k, v))
    s = q @ k.transpose(0, 2, 1) / np.sqrt(dh)
    s[:, mask] = -np.inf
    out = softmax(s) @ v                          # (heads, T, dh)
    return out.transpose(1, 0, 2).reshape(T, d) @ Wo

def feed_forward(x):
    return np.maximum(0, x @ W1) @ W2

def block(x):
    x = x + attention(layer_norm(x))      # residual around attention
    x = x + feed_forward(layer_norm(x))   # residual around the MLP
    return x

x = rng.normal(size=(T, d))
y = block(x)
params = Wqkv.size + Wo.size + W1.size + W2.size
print("in:", x.shape, " out:", y.shape)
print("parameters in one block:", params)
print("how far the block moved x:", round(float(np.abs(y - x).mean() / np.abs(x).mean()), 3))

The shape is unchanged, (6, 16) to (6, 16), which is why you can stack any number of blocks. And 3,072 parameters for a width-16 block is about 12 × d², the rule you will use to estimate the size of any transformer. Notice the block moved x by only about 1%: with small initial weights the block is nearly "nothing" thanks to the residual, which is what makes the start stable.

How It Works, Step by Step: how one token passes through a block

Going deeper

The order of the pieces matters. The original transformer put the norm after the addition (post-norm), and all modern models put it before attention and the feed-forward (pre-norm), because the residual path stays clean from input to output, which makes deep models train more stably with less learning-rate warm-up.

Worked example

The simplest piece first: layer normalisation. It takes each token's vector on its own and makes its mean zero and its spread one, whatever its original range. This stops numbers exploding or vanishing as they pass through dozens of blocks.

2. The difference from batch norm

Batch norm (the layers lesson) averages across the examples in a batch, so it behaves differently in training and in use and depends on batch size. Layer norm averages across one token's dimensions, so every token is independent, which is what suits sequences of different lengths.

Example: The ready-made block in PyTorch

import torch
import torch.nn as nn

block = nn.TransformerEncoderLayer(d_model=512, nhead=8, dim_feedforward=2048,
                                   batch_first=True, norm_first=True)
x = torch.randn(2, 10, 512)
mask = nn.Transformer.generate_square_subsequent_mask(10)
print(block(x, src_mask=mask, is_causal=True).shape)   # torch.Size([2, 10, 512])
print(sum(p.numel() for p in block.parameters()))       # 3,152,384

norm_first=True is pre-norm. And 3,152,384 is close to 12 × 512² = 3,145,728; the difference is the biases and the norm parameters. In the lesson on building a small transformer we write the block ourselves instead of this ready-made one, so you see every line.

Try it without code

Myth vs Reality

Try it yourself

Quick review cards

Check Your Understanding