Training Deep Networks

At this point, the learning loop looks simple: make a prediction, compute the loss, run backpropagation, update the weights, and repeat.

If backpropagation tells every weight which direction to move, why is training deep networks still difficult?

Because knowing the direction is not enough. The step can be too large, too small, too noisy, badly scaled, or pointed through a network whose activations are already unstable. The model can also memorize the training set while getting worse on new examples.

This chapter is about the practical techniques that make the training loop stable, efficient, and useful.

The Full Training Loop

Earlier chapters often said gradient descent updates the weights. In modern deep learning, an optimizer usually performs the update. The optimizer decides how to use the gradients.

Backpropagation is like a compass. It points downhill on the loss surface. Training techniques decide how you move while following that compass.

A careful hiker takes reasonable steps, remembers the direction of previous motion, adapts when the terrain becomes steep, keeps their equipment balanced, and stops before walking past the best point. That is the role of learning rates, optimizers, initialization, normalization, dropout, and early stopping.

The Training Toolkit

This is why the chapter feels different. We are no longer asking what a neural network is. We are asking how to make one train reliably.

Learning Rate

A basic gradient update is:

The learning rate $\eta$ controls the size of the step. If it is too small, training crawls. If it is too large, training may oscillate or diverge.

For example, if $w=4$ and $\frac{\partial L}{\partial w}=2$:

Tuning the learning rate remains one of the first things practitioners do when a neural network fails to train.

Mini-Batches

Deep learning almost always uses mini-batches. Instead of computing the gradient on the entire dataset, we compute an average gradient on a small batch:

Mini-batches balance three needs: gradients are less noisy than single-example SGD, updates happen much faster than full-batch gradient descent, and GPUs can process many examples in parallel.

Momentum

Plain SGD can zigzag when the loss surface is narrow in one direction and shallow in another. Momentum keeps a running velocity:

The optimizer develops inertia. Directions that persist accumulate speed, while directions that alternate back and forth cancel out.

Adaptive Learning Rates And Adam

Different parameters can receive gradients at very different scales. Adaptive optimizers adjust each parameter's effective step size.

RMSProp divides by a running estimate of recent squared gradients. Parameters with consistently huge gradients get smaller effective steps; parameters with tiny gradients get relatively larger ones.

Adam combines momentum and adaptive scaling. It tracks both an average gradient and an average squared gradient, which is why it often works well as a strong default optimizer.

Weight Initialization

If every weight starts at zero, neurons in the same layer receive identical gradients and learn identical features. Random initialization breaks this symmetry.

The scale of the random weights also matters. Too large can make activations and gradients explode. Too small can make signals vanish. Xavier and He initialization choose scales that keep information flowing through layers.

Batch Normalization

As earlier layers change, the inputs to later layers also change. Batch Normalization stabilizes those inputs during training by normalizing activations within a mini-batch:

The learned scale $\gamma$ and shift $\beta$ let the network recover whatever distribution is useful after normalization.

Dropout

Dropout randomly disables a fraction of neurons during training. This prevents the network from relying too heavily on any one neuron.

At inference time, dropout is turned off and all neurons are used. This distinction matters: dropout is a training-time regularization technique, not a prediction-time behavior.

Early Stopping

Training loss often keeps falling. Validation loss may fall at first and then rise as the model begins to overfit. Early stopping saves the model at the best validation performance rather than the final epoch.

This is why validation curves are not decoration. They are one of the main instruments for deciding whether training is working.

A modern PyTorch training loop separates the model, optimizer, training batches, and validation monitoring. This sketch shows where the techniques fit.

import torch
from torch import nn

model = nn.Sequential(
    nn.Linear(30, 128),
    nn.BatchNorm1d(128),
    nn.ReLU(),
    nn.Dropout(p=0.3),
    nn.Linear(128, 64),
    nn.ReLU(),
    nn.Linear(64, 1),
)

criterion = nn.BCEWithLogitsLoss()
optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)

best_val_loss = float("inf")
patience = 5
epochs_without_improvement = 0

for epoch in range(100):
    model.train()

    for X_batch, y_batch in train_loader:
        optimizer.zero_grad()

        logits = model(X_batch)
        loss = criterion(logits, y_batch)

        loss.backward()
        optimizer.step()

    model.eval()
    val_losses = []

    with torch.no_grad():
        for X_val, y_val in val_loader:
            val_logits = model(X_val)
            val_loss = criterion(val_logits, y_val)
            val_losses.append(val_loss.item())

    mean_val_loss = sum(val_losses) / len(val_losses)

    if mean_val_loss < best_val_loss:
        best_val_loss = mean_val_loss
        best_state = model.state_dict()
        epochs_without_improvement = 0
    else:
        epochs_without_improvement += 1

    if epochs_without_improvement >= patience:
        break

model.load_state_dict(best_state)

Notice the mode switches. $\text{model.train()}$ turns on training-time behavior such as dropout and batch-normalization batch statistics. $\text{model.eval()}$ turns off dropout and uses evaluation behavior.