Module 06: Transformer

Introduction

The transformer decoder block is the building block of GPT-style language models. GPT-2, GPT-3, and LLaMA stack 12 to 96 of these blocks.

This module combines everything built so far:

  • Multi-head attention from Module 05
  • Feed-forward networks (mini neural networks for each token)
  • Layer normalization (stabilizes training)
  • Residual connections (enables deep networks)

Each block performs two operations:

  1. Multi-head attention: Tokens communicate with each other
  2. Feed-forward network: Each token is processed independently

Decoder-Only vs Encoder-Decoder

This module implements decoder-only transformers (GPT-style). The two main transformer architectures:

Architecture Examples Use Case Attention
Decoder-only GPT, LLaMA, Claude Text generation Causal (can’t see future)
Encoder-Decoder T5, BART, original Transformer Translation, summarization Bidirectional encoder + causal decoder

Most modern LLMs use decoder-only architecture because it is simpler; we follow that approach.

What You’ll Learn

After this module, you can:

  • Understand the complete GPT-style transformer architecture
  • Implement LayerNorm, GELU, and feed-forward networks from scratch
  • Build a full transformer block with residual connections
  • Assemble a complete language model from components
  • Calculate parameter counts for different model sizes
  • Tie the LM head to the token embedding (logits = h·Eᵀ) and ablate the parameter saving

Prerequisites

This module requires familiarity with:

Complete Model Architecture

TipInteractive Architecture Walkthrough

Use the slider above to step through the forward pass. Each stage shows how tensor shapes transform as data flows through the model: - Input: Raw token IDs (integers) - Embeddings: Dense vectors capturing meaning and position - Blocks: Iterative refinement through attention and FFN - Output: Probability distribution over vocabulary

Single Transformer Block (Pre-Norm)

The key innovation is the residual connections (the + nodes). Instead of y = f(x), we compute y = x + f(x). This:

  • Helps gradients flow through deep networks
  • Makes it easy to learn identity (just set f(x) = 0)
  • Enables training of 100+ layer networks

The Components

We build each component from scratch, then show the PyTorch equivalents. The pattern: understand the math, implement it simply, then see how PyTorch optimizes it.

LayerNorm from Scratch

The Idea: Activations drift to extreme values during training, causing gradients to explode or vanish. Layer normalization fixes this by normalizing each token’s embedding to zero mean and unit variance, then applying learnable scale and shift.

The Formula:

\text{LayerNorm}(x) = \gamma \times \frac{x - \mu}{\sqrt{\sigma^2 + \epsilon}} + \beta

where:

  • \mu = mean across the embedding dimension
  • \sigma^2 = variance across the embedding dimension
  • \gamma (gamma) = learnable scale parameter (initialized to 1)
  • \beta (beta) = learnable shift parameter (initialized to 0)
  • \epsilon = small constant for numerical stability (typically 1e-5)

From Scratch Implementation:

import numpy as np
import torch
import torch.nn as nn

class LayerNormScratch:
    """Layer normalization from scratch using NumPy-style operations."""

    def __init__(self, dim, eps=1e-5):
        # Learnable parameters
        self.gamma = np.ones((dim,))   # scale (initialized to 1)
        self.beta = np.zeros((dim,))   # shift (initialized to 0)
        self.eps = eps

    def __call__(self, x):
        """
        Args:
            x: input array of shape (..., dim)
        Returns:
            normalized array of same shape
        """
        # Step 1: Compute mean across last dimension
        mean = x.mean(axis=-1, keepdims=True)

        # Step 2: Compute variance across last dimension
        var = ((x - mean) ** 2).mean(axis=-1, keepdims=True)

        # Step 3: Normalize (the "norm" in LayerNorm)
        x_norm = (x - mean) / np.sqrt(var + self.eps)

        # Step 4: Scale and shift with learnable parameters
        return self.gamma * x_norm + self.beta


# Test our from-scratch implementation
x = np.array([[2.0, 4.0, 6.0, 8.0],
              [1.0, 2.0, 3.0, 4.0]])

ln_scratch = LayerNormScratch(dim=4)
out_scratch = ln_scratch(x)

print("LayerNorm from Scratch:")
print(f"  Input:\n{x}")
print(f"  Output:\n{np.round(out_scratch, 4)}")
print(f"  Output mean per row: {out_scratch.mean(axis=-1).round(6)}")
print(f"  Output std per row: {out_scratch.std(axis=-1).round(4)}")
LayerNorm from Scratch:
  Input:
[[2. 4. 6. 8.]
 [1. 2. 3. 4.]]
  Output:
[[-1.3416 -0.4472  0.4472  1.3416]
 [-1.3416 -0.4472  0.4472  1.3416]]
  Output mean per row: [0. 0.]
  Output std per row: [1. 1.]

PyTorch’s nn.LayerNorm:

# PyTorch's optimized implementation
ln_pytorch = nn.LayerNorm(4, elementwise_affine=True)

# Initialize to match our scratch version (gamma=1, beta=0)
nn.init.ones_(ln_pytorch.weight)
nn.init.zeros_(ln_pytorch.bias)

x_torch = torch.tensor(x, dtype=torch.float32)
out_pytorch = ln_pytorch(x_torch)

print("PyTorch LayerNorm:")
print(f"  Output:\n{out_pytorch.detach().numpy().round(4)}")
print(f"  Matches scratch: {np.allclose(out_scratch, out_pytorch.detach().numpy(), atol=1e-5)}")
PyTorch LayerNorm:
  Output:
[[-1.3416 -0.4472  0.4472  1.3416]
 [-1.3416 -0.4472  0.4472  1.3416]]
  Matches scratch: True

Key Insight: LayerNorm is just normalize-scale-shift. The learnable \gamma and \beta let the network undo the normalization if needed, but start from a stable baseline. Unlike BatchNorm, LayerNorm normalizes across features (embedding dimension) rather than across batch, making it suitable for variable-length sequences.

RMSNorm: The Modern Simplification

Every model since LLaMA — Mistral, Qwen, Gemma, DeepSeek — replaces LayerNorm with RMSNorm. It is LayerNorm with two pieces removed. Understanding which two, and why removing them is safe, is the whole lesson.

The Idea: LayerNorm does two jobs at once. It re-centers (subtracts the mean \mu so the vector sits at zero) and it re-scales (divides by the standard deviation \sigma so the vector has unit spread). Zhang & Sennrich (2019) asked a sharp question: which job actually stabilizes training? They found the answer is re-scaling. The re-centering step contributes almost nothing. So RMSNorm throws it away — no mean, no variance, no shift \beta — and keeps only the rescaling by the vector’s root mean square:

The Formula:

\text{RMSNorm}(x) = \gamma \odot \frac{x}{\sqrt{\frac{1}{d}\sum_{i=1}^{d} x_i^2 + \epsilon}}

Compare it term-by-term with LayerNorm:

LayerNorm RMSNorm
Re-center (subtract \mu) ✅ yes ❌ no
Re-scale (divide) by \sigma (std) by RMS
Learnable gain \gamma ✅ ✅
Learnable shift \beta ✅ ❌ no
Params (dim d) 2d d
Statistics per token mean and variance one sum of squares

When the input already has roughly zero mean — which residual streams do — \sigma \approx \text{RMS}, so the two normalizers produce nearly identical vectors. RMSNorm just gets there with one pass over the features instead of two, and half the parameters. At the scale of a modern LLM that shaves a real slice off every forward and backward pass.

NoteKey Insight

RMSNorm is not a different kind of normalization — it is LayerNorm with the re-centering deleted. It works because, empirically, only the re-scaling matters for gradient stability. Fewer operations, half the parameters, the same training behavior.

From Scratch Implementation:

class RMSNormScratch:
    """RMS normalization from scratch — LayerNorm minus mean-centering and beta."""

    def __init__(self, dim, eps=1e-5):
        self.gamma = np.ones((dim,))   # scale (initialized to 1); no beta
        self.eps = eps

    def __call__(self, x):
        # Mean of the SQUARES (not of x) — no mean subtraction happens.
        mean_square = (x ** 2).mean(axis=-1, keepdims=True)
        # Divide by the root mean square, then apply the learnable gain.
        return self.gamma * x / np.sqrt(mean_square + self.eps)


# Reuse the same input as the LayerNorm example above
rms_scratch = RMSNormScratch(dim=4)
out_rms = rms_scratch(x)

print("RMSNorm from Scratch:")
print(f"  Input:\n{x}")
print(f"  Output:\n{np.round(out_rms, 4)}")
print(f"  Output mean per row: {out_rms.mean(axis=-1).round(4)}  <- NOT forced to 0")
print(f"  Output RMS per row:  {np.sqrt((out_rms ** 2).mean(axis=-1)).round(4)}  <- forced to ~1")
RMSNorm from Scratch:
  Input:
[[2. 4. 6. 8.]
 [1. 2. 3. 4.]]
  Output:
[[0.3651 0.7303 1.0954 1.4606]
 [0.3651 0.7303 1.0954 1.4606]]
  Output mean per row: [0.9129 0.9129]  <- NOT forced to 0
  Output RMS per row:  [1. 1.]  <- forced to ~1

Notice the difference from LayerNorm: RMSNorm does not drive the mean to zero. It only fixes the magnitude (RMS ≈ 1). The mean is left wherever the data puts it — that is exactly the re-centering step we dropped.

How close is it to LayerNorm? For a residual stream, very. Let’s normalize the same near-zero-mean vector both ways and measure the gap:

# A realistic activation vector at embedding width 512
rng = np.random.default_rng(0)
v = rng.standard_normal((1, 512))

ln_out = ln_scratch.__class__(dim=512)(v)   # LayerNorm (gamma=1, beta=0)
rms_out = RMSNormScratch(dim=512)(v)         # RMSNorm  (gamma=1)

diff = np.abs(ln_out - rms_out).max()
print(f"Input sample mean:  {v.mean():.4f}  (already close to 0)")
print(f"Max abs difference: {diff:.4f}")
print("→ The gap between LayerNorm and RMSNorm is exactly the input mean.")
print("  With a wide, near-zero-mean vector, the two are nearly identical.")
Input sample mean:  -0.0240  (already close to 0)
Max abs difference: 0.0246
→ The gap between LayerNorm and RMSNorm is exactly the input mean.
  With a wide, near-zero-mean vector, the two are nearly identical.

The transformer.py module ships a production RMSNorm(nn.Module) and a make_norm(norm_type, dim) factory. Building a modern, LLaMA-style stack is a one-word change — norm_type="rmsnorm":

from transformer import RMSNorm, GPTModel

# Drop-in normalization layer
norm = RMSNorm(dim=8)
sample = torch.randn(2, 5, 8)
print(f"RMSNorm output shape: {tuple(norm(sample).shape)}")
print(f"RMSNorm params (dim=8): {sum(p.numel() for p in norm.parameters())}  (LayerNorm would be 16)")

# Swap every norm in the whole model with one argument
modern = GPTModel(vocab_size=100, embed_dim=32, num_heads=4, num_layers=2,
                  max_seq_len=64, norm_type="rmsnorm")
print(f"Block norm: {type(modern.blocks[0].ln1).__name__}")
print(f"Final norm: {type(modern.ln_final).__name__}")
RMSNorm output shape: (2, 5, 8)
RMSNorm params (dim=8): 8  (LayerNorm would be 16)
Block norm: RMSNorm
Final norm: RMSNorm
WarningDon’t forget: RMSNorm keeps the mean

The whole point is that RMSNorm is not invariant to a constant shift of the input — if you add a DC offset to a vector, LayerNorm removes it but RMSNorm passes it through (rescaled). That is fine inside a transformer because the residual stream stays near zero-mean, but it means RMSNorm is not a literal drop-in if some upstream code was relying on LayerNorm to re-center for it.

See the Difference: Shift the Input

Add a constant offset to a feature vector and watch what each normalizer does. LayerNorm’s output never moves — re-centering erases the offset. RMSNorm’s output slides with it, because it only fixes the magnitude, never the center.

TipTry This
  1. Slide the offset. The orange LayerNorm bars stay frozen no matter where you drag — re-centering deletes the offset. The pink RMSNorm bars grow and shrink because the offset changes the vector’s magnitude.
  2. Return to offset 0. With a near-zero-mean input the LayerNorm and RMSNorm bars are almost the same height — this is why the swap is nearly free inside a transformer.
  3. Push the offset to ±3. Now the two disagree strongly. That gap is exactly the re-centering term LayerNorm applies and RMSNorm skips.

Dropout: Regularization by Noise

The Idea: During training, dropout zeros a random subset of activations. The network learns redundant representations rather than depending on any single feature. The key trick: scale remaining values by \frac{1}{1-p} so the expected value stays the same.

Why it works:

  • Forces the network to learn redundant representations
  • Acts like training an ensemble of sub-networks
  • At inference time, use all neurons (no dropout)

From Scratch Implementation:

class DropoutScratch:
    """Dropout from scratch."""

    def __init__(self, p=0.1):
        """
        Args:
            p: probability of dropping each element (not keeping!)
        """
        self.p = p

    def __call__(self, x, training=True):
        """
        Args:
            x: input array
            training: if False, return x unchanged
        Returns:
            x with dropout applied (if training)
        """
        if not training or self.p == 0:
            return x

        # Create random mask: True where we KEEP the value
        keep_prob = 1 - self.p
        mask = np.random.random(x.shape) < keep_prob

        # Apply mask and scale by 1/(1-p)
        # This keeps the expected value the same:
        # E[x * mask / keep_prob] = x * keep_prob / keep_prob = x
        return x * mask / keep_prob


# Demonstrate dropout
np.random.seed(42)
x = np.ones((2, 8))

dropout = DropoutScratch(p=0.5)

print("Dropout from Scratch (p=0.5):")
print(f"  Input (all 1s): {x[0]}")

# Apply dropout multiple times to see the randomness
for i in range(3):
    np.random.seed(i)
    out = dropout(x, training=True)
    print(f"  Trial {i+1}: {out[0].round(2)}")
    print(f"    Mean: {out[0].mean():.2f} (should be ~1.0 on average)")
Dropout from Scratch (p=0.5):
  Input (all 1s): [1. 1. 1. 1. 1. 1. 1. 1.]
  Trial 1: [0. 0. 0. 0. 2. 0. 2. 0.]
    Mean: 0.50 (should be ~1.0 on average)
  Trial 2: [2. 0. 2. 2. 2. 2. 2. 2.]
    Mean: 1.75 (should be ~1.0 on average)
  Trial 3: [2. 2. 0. 2. 2. 2. 2. 0.]
    Mean: 1.50 (should be ~1.0 on average)

The Scaling Trick Explained:

# Why divide by (1-p)?
# Without scaling, dropout reduces expected output
# With scaling, expected output stays the same

p = 0.5
np.random.seed(0)
x = np.array([1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0])

# Without scaling
mask = np.random.random(x.shape) < (1-p)
out_no_scale = x * mask
print(f"Without scaling: {out_no_scale} -> mean = {out_no_scale.mean():.2f}")

# With scaling (divide by keep probability)
np.random.seed(0)
mask = np.random.random(x.shape) < (1-p)
out_scaled = x * mask / (1-p)
print(f"With scaling:    {out_scaled} -> mean = {out_scaled.mean():.2f}")
print(f"\nThe scaling keeps expected value at 1.0 despite dropping 50% of values")
Without scaling: [0. 0. 0. 0. 1. 0. 1. 0.] -> mean = 0.25
With scaling:    [0. 0. 0. 0. 2. 0. 2. 0.] -> mean = 0.50

The scaling keeps expected value at 1.0 despite dropping 50% of values

PyTorch’s nn.Dropout:

# PyTorch handles training/mode automatically
dropout_pytorch = nn.Dropout(p=0.5)

x_torch = torch.ones(2, 8)

# Training mode (dropout active)
dropout_pytorch.train()
torch.manual_seed(42)
out_train = dropout_pytorch(x_torch)
print(f"PyTorch Dropout (training): {out_train[0].numpy()}")

# Inference mode (dropout disabled)
dropout_pytorch.eval()
out_inference = dropout_pytorch(x_torch)
print(f"PyTorch Dropout (inference): {out_inference[0].numpy()}")
PyTorch Dropout (training): [2. 2. 2. 2. 0. 2. 0. 0.]
PyTorch Dropout (inference): [1. 1. 1. 1. 1. 1. 1. 1.]

Key Insight: Dropout is random masking with scaling. Scaling by \frac{1}{1-p} during training preserves the expected value, so inference requires no adjustment.

Residual Connections: The Highway for Gradients

The Idea: Replace y = f(x) with y = x + f(x). This skip connection eliminates the vanishing gradient problem by giving gradients a direct path.

Why it helps:

Pre-Norm vs Post-Norm:

The original Transformer used “post-norm”: normalize after the residual addition.

# Post-Norm (original Transformer)
x = LayerNorm(x + Attention(x))
x = LayerNorm(x + FFN(x))

Modern LLMs use “pre-norm”: normalize before each sublayer.

# Pre-Norm (GPT-2, LLaMA, modern LLMs)
x = x + Attention(LayerNorm(x))
x = x + FFN(LayerNorm(x))

Why Pre-Norm is Better:

# Demonstrate the stability difference

def simulate_forward_pass(num_layers, prenorm=True):
    """Simulate activation magnitudes through layers."""
    x = 1.0  # Starting activation magnitude

    for _ in range(num_layers):
        if prenorm:
            # Pre-norm: normalize first, then residual keeps things bounded
            normed = 1.0  # After LayerNorm, magnitude is ~1
            sublayer_out = normed * 0.5  # Sublayer output
            x = x + sublayer_out  # Residual addition
        else:
            # Post-norm: residual can grow, then we normalize
            sublayer_out = x * 0.5
            x = x + sublayer_out  # Can grow unboundedly before norm
            x = 1.0  # LayerNorm resets to ~1

    return x

print("Activation stability comparison:")
print(f"  Pre-norm after 24 layers:  ~{simulate_forward_pass(24, prenorm=True):.1f}")
print(f"  Post-norm after 24 layers: ~{simulate_forward_pass(24, prenorm=False):.1f}")
print("\nPre-norm has a cleaner gradient path because the skip connection")
print("bypasses normalization - gradients flow directly from output to input.")
Activation stability comparison:
  Pre-norm after 24 layers:  ~13.0
  Post-norm after 24 layers: ~1.0

Pre-norm has a cleaner gradient path because the skip connection
bypasses normalization - gradients flow directly from output to input.

Key Insight: Residual connections transform y = f(x) into y = x + f(x). The gradient of this is dy/dx = 1 + df/dx. That + 1 is crucial - it means gradients always have a direct path through the network, even if df/dx is tiny.

The Full Transformer Block from Scratch

The Idea: Now we assemble all the pieces into a complete transformer block:

  1. LayerNorm + Multi-Head Attention + Residual
  2. LayerNorm + Feed-Forward Network + Residual
class FeedForwardScratch:
    """Simple feed-forward network from scratch."""

    def __init__(self, embed_dim, ff_dim):
        # Initialize weights with small random values
        scale = 0.02
        self.w1 = np.random.randn(embed_dim, ff_dim) * scale
        self.b1 = np.zeros(ff_dim)
        self.w2 = np.random.randn(ff_dim, embed_dim) * scale
        self.b2 = np.zeros(embed_dim)

    def gelu(self, x):
        """GELU activation: x * Phi(x) where Phi is standard normal CDF."""
        return 0.5 * x * (1 + np.tanh(np.sqrt(2/np.pi) * (x + 0.044715 * x**3)))

    def __call__(self, x):
        # Up projection: embed_dim -> ff_dim
        h = x @ self.w1 + self.b1
        # Activation
        h = self.gelu(h)
        # Down projection: ff_dim -> embed_dim
        return h @ self.w2 + self.b2


# NOTE: This uses a SIMPLIFIED attention (just linear projection) to focus on
# the overall block structure. Real attention with Q, K, V is in attention.py
class TransformerBlockScratch:
    """
    A complete transformer block from scratch (with simplified attention).

    Architecture (Pre-Norm):
        x = x + Attention(LayerNorm(x))
        x = x + FeedForward(LayerNorm(x))

    WARNING: The attention here is simplified to a linear projection for
    demonstration purposes. See m05_attention for full attention implementation.
    """

    def __init__(self, embed_dim, num_heads, ff_dim, dropout_p=0.1):
        self.embed_dim = embed_dim
        self.num_heads = num_heads
        self.head_dim = embed_dim // num_heads

        # Layer norms
        self.ln1 = LayerNormScratch(embed_dim)
        self.ln2 = LayerNormScratch(embed_dim)

        # Attention projections (simplified: no actual attention computation)
        # In a full implementation, this would include Q, K, V projections
        scale = 0.02
        self.attn_proj = np.random.randn(embed_dim, embed_dim) * scale

        # Feed-forward network
        self.ff = FeedForwardScratch(embed_dim, ff_dim)

        # Dropout
        self.dropout = DropoutScratch(dropout_p)

    def __call__(self, x, training=True):
        """
        Args:
            x: input of shape (batch, seq, embed_dim)
            training: whether to apply dropout
        Returns:
            output of shape (batch, seq, embed_dim)
        """
        # === Attention sub-block ===
        # 1. Layer norm (pre-norm)
        normed = self.ln1(x)

        # 2. Attention (simplified: just a linear projection for demo)
        # Real implementation would compute Q, K, V and attention weights
        attn_out = normed @ self.attn_proj

        # 3. Dropout
        attn_out = self.dropout(attn_out, training=training)

        # 4. Residual connection
        x = x + attn_out

        # === Feed-forward sub-block ===
        # 1. Layer norm (pre-norm)
        normed = self.ln2(x)

        # 2. Feed-forward network
        ff_out = self.ff(normed)

        # 3. Dropout
        ff_out = self.dropout(ff_out, training=training)

        # 4. Residual connection
        x = x + ff_out

        return x


# Test the from-scratch transformer block
np.random.seed(42)
block_scratch = TransformerBlockScratch(
    embed_dim=64,
    num_heads=4,
    ff_dim=256,
    dropout_p=0.0  # Disable dropout for reproducibility
)

x = np.random.randn(2, 8, 64)  # batch=2, seq=8, embed=64
out = block_scratch(x, training=False)

print("Transformer Block from Scratch:")
print(f"  Input shape:  {x.shape}")
print(f"  Output shape: {out.shape}")
print(f"  Input mean:   {x.mean():.4f}")
print(f"  Output mean:  {out.mean():.4f}")
print(f"\nThe block transforms each token while preserving shape.")
print("Residual connections keep the output close to input initially.")
Transformer Block from Scratch:
  Input shape:  (2, 8, 64)
  Output shape: (2, 8, 64)
  Input mean:   -0.0469
  Output mean:  -0.0443

The block transforms each token while preserving shape.
Residual connections keep the output close to input initially.

The Complete Picture:

# Visualize the transformer block structure
print("""
Transformer Block (Pre-Norm Architecture):
==========================================

    Input x
        |
        +------------------+
        |                  |
        v                  |
    LayerNorm              |
        |                  |
        v                  |
    Multi-Head Attention   |
        |                  |
        v                  |
    Dropout                |
        |                  |
        +--------(+)-------+  <- Residual connection
                  |
        +------------------+
        |                  |
        v                  |
    LayerNorm              |
        |                  |
        v                  |
    Feed-Forward           |
        |                  |
        v                  |
    Dropout                |
        |                  |
        +--------(+)-------+  <- Residual connection
                  |
                  v
              Output
""")

Transformer Block (Pre-Norm Architecture):
==========================================

    Input x
        |
        +------------------+
        |                  |
        v                  |
    LayerNorm              |
        |                  |
        v                  |
    Multi-Head Attention   |
        |                  |
        v                  |
    Dropout                |
        |                  |
        +--------(+)-------+  <- Residual connection
                  |
        +------------------+
        |                  |
        v                  |
    LayerNorm              |
        |                  |
        v                  |
    Feed-Forward           |
        |                  |
        v                  |
    Dropout                |
        |                  |
        +--------(+)-------+  <- Residual connection
                  |
                  v
              Output

Key Insight: A Transformer block contains only attention, MLP, residuals, and norms. Stacking dozens of blocks and training on billions of tokens produces the performance.

PyTorch Transformer Modules

PyTorch provides optimized versions of everything we built from scratch.

Comparison Table:

Component From Scratch PyTorch
LayerNorm Manual mean/var nn.LayerNorm
Dropout Random mask + scale nn.Dropout
FFN Two linear layers + GELU Custom or nn.Sequential
Full Block Manual assembly nn.TransformerDecoderLayer 
# PyTorch's TransformerDecoderLayer
# Note: This is for encoder-decoder models; for decoder-only like GPT,
# we typically build our own (as in transformer.py)

from torch.nn import TransformerDecoderLayer

# Create a decoder layer similar to our scratch implementation
pytorch_block = TransformerDecoderLayer(
    d_model=64,
    nhead=4,
    dim_feedforward=256,
    dropout=0.1,
    activation='gelu',
    batch_first=True,
    norm_first=True  # Pre-norm architecture
)

x_torch = torch.randn(2, 8, 64)

# For decoder-only, we use self-attention (memory = x)
pytorch_block.eval()
out_pytorch = pytorch_block(x_torch, x_torch)

print("PyTorch TransformerDecoderLayer:")
print(f"  Input shape:  {tuple(x_torch.shape)}")
print(f"  Output shape: {tuple(out_pytorch.shape)}")
print(f"  Parameters:   {sum(p.numel() for p in pytorch_block.parameters()):,}")
PyTorch TransformerDecoderLayer:
  Input shape:  (2, 8, 64)
  Output shape: (2, 8, 64)
  Parameters:   66,752

When to Use What:

  • Learning: Build from scratch to understand every step
  • Production: Use PyTorch’s optimized modules
  • Custom architectures: Mix both - understand the components, then optimize
# Our module's TransformerBlock (production quality)
from transformer import TransformerBlock

our_block = TransformerBlock(
    embed_dim=64,
    num_heads=4,
    ff_dim=256,
    dropout=0.1
)

x_torch = torch.randn(2, 8, 64)
our_block.eval()
out_ours = our_block(x_torch)

print("Our TransformerBlock (from transformer.py):")
print(f"  Input shape:  {tuple(x_torch.shape)}")
print(f"  Output shape: {tuple(out_ours.shape)}")
print(f"  Parameters:   {sum(p.numel() for p in our_block.parameters()):,}")
print("\nThis is what we use for training - it includes proper")
print("causal attention, not the simplified version in scratch code.")
Our TransformerBlock (from transformer.py):
  Input shape:  (2, 8, 64)
  Output shape: (2, 8, 64)
  Parameters:   49,984

This is what we use for training - it includes proper
causal attention, not the simplified version in scratch code.

More Component Details

Layer Normalization (PyTorch Details)

Normalizes activations across the embedding dimension:

\text{LayerNorm}(x) = \gamma \times \frac{x - \mu}{\sqrt{\sigma^2 + \epsilon}} + \beta

where \mu and \sigma^2 are the mean and variance across the embedding dimension, and \gamma, \beta are learnable parameters.

Why it helps:

  • Stabilizes activations: Prevents values from exploding or vanishing
  • Faster training: More stable gradients
  • Independent per token: Each token normalized separately

Feed-Forward Network (FFN)

The FFN is a mini neural network applied to each token independently:

  • 4x expansion: More capacity to learn complex transformations
  • GELU activation: Smoother than ReLU, better gradients
  • Same for all tokens: Unlike attention, no mixing between positions

SwiGLU: The Gated Feed-Forward

The FFN is the single largest block of parameters in a transformer — roughly two thirds of every non-embedding weight lives here. So it is exactly where a small architectural upgrade pays off the most. Every major open model since LLaMA — PaLM, Mistral, Qwen, Gemma — replaces the plain Linear → GELU → Linear FFN above with a gated variant called SwiGLU. Understanding it is a two-step story: a new activation, and a gate.

The Idea — give the FFN a gate. The standard FFN pushes every expanded feature through the same activation. A gated FFN instead makes two up-projections from the input — a gate branch and a value branch — and lets the gate decide, element by element, how much of the value to let through. It is the same move as an LSTM gate, but learned per token, per feature:

\text{SwiGLU}(x) = W_{\text{down}}\,\Big(\underbrace{\text{SiLU}(x W_{\text{gate}})}_{\text{gate}} \;\odot\; \underbrace{x W_{\text{up}}}_{\text{value}}\Big)

where \odot is element-wise multiplication. The gate is squashed by SiLU (also called Swish), the smooth activation the GLU family settled on:

\text{SiLU}(x) = x \cdot \sigma(x), \qquad \sigma(x) = \frac{1}{1 + e^{-x}}

Compare the two feed-forwards term by term:

Standard FFN (GPT-2) SwiGLU (LLaMA)
Up-projections 1 (W_1) 2 (W_{\text{gate}}, W_{\text{up}})
Down-projection 1 (W_2) 1 (W_{\text{down}})
Nonlinearity GELU on the whole expansion SiLU on the gate, then a multiply
Weight matrices 2 3
Bias terms yes (GPT-2) none (LLaMA)
Hidden width for equal params 4d \tfrac{2}{3}\cdot 4d \approx 2.7d
NoteKey Insight

SwiGLU is not a bigger FFN — it is a smarter one. Splitting the expansion into a gate and a value, then multiplying, lets the network learn a data-dependent mask over its own hidden units. Shazeer (2020), GLU Variants Improve Transformer found gated variants consistently lower loss at the same parameter and compute budget; the cost is a third weight matrix, which we pay for by shrinking the hidden width (the ⅔ rule below). SiLU/Swish itself is from Ramachandran et al. (2017).

Activations up close: SiLU vs GELU vs ReLU

Before the gate, look at the gate’s nonlinearity. SiLU and GELU are nearly the same curve — both are smooth, and both dip slightly below zero for small negatives (they are non-monotonic) before recovering. That gentle negative lobe is what the GLU family wanted; ReLU’s hard corner throws it away. Drag the marker to read each function’s value.

TipTry This
  1. Drag into the negatives. Around x=-1.5 both SiLU and GELU sit below zero — the non-monotonic dip. ReLU is flat at 0. That small negative signal is information ReLU discards.
  2. Overlay the smooth pair. SiLU and GELU nearly coincide everywhere; this is why SwiGLU uses the cheaper-to-compute SiLU and loses nothing.
  3. Go positive. For large x all three converge to the identity y=x — the differences live entirely near the origin.

The ⅔ rule: keeping the budget fixed

A gated FFN has three weight matrices where the standard one had two, so a naive swap at the same hidden width would balloon the FFN by 50%. LLaMA keeps the comparison fair by shrinking the gated hidden dimension to \tfrac{2}{3} of the standard 4d:

3 \cdot d \cdot \underbrace{\tfrac{2}{3}(4d)}_{\text{gated hidden}} \;=\; 8d^2 \;=\; 2 \cdot d \cdot \underbrace{4d}_{\text{standard hidden}}

so the gated FFN spends exactly the parameters of the standard one — it just spends them on a gate. (LLaMA then rounds the width up to a hardware-friendly multiple; for d=4096 that gives the famous hidden width of 11008.) Drag the width and watch the naive gated bar overshoot while the ⅔-scaled bar rides exactly on the standard FFN.

TipTry This
  1. Slide the width up. The SwiGLU (naive) bar always sits ~50% past the dashed Standard-FFN line — that is the extra matrix, unpaid for.
  2. Watch the ⅔ bar. The SwiGLU (⅔·4d) bar lands exactly on the dashed line at every width — same parameters, now spent on a gate.
  3. Set d = 4096. The ⅔ hidden dimension lands near LLaMA-1’s 11008 (before the multiple-of-256 rounding the real model applies).

From Scratch Implementation:

def silu(x):
    """SiLU / Swish-1 activation: x * sigmoid(x)."""
    return x / (1 + np.exp(-x))


class SwiGLUScratch:
    """Gated feed-forward from scratch — two up-projections and a gate."""

    def __init__(self, embed_dim, hidden_dim, seed=0):
        rng = np.random.default_rng(seed)
        s = 0.02
        # Three weight matrices, no biases (LLaMA-style).
        self.W_gate = rng.standard_normal((embed_dim, hidden_dim)) * s
        self.W_up = rng.standard_normal((embed_dim, hidden_dim)) * s
        self.W_down = rng.standard_normal((hidden_dim, embed_dim)) * s

    def __call__(self, x):
        gate = silu(x @ self.W_gate)   # squashed gate branch
        value = x @ self.W_up          # linear value branch
        return (gate * value) @ self.W_down   # gate modulates value, then project


x = np.random.default_rng(1).standard_normal((1, 3, 8))   # (batch, seq, embed_dim)
ffn = SwiGLUScratch(embed_dim=8, hidden_dim=16)
out = ffn(x)
print(f"Input shape:  {x.shape}")
print(f"Output shape: {out.shape}  <- same as input, ready for the residual")
Input shape:  (1, 3, 8)
Output shape: (1, 3, 8)  <- same as input, ready for the residual

The transformer.py module ships a production SwiGLU(nn.Module), a make_ffn(ffn_type, ...) factory, and the swiglu_hidden_dim budget helper. Building a modern LLaMA-style block — RMSNorm and SwiGLU — is a two-word change:

from transformer import SwiGLU, swiglu_hidden_dim, FeedForward, GPTModel

# The budget helper reproduces LLaMA-1's FFN width for d=4096
print(f"swiglu_hidden_dim(4*4096) = {swiglu_hidden_dim(4 * 4096)}  (LLaMA-1 uses 11008)")

# A budget-matched gated FFN has the same matrix params as a 4x GELU FFN
d = 768
std = FeedForward(d, 4 * d)
gated = SwiGLU(d, swiglu_hidden_dim(4 * d, multiple_of=1))
std_p = std.linear1.weight.numel() + std.linear2.weight.numel()
gated_p = sum(w.numel() for w in
              (gated.gate_proj.weight, gated.up_proj.weight, gated.down_proj.weight))
print(f"Standard FFN matrix params: {std_p:,}")
print(f"SwiGLU     matrix params:   {gated_p:,}  <- equal, by the 2/3 rule")

# Swap the whole model's FFN with one argument (ff_dim auto-scales for SwiGLU)
modern = GPTModel(vocab_size=100, embed_dim=32, num_heads=4, num_layers=2,
                  max_seq_len=64, norm_type="rmsnorm", ffn_type="swiglu")
print(f"Block FFN: {type(modern.blocks[0].ffn).__name__}")
print(f"Block norm: {type(modern.blocks[0].ln1).__name__}")
swiglu_hidden_dim(4*4096) = 11008  (LLaMA-1 uses 11008)
Standard FFN matrix params: 4,718,592
SwiGLU     matrix params:   4,718,592  <- equal, by the 2/3 rule
Block FFN: SwiGLU
Block norm: RMSNorm
WarningDon’t forget: three matrices means a narrower hidden layer

Swapping in SwiGLU at the same ff_dim you used for the GELU FFN quietly grows the model by ~50%, because you added a whole up-projection. To compare architectures fairly, scale the gated hidden width to swiglu_hidden_dim(4 * embed_dim). GPTModel(ffn_type="swiglu") does this for you when you leave ff_dim at its default; pass an explicit ff_dim only if you want the larger FFN.

Pre-Norm vs Post-Norm

We use Pre-Norm (LayerNorm before attention/FFN) rather than Post-Norm:

# Pre-Norm (GPT-2, LLaMA, modern LLMs)
x = x + Attention(LayerNorm(x))

# Post-Norm (original Transformer paper)
x = LayerNorm(x + Attention(x))

Why Pre-Norm is preferred:

  1. Cleaner gradient path: The residual connection bypasses normalization, so gradients flow directly
  2. More stable training: Especially important for deep networks (24+ layers)
  3. Requires final LayerNorm: Since the last block’s output isn’t normalized, we add a final LayerNorm before the output projection

Post-Norm achieves marginally better final performance with careful tuning, but Pre-Norm trains more robustly.

Code Walkthrough

Let’s build and explore transformer blocks:

import sys
import importlib.util
from pathlib import Path

import torch
import torch.nn as nn

print(f"PyTorch version: {torch.__version__}")
device = "cuda" if torch.cuda.is_available() else "mps" if torch.backends.mps.is_available() else "cpu"
print(f"Device: {device}")
PyTorch version: 2.9.1
Device: mps

GELU Activation

GPT uses GELU instead of ReLU. Let’s see why:

GELU formula: \text{GELU}(x) = x \cdot \Phi(x) where \Phi is the standard normal CDF.

Approximation used in practice: 0.5x(1 + \tanh(\sqrt{2/\pi}(x + 0.044715x^3)))

Activation function choices in modern LLMs:

Model Activation Notes
GPT-2, BERT GELU Smooth, good gradients
LLaMA, Mistral SwiGLU Gated variant, better performance
GPT-3 GELU Same as GPT-2

SwiGLU (used in LLaMA) is a gated linear unit: \text{SwiGLU}(x) = \text{SiLU}(xW_{\text{gate}}) \otimes xW_{\text{up}}. It adds a gate branch on top of GELU-style FFNs and usually improves performance. Our FeedForward defaults to GELU to match GPT-2, but the module also ships a from-scratch SwiGLU you can swap in with ffn_type="swiglu" — see SwiGLU: The Gated Feed-Forward below.

Layer Normalization Demo

# Manual LayerNorm demonstration
x = torch.tensor([[2.0, 4.0, 6.0, 8.0]])

mean = x.mean(dim=-1, keepdim=True)
var = x.var(dim=-1, unbiased=False, keepdim=True)
normalized = (x - mean) / torch.sqrt(var + 1e-5)

print("Manual LayerNorm:")
print(f"  Input: {x.numpy().tolist()}")
print(f"  Mean: {mean.item():.2f}")
print(f"  Variance: {var.item():.2f}")
print(f"  Normalized: {normalized.numpy().round(2).tolist()}")
print(f"  New mean: {normalized.mean().item():.4f}")
print(f"  New std: {normalized.std().item():.4f}")
Manual LayerNorm:
  Input: [[2.0, 4.0, 6.0, 8.0]]
  Mean: 5.00
  Variance: 5.00
  Normalized: [[-1.340000033378601, -0.44999998807907104, 0.44999998807907104, 1.340000033378601]]
  New mean: 0.0000
  New std: 1.1547
# PyTorch LayerNorm (with learnable gamma and beta)
ln = nn.LayerNorm(4)
pytorch_normalized = ln(x)

print(f"PyTorch LayerNorm output: {pytorch_normalized.detach().numpy().round(2).tolist()}")
print("(gamma and beta are learnable parameters)")
PyTorch LayerNorm output: [[-1.340000033378601, -0.44999998807907104, 0.44999998807907104, 1.340000033378601]]
(gamma and beta are learnable parameters)

Residual Connections Demo

Weight Initialization

Deep networks require proper initialization. Our implementation uses:

  • Embeddings: Normal distribution with std=0.02
  • Linear layers in FFN: Normal distribution with std=0.02, biases initialized to 0
  • Attention projections: Xavier uniform initialization
# Demonstrate the importance of initialization
import torch.nn as nn

# Bad initialization - too large
bad_linear = nn.Linear(768, 768)
nn.init.normal_(bad_linear.weight, std=1.0)  # Too large!

# Good initialization - small weights
good_linear = nn.Linear(768, 768)
nn.init.normal_(good_linear.weight, std=0.02)  # GPT-2 style

x = torch.randn(1, 10, 768)
bad_out = bad_linear(x)
good_out = good_linear(x)

print("Effect of initialization on output magnitude:")
print(f"  Bad init (std=1.0):  output std = {bad_out.std().item():.2f}")
print(f"  Good init (std=0.02): output std = {good_out.std().item():.2f}")
print("\nLarge outputs can cause exploding gradients and NaN losses!")
Effect of initialization on output magnitude:
  Bad init (std=1.0):  output std = 27.87
  Good init (std=0.02): output std = 0.55

Large outputs can cause exploding gradients and NaN losses!

GPT-2’s initialization trick: Scale the final projection in each residual block by 1/\sqrt{2N} where N is the number of layers. This keeps the variance stable as depth increases.

Building a Transformer Block

from transformer import (
    FeedForward,
    TransformerBlock,
    GPTModel,
    create_gpt_tiny,
    create_gpt_small,
)

# Create a transformer block
embed_dim = 64
num_heads = 4
ff_dim = 256

block = TransformerBlock(
    embed_dim=embed_dim,
    num_heads=num_heads,
    ff_dim=ff_dim,
    dropout=0.0
)

print(f"Transformer Block:")
print(f"  Embed dim: {embed_dim}")
print(f"  Num heads: {num_heads}")
print(f"  Head dim: {embed_dim // num_heads}")
print(f"  FF dim: {ff_dim}")
print(f"\nTotal parameters: {sum(p.numel() for p in block.parameters()):,}")
Transformer Block:
  Embed dim: 64
  Num heads: 4
  Head dim: 16
  FF dim: 256

Total parameters: 49,984
# Forward pass
x = torch.randn(1, 8, embed_dim)  # batch=1, seq=8
output, attention = block(x, return_attention=True)

print(f"Input shape: {x.shape}")
print(f"Output shape: {output.shape}")
print(f"Attention shape: {attention.shape}")
Input shape: torch.Size([1, 8, 64])
Output shape: torch.Size([1, 8, 64])
Attention shape: torch.Size([1, 4, 8, 8])
# Pass attention weights to OJS for visualization
attention_list = [attention[0, h].detach().tolist() for h in range(4)]
ojs_define(attentionWeightsData=attention_list)

Complete GPT Model

# Create a tiny GPT model
model = create_gpt_tiny(vocab_size=1000)

print("GPT Tiny Model:")
print(f"  Vocab size: {model.vocab_size}")
print(f"  Embed dim: {model.embed_dim}")
print(f"  Num layers: {len(model.blocks)}")
print(f"  Max seq len: {model.max_seq_len}")
print(f"\nTotal parameters: {model.num_params:,}")
GPT Tiny Model:
  Vocab size: 1000
  Embed dim: 128
  Num layers: 4
  Max seq len: 256

Total parameters: 954,112
# Parameter breakdown
counts = model.count_parameters()

print("Parameter breakdown:")
for name, count in counts.items():
    if count > 0:
        pct = 100 * count / counts['total']
        print(f"  {name}: {count:,} ({pct:.1f}%)")
Parameter breakdown:
  token_embedding: 128,000 (13.4%)
  position_embedding: 32,768 (3.4%)
  transformer_blocks: 793,088 (83.1%)
  final_layer_norm: 256 (0.0%)
  total: 954,112 (100.0%)
# Pass parameter counts to OJS for visualization
param_data = [
    {"category": name, "value": count}
    for name, count in counts.items()
    if count > 0 and name != 'total'
]
ojs_define(parameterData=param_data, totalParams=counts['total'])
# Forward pass
token_ids = torch.randint(0, 1000, (2, 32))  # batch=2, seq=32
logits = model(token_ids)

print(f"Input token IDs: {token_ids.shape}")
print(f"Output logits: {logits.shape}")
print(f"  (batch=2, seq=32, vocab=1000)")

# Get predictions
probs = torch.softmax(logits[0, -1], dim=-1)
top_5 = torch.topk(probs, 5)

print("\nTop 5 predicted next tokens (untrained, so random):")
for i, (idx, prob) in enumerate(zip(top_5.indices, top_5.values)):
    print(f"  {i+1}. Token {idx.item()}: {prob.item()*100:.2f}%")
Input token IDs: torch.Size([2, 32])
Output logits: torch.Size([2, 32, 1000])
  (batch=2, seq=32, vocab=1000)

Top 5 predicted next tokens (untrained, so random):
  1. Token 345: 0.20%
  2. Token 336: 0.20%
  3. Token 90: 0.18%
  4. Token 807: 0.18%
  5. Token 157: 0.18%

Hidden States Through Layers

# Get hidden states from all layers
logits, hidden_states = model(token_ids, return_hidden_states=True)

print(f"Number of hidden states: {len(hidden_states)}")
print(f"  (1 after embedding + {len(model.blocks)} after each block)")

# Show how representations change through layers
norms = [h.norm(dim=-1).mean().item() for h in hidden_states]
layer_names = ['Embed'] + [f'Block {i}' for i in range(len(model.blocks))]
Number of hidden states: 5
  (1 after embedding + 4 after each block)
# Pass data to OJS for visualization
layer_norms_data = [{"layer": name, "norm": norm} for name, norm in zip(layer_names, norms)]
ojs_define(layerNormsData=layer_norms_data)

Weight Tying: One Matrix, Two Jobs

Look at the two ends of the model. At the input, the token embedding is a matrix E of shape (V, d): row i is the vector for token i. At the output, the LM head is a matrix of shape (V, d) too — because nn.Linear(d, V) stores its weight as (V, d), and the logit for token j is the dot product of the hidden vector h with row j.

Same shape, and — it turns out — they want to learn the same thing. So GPT ties them: one matrix does both jobs.

  • Input job (read in): a token ID selects its row, E[i], as an embedding.
  • Output job (score out): the hidden state h is compared against every row of E to score every token.

Step through the two directions of the single shared matrix:

The Math

The pre-softmax projection is literally the transpose of the embedding lookup:

\text{logits} = h\,E^\top, \qquad E \in \mathbb{R}^{V \times d},\; h \in \mathbb{R}^{d}

Because both matrices are (V, d), sharing them is natural, not a hack: the same geometry that places similar tokens near each other on the way in is what you want for scoring similar tokens together on the way out.

NoteKey Insight

Weight tying makes the token-embedding matrix do double duty. It is the biggest single weight in a small model — for GPT-2 small it is 50257 \times 768 \approx 38.6\text{M} parameters, roughly a third of the whole model — so tying it away is the largest parameter saving in the architecture, and it tends to lower perplexity because the one matrix now gets V\times more gradient signal.

Code: the tie_weights toggle

GPTModel ties by default (tie_weights=True, the GPT-2 choice). Flip it off to get the untied baseline — the only difference is whether the head owns a separate weight. Everything else about the two models is identical.

from transformer import GPTModel, tied_lm_logits, weight_tying_savings

tied   = GPTModel(vocab_size=1000, embed_dim=128, num_heads=4, num_layers=4, max_seq_len=256, tie_weights=True)
untied = GPTModel(vocab_size=1000, embed_dim=128, num_heads=4, num_layers=4, max_seq_len=256, tie_weights=False)

print("Tied   → head IS the embedding:", tied.token_embedding.weight is tied.lm_head.weight)
print("Untied → head is separate:     ", untied.token_embedding.weight is untied.lm_head.weight)
print(f"\nParameters — tied:   {tied.num_params:,}")
print(f"Parameters — untied: {untied.num_params:,}")
print(f"Difference:          {untied.num_params - tied.num_params:,}  (= V×d = {1000*128:,})")
Tied   → head IS the embedding: True
Untied → head is separate:      False

Parameters — tied:   954,112
Parameters — untied: 1,082,112
Difference:          128,000  (= V×d = 128,000)

The tied model computes its logits as exactly h\,E^\top. We can reproduce the head by hand with tied_lm_logits and check it matches:

import torch

tied.eval()
ids = torch.randint(0, 1000, (1, 8))
logits, hidden_states = tied(ids, return_hidden_states=True)

# The head's input is the final-norm output of the last block.
h = tied.ln_final(hidden_states[-1])
manual = tied_lm_logits(h, tied.token_embedding.weight)   # h @ Eᵀ

print("logits == h @ Eᵀ ?", torch.allclose(logits, manual, atol=1e-5))
print("max abs difference:", (logits - manual).abs().max().item())
logits == h @ Eᵀ ? True
max abs difference: 0.0

The ablation: what tying buys

The saving is a single V\times d slab, so it grows with the vocabulary. Drive the vocabulary and width below and watch the untied bar sprout an extra block:

# A trustworthy reference point for the live calculator below.
ref = weight_tying_savings(vocab_size=50257, embed_dim=768, num_layers=12, max_seq_len=1024)
ojs_define(tyingRefSaved = ref["saved"])
TipTry This
  1. Grow the vocabulary. Push V from 1000 to 128000 at fixed d. The saved slab scales linearly with V, so tying matters more for large-vocabulary models.
  2. Grow the depth. Crank N up. The backbone (12Nd^2) swamps the embedding, so tying’s share shrinks even as the absolute saving stays fixed — tying is biggest, relatively, for shallow, wide-vocab models.

Why it also improves quality

Tying is not only cheaper. Because one matrix now sits on both the input and output paths, a single backward() deposits gradient into it from both roles — you can see they land on the same .grad tensor:

tied.zero_grad()
tied(ids).sum().backward()

g_embed = tied.token_embedding.weight.grad
g_head  = tied.lm_head.weight.grad
print("Same gradient tensor for both jobs?", g_embed is g_head)   # True — coupled

# In the untied model the two roles learn from separate gradients:
untied.zero_grad()
untied(ids).sum().backward()
print("Untied — separate gradients?      ",
      untied.token_embedding.weight.grad is not untied.lm_head.weight.grad)
Same gradient tensor for both jobs? True
Untied — separate gradients?       True

Press & Wolf (2016) showed this coupling lowers perplexity: the output side carries a strong learning signal (every position contributes a gradient to every row), and sharing lets the input embedding inherit it.

WarningTie carefully
  • The original Transformer scales the embedding by \sqrt{d} before adding positions (Vaswani et al. §3.4) precisely because the same weights serve the softmax; GPT-2 (and this model) drop that scaling and initialize at std=0.02. Match your init to your choice — don’t blindly copy one recipe’s scaling into the other.
  • Only tie when the input and output vocabularies are the same. They usually are for a decoder-only LM, but an encoder-decoder with different source/target vocabularies cannot share one matrix across both.
  • nn.Linear(d, V) stores a (V, d) weight, already the embedding’s shape — so tying is the one-liner lm_head.weight = token_embedding.weight. If you ever store the head transposed, you must transpose before assigning.

Model Sizes Comparison

Model Layers Heads Embed Dim Params
Tiny (ours) 4 4 128 ~1M
Small (ours) 6 6 384 ~10M
GPT-2 Small 12 12 768 117M
GPT-2 Medium 24 16 1024 345M
GPT-2 Large 36 20 1280 774M
GPT-2 XL 48 25 1600 1.5B

Parameter Counting Formulas

Knowing where parameters come from aids model sizing:

Per Transformer Block:

  • Attention Q, K, V projections: 3 \times d \times d (where d = embed_dim)
  • Attention output projection: d \times d
  • FFN first linear: d \times 4d
  • FFN second linear: 4d \times d
  • LayerNorm (x2): 2 \times 2d (gamma and beta for each)

Total per block: \approx 12d^2 parameters

Full Model:

  • Token embedding: V \times d (V = vocab size)
  • Position embedding: L \times d (L = max sequence length)
  • N transformer blocks: N \times 12d^2
  • Final LayerNorm: 2d
  • LM head: 0 (weight-tied with token embedding)

Approximate formula: \text{Params} \approx V \times d + 12Nd^2

For GPT-2 Small (V=50257, d=768, N=12): 50257 \times 768 + 12 \times 12 \times 768^2 \approx 117M

Scaling laws: more layers, heads, and dimensions lead to better performance (but diminishing returns and higher compute cost).

Architectural Variations

Modern LLMs have evolved beyond the original GPT-2 architecture. Here are key variations:

Normalization

Variant Used By Description
LayerNorm GPT-2, GPT-3 Normalize across embedding dimension
RMSNorm LLaMA, Mistral Simpler: just divide by RMS, no mean subtraction
Pre-Norm Most modern LLMs Normalize before sublayer (more stable)

Position Embeddings

Variant Used By Description
Learned absolute GPT-2 Separate embedding for each position
Rotary (RoPE) LLaMA, Mistral Encode position in attention via rotation
ALiBi BLOOM, MPT Add a per-head linear distance penalty to attention scores

Our implementation uses learned absolute position embeddings (GPT-2 style), which are simple but limit the model to the maximum trained sequence length. Both alternatives are built from scratch in Module 04: RoPE (rotate q/k so scores are offset-only) and ALiBi (a formula-based score penalty that extrapolates past the training length).

Feed-Forward Networks

Variant Used By Expansion Activation
Standard GPT-2 4x GELU
SwiGLU LLaMA, Mistral 8/3x (after gating) SiLU (Swish)

We build both — the GELU FFN by default and the gated SwiGLU you can swap in — in SwiGLU: The Gated Feed-Forward.

Common Pitfalls

When implementing or training transformers, watch out for:

  1. Forgetting the causal mask: Without it, the model “cheats” by seeing future tokens during training, which cripples generation at inference.

  2. Wrong normalization axis: LayerNorm should normalize across the embedding dimension (last axis), not the sequence or batch dimensions.

  3. Residual connection placement: Make sure to add the residual after dropout but before the next LayerNorm in Pre-Norm architecture.

  4. Large learning rates: Transformers are sensitive to learning rate. Start with 1e-4 to 3e-4 for Adam, use warmup.

  5. Numerical instability: Use float32 for training initially. Half precision (fp16/bf16) requires careful scaling.

  6. Forgetting final LayerNorm: In Pre-Norm, the output of the last block isn’t normalized. The final LayerNorm before the LM head is essential.

Interactive Exploration

Experiment with transformer architecture choices to understand where parameters come from and how they scale.

TipTry This

  1. FFN dominates: Set embed_dim=768, layers=12. Notice the Feed-Forward bars are ~2x the Attention bars (because FFN has 8d² params vs Attention’s 4d²).

  2. Embedding cost at small scale: With vocab=50257 and embed_dim=768, token embeddings are ~38M params - a large fraction for small models.

  3. Scaling law: Double embed_dim from 512 to 1024. Total params roughly quadruple (because most params scale with d²).

  4. Load GPT-2 presets and see how the 117M, 345M, 774M models break down.

  5. Head dimension check: Try numHeads that doesn’t divide embedDim evenly - you’ll see a warning.

Exercises

Exercise 1: Build a Custom Block

# Create a transformer block with different configurations
custom_block = TransformerBlock(
    embed_dim=128,
    num_heads=8,
    ff_dim=512,  # 4x expansion
    dropout=0.1
)

# Test it
x = torch.randn(4, 16, 128)  # batch=4, seq=16
output = custom_block(x)
print(f"Custom block: {x.shape} -> {output.shape}")
print(f"Parameters: {sum(p.numel() for p in custom_block.parameters()):,}")
Custom block: torch.Size([4, 16, 128]) -> torch.Size([4, 16, 128])
Parameters: 198,272

Exercise 2: Compare Model Scales

# Compare tiny vs small model
tiny = create_gpt_tiny(vocab_size=10000)
small = create_gpt_small(vocab_size=10000)

print(f"{'Model':<10} {'Embed':<8} {'Layers':<8} {'Heads':<8} {'Params':<15}")
print("-" * 50)
print(f"{'Tiny':<10} {tiny.embed_dim:<8} {len(tiny.blocks):<8} {tiny.blocks[0].attention.mha.num_heads:<8} {tiny.num_params:,}")
print(f"{'Small':<10} {small.embed_dim:<8} {len(small.blocks):<8} {small.blocks[0].attention.mha.num_heads:<8} {small.num_params:,}")
Model      Embed    Layers   Heads    Params         
--------------------------------------------------
Tiny       128      4        4        2,106,112
Small      384      6        6        14,684,160

Exercise 3: Information Flow

# See how a single token's representation changes through layers
model = create_gpt_tiny(vocab_size=100)
token_ids = torch.randint(0, 100, (1, 8))

_, hidden = model(token_ids, return_hidden_states=True)

# Track first token through layers
first_token_norms = [h[0, 0].norm().item() for h in hidden]
ex3_layer_names = ['Embed'] + [f'Block {i}' for i in range(len(model.blocks))]

print(f"First token embedding norm through {len(first_token_norms)} layers:")
for name, norm in zip(ex3_layer_names, first_token_norms):
    print(f"  {name}: {norm:.3f}")
First token embedding norm through 5 layers:
  Embed: 0.334
  Block 0: 10.669
  Block 1: 15.899
  Block 2: 19.688
  Block 3: 22.203
# Pass data to OJS for visualization
first_token_data = [{"layer": name, "norm": norm} for name, norm in zip(ex3_layer_names, first_token_norms)]
ojs_define(firstTokenNormsData=first_token_data)

Exercise 4: Untie and Re-count

Build the same model tied and untied, confirm the difference is exactly one V × d matrix, and check the tied model’s head is literally the embedding.

from transformer import GPTModel

V, d = 500, 96
tied   = GPTModel(V, d, num_heads=4, num_layers=3, max_seq_len=128, tie_weights=True)
untied = GPTModel(V, d, num_heads=4, num_layers=3, max_seq_len=128, tie_weights=False)

diff = untied.num_params - tied.num_params
print(f"Untied − tied = {diff:,}   (expected V×d = {V*d:,})")
assert diff == V * d
assert tied.lm_head.weight is tied.token_embedding.weight
assert untied.lm_head.weight is not untied.token_embedding.weight
print("Weight tying verified ✓")
Untied − tied = 48,000   (expected V×d = 48,000)
Weight tying verified ✓

Exercise 5: See the Gradient Coupling

One backward() on the tied model should deposit gradient into a single shared matrix from both the embedding and the head. Verify it, then confirm the untied model keeps two independent gradients.

ids = torch.randint(0, V, (2, 12))

tied.zero_grad()
tied(ids).sum().backward()
print("Tied — one shared .grad tensor?  ",
      tied.token_embedding.weight.grad is tied.lm_head.weight.grad)

untied.zero_grad()
untied(ids).sum().backward()
print("Untied — two separate .grads?    ",
      untied.token_embedding.weight.grad is not untied.lm_head.weight.grad)
Tied — one shared .grad tensor?   True
Untied — two separate .grads?     True

Summary

Key takeaways:

  1. Transformer architecture: Input embeddings -> N transformer blocks -> Final LayerNorm -> Output projection

  2. Each block has two sublayers:

    • Multi-head attention (tokens communicate)
    • Feed-forward network (tokens processed independently)

  3. Pre-Norm architecture: LayerNorm before each sublayer, with a “clean” residual path for stable gradients

  4. Layer normalization: Normalizes across the embedding dimension, keeping activations in a stable range. RMSNorm (LLaMA, Mistral) simplifies it — re-scale only, no mean-centering or shift — for half the parameters at the same stability

  5. Residual connections: x + f(x) enables gradient flow through very deep networks (100+ layers)

  6. Feed-forward networks: 4x expansion with GELU activation provides computational capacity. SwiGLU (LLaMA, Mistral) upgrades it with a learned gate — two up-projections whose SiLU-squashed gate modulates the value — at the same parameter budget via the ⅔ rule

  7. Weight tying: One V × d matrix does both jobs — the input lookup and the output scoring (logits = h·Eᵀ). It removes a full V × d weight (~38.6M params for GPT-2 small, about a third of the model) and lowers perplexity, because the shared matrix now gets gradient from both roles. tie_weights=True is the GPT-2 default; flip it off for the untied baseline

  8. Initialization matters: Small initial weights (std=0.02) prevent exploding activations

  9. Parameter scaling: Total params \approx V \times d + 12Nd^2 (dominated by FFN for large models)

  10. Architectural variations: Modern LLMs (LLaMA, Mistral) use RMSNorm, RoPE, and SwiGLU for better efficiency — and this module now builds RMSNorm and SwiGLU from scratch, swappable with norm_type and ffn_type

What’s Next

Module 07: Training trains our transformer on actual data using cross-entropy loss, learning rate scheduling, and gradient accumulation.