GenAI Course · Module 3 · Session 13

Optimizers & Training in Practice

⏱ 14 min read🎯 Intermediate🖼 2 diagrams🧪 2 live demos🐍 Python code

Plain gradient descent works, but it can be slow and fussy about the learning rate. Optimizers like Momentum and Adam are smarter ways to take each step. In this session we race four optimizers, schedule the learning rate, and learn the practical tricks that stop a model from memorising its training data.

In this session

SGD, Momentum, RMSProp and Adam: the idea behind each
A fair race between them on a tricky valley
Learning-rate schedules: step, exponential, cosine
Overfitting: spotting it with a validation set
Dropout and early stopping

01SGD, Momentum, RMSProp and Adam

Optimizer	Idea	Analogy
SGD	step = −lr × gradient	walking downhill, one careful step at a time
Momentum	keep a running average of past gradients (velocity)	a ball rolling downhill that builds speed
RMSProp	divide each weight’s step by its recent gradient size	smaller steps where the ground is steep, bigger where flat
Adam	momentum + RMSProp, with bias correction	a rolling ball that also adjusts its step for each direction

We implement all four from scratch and race them down a long, narrow valley (f = x² + 100y²), which is hard for plain SGD. To be fair, each optimizer gets its best learning rate from the same list of 10:

optimizers.py

import numpy as np

# A long, narrow valley: f(x, y) = x^2 + 100 y^2  (minimum at 0, 0)
grad = lambda p: np.array([2 * p[0], 200 * p[1]])
loss = lambda p: p[0] ** 2 + 100 * p[1] ** 2

def sgd(p, g, s, lr): return p - lr * g
def momentum(p, g, s, lr, beta=0.9):
    s["v"] = beta * s.get("v", 0) + g; return p - lr * s["v"]
def rmsprop(p, g, s, lr, beta=0.9, eps=1e-8):
    s["sq"] = beta * s.get("sq", 0) + (1 - beta) * g ** 2; return p - lr * g / (np.sqrt(s["sq"]) + eps)
def adam(p, g, s, lr, b1=0.9, b2=0.999, eps=1e-8):
    s["t"] = s.get("t", 0) + 1
    s["m"] = b1 * s.get("m", 0) + (1 - b1) * g
    s["v"] = b2 * s.get("v", 0) + (1 - b2) * g ** 2
    m_hat, v_hat = s["m"] / (1 - b1 ** s["t"]), s["v"] / (1 - b2 ** s["t"])
    return p - lr * m_hat / (np.sqrt(v_hat) + eps)

def steps_needed(step, lr, max_steps=2000):
    p, s = np.array([-4.0, 1.5]), {}
    for i in range(1, max_steps + 1):
        p = step(p, grad(p), s, lr)
        if not np.isfinite(loss(p)) or loss(p) > 1e9: return None     # diverged
        if loss(p) < 1e-3: return i
    return None

lrs = [0.001, 0.002, 0.005, 0.009, 0.01, 0.02, 0.05, 0.1, 0.2, 0.5]
for name, step in [("SGD", sgd), ("Momentum", momentum), ("RMSProp", rmsprop), ("Adam", adam)]:
    results = {lr: steps_needed(step, lr) for lr in lrs}
    ok = {lr: n for lr, n in results.items() if n}
    best = min(ok, key=ok.get)
    print(f"{name:9s} best lr={best:<6} -> {ok[best]:4d} steps   (works for {len(ok)} of {len(lrs)} learning rates)")

SGD best lr=0.009 -> 267 steps (works for 3 of 10 learning rates) Momentum best lr=0.002 -> 64 steps (works for 5 of 10 learning rates) RMSProp best lr=0.5 -> 10 steps (works for 5 of 10 learning rates) Adam best lr=0.5 -> 58 steps (works for 8 of 10 learning rates)

Demo · Steps to reach the bottom (best learning rate each)

Two lessons: adaptive methods (RMSProp, Adam) reach the bottom much faster than SGD, and Adam works for the widest range of learning rates. That robustness is why Adam (and its variant AdamW) is the default for training transformers and fine-tuning LLMs.

pytorch_optimizers.py

import torch
opt = torch.optim.SGD(model.parameters(), lr=0.01, momentum=0.9)
opt = torch.optim.Adam(model.parameters(), lr=1e-3)                      # common default
opt = torch.optim.AdamW(model.parameters(), lr=2e-5, weight_decay=0.01)  # transformers / LLMs

02Learning-rate schedules

A big learning rate helps early (move fast), a small one helps late (settle precisely). A schedule lowers it during training:

lr_schedule.py

import torch
torch.set_num_threads(1)                    # same results on every run
w = torch.nn.Parameter(torch.zeros(1))

def lrs(make):
    opt = torch.optim.SGD([w], lr=0.1); sch = make(opt); out = []
    for epoch in range(10):
        out.append(opt.param_groups[0]["lr"]); opt.step(); sch.step()
    return out

schedules = {
    "constant":         lambda o: torch.optim.lr_scheduler.LambdaLR(o, lambda e: 1.0),
    "step (x0.5 / 3)":  lambda o: torch.optim.lr_scheduler.StepLR(o, step_size=3, gamma=0.5),
    "exponential 0.8":  lambda o: torch.optim.lr_scheduler.ExponentialLR(o, gamma=0.8),
    "cosine":           lambda o: torch.optim.lr_scheduler.CosineAnnealingLR(o, T_max=10),
}
for name, make in schedules.items():
    print(f"{name:17s}", " ".join(f"{v:.3f}" for v in lrs(make)))

constant 0.100 0.100 0.100 0.100 0.100 0.100 0.100 0.100 0.100 0.100 step (x0.5 / 3) 0.100 0.100 0.100 0.050 0.050 0.050 0.025 0.025 0.025 0.013 exponential 0.8 0.100 0.080 0.064 0.051 0.041 0.033 0.026 0.021 0.017 0.013 cosine 0.100 0.098 0.090 0.079 0.065 0.050 0.035 0.021 0.010 0.002

Demo · Learning rate over 10 epochs

Warm-up + cosine is the standard for transformers: start tiny for a few hundred steps, rise to the peak learning rate, then follow a cosine curve down.

03Overfitting

A model overfits when it memorises the training data, including its noise, and does badly on new data. You only notice it if you hold back a validation set. Here we fit polynomials of growing degree to 15 noisy points:

overfit_poly.py

import numpy as np
rng = np.random.default_rng(0)
x = np.sort(rng.uniform(0, 1, 15)); y = np.sin(2 * np.pi * x) + rng.normal(0, 0.25, 15)
xv = rng.uniform(0, 1, 200); yv = np.sin(2 * np.pi * xv) + rng.normal(0, 0.25, 200)

for degree in [1, 3, 6, 12]:
    coef = np.polyfit(x, y, degree)
    train = np.mean((np.polyval(coef, x) - y) ** 2); val = np.mean((np.polyval(coef, xv) - yv) ** 2)
    print(f"degree {degree:2d}: train error {train:.3f}   validation error {val:.3f}")

degree 1: train error 0.193 validation error 0.393 degree 3: train error 0.017 validation error 0.089 degree 6: train error 0.013 validation error 0.134 degree 12: train error 0.002 validation error 102.817

Degree 3 is the sweet spot here: low error on both.

📚

Like a student who memorises last year’s paper. Perfect marks on that paper (training data), but lost on new questions (validation). Understanding the concepts is what we want.

04Dropout and early stopping

Dropout randomly switches off a fraction of neurons at every training step, so the network cannot rely on any single neuron. At evaluation time all neurons are on:

early_stop.py

import torch, torch.nn as nn
torch.set_num_threads(1)                    # same results on every run
torch.manual_seed(0)
x = torch.rand(40, 1); y = torch.sin(6 * x) + 0.3 * torch.randn(40, 1)          # small, noisy data
xv = torch.rand(400, 1); yv = torch.sin(6 * xv) + 0.3 * torch.randn(400, 1)

def train(p_drop, epochs=3000):
    torch.manual_seed(1)
    net = nn.Sequential(nn.Linear(1, 256), nn.ReLU(), nn.Dropout(p_drop),
                        nn.Linear(256, 256), nn.ReLU(), nn.Dropout(p_drop), nn.Linear(256, 1))
    opt = torch.optim.Adam(net.parameters(), lr=1e-2)
    history = []
    for epoch in range(epochs):
        net.train(); opt.zero_grad()
        loss = nn.functional.mse_loss(net(x), y); loss.backward(); opt.step()
        net.eval()
        with torch.no_grad(): history.append(nn.functional.mse_loss(net(xv), yv).item())
    best = min(range(epochs), key=history.__getitem__)
    return loss.item(), history, best

tr, h, best = train(p_drop=0.0)
for e in [50, 200, best, 1500, 2999]:
    print(f"epoch {e:4d}: validation loss {h[e]:.3f}" + ("   <- best: early stopping keeps this model" if e == best else ""))
print(f"at the end: training loss {tr:.3f} vs validation {h[-1]:.3f}. The growing gap is overfitting.")

torch.manual_seed(3); d = nn.Dropout(0.5); t = torch.ones(8)
d.train(); print("\ndropout in training:", d(t).tolist())
d.eval();  print("dropout in eval:    ", d(t).tolist())

epoch 50: validation loss 0.160 epoch 200: validation loss 0.109 epoch 610: validation loss 0.107 <- best: early stopping keeps this model epoch 1500: validation loss 0.110 epoch 2999: validation loss 0.114 at the end: training loss 0.046 vs validation 0.114. The growing gap is overfitting. dropout in training: [2.0, 2.0, 0.0, 2.0, 0.0, 0.0, 2.0, 2.0] dropout in eval: [1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0]

In training, dropout zeroed about half the values and doubled the rest (2.0) so the average stays the same. In eval() mode it does nothing. Always call model.eval() before validation and inference.
Early stopping: watch the validation loss and keep the best model. Here the best was at epoch 610; training to epoch 3,000 only made validation worse.
Other regularisers: weight decay (AdamW), data augmentation, and simply more data.

Early stopping: keep the model from the best validation epoch.

05Summary

Momentum builds speed; RMSProp scales steps per weight; Adam combines both.
Adam/AdamW converge fast and tolerate many learning rates: the default for transformers.
Schedules lower the learning rate over time; warm-up + cosine is standard for LLMs.
Overfitting = low training error but high validation error. Always keep a validation set.
Fight it with dropout, early stopping, weight decay and more data.

06Quick quiz

1. Adam combines which two ideas?

2. Why use a learning-rate schedule?

3. Train error 0.002, validation error 102. This is…

4. What does dropout do at evaluation time?

5. Early stopping keeps…

07References

Adam: A Method for Stochastic Optimization (Kingma & Ba, 2014)
Dropout: A Simple Way to Prevent Overfitting (Srivastava et al., 2014)
PyTorch: torch.optim
Decoupled Weight Decay Regularization (AdamW) (Loshchilov & Hutter, 2017)

Created by Ajeetkumar