Lesson 32: Optimization¶
Lessons 30-31 used plain gradient descent: take the gradient, take a fixed-size step against it, repeat. That's enough to prove the idea works, but it's rarely how real networks are trained. This lesson covers four practical upgrades — momentum, adaptive step sizes (Adam), weight initialization schemes, and learning rate schedules — each fixing a specific, concrete failure mode of plain gradient descent.
import numpy as np
import torch
import torch.nn as nn
import matplotlib.pyplot as plt
A loss landscape where plain gradient descent struggles¶
$f(x,y) = 0.05x^2 + 5y^2$ is a bowl that's steep in $y$ and shallow in $x$. A step size large enough to make progress along $x$ overshoots along $y$, causing the classic zigzag.
def f(p):
return 0.05 * p[0]**2 + 5 * p[1]**2
def grad_f(p):
return np.array([0.1 * p[0], 10 * p[1]])
start = np.array([-4.0, 1.0])
lr = 0.05
p = start.copy()
sgd_path = [p.copy()]
for _ in range(30):
p = p - lr * grad_f(p)
sgd_path.append(p.copy())
sgd_path = np.array(sgd_path)
print(f'plain SGD, 30 steps: final loss = {f(sgd_path[-1]):.4f}')
Momentum: remember where you were heading¶
Instead of stepping purely along the current gradient, accumulate a running velocity — a weighted average of past gradients — and step along that instead:
$$v \leftarrow \beta v - \eta\,\nabla f(p), \qquad p \leftarrow p + v$$
Consistent gradient directions (like the shallow $x$ direction here) reinforce each other and accelerate; oscillating directions (the steep $y$ direction, flipping sign every step) partially cancel out and damp down.
momentum = 0.9
p = start.copy()
v = np.zeros(2)
mom_path = [p.copy()]
for _ in range(30):
g = grad_f(p)
v = momentum * v - lr * g
p = p + v
mom_path.append(p.copy())
mom_path = np.array(mom_path)
print(f'SGD + momentum, 30 steps: final loss = {f(mom_path[-1]):.4f}')
Sanity check against PyTorch¶
p_t = torch.tensor(start.copy(), requires_grad=True)
opt = torch.optim.SGD([p_t], lr=lr, momentum=momentum)
for _ in range(30):
opt.zero_grad()
loss = 0.05 * p_t[0]**2 + 5 * p_t[1]**2
loss.backward()
opt.step()
print(f'our result: {mom_path[-1]}')
print(f'torch result: {p_t.detach().numpy()}')
print(f'max diff: {np.abs(p_t.detach().numpy() - mom_path[-1]).max():.2e}')
Adam: per-parameter adaptive step sizes¶
Adam (Kingma & Ba, 2014★) tracks both a momentum-like running mean of the gradient ($m$) and a running mean of the squared gradient ($v$), then divides the step by $\sqrt{v}$ — automatically shrinking the step size for parameters with consistently large gradients (like the steep $y$ direction here) and boosting it for parameters with small ones:
$$m \leftarrow \beta_1 m + (1-\beta_1)g, \qquad v \leftarrow \beta_2 v + (1-\beta_2)g^2, \qquad p \leftarrow p - \eta\frac{\hat{m}}{\sqrt{\hat{v}}+\epsilon}$$
($\hat{m}, \hat{v}$ are bias-corrected versions of $m, v$, which matter mainly in the first few steps.)
beta1, beta2, adam_eps = 0.9, 0.999, 1e-8
adam_lr = 0.3
p = start.copy()
m, v_sq = np.zeros(2), np.zeros(2)
adam_path = [p.copy()]
for t in range(1, 31):
g = grad_f(p)
m = beta1 * m + (1 - beta1) * g
v_sq = beta2 * v_sq + (1 - beta2) * g**2
m_hat = m / (1 - beta1**t)
v_hat = v_sq / (1 - beta2**t)
p = p - adam_lr * m_hat / (np.sqrt(v_hat) + adam_eps)
adam_path.append(p.copy())
adam_path = np.array(adam_path)
p_t2 = torch.tensor(start.copy(), requires_grad=True)
opt2 = torch.optim.Adam([p_t2], lr=adam_lr, betas=(beta1, beta2), eps=adam_eps)
for _ in range(30):
opt2.zero_grad()
loss = 0.05 * p_t2[0]**2 + 5 * p_t2[1]**2
loss.backward()
opt2.step()
print(f'Adam, 30 steps: final loss = {f(adam_path[-1]):.4f}')
print(f'max diff vs torch.optim.Adam: {np.abs(p_t2.detach().numpy() - adam_path[-1]).max():.2e}')
All three paths, visualized¶
xs = np.linspace(-4.5, 1, 200)
ys = np.linspace(-1.5, 1.5, 200)
XX, YY = np.meshgrid(xs, ys)
ZZ = 0.05 * XX**2 + 5 * YY**2
plt.contour(XX, YY, ZZ, levels=20, colors='lightgray', linewidths=0.7)
plt.plot(*sgd_path.T, '-o', markersize=3, label=f'plain SGD (loss={f(sgd_path[-1]):.3f})')
plt.plot(*mom_path.T, '-o', markersize=3, label=f'momentum (loss={f(mom_path[-1]):.3f})')
plt.plot(*adam_path.T, '-o', markersize=3, label=f'Adam (loss={f(adam_path[-1]):.3f})')
plt.scatter([0], [0], marker='*', s=150, color='black', zorder=5, label='minimum')
plt.legend(fontsize=8)
plt.title('30 steps of each optimizer on the same narrow bowl')
plt.show()
Plain SGD is stuck oscillating across the narrow valley, barely progressing along the shallow direction. Momentum smooths out the oscillation and travels farther. Adam, which independently rescales each direction, damps the steep axis and races along the shallow one, reaching the lowest loss of the three in the same 30 steps.
Weight initialization: why zero doesn't work¶
It's tempting to initialize all weights to zero — it seems like the most "neutral" starting point. It's actually catastrophic for any layer with more than one unit: every hidden unit computes the exact same function of the input, gets the exact same gradient, and gets updated by the exact same amount, forever. This symmetry never breaks on its own.
rng = np.random.default_rng(0)
theta_in = rng.uniform(0, 2 * np.pi, 60)
inner = np.stack([0.5 * np.cos(theta_in), 0.5 * np.sin(theta_in)], axis=1) + rng.normal(0, 0.1, (60, 2))
theta_out = rng.uniform(0, 2 * np.pi, 60)
outer = np.stack([2.0 * np.cos(theta_out), 2.0 * np.sin(theta_out)], axis=1) + rng.normal(0, 0.15, (60, 2))
X = np.vstack([inner, outer])
y = np.concatenate([np.zeros(60), np.ones(60)])
def sigmoid(z):
return 1 / (1 + np.exp(-z))
def relu(z):
return np.maximum(0, z)
def relu_deriv(z):
return (z > 0).astype(np.float64)
H = 4
def train_mlp(W1, b1, W2, b2, n_epochs=3000, lr=0.1):
for _ in range(n_epochs):
z1 = X @ W1 + b1
a1 = relu(z1)
z2 = (a1 @ W2 + b2).ravel()
p = sigmoid(z2)
n = len(y)
grad_z2 = ((p - y) / n).reshape(-1, 1)
grad_W2 = a1.T @ grad_z2
grad_b2 = grad_z2.sum(axis=0)
grad_a1 = grad_z2 @ W2.T
grad_z1 = grad_a1 * relu_deriv(z1)
grad_W1 = X.T @ grad_z1
grad_b1 = grad_z1.sum(axis=0)
W2 -= lr * grad_W2; b2 -= lr * grad_b2
W1 -= lr * grad_W1; b1 -= lr * grad_b1
return W1, b1, W2, b2, p
W1_zero, b1_zero = np.zeros((2, H)), np.zeros(H)
W2_zero, b2_zero = np.zeros((H, 1)), np.zeros(1)
W1_zero, b1_zero, W2_zero, b2_zero, p_zero = train_mlp(W1_zero, b1_zero, W2_zero, b2_zero)
acc_zero = ((p_zero > 0.5).astype(float) == y).mean()
rng2 = np.random.default_rng(8)
W1_rand = rng2.normal(size=(2, H)) * 0.7
W2_rand = rng2.normal(size=(H, 1)) * 0.7
b1_rand, b2_rand = np.zeros(H), np.zeros(1)
_, _, _, _, p_rand = train_mlp(W1_rand, b1_rand, W2_rand, b2_rand)
acc_rand = ((p_rand > 0.5).astype(float) == y).mean()
print(f'zero-initialized: accuracy = {acc_zero:.1%}')
print(f'randomly initialized: accuracy = {acc_rand:.1%}')
print()
print('all 4 hidden units still have identical weight vectors after 3000 steps of zero-init training:')
print(np.round(W1_zero, 4))
With zero initialization, all four hidden units stay locked at exactly zero forever — the model never escapes the trivial (chance-level) solution, no matter how long it trains. Small random initialization breaks the symmetry: every unit starts out computing something slightly different, so gradient descent can push them apart and let them specialize. This is why every framework's default layer initialization uses small random values, never zeros.
Weight initialization schemes: how large should "small random" be?¶
Zero-init fails completely; "small random" fixes it. But how small? Pick a fixed standard deviation and it works for one layer width, then quietly fails at another. Track the variance of activations after passing random data through a deep stack of tanh-activated linear layers, using the same fixed weight std at every width:
def forward_variance(width, std, depth=20):
x = torch.randn(100, width)
for _ in range(depth):
W = torch.randn(width, width) * std
x = torch.tanh(x @ W)
return x.var().item()
print('fixed weight std = 0.05, activation variance after 20 tanh layers:')
for width in [16, 64, 256]:
print(f' width={width:4d}: variance = {forward_variance(width, std=0.05):.2e}')
A fixed std of 0.05 vanishes to essentially zero after 20 layers, and gets worse as the layer gets wider — more incoming connections means more terms summed into each output, so a fixed per-weight std pushes the pre-activation further from zero even as the individual weights stay the same size. Xavier/Glorot initialization (Glorot & Bengio, 2010) fixes this by scaling the std with the layer's fan-in (number of inputs): std = sqrt(1 / fan_in). Repeat the same experiment with Xavier-scaled weights:
def forward_variance_xavier(width, depth=20):
x = torch.randn(100, width)
for _ in range(depth):
std = (1.0 / width) ** 0.5
W = torch.randn(width, width) * std
x = torch.tanh(x @ W)
return x.var().item()
print('Xavier std = sqrt(1/fan_in), activation variance after 20 tanh layers:')
for width in [16, 64, 256]:
print(f' width={width:4d}: variance = {forward_variance_xavier(width):.3f}')
# check the formula against PyTorch's own implementation
layer = nn.Linear(256, 256, bias=False)
nn.init.xavier_normal_(layer.weight)
expected_std = (2.0 / (256 + 256)) ** 0.5 # nn.init's xavier_normal_ uses fan_in AND fan_out
print(f'\nnn.init.xavier_normal_ weight std: {layer.weight.std().item():.4f} (formula predicts {expected_std:.4f})')
Xavier keeps the variance in the same ballpark regardless of width, instead of vanishing catastrophically. But Xavier's derivation assumes a symmetric activation like tanh; Kaiming/He initialization (He et al., 2015) is the same idea adjusted for ReLU, which zeros out roughly half its inputs and so needs twice the variance to compensate: std = sqrt(2 / fan_in). Using Xavier's tanh-derived scale on a ReLU network still vanishes:
def forward_variance_relu(width, depth=20, kaiming=True):
x = torch.randn(100, width)
for _ in range(depth):
std = (2.0 / width) ** 0.5 if kaiming else (1.0 / width) ** 0.5
W = torch.randn(width, width) * std
x = torch.relu(x @ W)
return x.var().item()
v_xavier_on_relu = forward_variance_relu(256, kaiming=False)
v_kaiming = forward_variance_relu(256, kaiming=True)
print(f'ReLU net, Xavier scale (sqrt(1/fan_in), no ReLU correction): variance after 20 layers = {v_xavier_on_relu:.2e}')
print(f'ReLU net, Kaiming scale (sqrt(2/fan_in)): variance after 20 layers = {v_kaiming:.3f}')
layer2 = nn.Linear(256, 256, bias=False)
nn.init.kaiming_normal_(layer2.weight, nonlinearity='relu')
expected_std2 = (2.0 / 256) ** 0.5
print(f'\nnn.init.kaiming_normal_ weight std: {layer2.weight.std().item():.4f} (formula predicts {expected_std2:.4f})')
The rule of thumb that follows: use Kaiming init for ReLU-family networks (the default for nn.Conv2d/nn.Linear is actually already a variant of this), and Xavier for tanh/sigmoid. Both are strictly better than picking a fixed std and hoping — they're derived, not tuned, from a simple requirement: keep activation variance roughly constant from layer to layer, so a network can be made arbitrarily deep without silently losing its signal before training even starts. This is the same instinct behind Lesson 35's xavier_normal_ call in the vanishing-gradient demo, and behind batch normalization (also Lesson 35) — normalizing activation statistics is a recurring fix for the same underlying problem, applied at different points (once at initialization, continuously during training).
Learning rate schedules¶
A single fixed learning rate for an entire training run is another thing that's easy to reach for and often not quite right: a rate large enough to make fast early progress is often too large to settle precisely once training gets close to a minimum. Learning rate schedules change the rate over time. Three common ones, validated against torch.optim.lr_scheduler:
- Step decay: multiply the rate by a fixed factor every
Nsteps. - Cosine annealing: smoothly decay along a cosine curve from the initial rate down to (near) zero.
- Linear warmup: ramp up from zero over the first few steps, before applying the main schedule — used because early updates, before a model's gradient statistics have stabilized, can be unreliable enough that a large rate from step one causes damage a few steps of gradual ramp-up would have avoided.
base_lr = 0.5
def check_schedule(name, manual_fn, torch_scheduler_factory, steps):
p = torch.tensor([1.0], requires_grad=True)
opt = torch.optim.SGD([p], lr=base_lr)
sched = torch_scheduler_factory(opt)
torch_lrs = []
for step in range(steps):
torch_lrs.append(opt.param_groups[0]['lr'])
opt.step()
sched.step()
manual_lrs = [manual_fn(step) for step in range(steps)]
max_diff = max(abs(a - b) for a, b in zip(torch_lrs, manual_lrs))
print(f'{name}: max diff vs torch.optim.lr_scheduler = {max_diff:.2e}')
return manual_lrs
step_lrs = check_schedule(
'step decay',
lambda step: base_lr * (0.5 ** (step // 5)),
lambda opt: torch.optim.lr_scheduler.StepLR(opt, step_size=5, gamma=0.5),
steps=20)
cos_lrs = check_schedule(
'cosine annealing',
lambda step: 0.5 * base_lr * (1 + np.cos(np.pi * step / 20)),
lambda opt: torch.optim.lr_scheduler.CosineAnnealingLR(opt, T_max=20),
steps=20)
warmup_steps = 5
warmup_lrs = check_schedule(
'linear warmup',
lambda step: base_lr * min(1.0, (step + 1) / warmup_steps),
lambda opt: torch.optim.lr_scheduler.LambdaLR(opt, lr_lambda=lambda step: min(1.0, (step + 1) / warmup_steps)),
steps=10)
plt.figure(figsize=(6, 4))
plt.plot(step_lrs, '-o', markersize=3, label='step decay')
plt.plot(cos_lrs, '-o', markersize=3, label='cosine annealing')
plt.plot(warmup_lrs, '-o', markersize=3, label='linear warmup')
plt.xlabel('step'); plt.ylabel('learning rate'); plt.legend(fontsize=8)
plt.title('Three learning rate schedules')
plt.show()
Why decay actually helps: noisy gradients¶
On the exact, noise-free narrow bowl from the top of this lesson, a well-chosen constant learning rate converges just fine — there's nothing to decay away from. Real training almost never sees the exact gradient, though: each step uses a mini-batch estimate, which is noisy. Simulate that by adding random noise to the gradient every step, and compare a constant rate against step decay.
def run_noisy(schedule_fn, steps=400, seed=0, noise_std=0.6):
rng = np.random.default_rng(seed)
p = np.array([-4.0, 1.0])
losses = []
for t in range(steps):
noisy_g = grad_f(p) + rng.normal(0, noise_std, size=2) # stand-in for mini-batch noise
p = p - schedule_fn(t) * noisy_g
losses.append(f(p))
return losses
noisy_lr = 0.18
const_losses = run_noisy(lambda t: noisy_lr)
decay_losses = run_noisy(lambda t: noisy_lr * (0.3 ** (t // 80)))
print(f'constant lr: mean loss over final 50 steps = {np.mean(const_losses[-50:]):.4f}')
print(f'step decay: mean loss over final 50 steps = {np.mean(decay_losses[-50:]):.4f}')
plt.figure(figsize=(6, 4))
plt.semilogy(const_losses, label='constant lr', alpha=0.8)
plt.semilogy(decay_losses, label='step decay', alpha=0.8)
plt.xlabel('step'); plt.ylabel('loss (log scale)'); plt.legend(fontsize=8)
plt.title('Noisy gradients: constant lr vs. step decay')
plt.show()
With noisy gradients, a constant learning rate large enough to make fast early progress never actually settles — it keeps bouncing around the minimum by an amount proportional to the learning rate itself, forever. Step decay makes the same fast early progress, then shrinks the rate as training continues, tightening that bounce and landing far closer to the true minimum. This is the real justification for learning rate schedules: not that a fixed rate is "wrong" on some idealized noise-free problem, but that it can't simultaneously be large (for speed) and small (for precision) when every gradient is a noisy estimate, which is the normal situation for real mini-batch training.
Exercise¶
- Increase
momentumfrom 0.9 to 0.99 on the narrow-bowl problem. Does it converge faster, or does it start to overshoot the minimum and oscillate on the shallow axis now? - Try initializing the MLP's weights identically but nonzero (e.g.
np.full((2, H), 0.3)for both layers) instead of all-zero. Does symmetry still fail to break? Why would you expect the same failure mode from any initialization where every hidden unit starts identical, not just an all-zero one? - Adam's
adam_lr(0.3) is much larger than plain SGD'slr(0.05) in this notebook, yet Adam remains stable while a plain SGD run withlr=0.3would diverge wildly on the steep axis. Try it and confirm. What does that suggest about why Adam is often described as being more forgiving of the learning-rate choice? forward_variance_relucompares Xavier vs. Kaiming scale atwidth=256. Rerun it atwidth=16andwidth=1024. Does the gap between the two get bigger or smaller as width grows, and does that match the fan-in scaling argument used to derive both formulas?- In the noisy-gradient learning-rate demo, try
noise_std=0.0(no noise at all). Does step decay still help, hurt, or make no difference relative to a constant rate — and does that match the claim that schedules matter because of gradient noise, not despite the loss surface itself?