Gradient Descent

01 — Why Gradients Drive Learning

Understand the geometry the optimizer sees: loss surface, gradient as steepest descent, and why following it minimizes risk.

Learning is optimization. Given data \( \{(x_i, y_i)\}_{i=1}^n \) and a parametric model \( f(x;\theta) \), define a per-example loss \( \ell(f(x_i;\theta), y_i) \) (squared error, cross-entropy, …) and the empirical risk

\[ L(\theta) = \frac{1}{n}\sum_{i=1}^n \ell(f(x_i;\theta), y_i). \]

Training means moving \( \theta \) downhill on the surface \( z = L(\theta) \). At any \( \theta \), the gradient \( \nabla L(\theta) = (\partial L/\partial \theta_1, \dots, \partial L/\partial \theta_d)^\top \) points in the direction of steepest ascent; \( -\nabla L \) is steepest descent. Taylor's theorem justifies it: \( L(\theta + \Delta) \approx L(\theta) + \nabla L(\theta)^\top \Delta \), so the biggest decrease for small \( \|\Delta\| \) is \( \Delta \propto -\nabla L \).

For smooth \( L \) with Lipschitz gradient \( \|\nabla L(x)-\nabla L(y)\| \le \beta\|x-y\| \), the descent lemma gives

\[ L(y) \le L(x) + \nabla L(x)^\top (y-x) + \frac{\beta}{2}\|y-x\|^2, \]

which will bound how large a step we can safely take (Lesson 06). In non-convex deep learning, we still follow \( -\nabla L \) — not because the landscape is convex, but because it is the only locally computable descent direction and, iteratively, it finds useful minima in practice.

Try it yourself: for \( L(w)= \tfrac12 (w-3)^2 \), hand-compute \( \nabla L(0) = -3 \) and verify \( -\nabla L \) points toward the minimizer \( w^*=3 \).

02 — The Optimization Problem

State empirical vs expected risk, stationary points, and conditions for convergence.

We minimize empirical risk \( L(\theta) \); the true goal is expected risk \( R(\theta)=\mathbb{E}_{(x,y)}[\ell(f(x;\theta),y)] \). Minimizing \( L \) approximates minimizing \( R \) subject to generalization. A stationary point satisfies \( \nabla L(\theta^*)=0 \); it may be a global minimum, local minimum, saddle, or (rarely) flat maximizer. In high dimensions, strict saddles are abundant.

Classical convergence theory assumes convexity or smoothness:

• Smoothness: \( \|\nabla L(x)-\nabla L(y)\|\le \beta\|x-y\| \). Steps up to \( 2/\beta \) can converge; \( \eta = 1/\beta \) is a safe heuristic.
• Strong convexity (\( \mu I \preceq \nabla^2 L \preceq \beta I \)): condition number \( \kappa=\beta/\mu \) controls speed; larger \( \kappa \) → slower (narrow valleys).
• Non-convex: gradient descent converges to stationary points; escaping saddles requires noise or momentum.

Batch gradients cost \( O(nd) \) per step. That cost motivates stochastic approximations — the core trade-off of the next three lessons.

03 — Batch Gradient Descent

Derive the exact update, cost, and when it is the right choice.

Batch (full-batch) GD uses the exact gradient over the whole dataset:

\[ \theta_{t+1} = \theta_t - \eta \, \nabla L(\theta_t) = \theta_t - \eta \cdot \frac{1}{n}\sum_{i=1}^n \nabla \ell_i(\theta_t). \]

for t in range(T):
    g = grad(L, theta)          # over all n examples
    theta -= eta * g
    if norm(g) < tol: break

Properties: deterministic, same trajectory each run, monotonic decrease if \( \eta \) small enough (by descent lemma). For smooth convex \( L \), error after \( T \) steps is \( O(1/T) \); for strongly convex, linear \( O((1-\mu/\beta)^T) \).

Cost: \( n \) forward+backward passes per step — prohibitive at \( n \gg 10^4 \). Memory for gradient is \( O(d) \), but data must be visited entirely.

When to use: small datasets, convex problems, or when exact gradients are needed for debugging/line search. In deep learning, almost never — mini-batch dominates.

Gotchas

• Exactness is not always helpful: in non-convex landscapes, determinism can trap you in sharp local minima that noisy updates would escape.

04 — Stochastic Gradient Descent

Treat the gradient as a random estimator; understand unbiasedness, variance, and the price of one example.

SGD replaces the full average with one random example \( i_t \sim \mathrm{Uniform}\{1..n\} \):

\[ g_t = \nabla \ell_{i_t}(\theta_t), \qquad \mathbb{E}_{i_t}[g_t \mid \theta_t] = \nabla L(\theta_t) \; (\text{unbiased}), \]

\[ \theta_{t+1} = \theta_t - \eta_t \, g_t. \]

Unbiasedness follows because \( \frac1n\sum_i \nabla \ell_i = \nabla L \). Variance is

\[ \mathrm{Var}(g_t) = \frac1n\sum_i \|\nabla \ell_i(\theta) - \nabla L(\theta)\|^2, \]

and it does not vanish even at the optimum unless all gradients agree (overparameterized interpolation). Hence fixed \( \eta \) leaves SGD hovering in a ball of radius \( O(\eta \cdot \text{Var}) \) around the minimizer.

Convergence: with decaying \( \eta_t = \eta_0/\sqrt{t} \) (convex) or \( \eta_t \propto 1/t \) (strongly convex), SGD converges in expectation at \( O(1/\sqrt{T}) \) and \( O(1/T) \) respectively — slower than GD, but each step is \( n\times \) cheaper, so wall-clock wins for large \( n \).

for epoch in range(epochs):
    perm = torch.randperm(n)
    for i in perm:               # single-example loop
        g = grad(ell[i], theta)
        theta -= eta_t * g

Gotchas: single-example noise makes loss non-monotonic; gradient clipping often helps; learning rate must decay or schedule, otherwise no convergence.

Try it yourself: estimate variance empirically: sample 100 stochastic gradients at a fixed \( \theta \) and plot their spread versus the batch gradient.

Primary sources: Bottou et al., Optimization Methods for Large-Scale ML; Robbins & Monro 1951.

05 — Mini-Batch Gradient Descent

Balance variance and throughput; see why powers of two dominate in practice.

Mini-batch samples \( B_t \subset \{1..n\}, |B_t|=B \) (typically 32–512):

\[ g_t^{(B)} = \frac{1}{B}\sum_{i\in B_t} \nabla \ell_i(\theta_t), \qquad \mathbb{E}[g_t^{(B)}]=\nabla L(\theta_t), \]

\[ \mathrm{Var}(g_t^{(B)}) = \frac{1}{B}\,\mathrm{Var}(g_t^{(1)}). \]

Variance shrinks as \( 1/B \) (sampling without replacement even better). Update is

\[ \theta_{t+1} = \theta_t - \eta_t \, g_t^{(B)}. \]

Hardware favors \( B \) that fills GPU/TPU lanes (powers of two, multiples of 32). Larger \( B \) allows larger \( \eta \) (linear scaling rule: double \( B \) → double \( \eta \) up to a limit) and reduces steps per epoch, but too large hurts generalization (sharp minima hypothesis) and wastes samples. Sweet spot is empirical — tune it.

for x_batch, y_batch in DataLoader(dataset, batch_size=128, shuffle=True):
    loss = criterion(model(x_batch), y_batch)
    loss.backward()
    optimizer.step(); optimizer.zero_grad()

Primary sources: Goyal et al. 2017 (large-batch training); Keskar et al. 2017.

06 — Problems of Vanilla GD

Diagnose why a single global \(\eta\) fails on realistic loss surfaces.

Vanilla GD \( \theta_{t+1}=\theta_t - \eta g_t \) has four fundamental weaknesses:

Learning-rate effects:

• \( \eta > 2/\beta \) → overshoot, may diverge.
• \( \eta \approx 1/\beta \) → conservative, monotone but slow.
• \( \eta_t \) decaying \( \propto 1/t \) or cosine/step schedule needed for SGD to converge.

Geometry intuition: for a 2-D quadratic \( L=\tfrac12(\kappa x^2 + y^2) \), gradients point across the valley (steep \( x \)), not along it. Vanilla GD bounces wall-to-wall. Preconditioning by curvature \( H^{-1} \) would fix it — but \( H^{-1} \) is \( O(d^2) \); adaptive methods approximate its diagonal.

Other failure modes: vanishing/exploding gradients in deep nets (chain-rule product), sharp minima that generalize poorly, dependence on initialization.

07 — Momentum & First Moment

Derive exponential moving average of gradients; see how it dampens oscillation and accelerates through ravines.

Idea: average past gradients so the step remembers direction. Define the first moment (exponential moving average) with decay \( \beta_1 \in [0,1) \) (often 0.9):

\[ m_t = \beta_1 m_{t-1} + (1-\beta_1) g_t, \qquad m_0 = 0, \]

which expands to

\[ m_t = (1-\beta_1)\sum_{k=1}^t \beta_1^{t-k} g_k. \]

Classical (Polyak) momentum uses the equivalent form with velocity \( v_t \):

\[ v_t = \mu v_{t-1} - \eta g_t, \qquad \theta_{t+1} = \theta_t + v_t, \quad \mu \approx 0.9. \]

With bias correction (Adam-style) the unbiased estimate is \( \hat m_t = m_t/(1-\beta_1^t) \); early steps are otherwise dampened because \( m_0=0 \).

Why it helps:

• Along consistent descent direction, \( m_t \) accumulates → larger effective step (≈ \( \eta/(1-\mu) \) in steady state).
• Across noisy/sign-flipping directions, averaging cancels → dampened zigzag, lower variance by factor \( (1-\beta_1)/(1+\beta_1) \) for i.i.d. noise.
• Through saddle/plateau, inertia carries the iterate forward where \( \|g\| \approx 0 \).

Nesterov accelerated gradient (NAG): look-ahead variant

\[ g_t = \nabla L(\theta_t + \mu v_{t-1}), \quad v_t = \mu v_{t-1} - \eta g_t, \quad \theta_{t+1}=\theta_t + v_t, \]

which anticipates where momentum will carry us and corrects, achieving \( O(1/T^2) \) on smooth convex problems — optimal among first-order methods.

Effective memory: half-life \( \approx 1/(1-\beta_1) \) steps; \( \beta_1=0.9 \) → ~10 steps, 0.99 → ~100.

# PyTorch: torch.optim.SGD(model.parameters(), lr=0.01, momentum=0.9, nesterov=True)

Primary sources: Polyak 1964; Nesterov 1983; Sutskever et al. 2013.

08 — Second Moment & Adaptive Rates

Make the step size per-parameter: large where gradients are sparse, small where they are large.

Observation: different parameters need different scales. Embedding weights see sparse gradients; biases see dense. Second moment estimates track per-coordinate magnitude:

\[ v_t = \beta_2 v_{t-1} + (1-\beta_2) \, g_t^2, \qquad v_t \in \mathbb{R}^d \; (\text{elementwise}). \]

Its expectation approximates \( \mathbb{E}[g^2] \); dividing by \( \sqrt{v_t} \) rescales steps inversely to typical magnitude.

AdaGrad

Accumulates all past squared gradients (no decay):

\[ G_t = G_{t-1} + g_t^2,\quad G_0=0, \qquad \theta_{t+1}= \theta_t - \frac{\eta}{\sqrt{G_t+\varepsilon}}\odot g_t. \]

Per-coordinate \( \eta_{t,j} = \eta / \sqrt{\sum_{k\le t} g_{k,j}^2 + \varepsilon} \). Infrequent features keep large steps; frequent ones anneal. Problem: \( G_t \) grows monotonically → \( \eta_{t}\to 0 \), premature stop on non-convex problems.

RMSProp

Fixes AdaGrad by exponential decay (Hinton, 2012; \( \beta_2 \approx 0.99 \), also called \( \rho \) or \( \gamma \)):

\[ v_t = \beta_2 v_{t-1} + (1-\beta_2) g_t^2, \qquad \theta_{t+1}= \theta_t - \frac{\eta}{\sqrt{v_t}+\varepsilon}\odot g_t. \]

Now the denominator forgets old magnitudes → non-vanishing steps, stable on non-convex. \( \varepsilon \sim 10^{-8} \) prevents division by zero.

Interpretation: \( v_t \) approximates diagonal of \( \mathbb{E}[gg^\top] \) → diagonal preconditioner \( \mathrm{diag}(v_t)^{-1/2} \) that rescales the loss surface toward isotropy.

Primary sources: Duchi et al. 2011 (AdaGrad); Tieleman & Hinton 2012 (RMSProp).

09 — Adam Optimizer

Combine first and second moments with bias correction; understand every symbol in the update.

Adam (Kingma & Ba, 2015) = momentum + RMSProp + bias correction. Hyperparameters: \( \beta_1=0.9, \beta_2=0.999, \varepsilon=10^{-8}, \eta\) (often \(10^{-3}\)). State: \( m_0=0, v_0=0 \in \mathbb{R}^d \).

At step \( t \) with stochastic gradient \( g_t \):

\[ m_t = \beta_1 m_{t-1} + (1-\beta_1) g_t \qquad \text{(first moment)} \]

\[ v_t = \beta_2 v_{t-1} + (1-\beta_2) \, g_t^2 \qquad \text{(second moment, elementwise)} \]

Bias-corrected (unbiased if \( g_t \) stationary):

\[ \hat m_t = \frac{m_t}{1-\beta_1^t}, \qquad \hat v_t = \frac{v_t}{1-\beta_2^t}. \]

Update:

\[ \theta_{t+1} = \theta_t - \eta \cdot \frac{\hat m_t}{\sqrt{\hat v_t} + \varepsilon} \quad (\text{elementwise}). \]

Why bias correction matters: with \( m_0=v_0=0 \), early \( m_t, v_t \) are pulled toward zero by \( \beta^t \) weight. Dividing by \( 1-\beta^t \) removes that initialization bias so step size is correct from \( t=1 \). Without it, Adam would crawl at start.

Reading the formula:

• \( \hat m_t \) = smoothed descent direction (momentum).
• \( \sqrt{\hat v_t} \) = per-parameter RMS gradient magnitude → adaptive denominator.
• Ratio \( \hat m_t / \sqrt{\hat v_t} \) ≈ sign-normalized, curvature-normalized step (≈ ±1 in each coordinate when gradients are consistent).
• \( \varepsilon \) inside denominator guards against \( v\approx 0 \) explosion and sets a max step \( \approx \eta/\varepsilon \) (numerical, not theoretical).

Invariance: rescaling any coordinate's gradient by \( c \) rescales numerator and denominator by \( c \) → Adam's step is invariant to diagonal rescaling, unlike SGD. This is why it tolerates mixed-magnitude layers.

Convergence note: Adam converges on convex problems with appropriate \( \eta_t \) but can fail to converge with constant \( \eta \) on some non-convex examples (Reddi et al. 2018); AMSGrad fixes by \( \hat v_t \leftarrow \max(\hat v_t, v_t) \).

opt = torch.optim.Adam(model.parameters(), lr=1e-3, betas=(0.9, 0.999), eps=1e-8)
# AMSGrad variant
opt = torch.optim.Adam(model.parameters(), lr=1e-3, amsgrad=True)

Gotchas: \( \beta_2 \) very close to 1 → slow adaptation; too large \( \varepsilon \) kills adaptivity; weight decay must be decoupled (Lesson 10).

Primary sources: Kingma & Ba 2015 (Adam); Reddi et al. 2018 (Adam non-convergence).

10 — AdamW, Schedules & Convergence

Fix Adam's weight decay, schedule the learning rate, and choose an optimizer deliberately.

AdamW: Decoupled Weight Decay

\( \ell_2 \) regularization \( \frac{\lambda}{2}\|\theta\|^2 \) naively adds \( \lambda\theta \) to \( g_t \); under Adam's adaptive denominator this gets rescaled incorrectly (strong regularization where \( v \) is large). Loshchilov & Hutter (2019) decouple it:

\[ \theta_{t+1} = \theta_t - \eta \frac{\hat m_t}{\sqrt{\hat v_t}+\varepsilon} - \eta \lambda \theta_t. \]

Decay is now applied after the adaptive step, independent of \( v_t \). In practice this is what weight_decay should mean; use AdamW for transformers/LLMs, not Adam+ \( \ell_2 \).

opt = torch.optim.AdamW(model.parameters(), lr=1e-4, betas=(0.9, 0.999), weight_decay=0.01)

Learning-Rate Schedules

When to Use What

Convergence Cheat Sheet

• Convex smooth: GD \(O(1/T)\), SGD \(O(1/\sqrt{T})\), Nesterov \(O(1/T^2)\) optimal.
• Non-convex: all first-order converge to stationary points at \( O(1/\sqrt{T}) \) in \(\min_t \|\nabla L\|^2\); noise/momentum help escape saddles.
• Effective step of Adam is \( \approx \eta \cdot \mathrm{sign}(m)/\sqrt{v} \) — clips extremes automatically.

Primary sources: Loshchilov & Hutter 2019 (AdamW); Loshchilov & Hutter SGDR; Smith 2017 (warmup).

11 — Practical Playbook

Translate theory into a reproducible training recipe.

Defaults that work: AdamW (\( \beta_1=0.9, \beta_2=0.999 \)), warmup 1–5% of steps, cosine decay to 10% of peak, weight decay 0.01 (transformers) / 1e-4 (CV), gradient clip 1.0, batch size 128–512 tuned to hardware.

Diagnostics: loss NaN → reduce \( \eta \), clip gradients; loss plateau → increase \( \eta \) or schedule; train \(\gg\) val → weak regularization / smaller batch; slow in valley → add momentum; diverged AdaGrad → switch to RMSProp/Adam.

Invariance check: multiply one layer's loss scale by 10; SGD diverges, Adam unchanged — proves second moment adaptation.

# Full recipe
model = MyModel()
opt = torch.optim.AdamW(model.parameters(), lr=3e-4, betas=(0.9, 0.999), weight_decay=0.01)
sched = torch.optim.lr_scheduler.CosineAnnealingLR(opt, T_max=epochs)
for epoch in range(epochs):
    for x, y in loader:
        opt.zero_grad()
        loss = criterion(model(x), y)
        loss.backward()
        torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
        opt.step()
    sched.step()
    print(f"epoch {epoch}: loss={loss.item():.4f} lr={sched.get_last_lr()[0]:.2e}")

Gotchas

• Weight decay placement: Adam(..., weight_decay=...) in old PyTorch is \( \ell_2 \), not AdamW — use AdamW.
• Epsilon matters: too large kills per-param adaptivity; too small causes \( \eta/\sqrt{v} \) blow-up on sparse params.
• Batch size vs LR: linear scaling holds to a point; beyond, keep LR and train longer or use LAMB/LARS.
• Seed the noise: SGD trajectory is random — fix seed for reproducibility, but test multiple seeds before claiming an optimizer "wins".

Try it yourself: train the same MLP twice — once with SGD+momentum (lr 0.05), once with AdamW (lr 3e-4) — and compare wall-clock to 95% accuracy and final validation gap.