M3.4 — Transformer Encoders and Decoders

Learning objectives

After this lesson, you should be able to:

• Describe the two steps of a pre-norm transformer block.
• Explain the roles of LayerNorm and residual connections.
• Explain why positional information has to be added.
• Read a transformer block’s shapes and count its parameters.
• Distinguish encoder, decoder and encoder-decoder transformers.
• Explain what a causal mask guarantees.

The one-sentence idea

A transformer repeats one block, attention then a per-position MLP, each wrapped in normalisation and a residual connection, so that shape is preserved and blocks can be stacked.

Same shape in, same shape out, so you can stack it.

The transformer block

M3.3 gave you attention, which mixes information between positions. A transformer is built by repeating one block, and the block contains attention plus a few supporting parts. In the form used by many modern models (called pre-norm), a block does two things in turn:

Each line has the same pattern: normalise, apply a layer, then add the result back to the input. The input and output of the block have the same shape, which is what allows blocks to be stacked as many times as you like.

The supporting parts

Layer normalisation

LayerNorm rescales each position’s vector to have mean 0 and standard deviation 1 across its features (then applies a learned scale and shift). It keeps the numbers flowing through a deep network at a stable size, which makes training much steadier.

import numpy as np

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

x = np.array([[1.0, 2.0, 3.0, 4.0],
              [10.0, 0.0, 0.0, 0.0]])
print(layer_norm(x).round(3))
print(np.allclose(layer_norm(x).mean(axis=-1), 0), layer_norm(x).std(axis=-1).round(3))

[[-1.342 -0.447  0.447  1.342]
 [ 1.732 -0.577 -0.577 -0.577]]
True [1. 1.]

Residual connections

The “x = x + …” part is a residual connection. Each layer only has to learn a change to its input, and the original signal, and the gradient, can flow straight through the addition. This is what makes networks of dozens of blocks trainable (M3.1 mentioned it). If a layer’s contribution is zero, the block is exactly the identity.

The MLP (feed-forward) layer

After attention has mixed positions, a small two-layer network from M3.1 is applied to each position separately, usually expanding to a hidden size several times larger and back. It is where much of a transformer’s capacity sits.

Position information

Attention treats its input as a set: shuffle the tokens and each token’s output is just shuffled with them. To tell the model where each token is, a positional encoding is added to the embeddings. One classic choice builds a vector of sines and cosines of different frequencies for each position; many models instead learn the position vectors or use rotary variants.

def positional_encoding(T, d):
    pos = np.arange(T)[:, None]
    i = np.arange(d // 2)[None, :]
    angle = pos / 10000 ** (2 * i / d)
    P = np.zeros((T, d))
    P[:, 0::2] = np.sin(angle)
    P[:, 1::2] = np.cos(angle)
    return P

print(positional_encoding(3, 4).round(3))

[[ 0.     1.     0.     1.   ]
 [ 0.841  0.54   0.01   1.   ]
 [ 0.909 -0.416  0.02   1.   ]]

Each row is a distinct fingerprint for one position. Adding it to a token’s embedding lets the same token mean something different at position 0 and position 2.

A block in code

This single-head block follows the two lines above exactly. The weights are random, so the numbers are meaningless, but the shapes and the residual structure are real:

rng = np.random.default_rng(0)
T_, d, d_ff = 4, 8, 32

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

def block(x, Wq, Wk, Wv, Wo, W1, W2):
    h = layer_norm(x)
    q, k, v = h @ Wq, h @ Wk, h @ Wv
    attn = softmax(q @ k.T / np.sqrt(d)) @ v
    x = x + attn @ Wo                              # attention, added back
    h = layer_norm(x)
    return x + np.maximum(0, h @ W1) @ W2          # MLP, added back

W = [rng.normal(size=s) * 0.3 for s in [(d, d)] * 4 + [(d, d_ff), (d_ff, d)]]
x = rng.normal(size=(T_, d))
print(block(x, *W).shape)

W_off = W[:3] + [np.zeros((d, d))] + [W[4], np.zeros((d_ff, d))]    # zero both outputs
print(np.allclose(block(x, *W_off), x))

(4, 8)
True

The output has the same shape as the input. With the two output matrices set to zero, the block returns its input unchanged: the residual path is an identity, and the layers only add corrections on top of it.

In PyTorch the block is provided ready-made. Counting its parameters is a good check that you can read the structure: with d_model = 8 and a hidden size of 32, attention holds 4 × (8 × 8) weights plus biases, the MLP holds 8 × 32 and 32 × 8 weights plus biases, and two LayerNorms hold a scale and shift each.

import torch
from torch import nn

layer = nn.TransformerEncoderLayer(d_model=8, nhead=2, dim_feedforward=32,
                                   dropout=0.0, batch_first=True, norm_first=True)
x = torch.randn(2, 5, 8)                       # (batch, tokens, features)
print(layer(x).shape)
print(sum(p.numel() for p in layer.parameters()))

stack = nn.TransformerEncoder(layer, num_layers=3, enable_nested_tensor=False)    # three independent copies
print(stack(x).shape, sum(p.numel() for p in stack.parameters()))

torch.Size([2, 5, 8])
872
torch.Size([2, 5, 8]) 2616

The 872 parameters add up as 216 (query, key and value) + 72 (output projection) + 288 + 264 (the two MLP layers) + 32 (the two LayerNorms). Three stacked layers have three times as many, and the shape never changes. Large models are this same block repeated tens of times, with a much larger d_model.

Encoders, decoders and encoder-decoders

Blocks can be arranged in three ways, and which one a model uses depends on its job:

Encoder

Every position can attend to every other (no mask). It reads a whole input and produces a rich representation of it. Vision Transformers and the text and image towers of CLIP are encoders.

Decoder

A causal mask lets each position attend only to earlier ones, so it can generate a sequence one token at a time. Most large language models are decoder-only.

Encoder-decoder

An encoder reads the input, and a decoder generates the output while using cross-attention (M3.3) to read the encoder’s result.

Many VLA models are built from a pretrained decoder-style language model that is also given image tokens, so one transformer reads the camera tokens and the instruction and then produces action tokens. The roadmap’s later phases look at specific designs. Here, what matters is that the causal mask really does stop information flowing backwards, which you can check directly: change the last token and confirm that earlier outputs do not move.

mask = nn.Transformer.generate_square_subsequent_mask(5)
y1 = layer(x, src_mask=mask)
x2 = x.clone()
x2[:, -1] += 1.0                                # change only the last token
y2 = layer(x2, src_mask=mask)
print(torch.allclose(y1[:, :-1], y2[:, :-1], atol=1e-6), torch.allclose(y1[:, -1], y2[:, -1]))

True False

Earlier positions are unchanged and the last position responds, which is exactly the behaviour a generating model needs.

Common misconceptions

Engineering takeaways

1. A block is normalise, attend, add back, then normalise, MLP, add back.
2. LayerNorm keeps each position’s features at a stable scale.
3. Residual connections add a layer’s result to its input, so a zero contribution is the identity.
4. The MLP is applied to each position independently and holds much of the capacity.
5. Positional encodings tell the model where each token is.
6. Input and output shapes match, so blocks stack.
7. Encoders attend freely, decoders use a causal mask, and encoder-decoders add cross-attention.
8. Large models are the same block repeated with larger dimensions.

Learn the one block, and you can read every transformer.

Knowledge check

Try each question first, then reveal the explanation.

What’s next

You now know what a transformer is made of. Next is how these models are trained in two stages: learning general features first, then adapting to a specific task.

Back to the VLA roadmap prerequisites