Week 9: Practical Training Dynamics

Learning Rate Schedules#

A constant learning rate is rarely optimal. The right schedule depends on the problem and the optimiser:

Step decay: Reduce $\eta$ by a factor $\gamma$ every $T$ epochs (or steps). Simple and effective for SGD: $\eta_t = \eta_0 \cdot \gamma^{\lfloor t / T \rfloor}$

Typical: $\gamma = 0.1$, $T = 30$ epochs (so $\eta$ drops to $0.1\eta_0$, then $0.01\eta_0$, etc.).

Cosine annealing: Smoothly decays $\eta$ following a cosine curve from $\eta_0$ to $\eta_{\min}$ over $T$ iterations: $\eta_t = \eta_{\min} + \frac{\eta_0 - \eta_{\min}}{2}\left(1 + \cos\frac{\pi t}{T}\right)$

Cosine schedules are popular in vision and generative model training because they avoid sudden drops that can destabilise training.

Exponential decay: $\eta_t = \eta_0 \cdot \gamma^t$. Continuous decay — good when you need a smooth schedule without discrete drops. Rarely used as the primary schedule; more common as a component of warmup.

Cyclical and warm restarts: $\eta$ oscillates between $\eta_{\min}$ and $\eta_{\max}$ with period $T$. The oscillations can help escape saddle points in nonconvex landscapes.

When each helps:

• Step decay: well-understood problems where you know when loss plateaus
• Cosine annealing: modern default for vision and generative models — smooth, no tuning of drop epochs
• Cyclical: nonconvex problems where exploration matters
• Constant: debugging and early experimentation (remove one variable)

Warmup#

In the first few hundred training steps, the model parameters are poorly initialised and the gradient estimates are noisy. Large steps early on can send parameters to irrecoverably bad regions of the loss landscape.

Linear warmup increases $\eta$ from a small value to the target $\eta_0$ over $N_{\text{warmup}}$ steps:

$\eta_t = \eta_0 \cdot \frac{t}{N_{\text{warmup}}} \quad \text{for } t \leq N_{\text{warmup}}$

Typical: $N_{\text{warmup}} = 2000$–$4000$ steps for transformer training. Less critical for small models trained with SGD on simple objectives.

Why warmup is especially important for Adam: Adam's adaptive denominator $\hat v_k$ is unstable early on (the EMA estimates are based on very few samples). Warmup gives Adam time to build reliable $m$ and $v$ estimates before taking large steps.

Gradient Clipping#

In recurrent networks, transformers, and unstable training setups, gradient norms can occasionally explode — a single bad mini-batch produces gradients with norms in the thousands. Without clipping, the parameter update can be catastrophic.

Clip by global norm: If the total gradient norm exceeds a threshold $C$, rescale the entire gradient vector:

$\text{if } \|g\|_2 > C, \quad g \leftarrow \frac{C}{\|g\|_2} \cdot g$

This preserves the gradient direction while capping the magnitude. Typical $C = 1.0$ or $5.0$ for RNNs/transformers.

Clip by value: Clamp each coordinate individually to $[-v, v]$:

$g_i \leftarrow \max(-v, \min(v, g_i))$

Less common in deep learning; more useful when individual gradient components can be pathologically large.

Diagnosing from Curves#

Your training loss curve is a diagnostic instrument. Here's what to look for:

Healthy training:

• Training loss decreases smoothly and monotonically
• Validation loss tracks training loss with a small gap
• Gradient norm decreases gradually

Plateau: Loss is flat, gradients are small but nonzero.

• Possible causes: $\eta$ too small; stuck in a saddle point (nonconvex); poor initialisation
• Try: increase $\eta$, add momentum, or switch to Adam

Oscillation: Loss bounces up and down significantly between steps.

• Possible causes: $\eta$ too large; batch size too small (noisy gradients)
• Try: reduce $\eta$, increase batch size, or add momentum to smooth

Divergence: Loss rises monotonically, often to NaN.

• Likely cause: $\eta$ far too large; gradient explosion
• Try: immediately reduce $\eta$ by 10×, add gradient clipping

Overfitting signal: Training loss continues to drop while validation loss rises.

• Not an optimisation problem per se — you're minimising the wrong thing (ERM gap)
• Solutions: regularisation, data augmentation, early stopping

The gradient norm curve complements the loss curve. A healthy run shows:

• Gradient norm near zero at convergence
• No sudden spikes (if so, clipping is needed)
• Gradual decay proportional to the loss decrease

When to Stop#

Early stopping: Halt training when validation loss stops improving for $p$ consecutive epochs (patience). This prevents overfitting and is an implicit form of regularisation — stopping GD early on an overparameterised model biases the solution toward the initialisation, which often helps generalisation.

Practical rule: Track the best validation loss seen so far, and stop if it hasn't improved in 5–10 epochs. Restore the parameters from the epoch with the best validation loss.

Browser lab#

Train a small linear regression model with different LR schedules (constant, step decay, cosine) and plot loss vs epoch on the same axes.

import numpy as np
import matplotlib.pyplot as plt

rng = np.random.default_rng(42)
n, d = 500, 5
w_true = rng.normal(size=d)
X = rng.normal(size=(n, d))
y = X @ w_true + rng.normal(scale=0.5, size=n)

def train(X, y, w0, eta_fn, K):
    w = w0.copy(); losses = []
    for k in range(K):
        g = X.T @ (X @ w - y) / n
        w = w - eta_fn(k) * g
        losses.append(0.5 * np.mean((X @ w - y)**2))
    return losses

w0 = rng.normal(scale=0.1, size=d)
K = 80

constant = train(X, y, w0, lambda k: 0.1, K)
step_decay = train(X, y, w0, lambda k: 0.1 * (0.5 ** (k // 20)), K)
cosine = train(X, y, w0, lambda k: 0.01 + 0.5 * 0.09 * (1 + np.cos(np.pi * k / K)), K)

plt.figure(figsize=(9, 5))
plt.plot(constant, label="Constant η=0.1")
plt.plot(step_decay, label="Step decay (×0.5 every 20)")
plt.plot(cosine, label="Cosine annealing")
plt.yscale("log"); plt.xlabel("Iteration"); plt.ylabel("Loss")
plt.title("Loss vs Iteration: Three LR Schedules")
plt.legend(); plt.grid(True, alpha=0.3); plt.show()