21. Optimization & Backpropagation

Level: Expert You'll use this for: actually understanding how a neural network trains, instead of treating .fit()/loss.backward() as magic.

Gradient descent, from the ground up

To minimize a loss function L(w), repeatedly step in the direction opposite the gradient (the direction of steepest ascent, Ch. 13):

w_new = w_old - η · ∇L(w_old)

η (eta) is the learning rate — too large and you overshoot/diverge, too small and training crawls.

def gradient_descent(grad_fn, w0, lr=0.1, steps=50):
    w = w0
    for _ in range(steps):
        w = w - lr * grad_fn(w)
    return w

# Minimize f(w) = (w-3)^2, whose derivative is 2(w-3)
w_min = gradient_descent(grad_fn=lambda w: 2*(w-3), w0=0.0)
print(w_min)   # converges to ~3.0

Full worked backpropagation example (2 layers, by hand)

Network: x → [w1] → h → relu → [w2] → y_pred, loss L = (y_pred - target)².

# Forward pass
x, target = 2.0, 10.0
w1, w2 = 1.5, 2.0

h = w1 * x                 # h = 3.0
h_relu = max(0, h)         # h_relu = 3.0
y_pred = w2 * h_relu       # y_pred = 6.0
loss = (y_pred - target)**2   # loss = 16.0

# Backward pass (chain rule, Ch. 13, applied link by link)
dL_dy = 2 * (y_pred - target)          # dL/dy_pred = 2*(6-10) = -8.0
dy_dw2 = h_relu                          # dy_pred/dw2 = h_relu = 3.0
dL_dw2 = dL_dy * dy_dw2                  # -8.0 * 3.0 = -24.0

dy_dh_relu = w2                          # dy_pred/dh_relu = w2 = 2.0
dh_relu_dh = 1.0 if h > 0 else 0.0       # ReLU derivative = 1 (since h=3>0)
dL_dh = dL_dy * dy_dh_relu * dh_relu_dh  # -8.0 * 2.0 * 1.0 = -16.0

dh_dw1 = x                               # dh/dw1 = x = 2.0
dL_dw1 = dL_dh * dh_dw1                  # -16.0 * 2.0 = -32.0

print("dL/dw1 =", dL_dw1, " dL/dw2 =", dL_dw2)
# Update: w1_new = w1 - lr*dL_dw1, w2_new = w2 - lr*dL_dw2

This hand-computed example is literally what loss.backward() in PyTorch does automatically, for every parameter, across every layer, via a computation graph — reverse-mode automatic differentiation is just this chain-rule bookkeeping, done exactly, at scale.

Why deep networks need more than plain gradient descent

• Vanishing/exploding gradients: chaining many derivatives through the chain rule means multiplying many terms together — if each is < 1, the product shrinks toward 0 exponentially with depth (vanishing); if each is > 1, it explodes. This motivates activation function choice (ReLU over sigmoid) and architectural fixes (residual/skip connections, normalization layers).
• Momentum: accumulates a running average of past gradients to smooth out noisy steps and accelerate through consistent directions.
• Adam: combines momentum with a per-parameter adaptive learning rate based on the running variance of gradients — the default optimizer for most modern deep learning.

Convexity, and why it matters

A function is convex if a straight line between any two points on it never dips below the curve — convex functions have exactly one global minimum, guaranteeing gradient descent finds the best answer. Most real deep learning loss landscapes are non-convex (many local minima and saddle points) — which is why training is empirically finicky and why techniques like momentum, learning-rate schedules, and good initialization matter so much in practice.

Lagrange multipliers (constrained optimization)

To minimize f(x) subject to a constraint g(x) = 0, form the Lagrangian L(x, λ) = f(x) - λ·g(x) and solve ∇L = 0. This is the mechanism behind Support Vector Machines (maximize margin subject to correct-classification constraints) and appears throughout constrained resource-allocation problems (e.g., "minimize latency subject to a fixed compute budget").

The learning-rate speed limit: η < 2/λ_max

Why does a learning rate that's too large diverge? Take the simplest loss, a 1-D bowl with curvature λ:

L(w) = ½ · λ · w²       L'(w) = λ·w       (minimum at w = 0)

update:  w_new = w - η·λ·w = (1 - η·λ) · w

After t steps, w_t = (1 - ηλ)ᵗ · w₀. That's a geometric sequence (Ch. 6), so it converges if and only if |1 - ηλ| < 1, that is, 0 < η < 2/λ. With λ = 10:

In a real network every Hessian eigenvalue imposes its own limit, and the sharpest direction (λ_max) sets the speed limit for all of them. But the flattest direction (λ_min) then crawls at a factor of 1 - ηλ_min ≈ 1. The ratio κ = λ_max/λ_min, the condition number, measures how hard the problem is. Plain gradient descent needs on the order of κ steps. Feature standardization (Ch. 9 z-scores), BatchNorm/LayerNorm, momentum and Adam are all, at heart, ways to reduce κ or work around it. Researchers have observed that training tends to drive λ_max up until it hovers right at 2/η ("the edge of stability", Cohen et al., 2021), so this formula is a live dynamic in real training runs.

Momentum and Adam, every symbol explained

Momentum (the "heavy ball"):

v ← β · v + ∇L(w)
w ← w - η · v

v is a velocity that accumulates past gradients, an EWMA from Ch. 9 with β ≈ 0.9, so about 10 steps of memory. Along consistent directions, steps add up and the effective learning rate is about η/(1-β) = 10η. Across oscillating directions the opposite signs cancel. That accelerates the slow, flat directions without speeding up the sharp ones, which is exactly the fix for high κ.

Adam (Kingma & Ba, 2014), the default optimizer for transformers:

m ← β₁·m + (1-β₁)·g                 1st moment: EWMA of gradients
v ← β₂·v + (1-β₂)·g²                2nd moment: EWMA of squared gradients
m̂ = m / (1 - β₁ᵗ)                   bias correction
v̂ = v / (1 - β₂ᵗ)
w ← w - η · m̂ / (√v̂ + ε)

Why bias correction? m and v start at 0, so early on they're dragged toward 0. At step 1, m = 0.1·g, which is 10× too small. Dividing by 1 - 0.9¹ = 0.1 fixes it exactly.

Worked step 1, gradient g = 0.5:

m = 0.1 · 0.5   = 0.05          m̂ = 0.05    / (1 - 0.9)   = 0.5
v = 0.001 · 0.25 = 0.00025      v̂ = 0.00025 / (1 - 0.999) = 0.25
step = η · 0.5 / √0.25 = η · 1.0

Try g = 0.0005 instead and the step is still η · 1.0. Adam's step size is roughly η whatever the gradient's scale. That's why one learning rate works across embedding tables, attention weights and layer norms whose gradients differ by orders of magnitude, and why η in Adam reads directly as "the maximum change per parameter per step."

Lagrange multipliers, worked: how should flows share a network?

A problem from networking that shows what a Lagrange multiplier means. Two links, each with capacity 10 Mbit/s. Flow 0 crosses both links. Flow 1 uses only link A, and flow 2 uses only link B:

flow 0:  ──[ link A ]──[ link B ]──
flow 1:  ──[ link A ]──
flow 2:              ──[ link B ]──

Frank Kelly (1998) proposed choosing rates xᵢ to maximize total "utility" Σ log(xᵢ). The log gives diminishing returns, so starving any flow is very costly. This is proportional fairness:

maximize    log x₀ + log x₁ + log x₂
subject to  x₀ + x₁ ≤ 10      (link A, multiplier λ_A)
            x₀ + x₂ ≤ 10      (link B, multiplier λ_B)

The Lagrangian folds each constraint into the objective, weighted by a multiplier:

ℒ = log x₀ + log x₁ + log x₂ - λ_A·(x₀ + x₁ - 10) - λ_B·(x₀ + x₂ - 10)

Set each partial derivative to zero:

∂ℒ/∂x₁ = 1/x₁ - λ_A = 0              ->  x₁ = 1/λ_A
∂ℒ/∂x₂ = 1/x₂ - λ_B = 0              ->  x₂ = 1/λ_B
∂ℒ/∂x₀ = 1/x₀ - λ_A - λ_B = 0        ->  x₀ = 1/(λ_A + λ_B)

Flow 0 pays the price of every link it crosses. By symmetry λ_A = λ_B = λ, both links are full, so x₀ + x₁ = 1/(2λ) + 1/λ = 10, giving λ = 0.15:

x₀ = 1/0.30 = 3.33 Mbit/s       (the 2-hop flow)
x₁ = x₂ = 1/0.15 = 6.67 Mbit/s  (the 1-hop flows)

Compare max-min fairness, where everyone gets 5. Proportional fairness gives the long flow less, because it uses twice the network resources. Kelly showed that TCP-style congestion control, where each flow backs off when it sees loss or delay (the "price") on its path, is a distributed algorithm solving this optimization. The multipliers are the congestion signals, and the links compute them without any central coordinator. This is also why long multi-hop TCP flows get less throughput than short ones in practice. ML uses the identical machinery: the SVM's support vectors are the training points whose Lagrange multipliers are non-zero.

Real-life engineering ties

• AI: this chapter is the actual mechanism of model.fit() — knowing it lets you debug training instability (exploding loss, NaN gradients, stuck training) as a first-principles problem instead of randomly changing hyperparameters.
• SWE/Network: gradient-based tuning shows up in auto-scaling controllers and adaptive rate limiters that continuously adjust a parameter to minimize an error signal (closely related to Ch. 29, control theory).
• Security: adversarial examples (inputs crafted to fool an ML model) are literally found via gradient ascent on the model's own loss function with respect to the input instead of the weights — understanding backprop is a prerequisite for understanding adversarial ML attacks (Ch. 28 connects this to game theory).

More real-life examples:

• The 2012 AlexNet result that kicked off the modern deep learning boom was, mathematically, nothing more than the gradient descent and backpropagation in this chapter run at then-unprecedented scale on GPUs — the math didn't change, the compute did.
• Modern large language model training runs use carefully engineered learning-rate schedules (a warmup period followed by cosine decay) specifically to avoid the too-high-divergence and too-low-slowness failure modes described in this chapter.
• ResNet (2015) introduced "skip connections" specifically to fix the vanishing-gradient problem in very deep networks — a direct, famous architectural response to the chain-rule multiplication problem explained above.
• Researchers have physically demonstrated adversarial attacks by placing small stickers on a stop sign that reliably fool a vision model into classifying it as a speed-limit sign (Eykholt et al., 2018) — a real-world instance of the gradient-ascent-on-the-input attack.

Common mistake

Setting the learning rate too high "to train faster" — this is the single most common cause of a loss curve that diverges or oscillates instead of converging; always look at the loss curve shape, not just the final number.