6. Optimization

Calculus tells you which way is downhill. Optimization is the craft of actually getting down a mountain with billions of dimensions, in fog, using noisy directions, on a limited budget.

Why it exists

The best parameters for a neural network have no formula; they must be searched for. Optimization studies how to search efficiently and reliably, and why some searches diverge, stall or land somewhere useless.

Core ideas

1. The problem: minimize a loss

You want θ that makes L(θ) as small as possible. A global minimum is the lowest point anywhere; a local minimum is lowest only nearby. A stationary point is anywhere the gradient is zero, which includes saddle points (Pillar 2).

2. Gradient descent, derived from first principles

Pillar 2’s first-order Taylor expansion says L(θ + h) ≈ L(θ) + ∇L · h. The step h that decreases L the most, for a given length, points straight against the gradient. So:

The learning rate η sets the step length. The linear approximation is only trustworthy nearby, so η must be small enough not to overshoot.

How small? Take the 1D bowl L(x) = ½λx², where λ is the curvature. One step gives x ← (1 − ηλ)x, which shrinks toward zero only if η < 2/λ. Above that, every step overshoots further and training diverges.

In many dimensions, each eigen-direction of the Hessian has its own curvature. The sharpest direction caps η, while the flattest direction sets how slowly you make progress. Their ratio, the condition number, explains the classic zig-zag down a narrow valley.

3. Convexity: when optimization is easy

A convex function is bowl-shaped everywhere: any line segment between two points on its graph lies above it. Every local minimum is then global, and gradient descent provably finds it.

Neural network losses are not convex. Yet gradient methods work remarkably well on them. Current understanding points to heavy over-parameterization, which creates many good minima, and to noise that helps escape saddle points.

4. Stochastic gradient descent: cheap, noisy, better

The true gradient averages over all training data, which is far too expensive. SGD estimates it from a random minibatch. That estimate is unbiased, its noise shrinks as the batch grows, and each step costs a tiny fraction of a full pass.

Many cheap, noisy steps beat a few exact ones. The noise even seems to help, by shaking the model out of saddles and sharp minima. Batch size and learning rate interact: larger batches mean less noise, which usually allows a larger learning rate.

5. Momentum: remember where you were going

Momentum keeps a running average of past gradients and steps along that instead. Directions that keep agreeing build up speed; directions that keep flipping sign cancel out. It is a heavy ball rolling downhill, and it damps the zig-zag in narrow valleys.

6. Adam: a step size per parameter

Adam combines momentum with a per-parameter scale. It tracks a running mean of gradients m and a running mean of squared gradients v, and divides one by the square root of the other:

The hats are bias corrections for the early steps, when m and v start at zero. Parameters with consistently large gradients get smaller steps, and rarely updated ones get larger steps. AdamW applies weight decay directly to the weights instead of through the gradient, and it is the default for training transformers.

Adam has a systems cost: it stores two extra numbers per parameter. For a model with billions of parameters, that optimizer state is a major share of training memory (Pillar 7).

7. Learning-rate schedules

Warmup starts with a tiny learning rate and ramps up over the first steps. Early on, gradients are large and Adam’s statistics are unreliable, so full-size steps can blow up training.

Decay lowers the learning rate toward the end, often along a cosine curve. Smaller steps average out gradient noise so the model settles into a low region instead of bouncing around it.

8. Regularization: preferring simple solutions

Regularization changes the problem so the optimizer prefers solutions that generalize. Explicit tools include weight decay (an L2 penalty), dropout (randomly zeroing activations), data augmentation and early stopping.

There is also implicit regularization: SGD’s noise and the architecture itself bias training toward certain solutions, even with no penalty term. Understanding it is an open research problem.

9. Constraints and Lagrange multipliers

Sometimes you must optimize under a constraint g(θ) = 0. At the best constrained point, you cannot improve f without violating g, so the two gradients must be parallel:

The multiplier λ measures how much the optimum would improve if you relaxed the constraint. This tool derives the softmax as the maximum-entropy distribution (Pillar 5).

It also explains modern preference tuning. Maximizing expected reward minus β times the KL divergence to a reference model has a closed-form optimum: the reference model’s distribution reweighted by exp(reward / β). Direct Preference Optimization (DPO) is built on exactly that solution.

10. Second-order methods and stabilizers

Newton’s method rescales the gradient by the inverse Hessian, which fixes bad conditioning in one move. But a Hessian for n parameters has n² entries, impossible at billions of parameters. Approximate second-order optimizers are an active research area.

Gradient clipping is the everyday stabilizer: if the gradient’s norm exceeds a threshold, scale it down. It prevents one bad batch from wrecking a long training run.

In AI

• A standard recipe: AdamW, a warmup, cosine decay and gradient clipping cover most transformer training runs.
• Fine-tuning usually uses a much smaller learning rate, so you adjust the model without erasing what it knows.
• Loss spikes often mean a learning rate too close to the stability limit, or a bad batch of data.
• Preference tuning (RLHF, DPO) is KL-constrained optimization.
• Reward hacking is optimization working too well on a flawed objective: Goodhart’s law with a gradient.

Questions to ask

• What does the loss landscape look like here: curvature, conditioning, noise?
• Is the learning rate on the right scale for the sharpest direction?
• When the loss spikes, is it the learning rate, the data, or numerics?
• What is the effective noise level, given the batch size and learning rate?
• What is being regularized, explicitly or implicitly?
• How much memory does this optimizer’s state cost?
• Is this objective what I actually want, or a proxy the model can exploit?

Fluency check

• Derive the stability limit η < 2/λ for a 1D quadratic.
• In NumPy, run gradient descent, momentum and Adam on an elongated 2D bowl and plot the three paths.
• Train logistic regression from scratch with minibatch SGD.
• Show that a minibatch gradient is an unbiased estimate of the full gradient.
• Use a Lagrange multiplier to show the maximum-entropy distribution with a fixed expected score is a softmax.
• Add warmup and cosine decay to your training loop and compare the loss curves.