DL S039 · The Full Transformer Block — Encoder, Decoder, and What Order Everything Goes In

Series: Deep Learning & LLMs From Scratch — 80 sessions · Session 39 / 80 · Module M07 · ~2 hours

The story

We have all the ingredients. Multi-head self-attention (S037). Positional encodings (S038). The FFN (a two-layer MLP, S013). LayerNorm (S016). Residual connections (S017). What's left is the recipe: which order do they go in?

The recipe matters more than beginners expect. "Attention Is All You Need" (2017) used post-LN: x + Sublayer(x) first, then LayerNorm. It works, but training is fragile — you need warmup, careful init, and it explodes at deeper depths (past ~12 layers) without special tricks. Every modern transformer (GPT-2, Llama, everything) uses pre-LN: x + Sublayer(LayerNorm(x)). Pre-LN trains stably at 100+ layers, needs less warmup, and is what you'll implement in S040.

This session we build the block twice — once as the paper wrote it (encoder + decoder), once as the modern decoder-only variant used by GPT/Llama. You'll see exactly what changed, why cross-attention exists only in the encoder-decoder version, and why decoder-only won for language modelling. By the end, you can draw the full "Attention Is All You Need" figure on a whiteboard and defend every line.

Prerequisites

• S036–S038 — attention, multi-head, positional.
• S013 — MLP forward.
• S016 — LayerNorm.
• S017 — residuals.

1 · The four ingredients

Every transformer block is built from these four:

1. Multi-head self-attention (MHA) — S037. Mixes information across positions.
2. Feed-forward network (FFN) — a two-layer MLP applied position-wise (same weights every position).
3. LayerNorm (LN) — normalises each token's features to unit mean/variance.
4. Residual connection (skip) — adds the input of a sublayer to its output.

The FFN is the compute workhorse. For d_model = 512, it's typically:

FFN(x) = Linear(d_ff, d_model)(GELU(Linear(d_model, d_ff)(x)))

with d_ff = 4 · d_model = 2048. That 4× expansion is universal — even Llama and GPT-4-ish models keep the 4× ratio (SwiGLU variants use ~2.67× · 2 = 5.33× in raw dims to hit the same effective compute, but conceptually still 4×).

Why 4×? The FFN carries ~2/3 of the parameters of a transformer block (attention is only ~1/3). Making it wider gives you more capacity per layer. Empirically, 4× is the sweet spot: bigger doesn't help much, smaller hurts.

2 · Post-LN — the original recipe (and why it's fragile)

The 2017 paper's block:

attn_out = MultiHeadAttention(x)
x = LayerNorm(x + attn_out)          # post-LN: add first, then norm
 
ffn_out = FFN(x)
x = LayerNorm(x + ffn_out)

In code:

class PostLNBlock(nn.Module):
    def __init__(self, d_model, h, d_ff):
        super().__init__()
        self.attn = MultiHeadAttention(d_model, h)
        self.ffn = nn.Sequential(
            nn.Linear(d_model, d_ff), nn.GELU(),
            nn.Linear(d_ff, d_model))
        self.ln1 = nn.LayerNorm(d_model)
        self.ln2 = nn.LayerNorm(d_model)
 
    def forward(self, x, mask=None):
        x = self.ln1(x + self.attn(x, mask))
        x = self.ln2(x + self.ffn(x))
        return x

Works fine for the paper's 6-layer encoder + 6-layer decoder. But scale to 24 layers and you'll find gradients at the input layer are wildly larger than at the output layer, training loss oscillates in early epochs, and the model needs a linear-warmup learning rate schedule to survive the first few thousand steps.

3 · Pre-LN — the modern recipe

Swap the order: normalise inside the residual, add outside:

x = x + MultiHeadAttention(LayerNorm(x))
x = x + FFN(LayerNorm(x))

class PreLNBlock(nn.Module):
    def __init__(self, d_model, h, d_ff):
        super().__init__()
        self.attn = MultiHeadAttention(d_model, h)
        self.ffn = nn.Sequential(
            nn.Linear(d_model, d_ff), nn.GELU(),
            nn.Linear(d_ff, d_model))
        self.ln1 = nn.LayerNorm(d_model)
        self.ln2 = nn.LayerNorm(d_model)
 
    def forward(self, x, mask=None):
        x = x + self.attn(self.ln1(x), mask)
        x = x + self.ffn(self.ln2(x))
        return x

Two lines swapped. Consequences: gradients through the residual go straight through (identity path), the sublayer's LN keeps sublayer inputs well-scaled, and depth scales to 100+ layers without divergence.

block = PreLNBlock(d_model=64, h=4, d_ff=256)
x = torch.randn(2, 10, 64)                    # (B, T, d_model)
y = block(x)
assert y.shape == x.shape
 
# Kill the FFN and attention entirely
with torch.no_grad():
    for p in block.attn.parameters(): p.zero_()
    for p in block.ffn.parameters():  p.zero_()
y_zero = block(x)
print((y_zero - x).abs().max().item())        # expect ~0 (float noise only)

4 · The causal mask (decoder-only or decoder side of enc-dec)

Language modelling is next-token prediction. When training on "the cat sat on the mat", the model must NOT be allowed to peek at "cat" when predicting the next token after "the". Every position i can only attend to positions 0…i.

Enforce this by adding -∞ to attention scores at "forbidden" positions before soft-max. After soft-max, those positions get weight zero.

# T = sequence length
mask = torch.tril(torch.ones(T, T))       # lower-triangular of 1s
# then inside attention:
scores = scores.masked_fill(mask == 0, float('-inf'))

The tril matrix looks like:

[[1, 0, 0, 0],
 [1, 1, 0, 0],
 [1, 1, 1, 0],
 [1, 1, 1, 1]]

Row i has 1s in columns 0…i and 0s beyond. masked_fill(mask == 0, -inf) blocks the future. After soft-max, row 0 attends only to position 0, row 1 splits between 0 and 1, etc.

5 · The encoder-decoder — original transformer

The 2017 paper is an encoder-decoder for machine translation. Left tower processes the source sentence; right tower generates the target token by token, attending to both its own previous tokens (self-attention) and the encoder's outputs (cross-attention).

Encoder block: [LN → MHA(self) → +] → [LN → FFN → +]. Bidirectional attention over source.

Decoder block: three sublayers:

1. [LN → MHA(self, causal)] — attend to previous target tokens.
2. [LN → MHA(cross, Q=decoder, K=V=encoder_out)] — pull from source.
3. [LN → FFN] — position-wise.

Cross-attention is the same math as self-attention but with Q from the decoder and K, V from the encoder. Shapes:

Q: (B, T_dec, d_model)  from LN(x_dec)
K, V: (B, T_enc, d_model)  from encoder_output (already normed)
scores: (B, h, T_dec, T_enc)

The cross-attention mask blocks padding tokens in the source but has no causal structure.

6 · Decoder-only — the modern LM

For pure language modelling (predicting the next token given the past), you don't need an encoder or cross-attention. Just stack N pre-LN decoder blocks with causal self-attention, plus a token+position embedding at the input and a linear head at the output.

class DecoderOnlyBlock(nn.Module):
    def __init__(self, d_model, h, d_ff, dropout=0.0):
        super().__init__()
        self.ln1 = nn.LayerNorm(d_model)
        self.attn = MultiHeadAttention(d_model, h, dropout)
        self.ln2 = nn.LayerNorm(d_model)
        self.ffn = nn.Sequential(
            nn.Linear(d_model, d_ff), nn.GELU(),
            nn.Linear(d_ff, d_model),
            nn.Dropout(dropout))
 
    def forward(self, x, mask):
        x = x + self.attn(self.ln1(x), mask)
        x = x + self.ffn(self.ln2(x))
        return x

And the full model:

class GPTLike(nn.Module):
    def __init__(self, vocab_size, max_len, d_model, h, d_ff, n_layers):
        super().__init__()
        self.tok_emb = nn.Embedding(vocab_size, d_model)
        self.pos_emb = nn.Embedding(max_len, d_model)         # learned PE
        self.blocks = nn.ModuleList([
            DecoderOnlyBlock(d_model, h, d_ff) for _ in range(n_layers)])
        self.ln_f = nn.LayerNorm(d_model)                     # final norm
        self.head = nn.Linear(d_model, vocab_size, bias=False)
        # tie weights (optional, saves params and improves quality)
        self.head.weight = self.tok_emb.weight
 
    def forward(self, idx):
        B, T = idx.shape
        pos = torch.arange(T, device=idx.device)
        x = self.tok_emb(idx) + self.pos_emb(pos)
        mask = torch.tril(torch.ones(T, T, device=idx.device))
        for blk in self.blocks:
            x = blk(x, mask)
        x = self.ln_f(x)
        return self.head(x)                                   # (B, T, vocab)

Roughly 50 lines. That is a GPT. We build the full training loop for this in S040.

7 · Why decoder-only won

For dedicated seq2seq tasks (machine translation, summarisation with strict input-output separation), encoder-decoder still wins on FLOPs-per-BLEU. But for general-purpose LMs and chatbots, decoder-only rules.

8 · Common bugs

9 · Parameter budget worksheet

For d_model = 512, h = 8, d_ff = 2048, n_layers = 12, vocab = 50257, max_len = 1024:

• Token embed: 50257 · 512 = 25.7M
• Position embed: 1024 · 512 = 0.5M
• Per block: 4 · d_model² (attn) + 2 · d_model · d_ff (FFN) = 4·262144 + 2·512·2048 = 1.05M + 2.10M = 3.15M. Times 12 = 37.8M.
• Head (tied with token embed) = 0.
• Final LN etc. = negligible.
• Total ≈ 64M.

That's about the size of GPT-2 small (117M — differs because GPT-2 uses d_ff = 4·d_model = 3072, n_layers = 12, d_model = 768; let's redo: 12 · (4·768² + 2·768·3072) = 12 · (2.36M + 4.72M) = 85M, plus 50257·768 = 38.6M embed = ~124M, matches "117M" once you account for weight tying reducing the effective count).

Do this arithmetic once. You'll never be intimidated by "how big is this model" again.

10 · Modern-2025 twist — the block that actually ships

The pre-LN block above is the GPT-2 (2019) shape. What every 2024–2026 frontier model actually ships is a small set of drop-in upgrades. Here they are, roughly in order of impact:

1. RMSNorm instead of LayerNorm (Zhang & Sennrich, 2019, arXiv:1910.07467). Drops the mean-subtraction and the bias term; just rescales by root-mean-square. Same-ish quality, ~7–15% faster, one less numerical failure mode. Used in Llama 1/2/3/4, Mistral, Qwen, DeepSeek, Gemma.

class RMSNorm(nn.Module):
    def __init__(self, d, eps=1e-6):
        super().__init__(); self.g = nn.Parameter(torch.ones(d)); self.eps = eps
    def forward(self, x):
        return x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps) * self.g

2. SwiGLU instead of GELU-FFN (Shazeer, 2020, arXiv:2002.05202). Replace the 2-layer MLP with a gated variant. Consistently ~1–2% lower loss at fixed compute. Universal in modern LMs.

class SwiGLU(nn.Module):
    def __init__(self, d, d_ff):
        super().__init__()
        self.w1 = nn.Linear(d, d_ff, bias=False)   # gate
        self.w2 = nn.Linear(d, d_ff, bias=False)   # value
        self.w3 = nn.Linear(d_ff, d, bias=False)   # out
    def forward(self, x):
        return self.w3(F.silu(self.w1(x)) * self.w2(x))

Because SwiGLU has 3 matrices instead of 2, Llama shrinks d_ff from 4d to ~2.67d so the parameter count matches a 4d GELU FFN.

3. No biases anywhere (Llama). Empirically neutral quality but simpler, slightly faster. All nn.Linear(..., bias=False).

4. RoPE instead of learned/sinusoidal positions — covered in S038 and S043. Baked into attention, not added to embeddings.

5. GQA / MQA / MLA — shrinks KV heads. Covered in S037/S042.

The "modern block" — all five upgrades — is essentially the Llama 3 block. Eight-line class:

class ModernBlock(nn.Module):
    def __init__(self, d, n_head, n_kv_head, d_ff):
        super().__init__()
        self.n1 = RMSNorm(d); self.attn = GQAAttention(d, n_head, n_kv_head)
        self.n2 = RMSNorm(d); self.ffn  = SwiGLU(d, d_ff)
    def forward(self, x, cos, sin, kv_cache=None):
        h, kv = self.attn(self.n1(x), cos, sin, kv_cache)
        x = x + h; x = x + self.ffn(self.n2(x)); return x, kv

That's the block behind every open weights release from mid-2023 on. Read Llama 3's model.py; this is essentially what you'll see.

2025 experimental additions worth knowing about:

• DeepNorm (post-LN variant, Wang et al., 2022, arXiv:2203.00555) — scale residual by α to stabilise 1000-layer training. Used in some Chinese-lab models.
• Value residual (modded-nanoGPT, 2024) — add a residual on V inside attention.
• QK-norm (Chameleon, Grok) — RMSNorm on Q and K separately before the dot product. Kills attention-logit explosions at scale.
• Tanh logit soft-capping (Gemma 2, Gemma 2 tech report, 2024) — logits = cap * tanh(logits / cap). Prevents softmax saturation.
• U-net-shaped residual streams (modded-nanoGPT speedrun 2025) — experimental, decoupled residual streams reduce loss further.

Further reading:

• Llama 2 paper (Touvron et al., 2023) §2.1 — the modern-block recipe stated cleanly.
• SwiGLU original (Shazeer, 2020) — the two-page paper that changed every FFN.
• modded-nanoGPT PR history — real-time architecture research.

import torch

def grad_at_input(depth, pre_ln):
    x = torch.randn(2, 16, 64, requires_grad=True)
    h = x
    for _ in range(depth):
        if pre_ln:
            h = h + branch(norm(h))     # residual stream untouched
        else:
            h = norm(h + branch(h))     # every add passes through norm
    h.sum().backward()
    return x.grad.norm().item()

for depth in (2, 6, 12, 24, 48):
    print(depth, grad_at_input(depth, True), grad_at_input(depth, False))
# pre-LN column: roughly stable as depth grows
# post-LN column: falls away with depth -- and that decay is
# exactly what the warmup schedule was compensating for

Retention scaffold

Recall questions

Stretch prompt

Implement the encoder-decoder transformer (six blocks each side) and train it on a small English→German dataset (like Multi30k). Compare BLEU to a decoder-only trained on the same data as a concatenated EN <sep> DE sequence. Decoder-only will be within a couple BLEU points despite being architecturally simpler — that's the modern LM story in miniature.

In your own words

"Explain to a friend, in 4–5 sentences, what makes a transformer 'a transformer' as opposed to just 'a stack of attention layers'."

Spaced review

• S016 — LayerNorm, the piece that makes deep transformers train.
• S017 — residuals, without which nothing above 4 layers would learn.
• S037 — MHA, the workhorse.

Next-session teaser

S040: Karpathy's nanoGPT, line by line. We take the ~300 lines of model.py and understand every symbol. If you followed today's session, tomorrow you'll be typing it from memory.

Bring back tomorrow

• Pre-LN order (LN inside the residual).
• FFN = 4× hidden, GELU or SwiGLU, position-wise.
• Causal mask via tril.
• Decoder-only is the default modern LM architecture.

Previous: ← DL S038 · Next: DL S040 →