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

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.

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

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.10.0+cu128
Device: cpu

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{Swish}(xW_1) \otimes xW_2\). It requires an extra linear layer but often improves performance. Our implementation uses standard GELU to match GPT-2.

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 = 28.06
  Good init (std=0.02): output std = 0.56

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 787: 0.20%
  2. Token 763: 0.18%
  3. Token 654: 0.18%
  4. Token 492: 0.17%
  5. Token 877: 0.17%

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

GPT shares weights between token embedding and output projection:

# Check weight tying
print("Weight Tying:")
print(f"  Token embedding weight id: {id(model.token_embedding.weight)}")
print(f"  LM head weight id: {id(model.lm_head.weight)}")
print(f"  Are they the same object? {model.token_embedding.weight is model.lm_head.weight}")

# This saves parameters!
vocab_size = 1000
embed_dim = 128
saved_params = vocab_size * embed_dim
print(f"\nParameters saved by weight tying: {saved_params:,}")
Weight Tying:
  Token embedding weight id: 140165962793680
  LM head weight id: 140165962793680
  Are they the same object? True

Parameters saved by weight tying: 128,000

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 Add position bias 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.

Feed-Forward Networks

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

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.394
  Block 0: 12.604
  Block 1: 18.714
  Block 2: 23.403
  Block 3: 25.180
# 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)

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

  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

  7. Weight tying: Sharing token embedding and output projection reduces parameters and improves performance

  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

What’s Next

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