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 problems.

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 surface.

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 are 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.

Knowledge check#

Browser lab: learning-rate schedules#

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()

print(f"Final loss  constant={constant[-1]:.4e}  step={step_decay[-1]:.4e}  cosine={cosine[-1]:.4e}")
print("Notice: all three schedules descend, but decaying η (step / cosine) settles to a lower loss than a constant rate.")

Key takeaways#

• The learning rate schedule is a first-class hyperparameter: step decay, cosine annealing, exponential decay, and cyclical schedules each fit different regimes.
• Warmup ramps $\eta$ up over the first steps precisely because Adam's moment estimates are unreliable early.
• Gradient clipping by global norm rescales the whole gradient to a bound $C$, preserving direction while stopping one bad batch from wrecking the parameters.
• Loss curves are diagnostic: a plateau means small gradients or a saddle, oscillation means $\eta$ is too large, divergence means it is far too large or the gradient exploded, and a rising validation loss means overfitting.
• The gradient-norm curve complements the loss curve: sudden spikes signal that clipping is needed.
• Early stopping is regularization: halting before convergence biases the solution toward its initialization, which often generalizes better.

Further reading#

• Ilya Loshchilov & Frank Hutter, "SGDR: Stochastic Gradient Descent with Warm Restarts" (ICLR 2017). Cosine annealing and warm restarts. Freely available at arXiv:1608.03983.
• Priya Goyal et al., "Accurate, Large Minibatch SGD: Training ImageNet in 1 Hour" (2017). Linear warmup and large-batch training at scale. Freely available at arXiv:1706.02677.
• Razvan Pascanu, Tomas Mikolov & Yoshua Bengio, "On the difficulty of training recurrent neural networks" (ICML 2013). Gradient clipping by global norm.