Optimization: from gradient descent to Adam

The problem

Training means solving $\min_{\theta}L(\theta)$ for a loss $L$ averaged over data. A point is a local minimum if no small move decreases $L$ and a global minimum if nothing anywhere does. A function is convex if the line segment between any two points on its graph lies above the graph — equivalently, by 1.5 Multivariable and matrix calculus, its Hessian is positive semidefinite everywhere — and for convex functions every local minimum is global, gradient descent provably converges, and the models of Phase 3 (linear and logistic regression, SVMs) are all convex. Deep networks are not: their landscapes have many minima and vast numbers of saddle points. Remarkably, it does not matter much; the surprising empirical fact is that first-order methods find excellent minima anyway, and 4.5 Optimizers and learning-rate schedules and 4.7 Regularization in deep nets explore why.

Gradient descent and the step size

The basic update is $\theta_{t+1}=\theta_t-\eta\,\nabla L(\theta_t)$, with learning rate $\eta>0$. To see what $\eta$ does, take the one-dimensional quadratic $L(\theta)=\tfrac{a}{2}\theta^2$ with curvature $a>0$. Then $\nabla L=a\theta$ and

The iterates converge if and only if $|1-\eta a|<1$, that is, $0<\eta<2/a$. For $\eta<1/a$ they slide smoothly to zero; for $1/a<\eta<2/a$ they overshoot and oscillate with shrinking amplitude; at $\eta=1/a$ they land exactly in one step (Newton's method); beyond $2/a$ they diverge. In many dimensions the quadratic $L=\tfrac12\theta^\top\mathbf{H}\theta$ decouples along the Hessian's eigenvectors (1.3 Eigen-things, SVD and low rank), each with its own curvature $\lambda_i$: the stable range is set by the largest curvature, $\eta<2/\lambda_{\max}$, but convergence along the smallest is then governed by $(1-\eta\lambda_{\min})\approx1-2\lambda_{\min}/\lambda_{\max}$. The ratio $\kappa=\lambda_{\max}/\lambda_{\min}$ — the condition number — determines the number of steps: a loss shaped like a long narrow valley forces tiny steps along the valley floor while bouncing between its walls. This is the zigzag you see in the demo. For a general smooth loss whose gradient changes at rate at most $\mathcal{L}$ (an "$\mathcal{L}$-smooth" function), $\eta\le1/\mathcal{L}$ guarantees the loss decreases at every step.

Intuition. The learning rate is a bet about curvature. Too small and you crawl; too large and you overshoot the walls of the steepest direction. Everything after this section is a way to make that bet smarter.

Stochastic gradient descent

The full gradient $\nabla L=\frac1n\sum_{i}\nabla\ell_i$ costs a pass over the entire dataset. Stochastic gradient descent uses a random mini-batch $B$ instead:

The estimate is unbiased, so on average each step points downhill, and its noise shrinks with batch size (the sample-mean variance of 1.8 Statistics essentials). One pass over the data is an epoch. The noise is a feature as well as a cost: it helps escape saddles and poor minima and acts as an implicit regularizer (4.7 Regularization in deep nets). Larger batches give cleaner gradients and better hardware utilization but fewer updates per epoch, so the learning rate is usually raised with the batch size (4.5 Optimizers and learning-rate schedules).

Momentum and Nesterov acceleration

Plain gradient descent forgets the past at every step. Momentum keeps a running velocity,

with $\beta\approx0.9$. Along a consistent downhill direction the velocity builds up to $\nabla L/(1-\beta)$ — an effective step ten times larger — while in a direction where the gradient flips sign every step the contributions cancel. Momentum therefore accelerates along the valley floor and damps the wall-to-wall oscillation, improving the dependence on the condition number from $\kappa$ to roughly $\sqrt{\kappa}$ steps. Nesterov momentum evaluates the gradient at the look-ahead point $\theta_t-\eta\beta\mathbf{v}_t$, which corrects the velocity a step earlier and is slightly more stable.

Adaptive methods: AdaGrad, RMSProp, Adam and AdamW

Different parameters see different gradient scales — embeddings for rare tokens get tiny, infrequent gradients; biases get large ones. Adaptive methods give each parameter its own step size by normalizing with a running estimate of its gradient magnitude. AdaGrad accumulates $\mathbf{s}_t=\mathbf{s}_{t-1}+\mathbf{g}_t^2$ and steps $\eta\,\mathbf{g}_t/(\sqrt{\mathbf{s}_t}+\epsilon)$; because $\mathbf{s}_t$ only grows, its steps shrink to nothing. RMSProp fixes that with an exponential moving average, $\mathbf{s}_t=\beta_2\mathbf{s}_{t-1}+(1-\beta_2)\mathbf{g}_t^2$. Adam combines RMSProp with momentum and corrects both moving averages for their start at zero:

with defaults $\beta_1=0.9$, $\beta_2=0.999$, $\epsilon=10^{-8}$. The update is a signal-to-noise ratio: the mean gradient divided by its root-mean-square, so each parameter moves at a scale set by $\eta$ regardless of how large or small its raw gradients are.

Math check (bias correction). Unrolling, $\mathbf{m}_t=(1-\beta_1)\sum_{k=1}^{t}\beta_1^{t-k}\mathbf{g}_k$. If the gradients have a stationary mean $\bar{\mathbf{g}}$, then $\mathbb{E}[\mathbf{m}_t]=(1-\beta_1)\bar{\mathbf{g}}\sum_{k=1}^{t}\beta_1^{t-k}=(1-\beta_1^t)\,\bar{\mathbf{g}}$ by the geometric series. Dividing by $1-\beta_1^t$ makes the estimate unbiased; the same argument gives $1-\beta_2^t$ for $\mathbf{v}_t$. Without it, early steps would be far too small for $\mathbf{m}$ and the ratio far too large.

AdamW separates weight decay from the gradient: $\theta_{t+1}=\theta_t-\eta\big(\hat{\mathbf{m}}_t/(\sqrt{\hat{\mathbf{v}}_t}+\epsilon)+\lambda\theta_t\big)$. Adding an $\ell_2$ term to the loss instead would get divided by $\sqrt{\hat{\mathbf{v}}_t}$ and shrink for parameters with large gradients — not what you want from regularization. AdamW is the default optimizer for transformers (4.5 Optimizers and learning-rate schedules, 7.8 Build GPT from scratch (nanoGPT walkthrough)).

Learning-rate schedules

Even the right $\eta$ should change over training. Warmup — increasing $\eta$ linearly from near zero over the first few hundred or thousand steps — prevents the huge early updates that occur while Adam's second-moment estimate is still unreliable and the loss landscape is at its steepest. Afterwards a cosine decay

shrinks the step to let the iterates settle into a minimum; step decays and the one-cycle policy are alternatives. Fine-tuning uses small $\eta$ (around $10^{-5}$ for full fine-tuning, $10^{-4}$ for LoRA, 8.2 Fine-tuning and PEFT); pretraining from scratch uses larger values with long warmup. The learning-rate range test of 4.8 Training dynamics and debugging finds the useful range empirically.

Constrained optimization and Lagrange multipliers

To minimize $f(\mathbf{x})$ subject to $g(\mathbf{x})=0$, note that at an optimum the gradient of $f$ must be perpendicular to the constraint surface — otherwise you could slide along the constraint and improve. That is, $\nabla f=\lambda\nabla g$ for some Lagrange multiplier $\lambda$, which you find by solving the stationarity conditions of the Lagrangian $\mathcal{L}(\mathbf{x},\lambda)=f(\mathbf{x})-\lambda g(\mathbf{x})$. Inequality constraints add sign conditions (the KKT conditions). Two places this matters on the path: the SVM's dual problem (3.10 Support vector machines and kernels) is built from the Lagrangian of the margin constraints, and a norm penalty $\lambda\|\mathbf{w}\|^2$ in a loss is the Lagrangian form of the constraint $\|\mathbf{w}\|\le c$ — regularization is constrained optimization (3.6 Regularization and generalization). The maximum-entropy characterization of the Gaussian in 1.6 Probability I: random variables and distributions is also a Lagrangian calculation.

Worked example: a narrow valley

Take $L(\theta)=\tfrac12(\theta_1^2+10\,\theta_2^2)$: curvatures $1$ and $10$, condition number $\kappa=10$. Gradient descent needs $\eta<2/10=0.2$. With $\eta=0.18$, starting at $(1,1)$: $\theta_2$ is multiplied by $1-1.8=-0.8$ each step (overshoot and oscillate) and $\theta_1$ by $0.82$, so after 10 steps $\theta_1=0.82^{10}\approx0.137$ and $\theta_2=(-0.8)^{10}\approx0.107$ — both directions crawl, for opposite reasons. With $\eta=0.05$, $\theta_2$ converges fast ($0.5^{10}\approx0.001$) but $\theta_1$ barely moves ($0.95^{10}\approx0.60$). Momentum with $\beta=0.9$ at $\eta=0.05$ reaches the same $\theta_1$ accuracy in a few steps. Run exactly this in the demo; then switch to Adam and notice that it moves both coordinates at similar speed because it normalizes each by its own gradient scale.

Where this shows up later

• Fitting linear and logistic regression by gradient descent — 3.2 Linear regression, 3.3 Logistic regression and softmax classification; the SVM dual — 3.10 Support vector machines and kernels; regularization as a constraint — 3.6 Regularization and generalization.
• Optimizers, schedules, warmup, gradient clipping and batch-size scaling in practice — 4.5 Optimizers and learning-rate schedules; diagnosing learning-rate problems — 4.8 Training dynamics and debugging; implicit regularization of SGD — 4.7 Regularization in deep nets.
• Pretraining and fine-tuning recipes — 7.8 Build GPT from scratch (nanoGPT walkthrough), 8.2 Fine-tuning and PEFT; policy-gradient optimization in RL — 4.11 Reinforcement learning essentials.