The transformer block

Read these first if they are new to you

Attention: queries, keys, values Advanced
The multilayer perceptron (MLP) Intermediate
Residual connections Intermediate
Layer normalization Intermediate

The idea in one paragraph

A transformer block updates a set of tokens in two steps. First, attention lets every token gather information from the other tokens. Second, an MLP processes each token on its own. Each step reads a layer-normalized copy of the tokens and adds its result back through a residual connection. A transformer stacks many such blocks, each with its own weights. The tokens come out of a block in the same shape they went in: n tokens of d numbers each.

Attention mixes across tokens; the MLP works within each token

A block has two sublayers, and they divide the work.

Attention is the only place where tokens exchange information. Each token's output is a weighted sum of values from all the tokens. Usually it is multi-head self-attention.
The MLP, also called the feed-forward network, applies the same small network to every token separately: a linear layer that widens the token, an activation function, and a linear layer that narrows it back. The hidden width is often 4d. The original transformer used ReLU; many newer models use a smoother activation called GELU.

Each sublayer is wrapped in two helpers. Layer normalization resets each token to a standard size before the sublayer reads it. A residual connection adds the sublayer's output to the token, so the sublayer makes an edit rather than a replacement.

This page draws the pre-norm block; the original transformer used post-norm

Two placements of the layer normalization are in common use. This page draws the pre-norm block, which normalizes the input of each sublayer. For tokens x (one row per token):

h = x + Attention( LN₁(x) )

y = h + MLP( LN₂(h) )

The original transformer of 2017 used the post-norm block, which normalizes after each residual addition:

h = LN₁( x + Attention(x) )

y = LN₂( h + MLP(h) )

Here LN₁ and LN₂ are two layer normalizations, each with its own γ and β. The running vector that the additions keep updating is called the residual path, or residual stream.

One transformer block, with its residual path in thick blue

Try it. Switch between the two versions and follow the thick blue path from "tokens in" to "tokens out". In which version does that path pass through no layer normalization at all? Why could that help a deep stack of blocks train?

Answer

In the pre-norm version. There, a token reaches the output through additions only, so each block's skip path contributes its plain "1" to the slope, as in the residual connections unit, and the training signal crosses many blocks intact. In the post-norm version every block rescales the path twice. A 2020 study (Xiong and others) found that post-norm stacks need a slow learning-rate warm-up to train, while pre-norm stacks train without one.

What to notice. Side branches leave the blue path, compute an edit, and add it back at a plus sign. The attention branch is the only place where tokens see each other. In the post-norm version, the orange boxes are layer normalizations that sit on the residual path itself.

Many recent models use pre-norm because deep pre-norm stacks train more stably. A pre-norm transformer usually adds one more layer normalization after its last block, so the final output is normalized too. Both placements appear in published models, so check which one a model uses before reading its code.

The learned weights do not depend on the number of tokens

Predict first. A block reads 1,000 tokens. Then it reads 2,000. Does it need more learned weights for the larger input?

Worked example: a block with d = 256, 8 heads and an MLP of width 1,024

Part	Learned numbers	Count
Attention	W_Q, W_K, W_V, W_O: each 256 × 256, plus a bias of 256	4 × 65,792 = 263,168
MLP	256 → 1,024: 262,144 weights and 1,024 biases; 1,024 → 256: 262,144 weights and 256 biases	525,568
Two layer norms	γ and β, 256 each, twice	1,024
Total		789,760

No entry mentions the number of tokens n, and none mentions the 8 heads: the heads split the same 256 × 256 matrices into 8 slices of 32. The block uses the same weights for 10 tokens or 10,000.

What does grow with n is the work. Each head computes a weight for every pair of tokens, so 1,000 tokens give 8 × 1,000 × 1,000 = 8,000,000 attention weights per block. These are computed from the input each time, not learned. Doubling n doubles the MLP's work but quadruples the attention weights.

block.py

import torch

# fix the random weights, tokens and shuffle, so each run prints the same
torch.manual_seed(0)
# The block from the worked example: d = 256, 8 heads, an MLP of width 1,024.
# dropout=0.0 switches off dropout, which zeroes values at random in training.
# batch_first=True makes the first number in x's shape count the examples.
block = torch.nn.TransformerEncoderLayer(
    d_model=256, nhead=8, dim_feedforward=1024,
    dropout=0.0, batch_first=True, norm_first=True)   # norm_first: pre-norm
print(sum([p.numel() for p in block.parameters()]))

x = torch.randn(1, 1000, 256)        # 1 example: 1000 tokens of 256 numbers
y = block(x)
print(y.shape)

perm = torch.randperm(1000)          # shuffle the tokens
# torch.allclose is True when the two agree in every entry, to within about 0.00001.
print(torch.allclose(block(x[:, perm]), y[:, perm], atol=1e-5))

789760
torch.Size([1, 1000, 256])
True

The last line shows that a block treats its tokens as a set: shuffling the input tokens shuffles the output tokens in the same way and changes nothing else. Attention has no notion of position, and the MLP and layer norms work on each token alone. A point cloud fits this directly. Models of sentences or images add position information to each token before the first block.

What this page leaves out

Real blocks often add dropout (randomly zeroing some numbers during training) and, in language models, an attention mask. Newer models swap in variants: RMS normalization instead of layer normalization, or a gated MLP. Generative models that must know a time step or a condition change the layer norms so that γ and β come from that input (see diffusion transformers). The overall pattern, normalize, sublayer, add back, stays the same.

Check yourself

Which sublayer lets a token's output depend on other tokens, and which treats each token alone?

Answer
Attention mixes information across tokens. The MLP, and the layer normalizations, work on each token by itself.
Work a similar problem. Count the learned numbers in a block with d = 128, 4 heads and an MLP of width 512.

Answer
Attention: 4 × (128 × 128 + 128) = 66,048. MLP: 128 × 512 + 512 + 512 × 128 + 128 = 131,712. Layer norms: 4 × 128 = 512. Total 198,272. PyTorch's TransformerEncoderLayer(128, 4, 512) reports the same.
Rewrite the post-norm block h = LN₁(x + Attention(x)) as pre-norm.

Answer
h = x + Attention(LN₁(x)). The normalization moves inside the branch, and the residual addition stays outside any normalization.
A block reads 2,000 tokens instead of 1,000. By how much do its learned weights grow, and by how much do its attention weights grow?

Answer
The learned weights do not grow at all. The attention weights grow four times, from 1,000 × 1,000 to 2,000 × 2,000 per head.

Go further

Optional background, for readers who want more than this page.

Transformers from scratch: "The transformer block" Peter Bloem Article with code Advanced

A short section with one diagram and the PyTorch code of a block. Its block is post-norm, the 2017 order, so compare it with the post-norm version of the figure above.
Transformer Explainer Georgia Tech Polo Club Interactive page Advanced

Open "Transformer Block" to see attention, the MLP, layer normalization and the residual connections of a live model, with the real numbers.
On Layer Normalization in the Transformer Architecture Xiong and others (2020) Paper Advanced

The study of pre-norm against post-norm. Its Figure 1 draws both blocks side by side.

Where this shows up

Units that build on this one

Diffusion transformers (DiT) Advanced