Neural networks (the parts that survived)

Why this matters

Every Transformer is, ultimately, a stack of fully-connected layers, residuals, and normalisations. If you understand a deep MLP — initialisation, optimisation, residual connections, regularisation — you understand 70% of a modern LLM block. This lesson rebuilds that foundation in a practical way (no derivatives by hand), aimed at someone who will spend the next year reading transformer papers.

Learning objectives

1. Build a feed-forward network in PyTorch.
2. Reason about activation functions, initialisation, and depth.
3. Understand residual connections and why they matter.
4. Apply dropout and weight decay correctly.
5. Diagnose common training pathologies (NaN, dead neurons, exploding loss).

1. The MLP and what each piece does

A feed-forward network is just Linear → activation → Linear → activation → ... → output.

import torch.nn as nn

class MLP(nn.Module):
    def __init__(self, d_in, d_hidden, d_out, depth=3):
        super().__init__()
        layers = [nn.Linear(d_in, d_hidden), nn.GELU()]
        for _ in range(depth - 2):
            layers += [nn.Linear(d_hidden, d_hidden), nn.GELU()]
        layers += [nn.Linear(d_hidden, d_out)]
        self.net = nn.Sequential(*layers)
    def forward(self, x):
        return self.net(x)

Even after 50 papers, the transformer FFN is just this: Linear → activation → Linear, with d_hidden = 4 × d_model.

Activations — the modern shortlist

You will reimplement SwiGLU in Phase 2.

Initialisation matters more than you think

Random init is not zero-or-one — the scale matters:

• Xavier (Glorot) — for tanh.
• Kaiming (He) — for ReLU/GELU. std = sqrt(2 / fan_in).
• Scaled init in transformers — std = 0.02 (GPT-2 convention) or scaled by 1/sqrt(2L) for residual layers (Megatron / GPT-J).

def init_linear(m):
    if isinstance(m, nn.Linear):
        nn.init.normal_(m.weight, mean=0, std=0.02)
        if m.bias is not None:
            nn.init.zeros_(m.bias)
model.apply(init_linear)

Get this wrong and your loss diverges in step 200.

2. Residual connections (the single most important deep-learning trick)

Without residuals, networks beyond ~10 layers stop training (the He et al. ResNet paper, 2015).

y = layer(x) + x      # residual

Two ideas this enables:

1. Gradient flow — the identity path lets gradients reach early layers without vanishing.
2. Iterative refinement — each layer can be a small correction to the running representation.

Every Transformer block is:

x = x + Attention(LayerNorm(x))     # pre-norm style
x = x + FFN(LayerNorm(x))

Without that x + ..., GPT-style models would not train past 6-12 layers.

Pre-norm vs post-norm

Original Transformer ("Attention Is All You Need") used post-norm: LayerNorm(x + Sublayer(x)). Modern LLMs use pre-norm: x + Sublayer(LayerNorm(x)). Pre-norm is much more stable at depth — every model after 2020 uses it.

3. Regularisation — when and why

Pretraining LLMs almost never uses dropout (the data is huge; dropout actively hurts at scale). But fine-tuning small models, training small heads, or building from scratch — you will use these tricks:

• Dropout p=0.1 — randomly zeros activations; classical regulariser.
• Weight decay 0.01-0.1 — L2 penalty on weights, decoupled in AdamW.
• Label smoothing — replace one-hot target with (1-ε) one_hot + ε/V uniform; less common in modern LLM SFT.
• Stochastic depth — randomly drop residual blocks during training (used in some vision transformers).
• Early stopping — irrelevant for big LLM pretraining (one pass over data).

For LLM SFT in 2026 you will typically run:

• dropout = 0.0
• weight_decay = 0.0 for embeddings/biases, 0.1 elsewhere
• 1-3 epochs (you usually overfit after that)

4. Optimisers under the hood (just enough)

AdamW keeps two running averages per parameter:

• m_t — exponential average of gradients (momentum).
• v_t — exponential average of squared gradients (variance).

Update step (simplified):

m_t = β1 m_{t-1} + (1-β1) g_t
v_t = β2 v_{t-1} + (1-β2) g_t^2
θ ← θ - lr * m_t / (sqrt(v_t) + ε) - lr * weight_decay * θ

The m / sqrt(v) term gives per-parameter adaptive learning rates: parameters with consistently large gradients get smaller effective lrs. That is why Adam works on noisy LLM gradients where vanilla SGD would diverge.

Memory cost: 2 extra tensors per parameter → roughly 3x the model's memory just for optimiser state. That is why 8-bit Adam (bitsandbytes) and FSDP optimizer sharding matter at scale.

5. Common training pathologies (read this twice)

• Loss = NaN at step 1 → bad init or learning rate too high. Halve the lr, retry.
• Loss explodes mid-training → gradient spike. Add clip_grad_norm_(1.0). If still bad, lower lr or use bf16 instead of fp16.
• Loss stuck at log(V) → model output is uniform. Check that the targets are shifted by 1 token (next-token prediction, not same-token).
• Loss decreases then plateaus high → too small a model, or learning rate too low, or data is too noisy.
• Train loss << val loss → overfitting. Reduce model size, add regularisation, get more data.
• Loss the same at any lr → masking bug; gradients are not flowing where you think.

The senior trick: plot per-layer parameter norms and gradient norms. If gradient norm is exploding in layer 1 only, you have a vanishing-gradient analogue from the top.

6. Putting it together — train an MLP on MNIST

import torch, torch.nn as nn, torch.nn.functional as F
from torchvision import datasets, transforms
from torch.utils.data import DataLoader

train = datasets.MNIST(".", train=True, download=True, transform=transforms.ToTensor())
test  = datasets.MNIST(".", train=False, download=True, transform=transforms.ToTensor())
tr = DataLoader(train, batch_size=256, shuffle=True)
te = DataLoader(test, batch_size=512)

class MLP(nn.Module):
    def __init__(self):
        super().__init__()
        self.fc1, self.fc2, self.fc3 = nn.Linear(784, 256), nn.Linear(256, 256), nn.Linear(256, 10)
    def forward(self, x):
        x = x.view(x.size(0), -1)
        x = F.gelu(self.fc1(x))
        x = F.gelu(self.fc2(x))
        return self.fc3(x)

m = MLP().cuda()
opt = torch.optim.AdamW(m.parameters(), lr=3e-4, weight_decay=0.01)

for epoch in range(3):
    m.train()
    for x, y in tr:
        x, y = x.cuda(), y.cuda()
        loss = F.cross_entropy(m(x), y)
        loss.backward()
        torch.nn.utils.clip_grad_norm_(m.parameters(), 1.0)
        opt.step(); opt.zero_grad()
    m.eval()
    with torch.no_grad():
        correct = sum((m(x.cuda()).argmax(-1) == y.cuda()).sum().item() for x, y in te)
        print(f"epoch {epoch}: acc={correct/len(test):.4f}")

You should hit ~98% in three epochs.

Hands-on lab (3 hours)

mlp_drills.ipynb:

1. Modify the MLP above to use ReLU instead of GELU. Train. Compare convergence.
2. Add a residual connection between fc1 and fc2 outputs. (Hint: project shapes if mismatched.) Show training curve.
3. Increase depth to 12 layers without residuals. Observe training fails or stagnates.
4. Add residuals to the same 12-layer net — show it now trains.
5. Add nn.LayerNorm(256) before each linear. Retrain.
6. Plot per-layer gradient norms during training. Discuss what you see.
7. Bonus: implement SwiGLU as a module: out = (SiLU(x W1) * (x W3)) @ W2. Replace GELU layer.

Common pitfalls

1. No init — PyTorch's default Linear init is fine for small nets but unstable for deep ones.
2. Forgetting LayerNorm before residuals in deep custom networks.
3. Calling .zero_grad() after .step() — fine, but easy to forget; loss seems "sticky."
4. Mixing data on CPU and GPU — silent data-copy slowdowns.
5. Using too-small batch_size with AdamW → noisy gradient stats (v_t) → instability.

Self-check

1. Why did residual connections enable deep networks?
2. What is SwiGLU and why is it preferred over GELU in modern LLMs?
3. Why is dropout typically off during LLM pretraining?
4. What is the memory cost of Adam vs SGD?
5. Pre-norm vs post-norm — which do modern LLMs use, and why?

References

• He et al. (2015), "Deep Residual Learning" (ResNet).
• Ba et al. (2016), "Layer Normalization."
• Loshchilov & Hutter (2019), "Decoupled Weight Decay Regularization" (AdamW).
• Shazeer (2020), "GLU Variants Improve Transformer" (SwiGLU/GeGLU origin).
• 3Blue1Brown, "Backpropagation, intuitively."