1 The core problem: gradient products
- By the chain rule, a gradient through the whole net is a product of many per-layer factors.
- Same architecture, data, and optimizer can either die (loss → NaN) or converge — the difference is purely how weights were initialized.
∂L/∂w₁ = (∂L/∂aₙ) · (∂aₙ/∂aₙ₋₁) · (∂aₙ₋₁/∂aₙ₋₂) · … · (∂a₂/∂w₁)
- Each factor < 1 → product → 0 (gradients vanish).
- Each factor > 1 → product → ∞ (gradients explode).
- The telephone game: even slightly off 1.0, repeated multiplication is exponential.
| Factor | After 10 layers | After 50 layers |
|---|---|---|
| 0.5 | 0.001 | ≈ 10⁻¹⁵ |
| 0.9 | 0.35 | 0.005 |
| 1.0 | 1.0 | 1.0 |
| 1.1 | 2.6 | 117 |
| 1.5 | 57.7 | ≈ 637,621 |
This is why 1990s networks couldn't go beyond 3–5 layers.
2 Vanishing vs exploding gradients
| Vanishing | Exploding |
|---|---|
| Early layers have near-zero gradients | Loss jumps to NaN |
| Loss plateaus (doesn't NaN) | Weights oscillate wildly |
| Only last few layers learn | Happens in first few epochs |
| Sneaky — no crash, just no learning | Dramatic — easy to detect |
- Vanishing is worse — you might not notice it. At least NaN tells you something is wrong.
- ReLU helps vanishing: its derivative is exactly 1 for positive inputs (no squashing, unlike sigmoid/tanh).
3 Preserve variance: the goal
- Too-small init (σ=0.01): activations shrink → vanish. Too-large (σ=1.0): activations grow → explode.
- Goal: keep activation variance constant across layers: Var(aₗ) ≈ Var(aₗ₋₁).
Linear layer y = Wx: Var(y) = nᵢₙ · Var(w) · Var(x) To keep Var(y) = Var(x) → Var(w) = 1 / nᵢₙ
4 Xavier & He initialization
Xavier / Glorot (2010) — accounts for both forward (activation variance) and backward (gradient variance) passes:
Var(w) = 2 / (nᵢₙ + nₒᵤₜ) w ~ N(0, 2/(nᵢₙ+nₒᵤₜ)) or w ~ U[ −√(6/(nᵢₙ+nₒᵤₜ)), √(6/(nᵢₙ+nₒᵤₜ)) ]
- Designed for symmetric activations: sigmoid, tanh (derivative ≈ 1 around 0).
- Breakthrough: enabled reliably training 5+ layer nets.
He / Kaiming (2015) — fixes Xavier+ReLU. ReLU zeros all negatives → kills half the activations → variance halves each layer. Double the variance to compensate:
Var(w) = 2 / nᵢₙ
- Activations stay stable across 20+ ReLU layers. Modern default for almost all deep nets.
| Activation | Init | Var(w) |
|---|---|---|
| Sigmoid / Tanh | Xavier (Glorot) | 2 / (nᵢₙ + nₒᵤₜ) |
| ReLU / LeakyReLU / ELU | He (Kaiming) | 2 / nᵢₙ |
- 99% of nets use ReLU variants → He init. Good init ≠ guaranteed convergence, but bad init guarantees failure.
- PyTorch defaults:
nn.Linearandnn.Conv2duse Kaiming (He) uniform — rarely need to change.
nn.init.kaiming_normal_(layer.weight, mode='fan_in') nn.init.xavier_normal_(layer.weight) nn.init.zeros_(layer.bias)
5 Batch Normalization (Ioffe & Szegedy, 2015)
For each mini-batch, normalize activations to zero mean / unit variance, then apply learned scale γ and shift β:
BN(x) = γ · (x − μ_B) / √(σ²_B + ε) + β
- Why? As earlier layers change, later layers see shifting input distributions ("internal covariate shift"). BN re-standardizes inputs before each layer.
- Computed independently per feature (column-wise across the batch).
- γ, β are learned: γ = learned std, β = learned mean. If γ=σ_B and β=μ_B, BN becomes a no-op → BN can never hurt (at worst learns to undo itself).
Worked example — 4 samples, 1 feature [2,4,6,8]:
μ_B = 5, σ²_B = 5 x̂ = (x−5)/√5 = [−1.34, −0.45, 0.45, 1.34] y = 2x̂ + 1 (with γ=2, β=1)
model.eval() makes BN use test-batch stats → batch-dependent, unreliable predictions.- Placement: Linear → BatchNorm → ReLU. Normalize before activation so ReLU gets well-distributed inputs.
- Tip: use
bias=Falsein the Linear before BN — BN's β replaces the bias.
nn.Sequential(nn.Linear(256,128,bias=False), nn.BatchNorm1d(128), nn.ReLU())
BN limitations: batch size = 1 (no meaningful stats), small batches (noisy → unstable), variable-length sequences (each timestep needs different stats), and it couples samples (output of A depends on its batch neighbors). These motivated LayerNorm.
6 Layer Normalization & BN vs LN (Ba et al., 2016)
For each sample, normalize across all features independently:
LN(x) = γ · (x − μ_features) / √(σ²_features + ε) + β
| BatchNorm | LayerNorm | |
|---|---|---|
| Normalizes | each feature across the batch | each sample across features |
| Needs | multiple samples | self-contained (1 sample OK) |
| Train vs test | different (running avgs) | same behavior |
| Best for | CNNs / vision, big batches | Transformers, sequences |
- LN works with batch size = 1, variable-length sequences, no sample coupling, identical train/test behavior.
- Used in virtually all Transformers (BERT, GPT, …) because they have variable-length seqs and often batch size 1 at inference — BN would break.
7 Gradient checking
Verify your analytical (backprop) gradient against a numerical estimate (central difference):
∂f/∂x ≈ (f(x+ε) − f(x−ε)) / (2ε) (ε ≈ 10⁻⁵) relative error = |g_analytical − g_numerical| / max(|g_analytical|, |g_numerical|)
| Relative error | Verdict |
|---|---|
| < 10⁻⁷ | Excellent |
| < 10⁻⁵ | Acceptable |
| > 10⁻³ | Something is wrong |
- Slow (2 forward passes per parameter) → use only for debugging custom backward passes, never in training. (PyTorch:
torch.autograd.gradcheck.)
8 Debugging checklist & failure modes
The checklist — follow in order, never skip ahead:
- 1. Overfit ONE batch (8 samples → zero loss). If this fails, stop.
- 2. Check tensor shapes at every layer — silent mismatches cause mysterious bugs.
- 3. Verify loss at init: for k classes, L₀ ≈ ln(k). Wrong value ⇒ bad init or loss.
- 4. Gradient check — numerical vs analytical, catches custom backward bugs.
- 5. Train on a small subset (~100 samples) — should overfit quickly, confirms pipeline.
- 6. Scale up — full data, regularization, tuning. Only after steps 1–5 pass.
| Symptom | Likely cause | Fix |
|---|---|---|
| Loss stuck at ln(classes) | Not learning at all | Check data pipeline, loss fn, LR |
| Loss → NaN | Exploding gradients | Lower LR, gradient clip, check init |
| Loss plateaus early | Vanishing gradients | Check activations, add skip connections |
| Train ↓, val ↑ | Overfitting | Regularize |
★ Likely exam questions
Q1. Why does Xavier init fail with ReLU, and what does He do differently?
Xavier assumes a symmetric activation with derivative ≈ 1 around 0. ReLU zeros all negatives → halves the variance each layer → activations still vanish. He init uses Var(w)=2/nᵢₙ (double Xavier's 1/nᵢₙ) to compensate for the halving.
Q2. All weights set to 0.01 in a 50-layer net — what happens to early-layer gradients?
They vanish. Each layer multiplies by a small factor, so the product shrinks exponentially toward 0 — early layers stop learning.
Q3. During inference, where does BatchNorm get μ and σ²? What if you forget model.eval()?
From running averages accumulated during training (NOT the current batch). Forgetting model.eval() makes BN use the test batch's stats → predictions become batch-dependent and noisy.
Q4. What do BN and LN normalize over, and why does it matter for Transformers?
BN normalizes each feature across the batch (needs multiple samples); LN normalizes each sample across its features (self-contained). Transformers use variable-length sequences and often batch size 1 at inference, where BN breaks — LN has no batch dependency.
Q5. Initial loss of a 10-class classifier is 4.5 instead of ≈ 2.3 — what's wrong?
Expected L₀ = ln(10) ≈ 2.3 for uniform random predictions. 4.5 suggests a wrong loss function, mislabeled data, or an architecture bug like applying softmax twice.
Q6. Loss goes to NaN after 3 epochs — give 3 causes and fixes.
(1) Exploding gradients → lower LR or gradient clipping. (2) Bad initialization → use He init. (3) Numerical overflow in loss e.g. log(0) → add epsilon / use a numerically stable implementation.
Q7. 30-layer ReLU net has dead gradients in early layers — first fix? And the first debugging step when a model won't learn at all?
Dead-gradient fix: switch to He initialization. Won't-learn-at-all first step: overfit one batch (checklist step 1).