Unit 09.00: Transformer blocks at a practical level

A transformer is one block repeated. Learn the block and the architecture stops being intimidating.

Four components, twice around

The block is: multi-head attention, add the input back and normalise, a small feed-forward network, add and normalise again.

The two "add the input back" steps are residual connections, and they are what makes depth trainable — gradients reach early layers through the addition even if the transformations in between are unhelpful.

import torch
from torch import nn

torch.manual_seed(0)
# One encoder block, built from the pieces you already know.
d_model, heads = 16, 4
block = nn.TransformerEncoderLayer(d_model=d_model, nhead=heads,
                                   dim_feedforward=32, batch_first=True)

x = torch.randn(2, 5, d_model)      # batch, tokens, features
out = block(x)
print("in :", tuple(x.shape))
print("out:", tuple(out.shape), "<- same shape, richer representation")

# A block is: attention -> add & norm -> feed-forward -> add & norm.
print("\ncomponents:", [n for n, _ in block.named_children()])
print("parameters:", sum(p.numel() for p in block.parameters()))

# Residual connections are why depth is trainable: the input has a path
# straight to the output, so gradients do not have to survive every layer.
identity_ish = torch.nn.functional.cosine_similarity(x.flatten(), out.flatten(), dim=0)
print("output still resembles input (residual):", round(identity_ish.item(), 3))

Input and output have the same shape. A block does not change the representation size; it enriches each token with context from the others.

The cosine similarity at the end shows the output still resembles the input, which is the residual connection working. Without it, stacking twelve blocks would destroy the signal.

The mistake this prevents

Assuming a transformer needs a huge dataset to be useful at all. A single small block used as a frozen feature extractor is often a reasonable starting point, and it trains in seconds.

Takeaway

One block: attention, add and norm, feed-forward, add and norm. Everything else is repetition and scale.

Mark complete