Module 1 · Reading 1.7 · read before S07: Regularization, Initialization, Normalization

The PDF posted in Canvas is the assigned version of this reading. It is built from this page, so the text is identical; anything here marked as deep reading is optional and is not in the PDF.

A sufficiently expressive network can fit a small training set closely, including its noise. Parameter count alone does not guarantee interpolation or poor generalization; architecture, data, and optimization also matter. Before training starts, an unsuitable weight scale can also cause activations to shrink or grow sharply through depth, making optimization difficult. Forward activation measurements diagnose signal scale; gradient measurements and training runs are needed to assess the resulting optimization problem. This section is about the three families of tools that address these failure modes: regularization (weight decay, dropout, early stopping) to control what the network learns, initialization (Xavier, Kaiming) so that training can start at all, and normalization (BatchNorm, LayerNorm) so that it keeps going in deep networks.

These three are the difference between the theory of the previous sections and networks that actually train. Every experiment below is small enough to rerun in seconds, and each one shows a real number going wrong before we fix it.

How to read this

Overfitting, weight decay, and early stopping

Take 30 noisy points from \(y = \sin(2x)\) and fit them with an MLP that has tens of thousands of parameters. The network can pass near every training point, noise included; the interesting question is what happens to error measured against the true function. One classical remedy is L2 regularization: add \(\tfrac\lambda2\lVert w\rVert^2\) to the objective. For plain SGD without momentum, its extra gradient \(\lambda w\) gives the multiplicative decay factor \(1-\eta\lambda\). AdamW applies this decay separately from the adaptive gradient update; coupled L2 in Adam is different, as explained in reading 1.6. The idea predates deep learning (ridge regression in statistics, Tikhonov regularization in inverse problems). Penalizing large weights can discourage sensitivity, but it does not guarantee a smoother learned function or better validation performance. The second classical remedy costs nothing: early stopping — track validation error during training and keep the weights from the epoch where it was lowest.

The experiment trains the same network with increasing weight decay (via AdamW; see optimizers for why AdamW rather than Adam), and also records the best validation error seen at any epoch, which is what early stopping would have captured:

import torch
import torch.nn as nn

torch.manual_seed(0)

# 30 noisy training points from y = sin(2x); a big MLP can memorize them.
x_train = torch.rand(30, 1) * 6 - 3
y_train = torch.sin(2 * x_train) + 0.3 * torch.randn(30, 1)
x_val = torch.linspace(-3, 3, 200).unsqueeze(1)
y_val = torch.sin(2 * x_val)            # noise-free target for evaluation

def train(weight_decay):
    torch.manual_seed(1)
    model = nn.Sequential(
        nn.Linear(1, 128), nn.Tanh(),
        nn.Linear(128, 128), nn.Tanh(),
        nn.Linear(128, 1),
    )
    opt = torch.optim.AdamW(model.parameters(), lr=1e-3, weight_decay=weight_decay)
    loss_fn = nn.MSELoss()
    best_val, best_epoch = float("inf"), 0
    for epoch in range(1, 4001):
        model.train()
        opt.zero_grad()
        loss = loss_fn(model(x_train), y_train)
        loss.backward()
        opt.step()
        model.eval()
        with torch.no_grad():
            val = loss_fn(model(x_val), y_val).item()
        if val < best_val:
            best_val, best_epoch = val, epoch
    # Both final metrics describe the model after the last optimizer update.
    with torch.no_grad():
        final_train = loss_fn(model(x_train), y_train).item()
    return final_train, val, best_val, best_epoch

print(f"{'weight_decay':>13s}{'train MSE':>11s}{'final val':>11s}{'best val':>11s}{'best epoch':>12s}")
for wd in (0.0, 1e-1, 1.0, 3.0):
    tr, val, bv, be = train(wd)
    print(f"{wd:13.0e}{tr:11.4f}{val:11.4f}{bv:11.4f}{be:12d}")

Running this prints:

 weight_decay  train MSE  final val   best val  best epoch
        0e+00     0.0008     0.2335     0.0275         400
        1e-01     0.0014     0.2265     0.0273         415
        1e+00     0.0196     0.0573     0.0303         465
        3e+00     0.0375     0.0420     0.0361        1472

Four observations follow from this run. Without decay, final training MSE is 0.0008 while error against the noise-free function is 0.2335. These evaluate different targets, so their ratio is not a generalization-gap measure on a common noisy distribution. The strongest evidence for early stopping is the same validation metric at two times: 0.0275 at epoch 400 versus 0.2335 at the end. Save the best checkpoint during training and restore it for use; recording its score alone does not restore its weights. This experiment knows the true function, an advantage real tasks usually lack; use held-out observations to choose checkpoints and reserve a test set for final evaluation.

Decay 0.1 changes final validation error only modestly here. Decay 1.0 reduces it to 0.0573, roughly a fourfold reduction from 0.2335. Decay 3.0 has an even lower final error (0.0420), but its best observed checkpoint (0.0361) is worse than the best checkpoints at the smaller settings. Thus the preferred setting depends on whether the comparison uses the last or the validation-selected checkpoint. The table does not prove that every update after epoch 400 only memorized noise, or that 0.0361 is the best error this architecture can ever achieve. Tune decay and stopping time together.

Two siblings of L2 deserve a sentence each before we move on. L1 regularization penalizes \(\sum \lvert w\rvert\) instead of \(\lVert w\rVert^2\); the L1 optimum often contains exact zeros, and proximal/soft-thresholding algorithms create them explicitly. Ordinary gradient or subgradient updates alone do not guarantee that an iterate lands exactly at zero. This is why L1 is a standard route to sparse solutions in classical statistics (the lasso), though it is rarer as the primary regularizer in deep learning. And the most effective regularizer in modern practice is often not a penalty at all: data augmentation (training on label-preserving transformations of the inputs) attacks overfitting at its source by making the training set effectively larger and by encoding invariances the network would otherwise have to discover.

Dropout

Weight decay penalizes where the weights are; dropout (introduced in a 2012 preprint by Hinton and colleagues, with the full JMLR paper by Srivastava et al. in 2014) attacks co-adaptation: during training, each activation is independently zeroed with probability p, so no unit can rely on any particular other unit being present. The surviving activations are scaled by 1 / (1 - p) so that the expected value seen by the next layer is unchanged — “inverted dropout.” At evaluation time, dropout does nothing at all. The mechanics are worth verifying once, directly:

import torch
import torch.nn as nn

torch.manual_seed(0)

drop = nn.Dropout(p=0.5)
x = torch.ones(10000)                 # activations, all equal to 1.0

drop.train()                          # training mode: random zeroing + rescale
y_train = drop(x)
kept = (y_train != 0).float().mean().item()
print(f"train mode: fraction kept = {kept:.4f}")
print(f"train mode: surviving values are   {y_train[y_train != 0][0].item():.1f}  (1 / (1-p) = 2.0)")
print(f"train mode: mean of output = {y_train.mean().item():.4f}  (expectation preserved)")

drop.eval()                           # eval mode: identity
y_eval = drop(x)
print(f"eval  mode: mean of output = {y_eval.mean().item():.4f}, zeros = {(y_eval == 0).sum().item()}")

Running this prints:

train mode: fraction kept = 0.4954
train mode: surviving values are   2.0  (1 / (1-p) = 2.0)
train mode: mean of output = 0.9908  (expectation preserved)
eval  mode: mean of output = 1.0000, zeros = 0

About half the units are zeroed, survivors are doubled, the mean is preserved on average, and in eval mode the layer is the identity. Dropout can be read as training an implicit ensemble of thinned subnetworks, or as injecting noise that discourages brittle solutions. In current practice its role has narrowed: it remains common in Transformer blocks and MLP heads at modest rates (around 0.1), while convolutional networks mostly rely on normalization, augmentation, and weight decay instead.

Where you place dropout matters as much as whether you use it. The standard position is after an activation, inside fully connected blocks; applying it to the output layer corrupts predictions directly, and applying strong dropout to low-dimensional inputs throws away data. Rates are architecture-dependent: the original work used up to 0.5 in the wide fully connected layers of 2012-era vision networks, while modern Transformer recipes use around 0.1; large language models at the biggest scales often set it to zero, since with vastly more data than parameters per pass, underfitting is the operative risk, not overfitting. The train/eval asymmetry also has a productive second life: leaving dropout on at inference and averaging many stochastic forward passes (“Monte Carlo dropout,” analyzed by Gal and Ghahramani in 2016) turns the noise into a rough uncertainty estimate — a reminder that the mode switch you must normally never forget is a modeling decision, not just a flag.

Pitfall. The number-one regularization bug is forgetting model.eval() at test time. With dropout still in training mode, your predictions are randomly corrupted — accuracy drops and, worse, changes between runs on identical inputs. With BatchNorm still in training mode, evaluation statistics get computed from your test batch, so your “predictions” depend on what else you batched together, and a batch size of 1 can crash outright (you will see this error below). If evaluation results look noisy or batch-size-dependent, check the mode first: model.training tells you.

Initialization: why the scale of random weights matters

Hidden units with identical incoming and outgoing parameters can remain symmetric and learn the same features. Random weight initialization is a practical way to break that symmetry; biases can commonly start at zero. We also need a suitable weight scale.

Track the second moment, \(q=\mathbb E[x^2]\), the average squared activation. Unlike variance, this does not subtract the mean. For one output \(z=\sum_{i=1}^{n}w_i x_i\) with zero bias, independent zero-mean weights of variance \(\sigma^2\), and weights independent of the inputs, the cross terms vanish in expectation. If the inputs have a common second moment \(q\), then \(\mathbb E[z^2]=n\sigma^2 q\). This is an initialization calculation over random weights and inputs, not a promise about each unit in one sampled network. When the inputs are centered, their second moment equals their variance.

Xavier (Glorot and Bengio, AISTATS 2010) balances forward and backward signal scales with weight variance \(2/(n_{\mathrm{in}}+n_{\mathrm{out}})\); at equal widths this is \(1/n\). The linearized argument is useful for tanh-like networks near their nonsaturated regime, and the activation’s gain matters in a concrete implementation.

Kaiming (He et al., ICCV 2015) accounts for ReLU. For a symmetric pre-activation distribution, \(\mathbb E[\operatorname{ReLU}(z)^2]=\tfrac12\mathbb E[z^2]\). It is the second moment that halves; ReLU usually makes the mean positive, so its variance does not simply halve. A hand calculation shows the distinction: values [-2, -1, 1, 2] have average square 2.5; after ReLU, [0, 0, 1, 2] has average square 1.25 and mean 0.75. The variance after ReLU is 1.25 - 0.75**2 = 0.6875, not 1.25.

Combining the linear layer and ReLU gives \(q_{\mathrm{next}}\approx\tfrac12 n\sigma^2 q\). Choosing \(\sigma^2=2/n\) keeps this approximate factor near one. Across many layers, a repeated factor below or above one can make the signal shrink or grow rapidly. Backward propagation has a related scale calculation involving fan-out and activation derivatives; this forward argument alone does not measure the gradients. The experiment below reports activation standard deviation, a useful observed diagnostic, across 40 ReLU layers. 

import math
import torch

torch.manual_seed(0)

depth, width, batch = 40, 256, 512
x0 = torch.randn(batch, width)          # inputs with std ~ 1

def forward_stats(weight_std, label):
    """Push x0 through `depth` ReLU layers; report activation std along the way."""
    torch.manual_seed(1)                # same weights-up-to-scale for each run
    x = x0
    stds = {}
    with torch.no_grad():
        for layer in range(1, depth + 1):
            W = torch.randn(width, width) * weight_std
            x = torch.relu(x @ W.T)
            if layer in (1, 5, 10, 20, 40):
                stds[layer] = x.std().item()
    row = "".join(f"{stds[l]:>12.2e}" for l in (1, 5, 10, 20, 40))
    print(f"{label:>22s}{row}")

print(f"{'activation std, layer:':>22s}" + "".join(f"{l:>12d}" for l in (1, 5, 10, 20, 40)))
forward_stats(0.01,                      "N(0, 0.01^2)")
forward_stats(math.sqrt(1.0 / width),    "Xavier  sqrt(1/n)")
forward_stats(math.sqrt(2.0 / width),    "Kaiming sqrt(2/n)")

Running this prints:

activation std, layer:           1           5          10          20          40
          N(0, 0.01^2)    9.40e-02    1.58e-05    2.82e-10    8.34e-20    8.72e-39
     Xavier  sqrt(1/n)    5.87e-01    1.51e-01    2.57e-02    6.90e-04    5.96e-07
     Kaiming sqrt(2/n)    8.30e-01    8.54e-01    8.21e-01    7.07e-01    6.25e-01

The small N(0, 0.01^2) initialization produces an activation standard deviation of 8.72e-39 at layer 40, in the float32 subnormal range. This diagnoses severe forward-signal collapse in this stack; the experiment does not test every learning rate or directly measure backward gradients. The equal-width Xavier scale also decays markedly under ReLU, from 0.587 at layer 1 to 5.96e-7 at layer 40. Kaiming keeps the observed standard deviation on a comparable scale (0.830 to 0.625), with finite-network drift rather than exact preservation. These results support matching initialization to the activation and architecture.

PyTorch’s nn.Linear calls kaiming_uniform_ with a=sqrt(5), equivalent to the historical uniform bound ±1/sqrt(fan_in). That default has variance 1/(3*fan_in) and differs from the 2/fan_in ReLU recipe tested here. For a deliberately initialized ReLU layer, nn.init.kaiming_normal_(layer.weight, mode="fan_in", nonlinearity="relu") expresses the intended choice. PyTorch stores a linear layer’s weight as (out_features, in_features) and computes with its transpose; the NumPy convention used in S05 stores (in_features, out_features). Identify which axis counts inputs before applying an initialization formula.

Variance-matching was the beginning of initialization research, not the end. A second line of work, from Saxe, McClelland, and Ganguli (ICLR 2014), analyzed deep linear networks exactly and showed that orthogonal weight matrices (matrices that rotate without stretching, so every singular value is 1) give learning dynamics whose time-to-converge does not grow with depth, where independent Gaussian entries merely preserve variance on average while individual directions still shrink or grow. That analysis is why torch.nn.init.orthogonal_ exists and why orthogonal-flavored inits became standard advice for recurrent networks, where the same matrix is applied dozens of times in a loop and any systematic stretch compounds. A third line adapts initialization to architecture: in a residual network, the skip connection already carries the signal forward, so the right question becomes how to initialize the residual branch — and the common answer is “start it at zero.” Goyal and colleagues’ 2017 large-batch recipe zero-initializes the scale of each residual block’s final normalization layer so every block begins as the identity, and Fixup (Zhang, Dauphin, and Ma, ICLR 2019) pushed the idea to its logical end, showing that with carefully rescaled and partially zeroed initialization, ten-thousand-layer residual networks can train without any normalization layers at all. The through-line of all three generations: understand how signal propagates through your particular architecture at step 0, then choose weights that make that propagation neutral.

Historical note. Before principled initialization, the accepted workaround for deep networks was greedy layer-wise pretraining: around 2006, Hinton and colleagues trained deep belief networks one layer at a time precisely because end-to-end training from random weights failed. Glorot and Bengio’s 2010 analysis showed much of that failure was an artifact of poor initialization (and saturating activations), and within a few years the combination of Kaiming init, ReLU, and BatchNorm made pretraining-for-depth unnecessary — an entire training paradigm displaced by getting a variance calculation right.

Normalization: BatchNorm and LayerNorm

Initialization chooses a starting signal scale; nothing forces that scale to remain appropriate as the weights move. Normalization layers re-standardize activations continuously during training. Batch normalization (Sergey Ioffe and Christian Szegedy, ICML 2015) centers and rescales each feature using a group of values. For BatchNorm1d on an (N, C) array in training mode, each group is one column across the batch; convolutional BatchNorm also includes spatial positions. With \(\hat x=(x-\mu)/\sqrt{s^2+\epsilon}\), the centered group has approximately unit variance when \(s^2\) is large relative to \(\epsilon\). A learned scale and shift (\(\gamma\), \(\beta\)) then let the model choose each feature’s output scale and offset. It made training dramatically faster and more forgiving of learning rates, and was adopted almost universally in convolutional networks within a year or two. (The original explanation, reducing “internal covariate shift,” has been challenged; Santurkar and colleagues argued in 2018 that BatchNorm’s benefit is better explained by a smoother optimization landscape. The empirical benefit is not in dispute; the mechanism is.)

During training, BatchNorm statistics depend on the batch (and on spatial positions for convolutional inputs). With the default track_running_stats=True, evaluation instead uses running statistics accumulated during training; with tracking disabled, batch statistics are used in both modes. Layer normalization (Jimmy Lei Ba, Jamie Ryan Kiros, and Geoffrey Hinton, 2016) uses the specified final feature dimensions of each item, independently of other batch items. LayerNorm(C) on (N, C) uses each row; on (N, T, C) it uses the C features of each token. The layer uses the same calculation in train and eval modes, although other layers such as dropout can still make the whole model behave differently. For a hand-sized (N, C) example, use [[1, 2], [3, 6]]. BatchNorm groups the columns: their means are 2 and 4, and their population variances are 1 and 4. LayerNorm groups the rows: their means are 1.5 and 4.5, and their population variances are 0.25 and 2.25. With scale 1, shift 0, and epsilon 1e-5, the normalized arrays round to [[-1, -1], [1, 1]] for BatchNorm and [[-1, 1], [-1, 1]] for LayerNorm. The following experiment shows which groups are normalized, and how training-mode BatchNorm behaves when its batch changes: 

import torch
import torch.nn as nn

torch.manual_seed(0)

# A batch of 8 samples, 4 features, deliberately different scales per feature.
x = torch.randn(8, 4) * torch.tensor([0.1, 1.0, 10.0, 100.0]) + torch.tensor([0.0, 1.0, -5.0, 50.0])

bn = nn.BatchNorm1d(4, affine=False)   # affine=False: just the normalization
ln = nn.LayerNorm(4, elementwise_affine=False)

bn.train()
y_bn = bn(x)
y_ln = ln(x)

def stats(t, dim):
    return t.mean(dim=dim), t.std(dim=dim, unbiased=False)

m, s = stats(y_bn, 0)
print("BatchNorm: per-FEATURE mean/std across the batch (dim=0)")
print("  mean:", " ".join(f"{v:8.4f}" for v in m))
print("  std :", " ".join(f"{v:8.4f}" for v in s))
m, s = stats(y_bn, 1)
print("BatchNorm: per-SAMPLE mean across features (dim=1) — NOT normalized")
print("  mean:", " ".join(f"{v:8.4f}" for v in m[:4]), "...")

m, s = stats(y_ln, 1)
print("LayerNorm: per-SAMPLE mean/std across features (dim=1)")
print("  mean:", " ".join(f"{v:8.4f}" for v in m[:4]), "...")
print("  std :", " ".join(f"{v:8.4f}" for v in s[:4]), "...")

# BatchNorm's weakness: the output for sample 0 depends on WHO ELSE is in the batch.
bn2 = nn.BatchNorm1d(4, affine=False)
bn2.train()
pair = bn2(x[[0, 7]])               # same sample 0, different batchmates
print("\nsample 0 normalized in a batch of 8 :", " ".join(f"{v:7.3f}" for v in y_bn[0]))
print("sample 0 normalized in a batch of 2 :", " ".join(f"{v:7.3f}" for v in pair[0]))
print("sample 0 under LayerNorm (batch-free):", " ".join(f"{v:7.3f}" for v in ln(x[:1])[0]))

# And with a batch of one, training-mode BatchNorm cannot even run:
try:
    bn2(x[:1])
except ValueError as e:
    print(f"\nBatchNorm1d on a single sample in train mode -> ValueError: {e}")

Running this prints:

BatchNorm: per-FEATURE mean/std across the batch (dim=0)
  mean:  -0.0000   0.0000   0.0000  -0.0000
  std :   0.9994   1.0000   1.0000   1.0000
BatchNorm: per-SAMPLE mean across features (dim=1) — NOT normalized
  mean:  -0.8359  -0.3273  -0.0156   0.6555 ...
LayerNorm: per-SAMPLE mean/std across features (dim=1)
  mean:   0.0000   0.0000  -0.0000  -0.0000 ...
  std :   1.0000   1.0000   1.0000   1.0000 ...

sample 0 normalized in a batch of 8 :  -1.465  -0.657  -0.882  -0.339
sample 0 normalized in a batch of 2 :  -1.000  -1.000  -1.000   1.000
sample 0 under LayerNorm (batch-free):   0.035   0.027  -1.445   1.382

BatchNorm1d on a single sample in train mode -> ValueError: Expected more than 1 value per channel when training, got input size torch.Size([1, 4])

The first blocks confirm the groups: BatchNorm standardizes each feature over this batch; LayerNorm standardizes each sample over its features. The same input sample has different BatchNorm outputs in batches of eight and two because its column statistics change. Its LayerNorm output stays the same because its own feature group stays the same. With default running-statistic tracking, bn.eval() uses stored estimates rather than the current batch. Poorly estimated running statistics or a distribution shift can still hurt evaluation. The final error is specific to training-mode BatchNorm1d on (1, C), which has only one value per channel; convolutional BatchNorm can use additional spatial values even when N=1.

Why Transformers use LayerNorm

LayerNorm(C) is a natural fit for token representations of shape (N, T, C): it normalizes the C features of each token independently, without estimating statistics from other tokens or batch items. The normalization layer uses the same group rule in training, evaluation, and single-token decoding; other model layers such as dropout can still behave differently across modes. This does not mean default BatchNorm evaluation fails for one token: with running-statistic tracking enabled, evaluation uses stored estimates and can process a single sample. The original Transformer used it from the start (“Attention Is All You Need,” 2017), and the placement question, after each sublayer (post-LN, as in the original) versus before (pre-LN, now more common because it stabilizes deep stacks), is one you will meet again in the Transformer. Between BatchNorm’s batch coupling in convolutional networks and LayerNorm’s per-token independence in Transformers, one habit generalizes: know which axis your normalization averages over, and what that implies at inference.

People. Kaiming He, whose name is on the ReLU-aware initialization above, is also lead author of the 2015 ResNet work that made 100-plus-layer networks trainable — initialization, normalization, and residual connections are in large part the story of a few researchers systematically removing the obstacles to depth. He, Glorot, and Bengio’s initialization schemes are now literally function names (kaiming_normal_, xavier_uniform_) in torch.nn.init, a rare case of researchers’ names becoming API.

The wider normalization zoo

BatchNorm and LayerNorm are two members of a family whose distinguishing feature is a single design decision: which axis do you average over? Two more members matter in practice. Group normalization (Yuxin Wu and Kaiming He, ECCV 2018) splits the channels of each sample into groups and normalizes within each group — per sample, like LayerNorm, but with finer granularity. Its motivation was concrete: object detection and segmentation models process large images and can only fit one or two samples per GPU, exactly the regime where BatchNorm’s batch statistics fall apart; the paper reports that on ResNet-50 with batch size 2, GroupNorm cuts error by over 10 percentage points relative to BatchNorm. (Instance normalization, common in style transfer, is the one-channel-per-group special case.) RMSNorm (Biao Zhang and Rico Sennrich, NeurIPS 2019) simplifies LayerNorm instead: drop the mean-subtraction entirely and just divide by the root-mean-square of the features, on the hypothesis that rescaling, not recentering, is what matters. The paper reports equal quality with 7–64% less normalization runtime across their models, and the hypothesis aged well: the Llama family and many recent large language models use RMSNorm in place of LayerNorm.

The axes are easier to keep straight in code than in prose:

import torch
import torch.nn as nn

torch.manual_seed(0)

# One batch: 4 samples, 6 features, wildly different scales and offsets.
x = torch.randn(4, 6) * 3.0 + 2.0

ln = nn.LayerNorm(6, elementwise_affine=False)
rms = nn.RMSNorm(6, elementwise_affine=False)
gn = nn.GroupNorm(num_groups=2, num_channels=6, affine=False)

y_ln, y_rms, y_gn = ln(x), rms(x), gn(x)

def row_stats(t):
    mean = t.mean(dim=-1)
    rms_val = t.pow(2).mean(dim=-1).sqrt()
    return mean, rms_val

for name, y in (("LayerNorm", y_ln), ("RMSNorm", y_rms)):
    m, r = row_stats(y)
    print(f"{name:9s} per-sample mean: " + " ".join(f"{v:7.4f}" for v in m))
    print(f"{name:9s} per-sample rms : " + " ".join(f"{v:7.4f}" for v in r))

# GroupNorm with 2 groups of 3 channels: zero mean WITHIN each group, per sample.
grp = y_gn.view(4, 2, 3)
print("GroupNorm per-(sample,group) means:")
for i in range(4):
    print("  sample", i, " ".join(f"{v:7.4f}" for v in grp[i].mean(dim=-1)))

# GroupNorm and RMSNorm are batch-size independent: batch of 1 is fine in train mode.
print(f"\nGroupNorm on a single sample, train mode: shape {tuple(gn(x[:1]).shape)} (no error)")
same = torch.allclose(gn(x[:1]), y_gn[:1])
print(f"same output as in the batch of 4: {same}")

Running this prints:

LayerNorm per-sample mean:  0.0000  0.0000  0.0000 -0.0000
LayerNorm per-sample rms :  1.0000  1.0000  1.0000  1.0000
RMSNorm   per-sample mean:  0.4796  0.4859  0.9046  0.3643
RMSNorm   per-sample rms :  1.0000  1.0000  1.0000  1.0000
GroupNorm per-(sample,group) means:
  sample 0  0.0000  0.0000
  sample 1 -0.0000  0.0000
  sample 2  0.0000 -0.0000
  sample 3 -0.0000  0.0000

GroupNorm on a single sample, train mode: shape (1, 6) (no error)
same output as in the batch of 4: True

The output is the family portrait. With the demonstrated affine settings and small epsilon, LayerNorm rows have zero mean and approximately unit RMS; RMSNorm rows have approximately unit RMS but generally nonzero means. The centering step is the difference between these two normalization calculations. GroupNorm’s means vanish within each (sample, group) pair rather than across the whole row. And the last two lines are the property BatchNorm cannot offer: a batch of one processes without error in training mode, and each sample’s output is identical whether it arrives alone or with company (compare the ValueError and the three different normalizations of sample 0 in the BatchNorm experiment above). A usable decision rule falls out directly: convolutional network with comfortable batch sizes, BatchNorm; small or wildly varying batch sizes, GroupNorm; Transformers and anything token-shaped, LayerNorm or RMSNorm.

When more parameters help: double descent

The weight-decay experiment at the top of this section told the classical story: a model with far more parameters than data will fit noise, so you must constrain it. But that story sits awkwardly next to the last decade of practice, where the largest, least-constrained models generalize best. The tension has a name, double descent, and you can reproduce it in a few seconds with nothing but least squares. As you grow a model’s capacity, test error first follows the classical U (down, then up as overfitting sets in), spikes at the interpolation threshold (the capacity where the model can first fit the training data exactly) and then, past the threshold, comes down again, often below the classical optimum. Belkin, Hsu, Ma, and Mandal assembled this picture for simple model families in an influential 2019 PNAS paper, and Nakkiran and colleagues (ICLR 2020) showed the same shape in modern deep networks, both as a function of width and, strikingly, of training time.

The demonstration below fits random ReLU features to 40 noisy training points from a linear teacher: draw a random projection, apply ReLU to get n_features nonlinear features, and solve the resulting linear regression with np.linalg.lstsq — which returns the ordinary least-squares fit while the system is overdetermined (fewer features than points) and the minimum-norm interpolating solution once it is underdetermined (more features than points).

import numpy as np

np.random.seed(0)

# Random-features regression around the interpolation threshold.
# Teacher: y = x @ w_true + noise, in d = 8 dimensions, n_train = 40 points.
d, n_train, n_test, n_trials = 8, 40, 2000, 40
X = np.random.randn(n_train, d)
X_test = np.random.randn(n_test, d)
w_true = np.random.randn(d)
y = X @ w_true + 0.5 * np.random.randn(n_train)
y_test = X_test @ w_true                      # noise-free targets for test MSE

def fit_random_features(n_feat, rng):
    V = rng.standard_normal((d, n_feat)) / np.sqrt(d)
    F = np.maximum(X @ V, 0.0)                # random ReLU features
    F_test = np.maximum(X_test @ V, 0.0)
    # lstsq returns the least-squares fit when overdetermined and the
    # MINIMUM-NORM interpolating solution when underdetermined.
    coef, *_ = np.linalg.lstsq(F, y, rcond=None)
    train_mse = np.mean((F @ coef - y) ** 2)
    test_mse = np.mean((F_test @ coef - y_test) ** 2)
    return train_mse, test_mse

print(f"n_train = {n_train}  (interpolation threshold at n_features ~ {n_train})")
print(f"{'n_features':>11s}{'train MSE':>12s}{'test MSE':>12s}")
for n_feat in (5, 10, 20, 30, 38, 40, 42, 50, 80, 160, 640):
    rng = np.random.default_rng(1)
    tr, te = np.mean([fit_random_features(n_feat, rng) for _ in range(n_trials)], axis=0)
    print(f"{n_feat:>11d}{tr:>12.4f}{te:>12.4f}")

Running this prints:

n_train = 40  (interpolation threshold at n_features ~ 40)
 n_features   train MSE    test MSE
          5      2.6903      3.2696
         10      1.4315      2.4529
         20      0.3905      1.3519
         30      0.1395      3.1679
         38      0.0183     16.3983
         40      0.0000  72956.6574
         42      0.0000     18.7203
         50      0.0000      3.3347
         80      0.0000      1.4523
        160      0.0000      0.9067
        640      0.0000      0.7029

Both descents and the spike are all here. From 5 to 20 features, test error falls: the classical regime where capacity helps. From 20 to 38, it rises: classical overfitting. At exactly 40 features (the interpolation threshold, one feature per training point), test error explodes: the averaged value of 73,000 is dominated by trials where the feature matrix is nearly singular, so the only interpolating solution requires enormous, delicately canceling coefficients whose off-data behavior is wild. And then the modern regime: with 640 features, test error is 0.70 — better than the best underparameterized model (1.35 at 20 features), despite train error being exactly zero. Memorizing the noise did not prevent generalization.

Why does adding features past the threshold help? Because beyond it there are infinitely many interpolating solutions, and lstsq picks the minimum-norm one; more features mean more room to interpolate the same points smoothly, so the chosen solution needs less contortion. Nothing in the loss asked for small norms — the solver’s default supplied that preference. This is the phenomenon called implicit regularization: the training algorithm itself biases which of the many zero-training-error solutions you get. Gradient descent has the same character (run from small initialization on least squares, it converges toward this same minimum-norm solution), and an influential 2017 paper by Chiyuan Zhang and colleagues (“Understanding deep learning requires rethinking generalization,” ICLR 2017) forced the issue for deep networks by showing that standard architectures can perfectly fit randomly labeled images: whatever makes them generalize on real labels, it cannot be a lack of capacity, so it must live in the interaction of the data, the architecture, and the optimizer’s preferences among fits. Making that interaction precise is a central open problem of deep learning theory.

For your own practice the message is calibration, not alarm. The spike lives near an interpolation threshold, and explicit regularization (ridge, weight decay, early stopping) can damp it. Many overparameterized models operate beyond such a threshold and sometimes exhibit a second descent, but whether the regime is benign depends on the data, architecture, optimizer, and noise. Retire the reflex that “train error zero” automatically means “broken”; held-out evidence decides.

Going deeper

Implicit regularization beyond least squares. The minimum-norm story has a provable analogue for classification: for linearly separable data, gradient descent on logistic-type losses drives the weights in the direction of the maximum-margin separator, even though nothing in the loss mentions margins (Soudry and colleagues, JMLR 2018). For deep nonlinear networks, characterizing the optimizer’s implicit preferences is an active research area rather than a settled theory; proposed ingredients include the bias of SGD noise toward “flatter” minima, architectural biases (convolutions prefer translation-respecting functions before a single gradient step), and norm-based capacity measures that behave better than raw parameter counts. Treat any single-sentence explanation of why deep nets generalize as a hypothesis; the honest current answer is that several biases stack, and disentangling them is ongoing work.

Regularizers interact — sometimes to the point of aliasing. Two examples of why you cannot tune regularizers independently. First, BatchNorm is itself a mild regularizer: its statistics are computed from a random minibatch, injecting noise much as dropout does, and the original paper already reported being able to reduce or remove dropout when using it. Second, weight decay behaves strangely on weights that feed a normalization layer: under suitable scale-invariant configurations, positively rescaling a pre-normalization signal approximately cancels in normalization when epsilon is negligible. Biases, other computation paths, and fixed running statistics can break this argument. In the scale-invariant case, shrinking the weight norm can change the effective learning rate of gradient steps on those weights (van Laarhoven spelled this out in 2017). Weight decay on normalized networks is thus doing real work, but not the work its name suggests. The practical corollary you will see in most Transformer training code: weight decay is applied to weight matrices but not to biases or normalization parameters.

Where normalization research went next. LayerNorm’s story continued in two directions you will meet in Chapter 3. Placement: the original Transformer put LayerNorm after each sublayer (post-LN), but pre-LN (normalizing at each sublayer’s entrance) turned out to stabilize very deep stacks and is now the common default, with published analyses attributing the difference to better-behaved gradients at initialization. Reduction: RMSNorm’s success suggests normalization’s essential service in Transformers is controlling activation scale, not centering — and lines of work like Fixup (and successors that remove normalization with initialization and reparameterization alone) probe whether explicit normalization is needed at all. That question is genuinely open at the research frontier; in this course, you should treat normalization layers as standard equipment and know precisely which axis each one averages over.

Common pitfalls

  • Symptom: great validation numbers during development, garbage predictions in deployment, varying run to run. Cause: model.eval() never called — dropout still active and/or BatchNorm using batch statistics. Fix: model.eval() plus torch.no_grad() for all evaluation; wrap it in a function so it cannot be forgotten.
  • Symptom: adding weight decay “does nothing.” Cause: the value is too small for the optimizer/duration (compare 1e-1 versus 1.0 in the table above), or it is being applied through torch.optim.Adam’s coupled L2, which the optimizers section shows can neutralize it for large-gradient weights. Fix: use AdamW and sweep decay over orders of magnitude.
  • Symptom: a custom deep network trains fine at 3 layers, not at all at 30; gradients at early layers are near zero. Possible cause: poor signal scaling from a manual initialization or an activation mismatch. Check: measure both activations and gradients by layer. Use an activation-appropriate initialization, such as kaiming_normal_ for ReLU or xavier_normal_ with an explicitly chosen tanh gain; the forward table alone does not diagnose every training failure.
  • Symptom: a BatchNorm model’s accuracy craters with small batches, or BatchNorm1d on (N, C) crashes when there is only one value per channel. Cause: train-mode statistics lack enough samples (or eval-mode running statistics came from a different regime). Convolutional BatchNorm can still have many values per channel when N=1 because spatial positions also contribute. Fix: evaluate in eval() mode; for inherently tiny effective batches or sequence settings, prefer LayerNorm or GroupNorm.
  • Symptom: training loss keeps improving for thousands of epochs while true performance quietly degrades. Cause: no validation-based stopping — the best epoch 400 versus epoch 4000 gap in the first experiment. Fix: track validation loss every epoch, keep the best checkpoint, stop after a patience window without improvement.

Try it yourself

  1. In the weight-decay experiment, add dropout (nn.Dropout(0.2) after each hidden layer) with weight_decay=0 and compare its final and best validation MSE against the wd=1.0 row. Remember model.train() / model.eval() — the experiment’s loop already switches modes.
  2. Extend the initialization experiment to tanh activations. Compare equal-width Xavier variance 1/n with variance 2/n, and record activation standard deviation and the fraction with abs(tanh(z)) > 0.95 by layer. Does the larger weight scale cause saturation in this run? Do not assume that either scale keeps the variance exactly constant. State any gain factor you use.
  3. Repeat the initialization experiment with gradient tracking. Backpropagate from the mean output and measure the gradient standard deviation for the first layer’s weights under each initialization. Compare that measurement with the forward activation statistics; distinguish an observed small gradient from an exact zero caused by floating-point underflow.
  4. Replace BatchNorm1d with LayerNorm inside a small MLP trained on the ring dataset from the optimizers section, and compare outputs for the same inputs at batch sizes 1 and 256 using torch.allclose with suitable floating-point tolerances. Evaluate both models in eval() mode first: default BatchNorm also uses fixed running statistics in this mode. Then switch the BatchNorm model to training mode and confirm that BatchNorm1d on (1, C) raises an error. Keep dropout disabled when isolating this normalization comparison.
  5. In the double-descent experiment, replace lstsq with ridge regression (solve \((F^\top F + \lambda I)\, w = F^\top y\) for a small regularizer \(\lambda\) (lam in code) such as 1e-3) and rerun the sweep. What happens to the spike at 40 features, and what does that tell you about the relationship between explicit regularization and the interpolation threshold?

Further reading

  • Glorot and Bengio, “Understanding the difficulty of training deep feedforward neural networks,” AISTATS 2010. proceedings.mlr.press/v9/glorot10a.html — the diagnosis of why pre-2010 deep nets failed to train, and the variance-balancing fix now called Xavier init.
  • He, Zhang, Ren, and Sun, “Delving Deep into Rectifiers,” ICCV 2015. arxiv.org/abs/1502.01852 — the ReLU-aware factor-of-2 correction demonstrated in this section’s initialization table.
  • Saxe, McClelland, and Ganguli, “Exact solutions to the nonlinear dynamics of learning in deep linear neural networks,” ICLR 2014. arxiv.org/abs/1312.6120 — the theory behind orthogonal initialization and depth-independent training dynamics.
  • Srivastava, Hinton, Krizhevsky, Sutskever, and Salakhutdinov, “Dropout: A Simple Way to Prevent Neural Networks from Overfitting,” JMLR 2014. jmlr.org/papers/v15/srivastava14a.html — the full dropout paper, including the ensemble interpretation and the experiments that set the default rates.
  • Ioffe and Szegedy, “Batch Normalization: Accelerating Deep Network Training by Reducing Internal Covariate Shift,” ICML 2015. arxiv.org/abs/1502.03167 — the BatchNorm paper; read it alongside Santurkar below to see a proposed mechanism and its challenge.
  • Ba, Kiros, and Hinton, “Layer Normalization,” 2016. arxiv.org/abs/1607.06450 — the batch-free alternative that Transformers adopted; see Zhang and Sennrich, “Root Mean Square Layer Normalization,” NeurIPS 2019 (arxiv.org/abs/1910.07467) for the centering-free simplification many current LLMs use.
  • Wu and He, “Group Normalization,” ECCV 2018. arxiv.org/abs/1803.08494 — normalization for the small-batch regime, with the clearest single figure comparing all the normalization axes.
  • Santurkar, Tsipras, Ilyas, and Madry, “How Does Batch Normalization Help Optimization?,” NeurIPS 2018. arxiv.org/abs/1805.11604 — evidence that BatchNorm’s benefit is landscape smoothing rather than covariate-shift reduction.
  • Belkin, Hsu, Ma, and Mandal, “Reconciling modern machine-learning practice and the classical bias-variance trade-off,” PNAS 2019. arxiv.org/abs/1812.11118 — the double-descent curve this section reproduced; see also Nakkiran et al., “Deep Double Descent,” ICLR 2020 (arxiv.org/abs/1912.02292) for the deep-network version.
  • Dive into Deep Learning, section 5.4, “Numerical Stability and Initialization.” d2l.ai/chapter_multilayer-perceptrons/numerical-stability-and-init.html — the variance analysis of this section with executable notebooks.

Before class you should be able to

  • Read a train/validation table and separate overfitting, ineffective decay, over-regularization, and what early stopping would have captured.
  • Explain inverted dropout’s rescaling and what the layer does in eval mode.
  • Explain why a linear layer multiplies the average squared signal by \(n\,\sigma^2\) under the initialization assumptions, and why ReLU contributes the factor of one half that motivates Kaiming initialization.
  • State which axis BatchNorm and LayerNorm each normalize over, and the consequences for small batches and inference.
  • Explain why Transformers standardized on LayerNorm, and give the one-line differences of RMSNorm and GroupNorm.
  • Sketch the double-descent curve and locate the interpolation threshold.
  • Diagnose the symptoms of a forgotten model.eval() (noisy, batch-dependent, or crashing evaluation).

Connects to

  • This session is covered by Quiz 1 (Module 1: S01–S08), given at S09.
  • The initialization variance analysis matters the moment you build a custom layer (in HW2’s convolution code, in your project, or in the exercises below), because a bad init visibly stalls training.
  • Weight decay, early stopping, and normalization return as everyday tools in the training-practice session (S12), and LayerNorm reappears inside every Transformer block from S16 onward.

← S06 · Optimizers · All readings · S08 · PyTorch in Practice →