1.3 Calculus, Optimization & Convexity¶

Training a neural network is fundamentally a calculus problem: given a scalar loss \(\mathcal{L}(\theta)\) that measures how wrong our model is, find the \(\theta\) that makes it as small as possible. Everything else — architectures, data pipelines, learning rate schedules — exists in service of that search. This chapter builds the mathematical scaffolding from first principles: how gradients generalize derivatives to high dimensions, what the Hessian tells you about the loss landscape, why the landscape is almost never convex in practice, and which algorithmic tricks let us navigate it efficiently anyway. We close with a worked numerical example, a full runnable implementation of gradient descent on a 2D loss surface, and an answer to one of the most satisfying questions in ML theory: why does SGD generalize?

Prerequisites: basic single-variable calculus and the matrix notation from Linear Algebra for Deep Learning. We will build on these ideas directly in Neural Networks From Scratch: MLPs & Backprop and Automatic Differentiation & PyTorch Internals.

Gradients, Jacobians, and Hessians¶

The Gradient: Steepest Ascent in Parameter Space¶

Let \(f : \mathbb{R}^n \to \mathbb{R}\) be a differentiable scalar-valued function. The gradient \(\nabla_\theta f\) is the vector of partial derivatives:

The gradient has a geometric meaning that makes it indispensable: at any point \(\theta\), \(\nabla_\theta f\) points in the direction of steepest increase of \(f\). Moving in the direction \(-\nabla_\theta f\) therefore decreases \(f\) most rapidly — that single insight is all of gradient descent.

The directional derivative of \(f\) along unit vector \(\mathbf{u}\) is \(D_\mathbf{u} f = \nabla f \cdot \mathbf{u} = \|\nabla f\| \cos\phi\), where \(\phi\) is the angle between \(\mathbf{u}\) and \(\nabla f\). It is maximized when \(\phi = 0\) (moving along \(\nabla f\)) and minimized when \(\phi = \pi\) (moving against \(\nabla f\)).

For a modern language model, \(\theta \in \mathbb{R}^n\) with \(n\) on the order of \(10^9\) to \(10^{12}\). The gradient is a vector of the same dimensionality — one scalar sensitivity per parameter.

The Jacobian: Gradient for Vector-Valued Maps¶

When the function maps \(\mathbb{R}^n \to \mathbb{R}^m\), the first-order derivative is the Jacobian matrix \(J \in \mathbb{R}^{m \times n}\):

Each row \(i\) of \(J\) is the gradient of the \(i\)-th output with respect to all inputs. The gradient is the special case \(m=1\), i.e., \(J \in \mathbb{R}^{1 \times n}\) (a row vector; transposing gives the usual column-vector gradient).

In neural network backpropagation, the backward pass through each layer is essentially a Jacobian-vector product (JVP) or vector-Jacobian product (VJP). PyTorch’s autograd computes VJPs by default (reverse-mode AD), which is efficient when \(m \ll n\) — exactly the situation when we have a scalar loss. See Automatic Differentiation & PyTorch Internals for the full derivation.

The Hessian: Second-Order Curvature¶

The Hessian \(H \in \mathbb{R}^{n \times n}\) collects all second-order partial derivatives:

The Hessian is symmetric (under mild regularity conditions, by Clairaut’s theorem). Its eigenvalues characterize the local curvature of the loss:

• Positive eigenvalues → the loss curves upward in that direction (locally convex).
• Negative eigenvalues → the loss curves downward (local maximum in that direction).
• Zero eigenvalues → flat directions.

For a network with \(n = 10^9\) parameters, storing \(H\) explicitly requires \(n^2 \approx 10^{18}\) floats — physically impossible. This is why second-order optimizers for large networks either approximate \(H\) (L-BFGS) or work with diagonal/structured approximations (Shampoo, K-FAC). We revisit these in Optimizers: SGD, Adam, Adafactor, Lion, Muon & Shampoo.

The Chain Rule and Backpropagation¶

The chain rule is the engine of all gradient computation in deep learning. For composite functions \(z = f(g(x))\):

In the multivariate case, let \(\mathbf{y} = g(\mathbf{x})\) and \(z = f(\mathbf{y})\). Then:

which in matrix form is the VJP \(\frac{\partial z}{\partial \mathbf{x}} = J_g^\top \frac{\partial z}{\partial \mathbf{y}}\).

For a deep network, the loss \(\mathcal{L}\) is a composition of \(L\) layer functions. Unwinding the chain rule layer by layer — from the output back to the input — is backpropagation. The key efficiency insight is that each intermediate Jacobian never needs to be materialized; we only need the VJP, which costs \(O(n)\) time per layer rather than \(O(n^2)\).

A Three-Layer Chain Rule Worked Out¶

Suppose \(\mathcal{L} = \text{MSE}(W_3 \sigma(W_2 \sigma(W_1 x)))\). Define:

Backward pass (pseudocode logic):

The backward pass mirrors the forward pass in structure — each layer’s contribution is gated by its local Jacobian — and costs roughly the same FLOPs as the forward pass.

Convex vs Non-Convex Landscapes¶

Convexity Defined¶

A set \(\mathcal{C}\) is convex if for any two points \(x, y \in \mathcal{C}\) and \(\lambda \in [0,1]\), the interpolant \(\lambda x + (1-\lambda)y \in \mathcal{C}\). A function \(f\) is convex if its domain is convex and:

Equivalently, \(f\) is convex if \(\nabla^2 f \succeq 0\) everywhere (positive semidefinite Hessian). For convex functions, any local minimum is a global minimum — gradient descent with a suitable step size is guaranteed to converge.

Logistic regression and linear regression with squared loss are convex. Almost every useful deep learning objective is not convex.

Why Neural Networks Are Non-Convex¶

Several structural reasons conspire to make neural loss landscapes non-convex:

1. Permutation symmetry. Swapping two hidden units in a layer with corresponding weight adjustments yields an identical function but a different parameter vector. This creates \(n!\) equivalent global optima for a layer with \(n\) units — the loss surface has many valleys, not one.

2. Composition of non-linearities. Even a simple two-layer ReLU network has regions where the loss surface is neither convex nor concave.

3. Saddle points at scale. In high-dimensional spaces, critical points (where \(\nabla \mathcal{L} = 0\)) are overwhelmingly likely to be saddle points rather than local minima. This was formalized by Dauphin et al. (2014) using random matrix theory: a random function of \(n\) variables has local minima that are exponentially rare; most critical points have some directions curving up and some curving down.

4. Flat regions and loss plateaus. Dead ReLU neurons, saturating sigmoids, and residual scaling can all create large flat regions where \(\nabla \mathcal{L} \approx 0\) but the current point is far from a minimum.

What Saves Us¶

Empirically, large neural networks trained with SGD find solutions that generalize well, despite the non-convexity. Several phenomena explain this:

• Overparameterization: When \(n \gg \text{data size}\), the loss landscape has many equivalent minima and gradient descent finds one efficiently.
• SGD noise (see the Interview Corner below) biases toward wider minima.
• The loss landscape is actually well-behaved locally for modern architectures: residual connections and layer normalization smooth the surface dramatically (Li et al., Visualizing the Loss Landscape of Neural Nets, 2018).

Convex landscape:          Non-convex landscape (schematic):

Loss                        Loss
 |                           |
 |     U-shaped              |  /\  saddle  /\
 |   /        \              | /  \___/\___/  \
 |  /          \             |/    local min   \
 +--x*-----------> θ        +--------------------> θ
   one global min            many local min + saddles

Gradient Descent: Variants and Analysis¶

Vanilla Gradient Descent (GD)¶

The update rule is simple: move opposite to the gradient, scaled by learning rate \(\eta\):

For a convex function with \(L\)-Lipschitz gradient (see section on Lipschitz constants below), GD with \(\eta \leq 1/L\) converges at rate \(O(1/t)\). For strongly convex functions (Hessian eigenvalues \(\geq \mu > 0\)), it converges exponentially: \(O(\exp(-\mu t / L))\).

Full-batch GD computes the gradient over the entire dataset, which is exact but absurdly expensive for large datasets (a billion-token corpus requires a full pass per step).

Stochastic Gradient Descent (SGD)¶

SGD replaces the full gradient with a noisy estimate computed on a single sample (or a small mini-batch of \(B\) samples):

This is an unbiased estimator of the true gradient: \(\mathbb{E}[\nabla_\theta \hat{\mathcal{L}}] = \nabla_\theta \mathcal{L}\). Variance scales as \(O(1/B)\).

The practical magic of SGD: - Each step costs \(O(B)\) rather than \(O(N)\), enabling many more updates per wall-clock second. - The gradient noise acts as implicit regularization (more on this in the Interview Corner). - It escapes sharp minima and saddle points more easily than full-batch GD.

SGD with Momentum¶

Pure SGD can oscillate and converge slowly in narrow valleys. Heavy-ball momentum (Polyak, 1964) introduces a velocity term:

With \(\beta = 0.9\), each update is an exponentially weighted moving average of past gradients. In the direction of consistent gradient signal, velocity accumulates; in oscillating directions, gradients cancel. The effective learning rate in a consistent direction is \(\eta / (1 - \beta)\) — with \(\beta = 0.9\), it is \(10\eta\).

Nesterov Accelerated Gradient (NAG) computes the gradient at the “lookahead” position rather than the current position:

This provides a corrective anticipation and achieves the optimal \(O(1/t^2)\) convergence rate for convex functions (versus \(O(1/t)\) for GD) — the first-order oracle lower bound.

Adaptive Learning Rate Methods¶

AdaGrad (Duchi et al., 2011) accumulates squared gradients to scale the learning rate per parameter:

Parameters with historically large gradients get smaller effective learning rates. This is powerful for sparse features but the accumulating denominator means the effective learning rate shrinks to zero over time.

RMSProp fixes this with an exponential moving average:

Adam (Kingma & Ba, 2015) combines momentum (first moment) with RMSProp (second moment):

Typical defaults: \(\beta_1 = 0.9\), \(\beta_2 = 0.999\), \(\epsilon = 10^{-8}\), \(\eta = 3 \times 10^{-4}\). The bias-correction terms (\(\hat{m}_t, \hat{v}_t\)) are crucial in early training when the moving averages have not yet warmed up.

For LLM pretraining specifically, the convention departs from those defaults in two ways: \(\beta_2\) is usually lowered to \(0.95\) (as in GPT-3 and the Llama models) so the second moment forgets stale curvature faster on a heavy-tailed, non-stationary gradient distribution, and plain Adam is replaced by AdamW, which applies weight decay as a separate shrinkage step \(\theta \leftarrow \theta - \eta\lambda\theta\) rather than folding an L2 term into \(g\) (where the \(1/\sqrt{\hat v_t}\) normalization would rescale it per parameter and destroy its meaning).

Adam is the workhorse optimizer for LLM pretraining. Its variants (AdamW, Adafactor, Lion) are covered in depth in Optimizers: SGD, Adam, Adafactor, Lion, Muon & Shampoo.

Summary of Variants¶

That memory column decides what you can fit on one GPU, so it is worth making concrete. Take the book’s capstone model, Stack-100M, at roughly \(n = 1.2 \times 10^8\) parameters trained in mixed precision. You pay 4 bytes/param for the fp32 master weights, 4 for the fp32 gradients, and 4 each for Adam’s \(m\) and \(v\) — 16 bytes/param, or about 1.9 GB — on top of the 0.24 GB of bf16 weights the forward pass actually uses. Optimizer state, not the model, is the dominant fixed cost; that single fact is why 8-bit optimizer states (bitsandbytes), Adafactor’s factored second moment, and ZeRO/FSDP optimizer-state sharding exist at all. See Optimizers: SGD, Adam, Adafactor, Lion, Muon & Shampoo for the full accounting and Optimizer & Schedule: Muon + MuonClip and Warmup-Stable-Decay for the concrete choice made for Stack-100M.

The Library Layer: torch.optim¶

Everything above is implemented for you. In practice you never hand-roll these update rules — you write them once from scratch to understand them (as we do below), then use torch.optim, whose implementations are fused into few CUDA kernels and battle-tested against edge cases you would not think of. The code below is the entire optimization layer of a real training loop: parameter groups, AdamW, a warmup + cosine schedule, gradient clipping, and a step.

"""
The library layer: every update rule above, as a few lines of torch.optim.
"""

import math
import torch
import torch.nn as nn

torch.manual_seed(0)
model = nn.Sequential(nn.Linear(64, 256), nn.GELU(), nn.Linear(256, 64))

# AdamW with the LLM-pretraining convention: decoupled weight decay applied to
# 2D matmul weights only -- never to biases or norm gains (1D tensors).
decay = [p for p in model.parameters() if p.ndim >= 2]
no_decay = [p for p in model.parameters() if p.ndim < 2]
opt = torch.optim.AdamW(
    [{"params": decay,    "weight_decay": 0.1},
     {"params": no_decay, "weight_decay": 0.0}],
    lr=3e-4,
    betas=(0.9, 0.95),                 # beta2=0.95 (not 0.999) is standard for LLMs
    eps=1e-8,
    fused=torch.cuda.is_available(),   # fused/foreach kernels: far fewer launches
)

# Warmup + cosine decay to 10% of peak, expressed as a multiplier on lr.
WARMUP, TOTAL = 100, 1000

def lr_multiplier(step: int) -> float:
    if step < WARMUP:
        return (step + 1) / WARMUP                       # linear ramp
    progress = (step - WARMUP) / max(1, TOTAL - WARMUP)  # in [0, 1]
    return 0.1 + 0.9 * 0.5 * (1.0 + math.cos(math.pi * progress))

sched = torch.optim.lr_scheduler.LambdaLR(opt, lr_multiplier)

x, y = torch.randn(32, 64), torch.randn(32, 64)
for step in range(3):
    loss = ((model(x) - y) ** 2).mean()
    opt.zero_grad(set_to_none=True)     # set_to_none frees the grad buffers
    loss.backward()                     # populates p.grad for every parameter
    gnorm = torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
    opt.step()                          # the update rule
    sched.step()                        # advance the schedule
    print(f"step {step}  loss {loss.item():.4f}  "
          f"pre-clip |g|={gnorm:.3f}  lr={sched.get_last_lr()[0]:.3e}")

# ---------------------------------------------------------------------------
# Our from-scratch heavy-ball update IS torch.optim.SGD(momentum=...).
# PyTorch buffers b <- mu*b + g, then theta <- theta - lr*b, which is
# algebraically identical to v <- beta*v - lr*g; theta <- theta + v at a
# CONSTANT lr (v_t = -lr * b_t). They diverge under a schedule: PyTorch
# applies today's lr to the whole accumulated buffer, while our form bakes
# each step's historical lr permanently into the velocity.
# ---------------------------------------------------------------------------
p = nn.Parameter(torch.tensor([4.0, 1.0]))
sgd = torch.optim.SGD([p], lr=0.005, momentum=0.9)    # nesterov=True gives NAG
theta, v = torch.tensor([4.0, 1.0]), torch.zeros(2)
for _ in range(10):
    sgd.zero_grad(set_to_none=True)
    (p[0] ** 2 + 10.0 * p[1] ** 2).backward()         # the quadratic from above
    sgd.step()
    g = torch.tensor([2.0 * theta[0], 20.0 * theta[1]])
    v = 0.9 * v - 0.005 * g                           # our chapter's form
    theta = theta + v
assert torch.allclose(p.detach(), theta, atol=1e-5)
print("from-scratch momentum matches torch.optim.SGD:", theta.tolist())

Three details in that snippet are load-bearing for real runs. Parameter groups are how you exclude 1D tensors from weight decay — decaying a LayerNorm gain toward zero is actively harmful. opt.state_dict() holds \(m\) and \(v\); a checkpoint without it silently restarts the moment estimates and produces a loss bump on resume (see Checkpointing, Fault Tolerance & Long-Running Jobs). And fused=True/foreach=True matter more than they look: with hundreds of parameter tensors, an unfused optimizer launches thousands of tiny kernels per step and can become a measurable fraction of step time.

Saddle Points and Escape Dynamics¶

In a \(n\)-dimensional landscape, a saddle point is a critical point (\(\nabla \mathcal{L} = 0\)) where the Hessian has both positive and negative eigenvalues. By random matrix theory arguments (Choromanska et al., 2015; Dauphin et al., 2014), the fraction of critical points that are true local minima decreases exponentially with \(n\). For \(n \sim 10^9\), essentially all critical points are saddle points.

How SGD Escapes Saddle Points¶

Near a strict saddle (one with a negative eigenvalue \(\lambda_{\min} < 0\)), the gradient noise in SGD will eventually kick the trajectory into the descent direction. The escape time scales as \(O(1/|\lambda_{\min}|)\) — fast escape from sharp saddles, slow from near-flat ones.

Gradient perturbation tricks add explicit noise: \(\theta \leftarrow \theta + \xi\) where \(\xi \sim \mathcal{N}(0, \sigma^2 I)\) every \(k\) steps. This is occasionally used in practice to escape persistent flat regions.

Negative curvature descent (used in some second-order methods) explicitly computes the minimum-curvature direction of the Hessian and steps along it. This escapes saddles in \(O(1)\) steps but requires Hessian-vector products, making it expensive.

Flat Regions vs Sharp Minima¶

Not all minima are equal. The sharpness of a minimum is captured by the maximum eigenvalue of the Hessian \(\lambda_{\max}(H)\) at that point. Sharp minima (large \(\lambda_{\max}\)) generalize worse than flat minima (small \(\lambda_{\max}\)) — intuitively, a flat basin is robust to small perturbations in \(\theta\), while a sharp spike requires precision.

This observation (Hochreiter & Schmidhuber, 1997; Keskar et al., 2017) motivates Sharpness-Aware Minimization (SAM), which explicitly penalizes sharp minima, and also explains why large-batch SGD (which has lower gradient noise) tends to find sharper minima than small-batch SGD.

Measuring Curvature Without Forming the Hessian¶

Sharpness is not just a theoretical quantity — you can measure \(\lambda_{\max}(H)\) for a real network in a few lines, and it is the diagnostic behind edge-of-stability plots and sharpness-aware methods. The trick is that although \(H\) is \(n \times n\), the Hessian-vector product \(Hv\) costs only about two backward passes: differentiate the scalar \(g^\top v\) (where \(g = \nabla\mathcal{L}\) is itself built with create_graph=True) once more, since \(\nabla_\theta (g^\top v) = Hv\) for constant \(v\). Power iteration on that operator then gives the top eigenvalue.

"""
Measuring curvature without ever forming the Hessian:
Hessian-vector products via double backward + power iteration for lambda_max.
"""

import torch
import torch.nn as nn

torch.manual_seed(0)
model = nn.Sequential(nn.Linear(8, 16), nn.Tanh(), nn.Linear(16, 1))
x, y = torch.randn(64, 8), torch.randn(64, 1)
params = [p for p in model.parameters()]


def hvp(vec):
    """Hessian-vector product H @ v. Cost ~2 backward passes; H never formed."""
    loss = ((model(x) - y) ** 2).mean()
    g = torch.autograd.grad(loss, params, create_graph=True)   # keep the graph
    gv = sum((gi * vi).sum() for gi, vi in zip(g, vec))        # scalar g . v
    return torch.autograd.grad(gv, params)                     # d/dtheta (g.v) = H v


# Power iteration: v <- Hv/||Hv|| converges to the eigenvector of largest |lambda|.
v = [torch.randn_like(p) for p in params]
nrm = torch.sqrt(sum((vi ** 2).sum() for vi in v))
v = [vi / nrm for vi in v]
for _ in range(60):
    Hv = hvp(v)
    nrm = torch.sqrt(sum((h ** 2).sum() for h in Hv)) + 1e-12
    v = [h / nrm for h in Hv]

# Rayleigh quotient v^T H v with ||v|| = 1 gives the eigenvalue itself, signed:
# a NEGATIVE value means the dominant direction is one of negative curvature,
# i.e. we are sitting at or near a saddle.
lam = sum((h * vi).sum() for h, vi in zip(hvp(v), v)).item()
print(f"lambda_max(H) ~= {lam:.4f}")
print(f"stability ceiling 2/lambda_max = {2.0 / lam:.4f}  (max usable GD lr)")

On this toy network it prints \(\lambda_{\max} \approx 4.09\) and a stability ceiling of \(2/\lambda_{\max} \approx 0.49\) — a learning rate above that will diverge, which you can verify by running the chapter’s run_gd on it. Two practical notes: power iteration finds the eigenvalue of largest magnitude, so a negative result is informative rather than a bug (it means negative curvature dominates); and for a real LLM you would estimate this on a fixed batch, since \(H\) depends on the data you evaluate it on.

PyTorch ships higher-level wrappers for the same primitives — torch.autograd.functional.hvp/vhp, and torch.func.jacrev/torch.func.hessian for functions small enough to materialize — and the PyHessian library builds on HVPs to estimate the top-\(k\) eigenvalues, the Hessian trace (via Hutchinson’s estimator), and the full eigenvalue density of a trained network.

Lipschitz Constants and the Condition Number¶

\(L\)-Smooth Functions¶

A function \(f\) has an \(L\)-Lipschitz gradient (is \(L\)-smooth) if:

Equivalently, \(\lambda_{\max}(H) \leq L\) everywhere. The gradient doesn’t change “too fast.” This is the condition required for gradient descent to make guaranteed progress: with \(\eta \leq 1/L\), each step decreases the loss by at least \(\|\nabla f\|^2 / (2L)\).

If \(\eta > 1/L\), GD can overshoot and diverge. A common rule of thumb: start with \(\eta = 1/L\) and use a learning rate finder or warmup (see Learning Rate Schedules, Warmup, Batch Size & Hyperparameters) to tune from there.

In practice, real neural network training does not stay safely inside this classical regime: Cohen et al. (2021) showed empirically that full-batch (and often mini-batch) gradient descent instead self-stabilizes at the “edge of stability,” with the Hessian’s largest eigenvalue hovering just above \(2/\eta\) and the loss oscillating non-monotonically step-to-step while still trending down over longer horizons. This doesn’t overturn the \(\eta \lesssim 1/L\) intuition for why a learning rate has a stability ceiling — it shows that real training runs right at that ceiling rather than staying comfortably beneath it.

The Condition Number¶

For strongly convex functions (Hessian eigenvalues bounded below by \(\mu > 0\)), the condition number \(\kappa = L / \mu\) governs convergence:

• GD converges at rate \(\left(1 - \frac{\mu}{L}\right)^t = \left(1 - \frac{1}{\kappa}\right)^t\).
• With Nesterov momentum, the rate improves to \(\left(1 - \frac{1}{\sqrt{\kappa}}\right)^t\).

A condition number of \(\kappa = 100\) means GD takes about \(100 \ln(1/\epsilon)\) steps while Nesterov takes \(10 \ln(1/\epsilon)\) — a \(10\times\) speedup.

For neural networks, the condition number of the loss Hessian can be in the thousands or higher, which is one reason adaptive methods like Adam work better in practice: they implicitly precondition the gradient by the square root of the second moment, effectively rescaling each parameter direction to have condition number closer to 1.

Layer-Wise Conditioning¶

A residual network with skip connections has a better-conditioned loss landscape than a plain deep network (He et al., 2016). The gradient signal passes directly through skip connections without squashing by weight matrices, keeping \(\lambda_{\min}(H)\) from collapsing to zero (vanishing gradients). This is a geometric explanation for why residual networks are easier to train: they reduce the condition number of the optimization problem.

Worked Numerical Example: GD on a Quadratic¶

Full Runnable Code: Gradient Descent on a Toy Loss Surface¶

The code below implements vanilla GD, SGD, momentum, and Adam from scratch on a 2D landscape, then generates “plots as data” (JSON-serializable trajectories you can feed into any visualization library or print directly).

"""
Gradient descent visualization on a 2D loss surface.
Implements GD, SGD (with noise), Momentum, and Adam from scratch.
No dependencies except numpy.
"""

import numpy as np
import json
from typing import Callable, Dict, List, Tuple


# ---------------------------------------------------------------------------
# Loss surface and gradient
# ---------------------------------------------------------------------------

def rosenbrock(theta: np.ndarray) -> float:
    """
    Rosenbrock function: f(x,y) = (1-x)^2 + 100*(y - x^2)^2
    Classic non-convex test: narrow curved valley, global min at (1, 1).
    """
    x, y = theta[0], theta[1]
    return (1.0 - x)**2 + 100.0 * (y - x**2)**2


def rosenbrock_grad(theta: np.ndarray) -> np.ndarray:
    """Analytical gradient of the Rosenbrock function."""
    x, y = theta[0], theta[1]
    dfdx = -2.0 * (1.0 - x) - 400.0 * x * (y - x**2)
    dfdy = 200.0 * (y - x**2)
    return np.array([dfdx, dfdy])


def quadratic(theta: np.ndarray) -> float:
    """Anisotropic quadratic: f(x,y) = x^2 + 10*y^2. Condition number 10."""
    return theta[0]**2 + 10.0 * theta[1]**2


def quadratic_grad(theta: np.ndarray) -> np.ndarray:
    """Gradient of the anisotropic quadratic."""
    return np.array([2.0 * theta[0], 20.0 * theta[1]])


# ---------------------------------------------------------------------------
# Optimizers (all from scratch, no libraries)
# ---------------------------------------------------------------------------

def run_gd(
    theta0: np.ndarray,
    grad_fn: Callable,
    lr: float = 0.005,
    n_steps: int = 500,
) -> List[np.ndarray]:
    """Vanilla gradient descent."""
    theta = theta0.copy()
    trajectory = [theta.copy()]
    for _ in range(n_steps):
        g = grad_fn(theta)
        theta = theta - lr * g
        trajectory.append(theta.copy())
    return trajectory


def run_sgd_with_noise(
    theta0: np.ndarray,
    grad_fn: Callable,
    lr: float = 0.005,
    noise_std: float = 0.1,
    n_steps: int = 500,
    rng_seed: int = 42,
) -> List[np.ndarray]:
    """
    SGD with simulated mini-batch noise.
    In practice the noise comes from random mini-batches; here we add
    Gaussian noise to the gradient to simulate the same effect.
    """
    rng = np.random.default_rng(rng_seed)
    theta = theta0.copy()
    trajectory = [theta.copy()]
    for _ in range(n_steps):
        g = grad_fn(theta) + rng.normal(0, noise_std, size=theta.shape)
        theta = theta - lr * g
        trajectory.append(theta.copy())
    return trajectory


def run_momentum(
    theta0: np.ndarray,
    grad_fn: Callable,
    lr: float = 0.005,
    beta: float = 0.9,
    n_steps: int = 500,
) -> List[np.ndarray]:
    """SGD with heavy-ball momentum."""
    theta = theta0.copy()
    velocity = np.zeros_like(theta)
    trajectory = [theta.copy()]
    for _ in range(n_steps):
        g = grad_fn(theta)
        velocity = beta * velocity - lr * g   # accumulate direction
        theta = theta + velocity
        trajectory.append(theta.copy())
    return trajectory


def run_adam(
    theta0: np.ndarray,
    grad_fn: Callable,
    lr: float = 0.01,
    beta1: float = 0.9,
    beta2: float = 0.999,
    eps: float = 1e-8,
    n_steps: int = 500,
) -> List[np.ndarray]:
    """
    Adam optimizer from scratch.
    Note the bias-correction terms: critical for the first ~1/(1-beta) steps.
    """
    theta = theta0.copy()
    m = np.zeros_like(theta)   # first moment (mean)
    v = np.zeros_like(theta)   # second moment (uncentered variance)
    trajectory = [theta.copy()]
    for t in range(1, n_steps + 1):
        g = grad_fn(theta)
        m = beta1 * m + (1.0 - beta1) * g          # update biased first moment
        v = beta2 * v + (1.0 - beta2) * g**2       # update biased second moment
        m_hat = m / (1.0 - beta1**t)               # bias correction
        v_hat = v / (1.0 - beta2**t)               # bias correction
        theta = theta - lr * m_hat / (np.sqrt(v_hat) + eps)
        trajectory.append(theta.copy())
    return trajectory


# ---------------------------------------------------------------------------
# Run experiments and emit results as JSON-serializable data
# ---------------------------------------------------------------------------

def trajectory_to_losses(
    traj: List[np.ndarray],
    loss_fn: Callable,
) -> List[float]:
    """Convert a list of parameter vectors into a list of loss values."""
    return [float(loss_fn(theta)) for theta in traj]


def main():
    theta0 = np.array([-1.5, 0.5])   # starting point for Rosenbrock

    print("=== Rosenbrock surface (non-convex) ===")
    print(f"  Starting loss: {rosenbrock(theta0):.4f}")
    print(f"  Global minimum at (1, 1), loss = 0\n")

    runs: Dict[str, List[np.ndarray]] = {
        "GD":       run_gd(theta0, rosenbrock_grad, lr=0.001, n_steps=2000),
        "Momentum": run_momentum(theta0, rosenbrock_grad, lr=0.001, beta=0.9, n_steps=2000),
        "Adam":     run_adam(theta0, rosenbrock_grad, lr=0.01, n_steps=2000),
        "SGD+Noise": run_sgd_with_noise(theta0, rosenbrock_grad, lr=0.001,
                                         noise_std=0.05, n_steps=2000),
    }

    results = {}
    for name, traj in runs.items():
        losses = trajectory_to_losses(traj, rosenbrock)
        final_theta = traj[-1]
        results[name] = {
            "final_loss": losses[-1],
            "final_theta": final_theta.tolist(),
            # Subsample trajectory for visualization (every 50 steps)
            "loss_curve": losses[::50],
            "path_x": [t[0] for t in traj[::50]],
            "path_y": [t[1] for t in traj[::50]],
        }
        print(f"  {name:12s}  final_loss={losses[-1]:.6f}  "
              f"theta=({final_theta[0]:.4f}, {final_theta[1]:.4f})")

    # Emit as JSON so downstream tools (e.g., matplotlib, Vega, Plotly) can render
    print("\n=== JSON output (for plotting) ===")
    print(json.dumps(results, indent=2)[:800], "...")   # truncate for display

    # -----------------------------------------------------------------------
    # Second experiment: condition number demonstration
    # -----------------------------------------------------------------------
    print("\n=== Quadratic surface (condition number κ=10) ===")
    theta0_q = np.array([4.0, 1.0])

    gd_traj    = run_gd(theta0_q, quadratic_grad, lr=0.04, n_steps=200)
    mom_traj   = run_momentum(theta0_q, quadratic_grad, lr=0.04, beta=0.9, n_steps=200)
    adam_traj  = run_adam(theta0_q, quadratic_grad, lr=0.1, n_steps=200)

    for name, traj in [("GD", gd_traj), ("Momentum", mom_traj), ("Adam", adam_traj)]:
        losses = trajectory_to_losses(traj, quadratic)
        # Find first step where loss < 0.01
        converge_step = next((i for i, l in enumerate(losses) if l < 0.01), None)
        print(f"  {name:10s}  steps to loss<0.01: "
              f"{'>' + str(len(losses)) if converge_step is None else converge_step}")


if __name__ == "__main__":
    main()

Running this script produces output like:

=== Rosenbrock surface (non-convex) ===
  Starting loss: 312.5000
  Global minimum at (1, 1), loss = 0

  GD            final_loss=0.053318  theta=(0.7693, 0.5908)
  Momentum      final_loss=0.000000  theta=(0.9999, 0.9998)
  Adam          final_loss=2.681244  theta=(-0.6368, 0.4102)
  SGD+Noise     final_loss=0.053240  theta=(0.7695, 0.5911)

=== Quadratic surface (condition number κ=10) ===
  GD          steps to loss<0.01: 45
  Momentum    steps to loss<0.01: 60
  Adam        steps to loss<0.01: 61

(Note: with \(\theta_0 = (-1.5, 0.5)\) the actual starting Rosenbrock loss is \((1-(-1.5))^2 + 100(0.5-(-1.5)^2)^2 = 6.25 + 306.25 = 312.5\), not \(6.25\) — the smaller term alone is only the \((1-x)^2\) piece.)

These numbers are worth sitting with, because they cut against the “Adam always wins” intuition. On the Rosenbrock surface, momentum — not Adam — reaches the global minimum almost exactly, because the curved valley rewards accumulating velocity along a consistent direction. Adam actually does worse than plain GD here: its per-coordinate normalization by \(\sqrt{\hat v_t}\) treats \(x\) and \(y\) independently, which fights against Rosenbrock’s tightly coupled, curved valley, and with these particular (untuned) hyperparameters it stalls well short of the minimum. On the quadratic, plain GD wins on both metrics with these hyperparameters: it not only reaches the loss \(<0.01\) threshold first (45 steps vs. 60–61), it also ends up at a lower final loss after 200 steps (\(5\times 10^{-14}\) vs. \(\sim 10^{-8}\) for momentum and Adam) — its learning rate happens to be unusually well-matched to this simple, low-dimensional bowl, while momentum and Adam overshoot and oscillate for the first several dozen steps (visible if you print the early loss curve) before settling down. The lesson: the textbook convergence-rate rankings (Nesterov beats GD by \(O(\sqrt{\kappa})\), Adam preconditions ill-conditioned directions) describe asymptotic, well-tuned behavior and are landscape-dependent — naive default learning rates do not automatically realize those speedups on every problem, and which optimizer “wins” for a given step budget and hyperparameter choice is an empirical question you have to check on your actual loss surface, not something to assume from the algorithm’s name.

Why SGD Generalizes: The Interview Question¶

Learning Rate, Convergence, and Practical Guidance¶

The Learning Rate as the Most Important Hyperparameter¶

The learning rate \(\eta\) controls the scale of each update. Too large: the loss diverges or oscillates. Too small: training is impossibly slow. For a quadratic with maximum curvature \(L\), the critical threshold is exactly \(\eta < 2/L\).

Learning rate warmup is standard for LLM training: start with \(\eta \approx 0\) and linearly ramp to the target over the first few hundred to few thousand steps. This matters because at initialization, gradients are large and noisy — a large learning rate early in training frequently causes loss spikes. Warmup effectively applies a smaller effective learning rate while the model is in its most sensitive initial state.

Cosine annealing then decays \(\eta\) to near zero following a cosine curve:

The gradual decay helps the optimizer settle into a flat minimum rather than oscillating around it. The full landscape of schedules (linear decay, polynomial, constant-then-cosine, etc.) is covered in Learning Rate Schedules, Warmup, Batch Size & Hyperparameters.

Gradient Clipping¶

Large gradient norms can destabilize training, especially early in training or after a data spike. Global gradient norm clipping rescales the gradient if its norm exceeds a threshold \(\tau\):

This is not a cure for bad initialization or bad data, but it prevents one catastrophically large step from derailing an otherwise healthy run. Typical values: \(\tau \in [0.5, 5.0]\), with \(\tau = 1.0\) common for LLM training.

In PyTorch this is the single line torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0), called after loss.backward() and before opt.step(). It returns the pre-clip total norm — always log that value, because a gradient norm that starts climbing is the earliest available warning of an impending loss spike. Two footguns: the norm must be computed over all shards, so under FSDP you call the wrapped module’s own clip_grad_norm_ method rather than the free function, and under gradient accumulation you clip once after the final micro-batch, not per micro-batch.

Connection to Second-Order Methods¶

The ideal update would be \(\theta \leftarrow \theta - H^{-1} \nabla \mathcal{L}\) (Newton’s method), which solves the local quadratic approximation exactly. This converges in \(O(\log(1/\epsilon))\) steps regardless of condition number — it completely handles ill-conditioning. The problem is cost: inverting \(H\) for \(n = 10^9\) is \(O(n^3)\), completely intractable.

Quasi-Newton methods (L-BFGS) approximate \(H^{-1}\) using the last \(k \approx 20\) gradient differences, costing \(O(kn)\) per step. They work well for small-to-medium networks but are rarely used for LLMs.

Distributed Shampoo (Anil et al., 2020) maintains full matrix preconditioners per layer using Kronecker products, achieving near-second-order convergence at manageable memory cost. It has been used in production at Google for pretraining runs. See Optimizers: SGD, Adam, Adafactor, Lion, Muon & Shampoo for details.

Putting It Together: The Optimization Landscape of an LLM¶

Training a large language model is not just applying Adam to a convex problem. The loss landscape has structure that the practitioner must respect:

1. Early training (steps 0–1000): High gradient variance, potential for large spikes. Learning rate warmup and gradient clipping are essential. The Hessian has many near-zero eigenvalues (the network is massively under-utilized early).

2. Mid training (bulk of steps): The loss landscape is well-behaved near a low-loss manifold. Adam efficiently navigates the per-parameter curvature variations. The condition number is high in some directions (embedding layers, which see many more gradient updates for common tokens) and low in others.

3. Late training / convergence: The learning rate decay helps the optimizer settle into a flat basin. The final loss is determined not by the optimization algorithm per se, but by the flatness of the basin found — a generalization property.

4. Loss spikes: Occasional sharp increases in loss, sometimes by \(2-10\times\), can result from bad batches, gradient accumulation bugs, or instability in normalization layers. Monitoring the gradient norm, the Adam second moment, and the per-layer update magnitudes is key. See Training Stability, Loss Spikes & Debugging Large Runs.

The full probabilistic picture of why the loss landscape is navigable — the statistical mechanics viewpoint, the role of overparameterization, and the relationship to the Neural Tangent Kernel (NTK) — goes beyond our scope here. For the practitioner, the takeaways from optimization theory are: use AdamW (with \(\beta_2 \approx 0.95\)), clip gradients, warm up the learning rate, and monitor training closely. AdamW remains the default, but matrix-aware optimizers that precondition 2D weight matrices with an orthogonalized update — Muon and its stabilized variant MuonClip — have moved from research curiosities to production pretraining runs since 2024–2025; they are covered in Optimizers: SGD, Adam, Adafactor, Lion, Muon & Shampoo, and applied to the book’s ~100M-parameter capstone model in Optimizer & Schedule: Muon + MuonClip and Warmup-Stable-Decay and The Pretraining Run: A Complete Single-GPU Training Loop.

Further Reading¶

• Ruder, Sebastian. “An Overview of Gradient Descent Optimization Algorithms.” arXiv, 2016. A comprehensive survey of SGD variants with intuition-building diagrams.
• Kingma, Diederik P. and Ba, Jimmy. “Adam: A Method for Stochastic Optimization.” ICLR 2015. The original Adam paper; essential reading.
• Dauphin, Yann N., et al. “Identifying and Attacking the Saddle Point Problem in High-Dimensional Non-Convex Optimization.” NeurIPS 2014. Establishes why saddle points dominate over local minima in deep networks.
• Keskar, Nitish Shirish, et al. “On Large-Batch Training for Deep Learning: Generalization Gap and Sharp Minima.” ICLR 2017. Empirical and theoretical case that large-batch SGD finds sharper minima.
• Li, Hao, et al. “Visualizing the Loss Landscape of Neural Nets.” NeurIPS 2018. Shows how skip connections smooth the loss landscape.
• Nesterov, Yurii. “A Method of Solving a Convex Programming Problem with Convergence Rate \(O(1/k^2)\).” Soviet Mathematics Doklady, 1983. The original Nesterov acceleration paper.
• Polyak, Boris T. “Some Methods of Speeding Up the Convergence of Iteration Methods.” USSR Computational Mathematics and Mathematical Physics, 1964. Origin of heavy-ball momentum.
• Boyd, Stephen and Vandenberghe, Lieven. Convex Optimization. Cambridge University Press, 2004. The definitive reference on convexity, duality, and gradient methods. Available free online.

Exercises¶

1. (Conceptual) The chapter states that moving in the direction \(-\nabla_\theta f\) decreases \(f\) most rapidly. Using the directional-derivative identity \(D_\mathbf{u} f = \nabla f \cdot \mathbf{u} = \|\nabla f\| \cos\phi\) (with \(\mathbf{u}\) a unit vector and \(\phi\) the angle between \(\mathbf{u}\) and \(\nabla f\)), prove that among all unit directions, \(\mathbf{u} = -\nabla f / \|\nabla f\|\) gives the most negative directional derivative. What is the value of \(D_\mathbf{u} f\) in that direction, and what does it equal when \(\nabla f = 0\)?

2. (Quantitative) Consider the anisotropic quadratic from the chapter, \(\mathcal{L}(\theta_1, \theta_2) = \theta_1^2 + 10\theta_2^2\), optimized with vanilla gradient descent \(\theta_{t+1} = \theta_t - \eta \nabla\mathcal{L}\).

(a) Derive the per-coordinate update maps and find the largest learning rate \(\eta\) for which GD converges. Relate your answer to the smoothness constant \(L\) and the threshold \(\eta < 2/L\) from the chapter.

(b) With \(\eta = 0.05\) and starting point \(\theta_0 = (2,\ 0.5)\), compute \(\theta_1\) (one step) and the loss before and after the step by hand.

3. (Quantitative) Adam’s bias-correction terms \(\hat m_t, \hat v_t\) are described as “crucial in early training.” Take the very first Adam step (\(t=1\)) with the chapter’s defaults \(\beta_1 = 0.9\), \(\beta_2 = 0.999\), \(\epsilon = 10^{-8}\), \(\eta = 0.01\), and moments initialized at \(m_0 = v_0 = 0\). For a gradient \(g_1 = (0.1,\ -0.2)\), compute \(m_1\), \(v_1\), \(\hat m_1\), \(\hat v_1\), and the resulting parameter update \(\Delta\theta = \theta_1 - \theta_0\). What would the update have been without bias correction (using \(m_1, v_1\) directly), and what general property of the corrected first step does this reveal?

4. (Conceptual/quantitative) The chapter claims that with momentum, “the effective learning rate in a consistent direction is \(\eta/(1-\beta)\) — with \(\beta = 0.9\), it is \(10\eta\).” Derive this. Assume the gradient is a constant vector \(g\) (a perfectly consistent direction) and use the heavy-ball recursion \(v_{t+1} = \beta v_t - \eta g\), \(\theta_{t+1} = \theta_t + v_{t+1}\), with \(v_0 = 0\). Find the steady-state velocity \(v_\infty\) and interpret it as an effective step size. Evaluate for \(\beta = 0.9\) and for \(\beta = 0.99\).

5. (Conceptual) The chapter lists permutation symmetry as one structural cause of non-convexity. Consider a single hidden layer with \(k\) units. (a) How many parameter vectors are functionally identical to a given one by permuting hidden units alone? Evaluate for \(k = 4\) and \(k = 512\). (b) Use the definition of a convex function to argue that a loss with two distinct, isolated global minimizers separated by a higher-loss barrier cannot be convex.

6. (Implementation) The chapter’s runnable code implements GD, SGD, momentum, and Adam but not Nesterov Accelerated Gradient (NAG), even though NAG achieves the optimal \(O(1/t^2)\) convex rate discussed in the text. Implement run_nesterov in the same style as run_momentum, using the chapter’s NAG update \(v_{t+1} = \beta v_t - \eta\nabla\mathcal{L}(\theta_t + \beta v_t)\), \(\theta_{t+1} = \theta_t + v_{t+1}\). Then explain, in one or two sentences, the single line that differs from heavy-ball momentum and why it matters.

def run_nesterov(
    theta0: np.ndarray,
    grad_fn: Callable,
    lr: float = 0.005,
    beta: float = 0.9,
    n_steps: int = 500,
) -> List[np.ndarray]:
    """SGD with Nesterov accelerated gradient (NAG).

    Differs from heavy-ball momentum only in evaluating the gradient at the
    look-ahead position theta + beta*velocity, giving a corrective
    anticipation and the optimal O(1/t^2) convex convergence rate.
    """
    theta = theta0.copy()
    velocity = np.zeros_like(theta)
    trajectory = [theta.copy()]
    for _ in range(n_steps):
        lookahead = theta + beta * velocity      # peek ahead along momentum
        g = grad_fn(lookahead)                   # <-- the one changed line
        velocity = beta * velocity - lr * g      # accumulate direction
        theta = theta + velocity
        trajectory.append(theta.copy())
    return trajectory