Inside the Transformer Block: Residuals and Normalisation

A transformer block is where the residual stream lives. The input X of shape [B, T, d] flows through two sub-layers — attention and MLP — each wrapped in a residual connection. The residual stream is the highway that lets depth grow without changing representation size. This lesson traces that highway with explicit shapes, compares the normalization and activation choices that define modern decoder-only models, and builds a complete runnable block.

The one-sentence version

A decoder block preserves the [B, T, d] residual stream through two pre-norm sub-layers — attention and SwiGLU MLP — each adding its output back to the stream via residual connections; LayerNorm/RMSNorm and GELU/SwiGLU are the normalization and activation choices that distinguish model families.

Everything else in this lesson makes that concrete.

Lab: Run the companion CPU exercise at /labs/track1/02-tiny-decoder/ (residual stream, RMSNorm, SwiGLU, causal safety).

1. The residual stream and sub-layer structure

A decoder block receives X of shape [B, T, d] — the residual stream — and processes it through two sub-layers:

Sub-layer 1 (Attention):
  X1 = X + Attention(Norm(X))

Sub-layer 2 (MLP):
  X2 = X1 + MLP(Norm(X1))

Key invariant: The residual stream X maintains shape [B, T, d] throughout. Every sub-layer input and output has this shape. The residual connections are what allow depth to increase without changing representation size.

This is pre-norm placement: normalization happens inside the residual branch, before the sub-layer. The original Transformer used post-norm: Norm(X + SubLayer(X)). Pre-norm gives gradients a direct path through the residual connection, improving training stability at depth.

Pre-norm vs post-norm

Pre-norm (GPT-2, GPT-3, Llama, PaLM, Gemma): X + SubLayer(Norm(X)). Gradients flow directly through the residual path. Better for deep models. Post-norm (original Transformer): Norm(X + SubLayer(X)). Gradients must pass through Norm. Can be unstable at depth. Modern decoder-only models almost universally use pre-norm.

2. LayerNorm vs RMSNorm

Both normalization layers are applied per-token (independent across batch and sequence). For input x of shape [d]:

LayerNorm (Ba et al., 2016)

μ = mean(x)
σ² = var(x)
x̂ = (x - μ) / sqrt(σ² + ε)
y = x̂ * γ + β

• Parameters: γ (gain, [d]), β (bias, [d])
• Operation: centers (subtracts mean) then scales
• Used in: GPT-2, BERT, original Transformer

RMSNorm (Zhang & Sennrich, 2019)

RMS = sqrt(mean(x²) + ε)
x̂ = x / RMS
y = x̂ * γ

• Parameters: γ (gain, [d]), no β
• Operation: scales only, no centering
• Used in: Llama 2/3, PaLM, Gemma, Mistral

LayerNorm vs RMSNorm

	LayerNorm	RMSNorm
Centers data	Yes (subtracts mean)	No
Scales data	Yes (gain γ)	Yes (gain γ)
Bias parameter β	Yes	No
Parameters per d	2d	d
Operations	mean, var, sqrt	mean(x²), sqrt
Used in	GPT-2, BERT, original Transformer	Llama 2/3, PaLM, Gemma, Mistral

Key differences:

• LayerNorm re-centers and re-scales; RMSNorm only re-scales
• RMSNorm saves d parameters (no β) and one reduction (no mean computation)
• Empirically similar performance; RMSNorm slightly faster
• Both maintain the [B, T, d] shape

Which to choose?

RMSNorm is the modern default for decoder-only models (Llama, PaLM, Gemma, Mistral). Fewer parameters, slightly faster, no centering needed. LayerNorm remains standard for encoder models (BERT) and GPT-2 compatibility. If building a new decoder-only model today: default to RMSNorm.

3. GELU vs SwiGLU in the MLP

The MLP (Feed-Forward Network) expands each position from d to a hidden width (typically 4d), applies a non-linearity, and projects back to d. Positions are transformed independently.

GELU (Gaussian Error Linear Unit)

GELU(x) = 0.5 * x * (1 + tanh(√(2/π) * (x + 0.044715 * x³)))

• Smooth approximation to ReLU
• Single projection: W₁: [d, 4d], then W₂: [4d, d]
• Used in GPT-2, GPT-3, BERT

SwiGLU (Shazeer, 2020) — Swish-Gated Linear Unit

SwiGLU(x) = (x @ W_G + b_G) ⊙ SiLU(x @ W_U + b_U)

Where SiLU(x) = x * sigmoid(x), and ⊙ is element-wise multiplication.

• Gated activation: splits expansion into gate and value branches
• Two projections: W_G: [d, 2d], W_U: [d, 2d], then W₂: [2d, d]
• In practice: combined projection W: [d, 4d], split → SiLU(gate) * val → W_out: [2d, d]
• Used in Llama 2/3, PaLM, Gemma, Mistral

GELU vs SwiGLU MLP expansion (d → 4d → d)

	GELU	SwiGLU
Expansion matrix	One: [d, 4d]	Two (gate/val): [d, 2d] each
Activation	GELU (single)	SiLU gate ⊙ value (gated)
Output of expansion	[B, T, 4d]	[B, T, 2d] (after gate*val)
Output projection	[4d, d]	[2d, d]
Total expansion params	4d²	4d² (same!)
Used in	GPT-2, GPT-3, BERT	Llama 2/3, PaLM, Gemma, Mistral

Key insight: Both use 4d² expansion parameters. SwiGLU adds a gating mechanism that empirically improves quality per parameter.

Which to choose?

SwiGLU is the modern default for decoder-only LLMs (Llama, PaLM, Gemma, Mistral). Gating improves quality per parameter. GELU remains standard for GPT-2/3 compatibility and encoder models. If building a new decoder-only model today: default to SwiGLU.

3. Fully worked numeric example (verified, seed 42)

Setup (toy decoder block):

• B = 1, T = 4, d = 8
• Pre-norm decoder block with single-head attention
• RMSNorm (no bias)
• SwiGLU MLP with expansion factor 4 (hidden = 32)

import torch
import torch.nn as nn
import torch.nn.functional as F

torch.manual_seed(42)
torch.set_default_dtype(torch.float32)

B, T, d = 1, 4, 8
X = torch.randn(B, T, d)  # [1, 4, 8] — residual stream input

# --- Sub-layer 1: Attention with pre-norm RMSNorm ---

# RMSNorm
eps = 1e-6
rms = torch.sqrt(X.pow(2).mean(dim=-1, keepdim=True) + 1e-6)  # [1, 4, 1]
γ_attn = torch.ones(d)  # gain only, no bias
X_norm = X / rms * γ_attn  # [1, 4, 8]

# Single-head attention (d_k = d_v = 8)
W_Q = torch.randn(d, d)
W_K = torch.randn(d, d)
W_V = torch.randn(d, d)

Q = X_norm @ W_Q  # [1, 4, 8]
K = X_norm @ W_K  # [1, 4, 8]
V = X_norm @ W_V  # [1, 4, 8]

scores = Q @ K.transpose(-2, -1) / (d ** 0.5)  # [1, 4, 4]
attn_weights = F.softmax(scores, dim=-1)  # [1, 4, 4]
attn_out = attn_weights @ V  # [1, 4, 8]

# Residual
X1 = X + attn_out  # [1, 4, 8]

# --- Sub-layer 2: MLP with pre-norm RMSNorm and SwiGLU ---

rms = torch.sqrt(X1.pow(2).mean(dim=-1, keepdim=True) + 1e-6)
γ_mlp = torch.ones(d)
X1_norm = X1 / rms * γ_mlp  # [1, 4, 8]

# SwiGLU: d=8 → 32 → 16 → 8
W_swiglu = torch.randn(d, 32)  # combined projection
W_mlp_out = torch.randn(16, d)  # output projection

x_proj = X1_norm @ W_swiglu  # [1, 4, 32]
x_gate, x_val = x_proj.chunk(2, dim=-1)  # each [1, 4, 16]
x_swish = F.silu(x_gate) * x_val  # [1, 4, 16]
mlp_out = x_swish @ W_mlp_out  # [1, 4, 8]

# Residual
X2 = X1 + mlp_out  # [1, 4, 8] — final block output

Shape trace summary:

Step	Tensor	Shape
Input	X	[1, 4, 8]
After RMSNorm (attn)	X_norm	[1, 4, 8]
Attention output	attn_out	[1, 4, 8]
After residual 1	X1	[1, 4, 8]
After RMSNorm (MLP)	X1_norm	[1, 4, 8]
SwiGLU projection	x_proj	[1, 4, 32]
After gate*val	x_swish	[1, 4, 16]
MLP output	mlp_out	[1, 4, 8]
Final output	X2	[1, 4, 8]

The residual stream shape [B, T, d] is preserved throughout.

4. Pre-norm vs Post-norm placement

Pre-norm vs Post-norm

Pre-norm places normalization inside the residual branch; post-norm places it after the residual addition. Pre-norm gives gradients a direct path through the residual connection.

Read this diagram as text

Text equivalent. Left panel (pre-norm): input X flows to Norm, then SubLayer, then adds to X. Right panel (post-norm): input X flows to SubLayer, adds to X, then Norm. Pre-norm gives the residual path a direct gradient route; post-norm forces gradients through the normalization layer.

Aspect	Pre-norm	Post-norm
Structure	X + SubLayer(Norm(X))	Norm(X + SubLayer(X))
Gradient path	Direct through residual	Through Norm
Depth stability	Excellent	Degrades at depth
Used in	GPT-2, GPT-3, Llama, PaLM, Gemma	Original Transformer

Why pre-norm wins for deep models: The residual connection X + ... provides a direct gradient highway. In pre-norm, this highway is unobstructed. In post-norm, gradients must pass through the normalization layer, which can amplify or diminish them unpredictably at depth.

5. GELU vs SwiGLU expansion detail

MLP expansion: GELU vs SwiGLU

GELU uses a single [d, 4d] projection; SwiGLU uses a [d, 4d] projection split into gate and value, then element-wise SiLU(gate) * value, then [2d, d] output projection. Both use 4d² expansion parameters.

Read this diagram as text

Text equivalent. Top path (GELU): input [B,T,d] projects to [B,T,4d] via W₁, applies GELU, projects to [B,T,d] via W₂. Bottom path (SwiGLU): input [B,T,d] projects to [B,T,4d] via W, splits into gate and value [B,T,2d] each, computes SiLU(gate)*value to get [B,T,2d], projects to [B,T,d] via W_out. Both paths preserve [B,T,d] input/output.

GELU expansion (single projection)

# GELU: d → 4d → d
W₁ = nn.Linear(d, 4*d)      # [d, 4d]
W₂ = nn.Linear(4*d, d)      # [4d, d]

x_proj = x @ W₁              # [B, T, 4d]
x_act = F.gelu(x_proj)       # [B, T, 4d]
out = x_act @ W₂             # [B, T, d]

SwiGLU expansion (gated, combined projection)

# SwiGLU: d → 4d (split) → 2d → d
W = nn.Linear(d, 4*d)       # [d, 4d] combined gate+value
W_out = nn.Linear(2*d, d)   # [2d, d]

x_proj = x @ W               # [B, T, 4d]
gate, val = x_proj.chunk(2, dim=-1)  # each [B, T, 2d]
x_swish = F.silu(gate) * val  # [B, T, 2d]
out = x_swish @ W_out        # [B, T, d]

Parameter count is identical (4d² for expansion), but SwiGLU’s gating empirically improves quality.

6. Runnable decoder block

import torch
import torch.nn as nn
import torch.nn.functional as F

class RMSNorm(nn.Module):
    def __init__(self, d, eps=1e-6):
        super().__init__()
        self.eps = eps
        self.weight = nn.Parameter(torch.ones(d))

    def forward(self, x):
        rms = torch.sqrt(x.pow(2).mean(dim=-1, keepdim=True) + self.eps)
        return x / rms * self.weight

class SwiGLU(nn.Module):
    def __init__(self, d, hidden_mult=4):
        super().__init__()
        hidden = d * hidden_mult
        self.w = nn.Linear(d, hidden, bias=False)  # combined gate+value
        self.w_out = nn.Linear(hidden // 2, d, bias=False)

    def forward(self, x):
        x_proj = self.w(x)              # [B, T, 4d]
        gate, val = x_proj.chunk(2, dim=-1)  # each [B, T, 2d]
        x_swish = F.silu(gate) * val    # [B, T, 2d]
        return self.w_out(x_swish)      # [B, T, d]

class DecoderBlock(nn.Module):
    def __init__(self, d, n_heads=4):
        super().__init__()
        self.norm1 = RMSNorm(d)
        self.attn = nn.MultiheadAttention(d, n_heads, batch_first=True)
        self.norm2 = RMSNorm(d)
        self.mlp = SwiGLU(d)

    def forward(self, x, attn_mask=None):
        # Sub-layer 1: Attention with pre-norm
        x_norm = self.norm1(x)
        T = x_norm.size(1)
        # is_causal=True applies a lower-triangular mask internally so that
        # position t can attend only to positions 0..t. This makes the block
        # causally safe for decoder-only (auto-regressive) use.
        # PyTorch requires attn_mask to be non-None when is_causal=True; pass
        # an all-zero additive mask and let is_causal do the masking.
        if attn_mask is None:
            attn_mask = torch.zeros(T, T, device=x.device, dtype=x.dtype)
        attn_out, _ = self.attn(x_norm, x_norm, x_norm,
                                attn_mask=attn_mask, need_weights=False,
                                is_causal=True)
        x = x + attn_out

        # Sub-layer 2: MLP with pre-norm
        x = x + self.mlp(self.norm2(x))
        return x

# Test 1: shape preservation
torch.manual_seed(42)
block = DecoderBlock(d=256, n_heads=8)
X = torch.randn(2, 128, 256)  # [B=2, T=128, d=256]
out = block(X)
print(f"Input: {X.shape}, Output: {out.shape}")  # Both [2, 128, 256]
assert out.shape == X.shape, "Residual stream shape not preserved"

# Test 2: causal safety — future tokens cannot affect earlier outputs
block2 = DecoderBlock(d=32, n_heads=4)
X_base = torch.randn(1, 4, 32)
X_corrupt = X_base.clone()
X_corrupt[0, 2, :] = 999.0  # Corrupt position 2 (future for outputs 0 and 1)
out_base = block2(X_base)
out_corrupt = block2(X_corrupt)
# Positions 0 and 1 must be identical (position 2 is future for them)
assert torch.allclose(out_base[0, 0], out_corrupt[0, 0], atol=1e-5), \
    "Position 0 affected by future position 2 — block not causally safe!"
assert torch.allclose(out_base[0, 1], out_corrupt[0, 1], atol=1e-5), \
    "Position 1 affected by future position 2 — block not causally safe!"
print("PASS: Future tokens cannot affect earlier outputs (causally safe).")
print(f"Input: {X.shape}, Output: {out.shape}")  # Both [2, 128, 256]

7. Engineering trade-offs

Normalization and activation choices for decoder-only models

Choice	When to prefer	Trade-off
Pre-norm	Default for decoder-only	Better gradient flow, stable at depth
Post-norm	Encoder-decoder, shallow models	Original Transformer compatibility
RMSNorm	Modern decoder-only (Llama, PaLM, Gemma)	Fewer params, faster, no centering
LayerNorm	BERT, GPT-2, encoder models	Centers data, more established
SwiGLU	Modern LLMs (Llama, PaLM, Gemma)	Gating improves quality/param
GELU	GPT-2/3, BERT, compatibility	Simpler, well-understood

Default for new decoder-only models (2024+): Pre-norm + RMSNorm + SwiGLU. This is the Llama/Gemma/PaLM/Mistral stack and offers the best quality/parameter/compute trade-off.

8. Failure modes

Failure mode	Symptom	Root cause	Test
Residual shape mismatch	Runtime error on x + sublayer(x)	Sub-layer output shape ≠ input shape	Assert output.shape == input.shape
Norm applied to wrong tensor	Training instability, different from reference	Norm after residual instead of before (or vice versa)	Compare against known-good implementation
Missing residual	Representation collapses, can’t train deep	Forgot x + in forward pass	Assert gradient flows to input
Norm params not learned	Norm has no effect	weight not registered as nn.Parameter	Check list(module.parameters())
SwiGLU gate/value split wrong	Wrong output shape, silent quality loss	Chunk dimension or size incorrect	Assert intermediate shapes
RMSNorm without eps	NaN with zero inputs	Division by zero	Test with all-zero input
Not causally safe	Future tokens affect earlier outputs	No causal mask / is_causal=True on attention	Corrupt future position, check earlier outputs unchanged

9. Related lessons

Related lessons

• LLM Internals & Architecture · Transformer architecture and model inputsTransformer Architecture Explained: Trace a Decoder-Only LLM14 min
• LLM Internals & Architecture · Attention mechanisms from scratchSelf-Attention: Queries, Keys, and Values14 min
• LLM Internals & Architecture · Attention mechanisms from scratchCausal and Padding Masks: Implementation and Tests15 min
• LLM Internals & Architecture · Attention mechanisms from scratchMulti-Head Attention: Splitting, Merging, and Projecting14 min
• LLM Internals & Architecture · Building and training a language modelPositional Information: Absolute Embeddings and RoPE13 min

9. Sources

Source	Relevance
Vaswani et al., “Attention Is All You Need” (arXiv:1706.03762), §3.1, §3.3, §5.1	Original Transformer: post-norm, LayerNorm, FFN with ReLU, residual structure
Ba et al., “Layer Normalization” (arXiv:1607.06450)	LayerNorm formulation: mean/var centering, gain/bias
Zhang & Sennrich, “Root Mean Square Layer Normalization” (arXiv:1910.07467)	RMSNorm formulation: RMS scaling, gain only, no bias
Shazeer, “GLU Variants Improve Transformer” (arXiv:2002.05202)	SwiGLU: gated linear unit with SiLU
Hendrycks & Gimpel, “Gaussian Error Linear Units (GELU)” (arXiv:1606.08415)	GELU formulation and tanh approximation
GPT-2 model source: model.py	Pre-norm, LayerNorm, GELU, residual structure
Llama 2 model source: model.py	Pre-norm, RMSNorm, SwiGLU, RoPE, residual structure
PyTorch nn.LayerNorm, nn.RMSNorm (v2.4+) docs	Official implementations
PyTorch F.silu, F.gelu docs	Official activation functions

10. Exercises

Try it yourself

1. 01Implement RMSNorm from scratch (without using nn.RMSNorm) and verify it matches PyTorch's nn.RMSNorm on random input [2, 16, 512].

   Show a hint

   Compute RMS = sqrt(mean(x²) + eps), then return x / RMS * weight. No bias parameter.

2. 02Given a SwiGLU with d=512, what are the shapes of the combined projection weight and output projection weight? How many parameters in total for the MLP (including output projection)?

   Show a hint

   Combined projection: [d, 4d]. Output projection: [2d, d]. Total = 4d² + 2d² = 6d².

3. 03Modify the DecoderBlock to use LayerNorm instead of RMSNorm, and GELU instead of SwiGLU. What changes in the forward pass?

   Show a hint

   Replace RMSNorm with nn.LayerNorm(d), SwiGLU with nn.Sequential(nn.Linear(d, 4*d), nn.GELU(), nn.Linear(4*d, d)). The residual structure stays identical.

---