Optimization for Machine Learning

Every trained model is the output of an optimizer. The architecture decides what a model could represent; the optimizer, its learning rate and its schedule decide what it actually ends up representing, how long that takes, and whether training diverges at step 40,000 after burning a week of GPU time. This lesson builds the optimizers that matter in practice, from plain gradient descent to AdamW, with their exact update equations, the geometry that motivates each one, the schedules used to train modern networks, and a from-scratch implementation you can run.

Intuition: hiking down in fog

You are on a mountainside in thick fog and want to reach the valley floor. You cannot see the valley; you can only feel the slope under your feet. The obvious strategy is to step downhill, re-check the slope, and repeat. That is gradient descent. The step length is the learning rate: too short and you will be walking all week; too long and you overshoot the valley and end up higher on the opposite slope. If the valley is a long, narrow ravine, the steepest direction points mostly across the ravine rather than along it, so you zig-zag; a heavy ball that builds up speed (momentum) or boots that adapt their stride per direction (Adam) do better. And if your slope readings come from a noisy instrument (a mini-batch), you need smaller steps toward the end.

Loss landscapes

Training minimizes an empirical risk

where $\boldsymbol\theta \in \mathbb{R}^d$ are the parameters, $f$ is the model, $\ell$ is the per-example loss (from the previous lesson) and $N$ is the number of training examples. The loss landscape is the surface of $L$ over parameter space. Near any point, the gradient $\nabla L$ gives the slope and the Hessian $H = \nabla^2 L$ (the $d \times d$ matrix of second derivatives) gives the curvature. The second-order Taylor expansion summarizes local geometry:

The eigenvalues of $H$ (from the linear algebra lesson) are the curvatures along its eigenvector directions. The ratio of largest to smallest, the condition number $\kappa = \lambda_{\max}/\lambda_{\min}$, measures how elongated the valley is.

Convexity

A function is convex if the line segment between any two points on its graph lies on or above the graph: $L(t\mathbf{a} + (1-t)\mathbf{b}) \le tL(\mathbf{a}) + (1-t)L(\mathbf{b})$ for $t \in [0,1]$. For twice-differentiable functions this is equivalent to the Hessian being positive semidefinite everywhere. Convexity's great gift: every local minimum is a global minimum, and for smooth convex functions gradient descent with a suitable learning rate is guaranteed to converge to the minimum value. Linear regression with squared error and logistic regression with cross-entropy are convex in their parameters.

Neural networks are not convex. Their landscapes contain many minima (often equivalent, because permuting hidden units gives the same function), flat regions, and saddle points: stationary points ($\nabla L = 0$) that curve up in some directions and down in others, like a horse saddle, e.g. $f(x, y) = x^2 - y^2$ at the origin. In very high dimensions, a random stationary point is far more likely to be a saddle than a minimum, because a minimum requires every one of millions of curvatures to be positive. In practice, noisy gradients and momentum usually carry training past saddles, and the minima that large networks reach typically have similar, good loss values, so non-convexity is less catastrophic than it sounds. It does mean that results depend on initialization, data order and hyperparameters, which is why you train several seeds.

Gradient descent and the learning rate

Gradient descent (GD) repeats

with learning rate $\eta > 0$.

Worked example: why the learning rate has a ceiling. Take $L(w) = w^2$, so $\nabla L = 2w$ and the update is $w_{t+1} = w_t - 2\eta w_t = (1 - 2\eta)\,w_t$. Starting at $w_0 = 1$:

In general, for a quadratic with curvature (second derivative) $\lambda$, GD is stable only if $\eta < 2/\lambda$. With many directions, the largest curvature sets the ceiling, $\eta < 2/\lambda_{\max}$, while the smallest curvature sets the speed: progress along that direction shrinks the error by a factor $|1 - \eta\lambda_{\min}|$ per step. With the best fixed rate, the number of steps needed grows roughly in proportion to $\kappa$. For $L(x, y) = \tfrac12(x^2 + 25y^2)$, the curvatures are 1 and 25, so $\eta < 0.08$, and along $x$ each step only shrinks the error by a factor of at least $0.92$. That is the narrow-ravine problem in numbers.

Stochastic gradient descent and mini-batches

Computing $\nabla L$ over all $N$ examples per step is expensive. Stochastic gradient descent (SGD) instead uses a mini-batch $\mathcal{B}$ of $B$ random examples:

$\mathbf{g}_t$ is an unbiased estimate of the full gradient, and by the central-limit reasoning of the probability lesson its noise standard deviation shrinks like $1/\sqrt{B}$. One pass over the data is an epoch. Trade-offs:

• Larger batches use hardware better and give smoother gradients, but each example contributes less progress, and beyond some batch size returns diminish.
• Smaller batches are noisier; that noise helps escape saddles and is often associated with better generalization, but requires smaller learning rates.
• A common heuristic when scaling batch size $k$-fold is to scale the learning rate $k$-fold too (the linear scaling rule), combined with warmup; Goyal et al. (2017) used this to train ResNet-50 on ImageNet with 8,192-image batches.
• When memory limits the batch, gradient accumulation sums gradients over several micro-batches before one update, emulating a larger batch.

Momentum

Momentum keeps a running velocity $\mathbf{v}$ that accumulates past gradients:

with $\beta$ typically 0.9. In directions where gradients consistently agree (along the ravine floor), the velocity builds up to about $1/(1 - \beta) = 10$ times a single gradient; in directions where they flip sign (across the ravine), they cancel. Momentum thus speeds up progress on ill-conditioned problems and averages out mini-batch noise. Nesterov momentum evaluates the gradient at the look-ahead point $\boldsymbol\theta_t - \eta\beta\mathbf{v}_t$, which often reduces overshoot.

RMSProp: per-parameter step sizes

Different parameters see gradients of wildly different magnitudes (embedding rows for rare tokens vs. common ones; first vs. last layers). RMSProp divides each coordinate's step by a running root-mean-square of its recent gradients:

where squares, square roots and division are elementwise, $\rho \approx 0.9$–$0.99$, and $\epsilon$ (e.g. $10^{-8}$) prevents division by zero. Coordinates with large gradients take smaller steps and vice versa, which partly undoes poor conditioning when it is aligned with the coordinate axes.

Adam and AdamW

Adam (Kingma and Ba, ICLR 2015) combines momentum (a first-moment estimate) with RMSProp-style scaling (a second-moment estimate) and corrects both for their initialization at zero. With $\mathbf{m}_0 = \mathbf{v}_0 = \mathbf{0}$ and step counter $t = 1, 2, \dots$:

Defaults from the paper: $\beta_1 = 0.9$, $\beta_2 = 0.999$, $\epsilon = 10^{-8}$. Many large-language-model recipes lower $\beta_2$ (values around 0.95 are common) so the second-moment estimate reacts faster to gradient spikes.

Worked example: Adam's first step. One parameter, $g_1 = 0.5$, $\eta = 10^{-3}$.

1. $m_1 = 0.9 \cdot 0 + 0.1 \cdot 0.5 = 0.05$; $v_1 = 0.999 \cdot 0 + 0.001 \cdot 0.25 = 0.00025$.
2. Bias correction: $\hat m_1 = 0.05/(1 - 0.9) = 0.5$; $\hat v_1 = 0.00025/(1 - 0.999) = 0.25$.
3. Step: $\eta \cdot 0.5/(\sqrt{0.25} + 10^{-8}) \approx 10^{-3} \cdot 1 = 10^{-3}$.

Without bias correction, the step would have been $10^{-3} \cdot 0.05 / \sqrt{0.00025} \approx 3.2 \times 10^{-3}$, distorted by the zero initialization. Note also that the step size was $\eta$ regardless of the gradient's magnitude: Adam's updates are roughly scale-invariant, so $\eta$ is approximately the maximum per-parameter change per step. That makes it much easier to tune than SGD across layers and architectures.

Weight decay and AdamW. L2 regularization adds $\frac{\lambda}{2}\|\boldsymbol\theta\|_2^2$ to the loss, contributing $\lambda\boldsymbol\theta$ to the gradient. For plain SGD this is identical to weight decay, i.e. shrinking the weights by a factor $(1 - \eta\lambda)$ each step. For Adam it is not: the $\lambda\boldsymbol\theta$ term gets divided by $\sqrt{\hat{\mathbf v}}$, so weights with large gradient history are barely regularized. Loshchilov and Hutter's AdamW ("Decoupled Weight Decay Regularization") applies the decay directly to the weights, outside the adaptive scaling:

where $\mathbf{g}_t$ inside $\mathbf{m}$ and $\mathbf{v}$ is the gradient of the unregularized loss. AdamW is the default optimizer for training transformers. Typical practice exempts biases and normalization-layer parameters from weight decay. (The formula above follows PyTorch's torch.optim.AdamW, where $\lambda$ is multiplied by the full learning rate $\eta$. In Loshchilov and Hutter's paper the decay is multiplied only by the schedule multiplier, not by the base step size, so the same nominal $\lambda$ means different things across implementations.)

Learning-rate schedules

A fixed learning rate is rarely best. Two ideas dominate:

• Warmup: start near zero and increase linearly over the first $W$ steps. Early in training, Adam's second-moment estimates are unreliable and the network is far from any reasonable region, so large steps can blow up; warmup is standard for transformers.
• Decay: reduce the rate later so the noisy SGD iterates settle into a minimum. Step decay (divide by 10 at fixed epochs) was common for CNNs; cosine decay (popularized by Loshchilov and Hutter's SGDR paper) is now the default for most large models.

Linear warmup followed by cosine decay from $\eta_{\max}$ to $\eta_{\min}$ over $S$ total steps:

Worked example. $\eta_{\max} = 3 \times 10^{-4}$, $\eta_{\min} = 0$, $W = 1{,}000$, $S = 10{,}000$.

• $t = 500$: warmup, $\eta = 3\times10^{-4} \cdot 0.5 = 1.5 \times 10^{-4}$.
• $t = 5{,}500$: progress $(5500 - 1000)/9000 = 0.5$, $\cos(\pi/2) = 0$, $\eta = \tfrac12 \cdot 3\times10^{-4} = 1.5\times 10^{-4}$.
• $t = 7{,}750$: progress $0.75$, $\cos(0.75\pi) \approx -0.7071$, $\eta \approx 1.5\times10^{-4} \times 0.2929 \approx 4.39\times10^{-5}$.

Variants include warmup-stable-decay (a long constant phase, then a short decay, convenient when you do not know the total step count in advance) and one-cycle schedules.

Regularization through the optimizer's eyes

Regularization reduces overfitting by preferring simpler parameter settings.

• L2 / weight decay penalizes $\|\boldsymbol\theta\|_2^2$; it shrinks all weights proportionally and, as a MAP estimate, corresponds to a Gaussian prior on the weights.
• L1 penalizes $\|\boldsymbol\theta\|_1$; its (sub)gradient $\lambda\,\mathrm{sign}(\theta)$ pushes every weight toward zero by a constant amount, so the penalized optimum has many weights exactly zero (sparsity, feature selection). Solvers that actually land on those zeros use soft-thresholding (proximal gradient) or coordinate descent; plain subgradient steps only jitter around zero. It corresponds to a Laplace prior.
• Implicit regularizers include early stopping (stop when validation loss stops improving), the noise of small-batch SGD, dropout, data augmentation, and the learning-rate schedule itself.

Implementing gradient descent for linear regression

Model: $\hat{\mathbf y} = X\mathbf w + b$ with $X \in \mathbb{R}^{N \times d}$. Loss: mean squared error $L = \frac{1}{N}\|X\mathbf w + b - \mathbf y\|^2$. Gradients (using the calculus lesson's matrix rules), with residual $\mathbf r = X\mathbf w + b - \mathbf y$:

import numpy as np

rng = np.random.default_rng(0)
N, d = 1000, 3
X = rng.normal(size=(N, d)) * np.array([1.0, 10.0, 0.1])   # badly scaled features
w_true, b_true = np.array([2.0, -0.3, 5.0]), 1.0
y = X @ w_true + b_true + rng.normal(scale=0.5, size=N)

# Standardize features: fixes conditioning (try removing this and see GD crawl)
mu, sd = X.mean(0), X.std(0)
Xs = (X - mu) / sd

def minibatch_sgd(X, y, lr=0.01, epochs=30, batch=32, wd=0.0):
    w, b = np.zeros(X.shape[1]), 0.0
    for epoch in range(epochs):
        perm = rng.permutation(len(y))
        for i in range(0, len(y), batch):
            idx = perm[i:i + batch]
            r = X[idx] @ w + b - y[idx]                 # residuals (B,)
            gw = 2 * X[idx].T @ r / len(idx)            # (d,)
            gb = 2 * r.mean()
            w -= lr * (gw + wd * w)                     # L2 / weight decay
            b -= lr * gb
        if epoch % 10 == 0:
            print(epoch, round(np.mean((X @ w + b - y) ** 2), 4))
    return w, b

w, b = minibatch_sgd(Xs, y)
w_orig = w / sd                                         # undo standardization
b_orig = b - (mu / sd) @ w
print("SGD   :", w_orig.round(3), round(b_orig, 3))
A = np.c_[X, np.ones(N)]                                # closed-form check
print("lstsq :", np.linalg.lstsq(A, y, rcond=None)[0].round(3))

The loop shuffles each epoch, computes residuals and gradients on a mini-batch, and updates. Standardizing features matters: in the raw data the feature scales differ by a factor of 100, so curvatures differ by about $10^4$ and any single learning rate is either unstable for one feature or glacial for another. After standardization the condition number is near 1. Running it, the first two weights and the bias agree with lstsq to about two decimals (1.976 vs 1.975, −0.297 vs −0.297, 1.019 vs 1.018). The third weight lands at about 5.24 against lstsq's 5.18: that feature has a tiny scale (standard deviation 0.1), so the data pin its coefficient down only loosely (even the exact least-squares answer is 0.18 away from the true 5.0), and SGD's constant-rate jitter in standardized units is multiplied by $1/0.1 = 10$ when converted back. Decaying the learning rate, or a few full-batch steps at the end, closes the gap. Try lr=0.05: training still converges, but the final jitter is visibly larger.

Comparing optimizers on a ravine

import numpy as np

curv = np.array([0.01, 1.0])                    # L = 0.5*(0.01 x^2 + y^2), kappa = 100
grad = lambda th: curv * th

def run(update, steps=100):
    th, state = np.array([10.0, 1.0]), {}
    for t in range(1, steps + 1):
        th = update(th, grad(th), state, t)
    return 0.5 * (curv * th**2).sum()

def gd(th, g, s, t, lr=1.9):                    # just under 2 / lambda_max
    return th - lr * g

def momentum(th, g, s, t, lr=0.5, beta=0.9):
    s["v"] = beta * s.get("v", 0) + g
    return th - lr * s["v"]

def adamw(th, g, s, t, lr=0.3, b1=0.9, b2=0.999, eps=1e-8, wd=0.0):
    s["m"] = b1 * s.get("m", 0) + (1 - b1) * g
    s["v"] = b2 * s.get("v", 0) + (1 - b2) * g**2
    mhat, vhat = s["m"] / (1 - b1**t), s["v"] / (1 - b2**t)
    return th - lr * (mhat / (np.sqrt(vhat) + eps) + wd * th)

for name, f in [("GD", gd), ("Momentum", momentum), ("Adam", adamw)]:
    print(f"{name:9s} final loss {run(f):.2e}")

The starting loss is 1.0 and the condition number is 100. GD runs just under its maximum stable rate ($1.9 < 2/1$) and still only reaches about$1.1 \times 10^{-2}$ after 100 steps, because along the shallow $x$ direction each step shrinks the error by just $1 - 1.9 \times 0.01 = 0.981$. Momentum reaches about $1.0 \times 10^{-5}$ and Adam about $1.5 \times 10^{-5}$: three orders of magnitude better. Momentum wins by accumulating velocity along the consistent shallow direction; Adam wins by rescaling each coordinate so both directions move at a similar rate. The learning rates were tuned by hand for each method, as you would in practice; results on a toy quadratic do not transfer one-to-one to deep networks, but the mechanisms are the same.

Common pitfalls

• Learning rate too high or too low. Loss that explodes or becomes nan within a few hundred steps usually means too high; a loss that falls painfully slowly usually means too low. Run a short learning-rate sweep on a log scale (e.g. $10^{-5}$ to $10^{-2}$) before long runs.
• Using Adam with L2 in the loss when you meant weight decay. Use AdamW (decoupled decay) and exclude biases and norm parameters.
• No warmup for transformers. Early instability or loss spikes often disappear with a few hundred to a few thousand warmup steps.
• Unscaled features for linear models and MLPs on tabular data: standardize first.
• Forgetting to shuffle or reusing the same batch order makes SGD gradients biased toward data ordering.
• Changing batch size without retuning the learning rate. The best learning rate depends on batch size.
• Judging by training loss alone. Track validation loss; use early stopping or checkpoint selection.

In practice

• Default starting point for deep networks: AdamW, $\beta_1 = 0.9$, $\beta_2$ between 0.95 and 0.999, weight decay around 0.01 to 0.1, linear warmup, cosine decay, and gradient clipping at a global norm of 1.0. Then sweep the peak learning rate, which matters most.
• SGD with momentum (often Nesterov) and step or cosine decay remains competitive for many vision models and is cheaper in memory: Adam stores two extra values per parameter, which adds up for large models.
• Log the learning rate, loss, gradient norm and update-to-weight ratio every step; most training failures are visible in these four curves before they become fatal.
• For fine-tuning pretrained models, use learning rates one or two orders of magnitude lower than for pretraining and short warmups; adapters such as LoRA tolerate somewhat higher rates because only small matrices are trained.
• Save optimizer state along with weights in checkpoints; resuming Adam without its moment estimates behaves like a fresh, unwarmed start.