Training Deep Networks Effectively
Batch Normalization & Layer Normalization
As a network trains, every layer's weights are constantly updating (Chapter 2's optimizers) — which means the distribution of values flowing into every subseque
Jr Codex Deep Learning Notes
Level: Intermediate Prerequisites: Chapter 2: Optimizers Time to complete: ~20 minutes
Table of Contents
- The Problem: Internal Covariate Shift
- Batch Normalization
- What BatchNorm Learns
- BatchNorm's Train vs Eval Behavior
- Layer Normalization
- BatchNorm vs LayerNorm — When to Use Which
- Normalization in PyTorch
- Summary & Next Steps
1. The Problem: Internal Covariate Shift
As a network trains, every layer's weights are constantly updating (Chapter 2's optimizers) — which means the distribution of values flowing into every subsequent layer is also constantly shifting. Each layer is effectively trying to learn from a "moving target" input distribution, which slows and destabilizes training. This phenomenon is called internal covariate shift.
Why This Compounds in DEEP Networks
─────────────────────────────────────────
A small shift in Layer 1's output distribution becomes a
LARGER shift by the time it reaches Layer 5, since each
layer's output depends on (and can amplify) the layer
before it. In a very deep network, later layers can be
chasing a constantly and substantially shifting target.
─────────────────────────────────────────
2. Batch Normalization
Batch Normalization (BatchNorm) directly addresses this by normalizing each layer's activations — for every mini-batch, forcing them to have zero mean and unit variance — before passing them to the next layer.
import numpy as np
def batch_norm(Z, epsilon=1e-8):
"""
Z: activations for a MINI-BATCH, shape (batch_size, num_features)
"""
batch_mean = np.mean(Z, axis=0) # mean ACROSS the batch, per feature
batch_var = np.var(Z, axis=0)
Z_normalized = (Z - batch_mean) / np.sqrt(batch_var + epsilon)
return Z_normalized
# A mini-batch of 4 examples, 3 features each
Z = np.array([
[1.0, 50.0, -3.0],
[2.0, 45.0, -1.0],
[1.5, 55.0, -2.0],
[0.5, 40.0, -4.0],
])
Z_norm = batch_norm(Z)
print(f"Mean per feature (should be ~0): {Z_norm.mean(axis=0)}")
print(f"Variance per feature (should be ~1): {Z_norm.var(axis=0)}")3. What BatchNorm Learns
Simply forcing zero mean and unit variance everywhere could actually reduce a network's representational power (e.g., ReLU's usefulness depends partly on the scale of its input). BatchNorm compensates by adding two learnable parameters per feature — gamma (scale) and beta (shift) — letting the network learn the optimal normalized scale for each feature, rather than being forced into a fixed one.
def batch_norm_with_learnable_params(Z, gamma, beta, epsilon=1e-8):
batch_mean = np.mean(Z, axis=0)
batch_var = np.var(Z, axis=0)
Z_normalized = (Z - batch_mean) / np.sqrt(batch_var + epsilon)
Z_scaled_shifted = gamma * Z_normalized + beta # gamma, beta are LEARNED via backpropagation
return Z_scaled_shifted
gamma = np.ones(3) # initialized to 1 (no scaling initially)
beta = np.zeros(3) # initialized to 0 (no shift initially)Why This Matters
─────────────────────────────────────────
If gamma and beta are LEARNED to exactly UNDO the
normalization (gamma = sqrt(batch_var), beta = batch_mean),
the network can recover the ORIGINAL, un-normalized
activations if that turns out to be optimal — BatchNorm
never REDUCES what the network is capable of representing,
it just gives training an easier, more stable STARTING point.
─────────────────────────────────────────
4. BatchNorm's Train vs Eval Behavior
A subtlety that causes real bugs if misunderstood: BatchNorm behaves differently during training versus inference.
Training Time vs Inference Time
─────────────────────────────────────────
TRAINING: normalize using the CURRENT mini-batch's own
mean and variance (Section 2) — but ALSO
maintain a running average of these
statistics across all batches seen so far
INFERENCE: use the RUNNING AVERAGE statistics computed
during training, NOT the current batch's
— because at inference time, you might
only have ONE example (no "batch" to
compute statistics from), and you want
CONSISTENT, deterministic behavior
regardless of what other examples happen
to be in the same request
─────────────────────────────────────────
Forgetting to switch a PyTorch model between these two modes (model.train() vs model.eval()) is one of the most common real-world bugs in deep learning code — a model evaluated in training mode, or deployed in training mode (Module 9), will silently produce inconsistent or degraded predictions.
5. Layer Normalization
Layer Normalization (LayerNorm) normalizes across the features of a single example, rather than across the batch. This makes it independent of batch size and other examples in the batch entirely.
def layer_norm(z, epsilon=1e-8):
"""
z: activations for a SINGLE example, shape (num_features,)
"""
feature_mean = np.mean(z) # mean ACROSS this example's OWN features
feature_var = np.var(z)
return (z - feature_mean) / np.sqrt(feature_var + epsilon)
z = np.array([1.0, 50.0, -3.0, 20.0])
z_norm = layer_norm(z)
print(f"Normalized: {z_norm}")
print(f"Mean (should be ~0): {z_norm.mean():.6f}")6. BatchNorm vs LayerNorm — When to Use Which
The Deciding Factor
─────────────────────────────────────────
BatchNorm: depends on OTHER examples in the same batch
— works well for CNNs (Module 4) with
reasonably large, consistent batch sizes,
but degrades with very small batches (fewer
examples to compute reliable statistics from)
LayerNorm: depends ONLY on the current example's own
features — unaffected by batch size,
and works naturally with VARIABLE-length
sequences, which is why it is the
near-universal standard in RNNs (Module
5) and Transformers (Module 6) instead
of BatchNorm
─────────────────────────────────────────
7. Normalization in PyTorch
import torch
import torch.nn as nn
# BatchNorm — typical usage in a CNN (Module 4)
model_cnn = nn.Sequential(
nn.Linear(784, 256),
nn.BatchNorm1d(256), # normalizes ACROSS the batch, for each of the 256 features
nn.ReLU(),
)
# LayerNorm — typical usage in a Transformer (Module 6)
model_transformer_block = nn.Sequential(
nn.Linear(512, 512),
nn.LayerNorm(512), # normalizes ACROSS the 512 features, for EACH example independently
nn.ReLU(),
)
# The train/eval switch from Section 4 — CRITICAL to get right
model_cnn.train() # use CURRENT batch statistics
# ... training loop ...
model_cnn.eval() # use STORED running average statistics
# ... inference / evaluation ...8. Summary & Next Steps
Key Takeaways
- Internal covariate shift — the constantly shifting distribution of activations during training — slows and destabilizes deep network training.
- Batch Normalization normalizes activations per mini-batch, then applies learnable scale (gamma) and shift (beta) parameters so the network never loses representational power.
- BatchNorm behaves differently during training (uses current batch statistics) versus inference (uses stored running averages) — forgetting to switch modes is a common real bug.
- Layer Normalization normalizes across a single example's own features instead of across the batch, making it independent of batch size and the standard choice for RNNs and Transformers.
Concept Check
- What problem does batch normalization solve, and why does it compound in deeper networks?
- Why does BatchNorm include learnable gamma and beta parameters instead of just normalizing to zero mean and unit variance?
- Why is LayerNorm generally preferred over BatchNorm for Transformer architectures?
Next Chapter
→ Chapter 4: Regularization — Dropout, Weight Decay, Early Stopping
Jr Codex — 1-on-1 Personalized Coaching | Back to Module Index | Back to DL Index