AI Engineering Foundations 08: Direction Is Free, Step Size Is Deadly

This is lesson 8 of the *AI Engineering Foundations* series, working through the 58k-star open-source course ai-engineering-from-scratch by Rohit Ghumare. Today's topic: optimization—the gradient descent family (Build track, 75 minutes). The derivatives and gradients from the previous two lessons finally take the stage: once you know which way is steeper, how do you get down the mountain?

The fog valley picture

Imagine standing in a fog-covered valley. You can feel the slope under your feet—that's the gradient, and it *free*ly tells you the steepest direction. Direction needs no thought: walk against the gradient. The deadly question is how big a step to take.

The course's test function is Rosenbrock—a banana-shaped narrow valley with its bottom at (1,1), starting point (-1,1). Running the course code verbatim:

• Learning rate 0.001: after 5,000 steps, loss = 0.0036, settled in the valley.
• Learning rate 0.005: by step ten, x explodes to 10^52 and Python throws OverflowError.

A 5x difference in step size: one side is the answer, the other is a corpse of astronomical numbers.

Why vanilla GD fails on Rosenbrock

Plain gradient descent is one line: weights -= learning_rate * gradient. It relies only on local slope, never looking at a map. Rosenbrock has a vicious geometric property: across the valley walls is extremely steep, along the valley floor is extremely flat—the curvature ratio is about 2,500x. A single scalar learning rate cannot serve two directions that differ by 2,500x. Every optimizer trick below starts from this complaint.

Trick one: momentum

A ball rolling downhill doesn't restart at every dip. Momentum accumulates past gradients into a velocity: velocity = 0.9 * velocity + gradient, and weights move by velocity. Oscillating components cancel; consistent components accelerate. The zigzag crawl in narrow valleys gets cured.

Trick two: Adam

Adam keeps two ledgers per weight: a moving average of gradients (direction) and a moving average of squared gradients (magnitude), then updates using the first divided by the square root of the second. Verified empirically: when the gradient grows from 0.28 to 4.4, step length stays nearly constant, tracking the learning rate 0.01. Steep areas get automatic small steps, flat areas large steps—more precisely, it removes "steps scaling with gradient" entirely, leaving only direction and constant step size.

The race

Starting at (-1,1), target loss 10^-6:

| Optimizer | Steps to converge | |---|---| | Adam | ~1,000 | | SGD + momentum | ~15,000 | | Vanilla GD (lr=0.0005) | Not reached in 20,000 (loss 0.041 at 5,000 steps) |

Passing customs: what checks out

• Rosenbrock's gradient formula verified by numerical differencing: errors below 10^-9.
• Step 5's expected outputs—Adam fastest, momentum smoother, vanilla slow—all confirmed.
• Adam's formula matches Kingma & Ba 2014 line by line, with epsilon outside the square root.
• Cosine annealing formula correct; one million weights means a one-million-and-one-dimensional loss landscape, arithmetic checks out.

What doesn't check out

1. "There's no formula for learning rates." There is—it's just unaffordable. For quadratic functions, gradient descent converges when the learning rate is less than 2 divided by the maximum curvature. At the course's own starting point (-1,1), max curvature is 1001.6, giving a threshold of 0.001997. The course's own exercise (sweeping 0.0001 to 0.01) shows 0.001 converges and 0.005 diverges—the boundary sits exactly on this formula. "No formula" is colloquial shorthand for "the formula is too expensive."

2. How divergence actually dies. With lr=0.005, x's trajectory: -1.0, -0.98, -1.04, -0.78, -1.50, 3.42, -61, 4.6e4, -2×10^17, 1.5×10^52. The first five steps genuinely bounce between walls; from step six it's not oscillation but geometric explosion—blowing through float64 in ten steps. Loss sequence: 4, 4.1, 5.5, 26, 274, 8982, 1.4×10^9, 4.6×10^24. Bonus bug: the course's diagram labels the "just right" case as lr=0.01—but on this course's own Rosenbrock, 0.01 explodes in nine steps; the survivor is 0.001. The diagram's number belongs to someone else's function.

3. Momentum's learning rate was silently cut 5x. Vanilla GD uses 0.0005; SGD+momentum uses 0.0001—with no explanation. The accounting: with β=0.9, steady-state velocity is 10x a single-step gradient, so effective step size is 10x the learning rate times gradient. 0.0001 times ten lands back at ~0.001. Control experiment: GD at 0.001 vs SGD+momentum at 0.0001—losses after 5,000 steps are 0.00363 vs 0.00356, essentially the same point. This isn't a tuning tip; it's a conservation law: adding momentum changes what the learning rate *means*.

4. Exercise 4's counterintuitive result. Adding 0.999 exponential decay to vanilla GD: without decay, loss 0.041 at 5,000 steps; with decay, 2.11—50x slower. Reason: 0.999^5000 ≈ 0.0067; the learning rate gets strangled mid-journey. Decay is a tool for "long haul done, only fine-tuning left," while Rosenbrock's 5,000 steps are all on the Long March. The exercise teaches, hands-on, what the course never wrote: schedules are late-stage tools, not whole-journey tools.

5. Slogan audit. The lesson opens with: "Training a neural network is just finding the bottom of a valley." This hides two things:

• You walk the *training loss* mountain; you care about the *test loss*—a different mountain, and the two don't coincide. The course's ending on sharp valleys generalizing worse and flat valleys generalizing better builds a bridge between the two mountains without saying it plainly: the function you optimize is not the function you care about.
• The lesson's own demo shows Adam winning decisively (15x faster); the ending claims SGD+momentum often wins on final accuracy due to noise preventing sharp minima. Both statements are true but live in different worlds—the demo uses deterministic gradients; the claim applies to real mini-batch training. The winning reason (noise) is precisely what the demo lacks. Side note: sharp-minima-poor-generalization is a hypothesis, not a theorem—a 2017 paper showed sharp minima can become flat under reparameterization; the cited Li 2018 paper's actual contribution is loss landscape visualization.

Exercise 2 spoiler

Momentum 0.99 with lr 0.0001 reaches loss 10^-8 in 1,957 steps; momentum 0.9 at the same lr takes 20,000. Momentum 0.99's effective step size is 10x that of 0.9—by intuition, more dangerous—yet it's both faster and stable: overshooting to (1.22, 1.48) at step 500, swinging back to (1.03, 1.05) at step 1,000, settling at the valley bottom by step 2,000 with loss 2.4×10^-9.