Lesson 33: Optimization¶
Lessons 31-32 used plain gradient descent: take the gradient, take a fixed-size step against it:
$$p \leftarrow p - \eta\,\nabla f(p)$$
where $p$ is whatever's being learned (a weight, a bias, or the whole parameter vector at once), $\eta$ is the learning rate, and $\nabla f(p)$ is the loss gradient at the current parameters. 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.
Note that in real training, the gradient is almost never computed over the entire training set at once — it wouldn't fit into memory. Rather, it's estimated from a small, randomly sampled mini-batch (say, 32 or 256 examples) at every step instead. The resulting noisy, sample-to-sample-varying estimate is what gives stochastic gradient descent (SGD) its name. This lesson mostly uses the exact, noise-free gradient of a fixed function (calling it "SGD" for convenience); the last section simulates mini-batch noise directly to show why it matters.
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.15
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}')
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, color='tab:blue')
plt.scatter([0], [0], marker='*', s=150, color='black', zorder=5, label='minimum')
plt.legend(fontsize=8)
plt.title('Plain gradient descent: the classic zigzag')
plt.show()
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$$
Momentum's real benefit is acceleration, not damping. Along a consistent direction (like the shallow $x$ axis here), every gradient points the same way, so the running velocity keeps growing step after step — the effective step size along that axis becomes much larger than $\eta$ alone would give, exactly like a ball picking up speed rolling downhill. Along an oscillating direction (the steep $y$ axis, flipping sign every step), consecutive gradients partially cancel in that same running average — but only partially: a single momentum coefficient $\beta$ is shared across every direction, so that cancellation isn't tuned to any one direction's curvature, and the steep axis can still swing substantially — sometimes even more than plain gradient descent's own swing — before it decays. The net effect, over a full run, is usually still a win: faster progress on the well-behaved axis outweighs the imperfectly-damped one. But it's a win from acceleration plus a rough, untuned brake, not a clean fix for oscillation — that's what Adam, next, is actually built for.
(Momentum has a second benefit that matters even more in real training: with noisy mini-batch gradients rather than the exact gradient used here, the running average acts as a low-pass filter, smoothing sample-to-sample noise into a more reliable direction. We'll see the noisy case directly near the end of this lesson.)
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¶
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's accumulated velocity clearly accelerates progress along that shallow direction, and it ends with a lower loss than plain SGD — but notice its path actually swings wider in $y$ than plain SGD's own, not narrower, since a single shared momentum coefficient can't damp each axis to the degree that axis specifically needs. Adam, which independently rescales each direction, is what actually tames the steep axis directly, while still racing 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. Although this 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 the amount it vanishes depends heavily on layer width. Each layer's pre-activation variance scales as fan_in * std^2, where fan-in is the layer's number of inputs: at width 16, that's 16 * 0.05^2 ≈ 0.04, so each layer keeps only about 4% of the previous layer's variance, and 20 compounding layers erase the signal almost completely. At width 256, the same fixed std keeps 256 * 0.05^2 ≈ 0.64 of the variance per layer — still shrinking, but far more slowly, which is exactly why the width-256 column above vanishes so much less severely than width-16's. A fixed std's effect depends entirely on fan-in, so the same number that badly vanishes a narrow layer would, at a wide enough layer, push fan_in * std^2 past 1 and start exploding instead.
Xavier initialization (Glorot & Bengio, 2010) — also known as Glorot initialization — fixes this by scaling the std with the layer's fan-in: std = sqrt(1 / fan_in), which keeps fan_in * std^2 pinned at exactly 1 regardless of width. 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 initialization (He et al., 2015) — also known as He initialization — 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})')
Rule of thumb: 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.
Learning rate schedules¶
A single fixed learning rate for the whole training run is tempting — it's the obvious default — but it's often the wrong choice: a rate large enough for fast early progress is usually too large to settle precisely once training nears 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 — the name is borrowed directly from simulated annealing, the classical optimization technique of starting "hot" (free to explore, even accepting worse moves) and gradually "cooling" toward a fixed solution.
- 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. As mentioned above, 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.
Exercises¶
- 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.15) 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?