21 / 27 · Neural Networks
Training deep networks
What makes a deep network trainable and keeps it from overfitting: initialization that preserves the variance of signals (Glorot and He), dropout as an ensemble, weight decay against an L2 penalty under Adam, learning-rate schedules with warm-up and early stopping, data augmentation, residual connections and layer normalization.
Backpropagation · Stochastic and adaptive optimizers · Regularization
On this page 14
The previous note computed exact gradients for a network of any depth. That does not make deep networks easy to train. In a stack of 30 layers the gradient that reaches the first layer is a product of 30 Jacobians; for the plain 30-layer network of this note’s Code section its norm at initialization is , and the early layers never move; and a network flexible enough to fit anything also fits the noise in its training labels. This note collects the standard remedies that are not covered elsewhere in the courses: how to initialize the weights, dropout, weight decay, learning-rate schedules and early stopping, data augmentation, residual connections and layer normalization. Three neighbouring topics have their home in the Generative Models course and are only linked here: batch normalization and instance normalization, the mechanism of vanishing and exploding gradients, and a longer walk-through of a training run; the essentials of that routine are in the default recipe at the end of this note.
Notation. A layer computes (biases are zero at initialization and are left out) and . and are its fan-in and fan-out, the input and output widths. The weights start as independent draws with mean 0 and variance . is the learning rate, the weight decay or L2 coefficient, the dropout probability.
Initialization that preserves the scale#
Why not start from zero? If all weights of a layer start equal, every unit of the layer computes the same function and receives the same gradient, so they stay identical forever: the layer behaves like a single unit. Random weights break this symmetry. Their scale is the question.
Forward. Take one pre-activation at initialization. The weights are independent of each other and of the inputs , with mean 0. The terms are uncorrelated, since for , so their variances add (joint distributions note). Each term has variance . Hence
Note the second moment , not the variance of : after a ReLU the activations are not centred (He et al., 2015, section 2.2).
- Linear regime (tanh near zero, ): , so , and the variance is preserved when .
- ReLU: if is symmetric around zero, the ReLU keeps half of its distribution and . Then , preserved when
Over layers the variance is multiplied by : any constant factor other than 1 per layer grows or shrinks geometrically.
Backward. The same argument applied to sums over the outputs, so the gradients keep their scale when in the linear regime, or for ReLU. Forward and backward agree only if . Glorot and Bengio (2010), section 4.2.1, eqs. 10–12 and 16, take the compromise and draw the weights from the uniform distribution on , whose variance is exactly that (a uniform on has variance ). For ReLU, He et al. (2015), eqs. 10 and 14, show that preserving either direction is enough in practice, because the two products differ only by the ratio of the first and last widths.
Worked example. For a square layer with , Glorot gives standard deviation and He . With ReLU and the Glorot variance , each layer multiplies the variance by : after 30 layers the output has of its first-layer variance. With He it stays put; with it grows to times.
The scale of the signal in a deep network under different initializations
The simulation follows the theory up to random fluctuations, and the Code section repeats it at width 256. He initialization is what lets a plain 30-layer ReLU stack train at all: in the Code section, the same stack with PyTorch’s default initialization has a gradient of at its first layer and stays at the loss of a coin flip, while with He initialization it fits the data. PyTorch’s default for nn.Linear and nn.Conv2d has variance , six times below He; the convolution section of the Generative Models course explains where that comes from and why normalization layers compensate for it. Initialization fixes the scale only at the start; normalization layers and residual connections, below, keep it under control during training.
Dropout#
Dropout (Srivastava et al., 2014) switches off each unit of a layer at random during training. In the form used today (“inverted” dropout), each activation is replaced by
with a fresh mask for every example and every step. The division by keeps the expectation, , so at test time the mask is simply removed (model.eval() in PyTorch; p in nn.Dropout(p) is the probability of dropping). The original paper keeps units with probability and instead multiplies the outgoing weights by at test time (Srivastava et al., 2014, section 1 and Figure 2; section 10 notes the equivalence): mind that its is a retention probability.
An ensemble of thinned networks. A network with droppable units contains “thinned” subnetworks, one per mask, all sharing the same weights. Each training step trains one of them, sampled at random. At test time, the full network with expectation-preserving scaling stands in for the average of all predictions. This is exact only in special cases, such as a single linear layer, but works well in practice: on MNIST, averaging the predictions of about 50 sampled thinned networks matches the scaled network (Srivastava et al., 2014, section 7.5). Like bagging, dropout averages many models, and like a random forest it decorrelates them by randomizing what each one sees; unlike bagging, the members share their weights.
Worked example: dropout as a penalty. For a linear model with dropout on its inputs, the expected squared error can be computed in closed form. Write with , so that and . Since ,
Dropout adds a ridge-like penalty whose weight for grows with the size of feature and with the dropout rate (Srivastava et al., 2014, section 9.1, in their parametrization). For , , and : , the penalty is , and the expected loss is . A Monte Carlo average over a million masks in the Code section gives . The regularization note develops the ridge penalty itself.
Typical rates are to on hidden layers. Dropout matters most for large fully connected layers trained on modest data; Transformers use on attention weights and residual branches, and convolutional networks with batch normalization often use little or none.
Weight decay is not an L2 penalty under Adam#
With SGD they are the same. Adding to the loss adds to the gradient, and an SGD step becomes
the weights shrink by the factor every step, which is weight decay. The penalty is the negative log of a Gaussian prior, so the minimizer is a MAP estimate.
With Adam they differ. Adam divides each coordinate of the step by a running estimate of the root mean square of that coordinate’s gradient. With an L2 penalty, the term enters the gradient and is divided too: a weight whose gradients are typically large is decayed less, one with small gradients more. Decoupled weight decay (AdamW) applies the shrinkage outside the adaptive step, as PyTorch’s AdamW does:
so that every weight decays at the same rate (Loshchilov and Hutter, 2019, section 2: Proposition 1 for the equivalence under SGD, Proposition 2 for its failure under adaptive methods, Algorithm 2 for AdamW, whose decay is multiplied by the schedule factor but not by the base step).
Worked example. With and , a weight shrinks by per step under AdamW, whatever its gradients. Under Adam with an L2 penalty the shrinkage is about (ignoring the penalty’s own small share of ): if the weight’s gradients have root mean square 10, and if it is 0.1. The L2 version regularizes most the weights that the data barely constrain and least the ones it pushes hardest, a hundredfold difference in this example.
In practice AdamW with between and is the default for Transformers, and biases and the scales of normalization layers are usually exempt from decay. The Generative Models course uses AdamW at its PyTorch defaults.
Learning-rate schedules, warm-up and early stopping#
A constant learning rate is a compromise. Large steps make fast progress, but with stochastic gradients they leave the iterate bouncing in a region around the minimum whose size grows with , as the optimizers note shows; small steps settle closer but take long to get there. A schedule starts large and decreases:
- Step decay: divide by 10 at fixed epochs.
- Cosine decay (Loshchilov and Hutter, 2017, section 3, eq. 5): , a smooth glide from to over epochs.
- Inverse square root, the original Transformer’s schedule (Vaswani et al., 2017, section 5.3, eq. 3): , linear growth for steps ( there) and then decay like .
Warm-up starts from a small rate and increases it linearly over the first steps. At the start the weights are random, the gradients large and inconsistent, and Adam’s estimate rests on very few steps; a full-size step then can throw the network into a bad region. Goyal et al. (2017), section 2.2, needed a gradual warm-up over 5 epochs to train ResNet-50 with large batches and a correspondingly large rate, and post-norm Transformers need it because of large gradients near the output at initialization, as discussed in the Transformer note.
Three learning-rate schedules
Worked example. With , and , the cosine schedule gives , and : it stays near the peak for longer than a linear decay and spends its last quarter at small rates.
Early stopping uses the number of steps as a regularizer. Training error keeps falling, while validation error falls and then rises once the network starts fitting noise; keep the checkpoint with the lowest validation loss, and stop when it has not improved for a set number of epochs (the patience). For a quadratic loss, steps of gradient descent with rate act approximately like an L2 penalty with (Goodfellow, Bengio and Courville, 2016, section 7.8; Bishop, 2006, section 5.5.2): fewer steps, stronger regularization. In the Code section a network is trained on 300 digits of which 30% have random labels. The validation loss is lowest at epoch 127 () and has risen to by epoch 300; the checkpoint from epoch 127 classifies of the test digits correctly, the last one only . The validation set is used for a decision here, so the reported accuracy must come from a separate test set, as the validation note insists.
Data augmentation#
When the label is known to be unchanged by some transformations of the input, we can train on transformed copies: random crops, flips and colour changes of photographs, time shifts and added noise for audio, word dropout or back-translation for text. The model sees a new variant of each example in every epoch and is pushed to be invariant to those changes. Formally, the training loss becomes an expectation over a random transformation ,
and each minibatch estimates it with one sampled per example, a Monte Carlo estimate. The transformations encode prior knowledge, and wrong knowledge hurts: a horizontally flipped “b” is a “d”, and a 6 rotated by 180° is a 9. AlexNet trained on random crops and their mirror images, which enlarged the training set by a factor of 2048, and averaged its predictions over ten crops at test time (Krizhevsky, Sutskever and Hinton, 2012, section 4.1). Mixup goes further and trains on convex combinations of pairs of examples and of their labels (Zhang et al., 2018, section 2).
Worked example. In the Code section a network is trained on 300 digits, once as they are and once with four extra copies shifted by one pixel. On the test digits, centred like the training ones, both classify correctly ( and ). On test digits shifted by one pixel in a random direction, the plain network drops to and the augmented one keeps . Augmentation buys robustness to the variation it covers; whether that raises the test score depends on whether the test data contain that variation.
Residual connections#
Deeper is not automatically better, even on the training set. A plain 56-layer convolutional network has higher training error on CIFAR-10 than a 20-layer one (He et al., 2016a, Figure 1). This degradation is not overfitting, since the deeper model could copy the shallower one and set its extra layers to the identity; the optimizer simply fails to find such solutions. A residual connection makes the identity the default:
where is a small stack of layers, the residual branch. With small weights, and the block passes its input through; the block only has to learn a correction.
Two consequences explain why this helps optimization (He et al., 2016b, eqs. 4–5). Unrolled, the output of a deep stack is the input plus a sum of corrections, , instead of a product. And by the chain rule
so the gradient reaches every layer through the identity term without passing through any weight matrix. That makes a vanishing gradient much less likely, but it is not a guarantee: the second term can cancel the identity. A branch outputs for every input, and the block’s Jacobian is zero. Equivalently, a stack of residual blocks contains paths of different lengths, depending on which blocks each path goes through, and most of the gradient of a 110-layer residual network flows along paths only 10 to 34 layers deep (Veit, Wilber and Belongie, 2016).
In the Code section, a stack of 30 layers of width 32 with PyTorch’s default initialization has a gradient of at its first layer and never leaves the loss , while the same layers with residual connections get a first-layer gradient of and fit the two-moons data perfectly within 200 Adam steps. Residual connections require the input and the output of a block to have the same shape (or a projection on the skip path); they are why ResNets reach hundreds of layers and why Transformers stack dozens of blocks. The Generative Models course lists them among the remedies for vanishing gradients, next to gated recurrent cells.
Layer normalization#
Normalization layers rescale activations during training, not only at initialization. Batch normalization standardizes each feature over the examples of a batch, which ties every output to the rest of the batch and requires running averages at inference (Generative Models course). Layer normalization (Ba, Kiros and Hinton, 2016) standardizes each example over its own features, with and the mean and variance of the entries of . It does not depend on the batch, behaves the same in training and inference, and works for sequences of any length, which is why every Transformer uses it (or its variant RMSNorm), as the Transformer note describes. The table of normalization layers in the Generative Models course compares which values share a mean and a variance in batch, instance and layer normalization.
A default recipe#
Put together, a reasonable starting point for a new network:
- He initialization for ReLU-family layers (or the library default plus normalization layers).
- Residual connections and a normalization layer once the network has more than a handful of layers.
- AdamW with weight decay of to , not applied to biases and normalization parameters.
- A peak learning rate found by a short sweep, a linear warm-up, then cosine decay.
- Augmentation for every invariance the task really has; dropout where the network still overfits.
- A validation set checked every epoch, keeping the best checkpoint; before any of this, check that the network can drive the loss to nearly zero on a single small batch, which catches most bugs in the model and the loss.
- Evaluate in evaluation mode:
model.eval()turns dropout off and makes batch normalization use its running statistics,torch.no_grad()records no graph, andmodel.train()switches back before the next step. A checkpoint stores the model and optimizer state and the epoch, so a run can resume or return to its best point. Fixed seeds (torch.manual_seedand the NumPy generator) make a run repeatable, and a few seeds show how much of a difference is noise. The test set is evaluated once, at the end, with the settings chosen on validation.
What changed#
| A plain deep network | With the remedies of this note | |
|---|---|---|
| Signal scale at the start | shrinks or grows geometrically with depth | preserved by He or Glorot initialization |
| Gradient path to early layers | a product of one Jacobian per layer | an identity path through residual connections |
| Regularization | none: the network fits label noise | dropout, decoupled weight decay, augmentation, early stopping |
| Learning rate | one constant | warm-up, then decay |
| Normalization | none | batch norm for convolutional networks, layer norm for Transformers |
Limitations#
- The derivations hold at initialization. Variance preservation assumes independent weights and inputs; after training starts, the weights become correlated with the data and only normalization keeps the scale in check.
- Dropout is not free. It slows convergence, interacts with batch normalization, and its test-time scaling is only an approximation of the ensemble average.
- Every remedy adds hyperparameters: the dropout rate, , the peak rate, the warm-up length, the patience. They interact, and the defaults above are starting points, not guarantees.
- Augmentation needs domain knowledge. A transformation that changes the label teaches the wrong invariance, and for tabular data there is usually no natural one.
- Theory lags practice. Why residual connections and normalization help is partly explained (signal propagation, smoother losses) and partly still debated.
What you should remember#
- Initialize so that stays constant: for ReLU (He), for tanh (Glorot); any other constant factor per layer grows or shrinks geometrically.
- Dropout multiplies activations by with during training only; it trains an ensemble of weight-sharing subnetworks and acts like a feature-scaled ridge penalty on a linear model.
- Under SGD, an L2 penalty is weight decay; under Adam it is not, and AdamW decays every weight at the rate .
- Warm up, then decay the learning rate (step, cosine, ); stop at the best validation checkpoint, which acts like an L2 penalty .
- Augmentation trains on label-preserving transformations and buys exactly the invariances it covers.
- Residual connections give the gradient an identity path; layer normalization standardizes each example over its features, independently of the batch.
Code#
The scale of activations under three initializations, dropout as a penalty, the shrinkage of AdamW against L2, early stopping on noisy labels, augmentation by shifts, and a 30-layer network with and without He initialization and residual connections:
import numpy as np, torch
from torch import nn
from sklearn.datasets import load_digits, make_moons
from sklearn.model_selection import train_test_split
T = lambda a: torch.tensor(a, dtype=torch.float32)
# 1. the scale of z_30 in a stack of 30 ReLU layers of width 256, for weight variances 1/n, 2/n, 4/n
# (theory: 2^-29 = 1.9e-9, 2, 4 * 2^29 = 2.1e9)
torch.manual_seed(0)
x = torch.randn(512, 256)
for var in (1 / 256, 2 / 256, 4 / 256):
h = x
for _ in range(30):
z = h @ (torch.randn(256, 256) * var**0.5)
h = torch.relu(z)
print(f'{z.pow(2).mean().item():.2g}', end=' ')
print() # 1e-09 1.1 7.5e+09
# 2. dropout on a linear model: E(y - sum m_j w_j x_j / (1-p))^2 = (y - w.x)^2 + p/(1-p) sum w_j^2 x_j^2
xv, w, yv, p = np.array([1.0, 2.0]), np.array([0.5, -0.25]), 1.0, 0.5
m = np.random.default_rng(0).random((1_000_000, 2)) > p # keep each input with probability 1 - p
print(round(((yv - (m / (1 - p) * w * xv).sum(1)) ** 2).mean(), 3), (yv - w @ xv) ** 2 + p / (1 - p) * (w**2 * xv**2).sum()) # 1.499 1.5
# 3. weight decay: the shrinkage per step of a weight theta = 1 under Adam with L2 (eta*lam*theta / RMS of its
# gradients) and under AdamW (eta*lam*theta), for gradient RMS 10 and 0.1
eta, lam = 1e-3, 1e-2
print([f'{rms}: {eta * lam / rms:.0e} vs {eta * lam:.0e}' for rms in (10.0, 0.1)]) # ['10.0: 1e-06 vs 1e-05', '0.1: 1e-04 vs 1e-05']
# 4. early stopping: 300 digits with 30% of the labels replaced at random, a 64 -> 256 -> 10 network, full-batch Adam
D = load_digits()
Xtr, Xrest, ytr, yrest = train_test_split(D.data / 16, D.target, train_size=300, random_state=0, stratify=D.target)
Xva, Xte, yva, yte = train_test_split(Xrest, yrest, train_size=300, random_state=0, stratify=yrest)
rng = np.random.default_rng(0); noisy = ytr.copy(); flip = rng.random(300) < 0.3; noisy[flip] = rng.integers(0, 10, flip.sum())
torch.manual_seed(0)
net = nn.Sequential(nn.Linear(64, 256), nn.ReLU(), nn.Linear(256, 10))
opt = torch.optim.Adam(net.parameters(), lr=1e-3)
best = (np.inf, 0, None)
for epoch in range(1, 301):
opt.zero_grad(); nn.functional.cross_entropy(net(T(Xtr)), torch.tensor(noisy)).backward(); opt.step()
with torch.no_grad():
val = nn.functional.cross_entropy(net(T(Xva)), torch.tensor(yva)).item()
if val < best[0]:
best = (val, epoch, {k: v.clone() for k, v in net.state_dict().items()})
with torch.no_grad():
last_acc = (net(T(Xte)).argmax(1).numpy() == yte).mean()
net.load_state_dict(best[2]); best_acc = (net(T(Xte)).argmax(1).numpy() == yte).mean()
print(best[1], round(best[0], 3), round(val, 3), round(best_acc, 3), round(last_acc, 3)) # 127 0.592 0.757 0.881 0.788
# 5. augmentation: 300 training digits, with and without copies shifted by one pixel; tested on centred
# images and on images shifted by one pixel in a random direction
def shift(img, dx, dy):
out = np.zeros_like(img)
out[max(dy, 0):8 + min(dy, 0), max(dx, 0):8 + min(dx, 0)] = img[max(-dy, 0):8 + min(-dy, 0), max(-dx, 0):8 + min(-dx, 0)]
return out
moves = [(1, 0), (-1, 0), (0, 1), (0, -1)]
Itr, Ite, ltr, lte = train_test_split(D.images / 16, D.target, train_size=300, random_state=0, stratify=D.target)
aug = np.concatenate([Itr] + [np.stack([shift(i, *mv) for i in Itr]) for mv in moves])
Ish = np.stack([shift(i, *moves[k]) for i, k in zip(Ite, np.random.default_rng(1).integers(0, 4, len(Ite)))])
def fit(X, y, epochs):
torch.manual_seed(0)
net = nn.Sequential(nn.Flatten(), nn.Linear(64, 128), nn.ReLU(), nn.Linear(128, 10))
opt = torch.optim.Adam(net.parameters(), lr=1e-3)
for _ in range(epochs):
for idx in torch.randperm(len(X)).split(64):
opt.zero_grad(); nn.functional.cross_entropy(net(T(X[idx])), torch.tensor(y[idx])).backward(); opt.step()
return lambda Z: (net(T(Z)).argmax(1).numpy() == lte).mean().round(3)
plain, augmented = fit(Itr, ltr, 100), fit(aug, np.tile(ltr, 5), 60)
print(plain(Ite), augmented(Ite), plain(Ish), augmented(Ish)) # 0.938 0.937 0.46 0.888
# 6. depth 30, width 32, two moons: a plain stack with PyTorch's default initialization, the same stack with
# He initialization, and the default stack with residual connections; gradient norm of the first layer at
# initialization, then the training loss and accuracy after 200 Adam steps
Xm, ym = make_moons(500, noise=0.1, random_state=0)
X, Y = T(Xm), T(ym)
class Deep(nn.Module):
def __init__(self, residual, he):
super().__init__()
self.inp, self.out, self.residual = nn.Linear(2, 32), nn.Linear(32, 1), residual
self.blocks = nn.ModuleList(nn.Linear(32, 32) for _ in range(30))
if he:
for layer in [self.inp, *self.blocks, self.out]:
nn.init.kaiming_normal_(layer.weight, nonlinearity='relu'); nn.init.zeros_(layer.bias)
def forward(self, x):
h = self.inp(x)
for block in self.blocks:
h = h + block(torch.relu(h)) if self.residual else block(torch.relu(h))
return self.out(torch.relu(h)).squeeze(1)
for residual, he in ((False, False), (False, True), (True, False)):
torch.manual_seed(0); net = Deep(residual, he)
nn.functional.binary_cross_entropy_with_logits(net(X), Y).backward()
g0 = net.inp.weight.grad.norm().item()
opt = torch.optim.Adam(net.parameters(), lr=1e-3)
for _ in range(200):
opt.zero_grad(); loss = nn.functional.binary_cross_entropy_with_logits(net(X), Y); loss.backward(); opt.step()
print(residual, he, f'{g0:.1e}', f'{loss.item():.3f}', ((net(X) > 0).float() == Y).float().mean().item())
# False False 3.2e-13 0.693 0.5 | False True 1.8e-01 0.000 1.0 | True False 2.4e+00 0.000 1.0
Next#
Every method so far, from nearest neighbours to deep networks, learned from examples with labels: a price, a class, a rating. Often there are no labels at all, only the inputs, and the question is what structure they have: groups of similar customers, typical and unusual transactions, a few directions that explain most of the variation. The next note opens the part of the course on learning without labels with clustering, starting from k-means and its loss.
References#
- X. Glorot, Y. Bengio. Understanding the difficulty of training deep feedforward neural networks. AISTATS, 2010. Section 4.2.1, eqs. 10–12 and 16 (normalized initialization).
- K. He, X. Zhang, S. Ren, J. Sun. Delving deep into rectifiers: surpassing human-level performance on ImageNet classification. ICCV, 2015. Section 2.2, eqs. 7–14 (initialization for rectifiers); Figure 3 (a 30-layer network that stalls with Glorot initialization).
- N. Srivastava, G. Hinton, A. Krizhevsky, I. Sutskever, R. Salakhutdinov. Dropout: a simple way to prevent neural networks from overfitting. Journal of Machine Learning Research, 15:1929–1958, 2014. Section 1 and Figure 2 (thinned networks, weight scaling), section 7.5 (Monte Carlo averaging against weight scaling), section 9.1 (linear regression), section 10 (scaling at training time).
- I. Loshchilov, F. Hutter. Decoupled weight decay regularization. ICLR, 2019. Section 2, Propositions 1–3, Algorithm 2.
- I. Loshchilov, F. Hutter. SGDR: stochastic gradient descent with warm restarts. ICLR, 2017. Section 3, eq. 5 (cosine annealing).
- P. Goyal et al. Accurate, large minibatch SGD: training ImageNet in 1 hour. arXiv:1706.02677, 2017. Section 2.2 (gradual warm-up).
- A. Vaswani et al. Attention is all you need. NeurIPS, 2017. Section 5.3, eq. 3 (the learning-rate schedule).
- A. Krizhevsky, I. Sutskever, G. E. Hinton. ImageNet classification with deep convolutional neural networks. NeurIPS, 2012. Section 4.1 (data augmentation).
- H. Zhang, M. Cisse, Y. N. Dauphin, D. Lopez-Paz. mixup: beyond empirical risk minimization. ICLR, 2018. Section 2.
- K. He, X. Zhang, S. Ren, J. Sun. Deep residual learning for image recognition. CVPR, 2016 (cited as 2016a). Figure 1 (degradation), eq. 1.
- K. He, X. Zhang, S. Ren, J. Sun. Identity mappings in deep residual networks. ECCV, 2016 (cited as 2016b). Eqs. 4–5.
- A. Veit, M. Wilber, S. Belongie. Residual networks behave like ensembles of relatively shallow networks. NeurIPS, 2016.
- J. L. Ba, J. R. Kiros, G. E. Hinton. Layer normalization. arXiv:1607.06450, 2016.
- I. Goodfellow, Y. Bengio, A. Courville. Deep Learning. MIT Press, 2016. Sections 7.4 (dataset augmentation), 7.8 (early stopping), 7.12 (dropout), 8.4 (initialization).
- C. M. Bishop. Pattern Recognition and Machine Learning. Springer, 2006. Section 5.5.2 (early stopping as regularization).
- E. Sokolov et al. Машинное обучение (HSE), lecture notes, lecture 7: introduction to deep learning (normalization, augmentation, residual networks), and the class materials on batch normalization, CNNs and ResNets.