Research Scientist
Optimization Theory & Training Dynamics
Training deep learning models is an optimization problem โ but not one that classical theory fully explains. You must understand why Adam converges faster than SGD, why SGD often generalizes better, what sharpness-aware minimization does, and how to debug training instabilities using optimization theory.
Core Theory
- 1Gradient descent convergence for L-smooth functions: if f is L-smooth (||โf(x) - โf(y)|| โค L||x-y||), then GD with step size 1/L converges at rate: f(x_T) - f* โค L||xโ-x*||ยฒ/(2T) = O(1/T). For ฮผ-strongly convex functions: linear convergence O((1-ฮผ/L)^T). Condition number ฮบ = L/ฮผ determines convergence speed. Adam approximately normalizes the condition number per parameter.
- 2SGD as implicit regularizer: the noise in SGD gradient estimates (from mini-batches) acts as a regularizer. Stochastic gradient noise prevents the optimizer from getting stuck in sharp minima โ it effectively explores the loss landscape. SGD converges to flatter minima than Adam, which often generalizes better. This is the sharp/flat minima hypothesis (Keskar et al. 2017, though contested).
- 3Adam optimizer: exponential moving average of first moment m_t = ฮฒโm_{t-1} + (1-ฮฒโ)g_t and second moment v_t = ฮฒโv_{t-1} + (1-ฮฒโ)g_tยฒ. Bias correction: mฬ = m/(1-ฮฒโ^t), vฬ = v/(1-ฮฒโ^t) โ critical in early training when moments are underestimated. Update: ฮธ -= lr ร mฬ/(โvฬ + ฮต). The per-parameter adaptive learning rate โvฬ effectively normalizes the curvature per parameter.
- 4AdamW vs Adam: L2 regularization added to the gradient (standard Adam with weight_decay) effectively reduces the adaptive learning rate for frequently updated parameters โ large gradient parameters get less weight decay. AdamW decouples weight decay: ฮธ -= lr ร (mฬ/(โvฬ + ฮต) + ฮปฮธ). This correctly applies weight decay regardless of gradient history. All modern LLM training uses AdamW.
- 5Learning rate schedules: (1) linear warmup โ critical for transformers. Without warmup, large initial gradients destabilize early training. Typical: linear increase over 100-2000 steps. (2) Cosine annealing โ gradually decrease LR to enable fine-grained convergence. LR_t = LR_min + (LR_max - LR_min) ร (1 + cos(ฯ ร t/T)) / 2. (3) Linear decay โ simpler than cosine, used in LLM training. (4) One-cycle policy (Super-Convergence) โ aggressive warmup then cosine decay, enables faster convergence with larger max LR.
- 6Sharpness-Aware Minimization (SAM): instead of minimizing f(ฮธ), minimize max_{||ฮต||โคฯ} f(ฮธ+ฮต) โ the loss in the worst perturbation neighborhood. This finds parameters in flat loss landscape regions that generalize well. Two-step: (1) compute ฮต* = ฯ ร โf/||โf||, (2) compute gradient at ฮธ+ฮต* and apply to ฮธ. SAM adds roughly 2ร compute per step but improves generalization by 1-3% on ImageNet.
- 7Gradient flow analysis: in deep networks, the gradient magnitude should be roughly constant through depth. ReLU dead neurons: if a neuron's pre-activation is always negative, its gradient is always 0 โ the neuron is permanently dead. Detection: monitor the fraction of activations that are positive at each layer. Fix: use Leaky ReLU (gradient=ฮฑ for negative input), initialize weights so pre-activations are positive, use batch normalization.
- 8Loss landscape visualization (Li et al. 2018): project the loss onto 2D by perturbing weights in two random directions with filter normalization (normalize per filter to remove scale invariance). Sharp minima: loss landscape peaks quickly. Flat minima: loss changes slowly with perturbation. Empirically: SGD finds flatter minima than Adam, which may explain generalization differences in some settings.
Key Formulas
| 1 | import numpy as np |
| 2 | import matplotlib |
| 3 | matplotlib.use('Agg') |
| 4 | import matplotlib.pyplot as plt |
| 5 | |
| 6 | np.random.seed(42) |
| 7 | |
| 8 | # โโโ 2D ill-conditioned quadratic loss โโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโ |
| 9 | # L(x,y) = 100*(x-1)^2 + (y-1)^2 (condition number = 100) |
| 10 | # True minimum at (1, 1) |
| 11 | |
| 12 | def loss(theta): |
| 13 | return 100 * (theta[0] - 1)**2 + (theta[1] - 1)**2 |
| 14 | |
| 15 | def grad_full(theta): |
| 16 | return np.array([200*(theta[0]-1), 2*(theta[1]-1)]) |
| 17 | |
| 18 | def grad_stochastic(theta, noise_std=0.5): |
| 19 | """Add gradient noise to simulate mini-batch stochasticity.""" |
| 20 | return grad_full(theta) + np.random.randn(2) * noise_std |
| 21 | |
| 22 | # โโโ Vanilla GD โโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโ |
| 23 | def gradient_descent(lr=0.005, n_steps=500): |
| 24 | theta = np.array([0.0, 0.0]) |
| 25 | losses = [] |
| 26 | for _ in range(n_steps): |
| 27 | g = grad_full(theta) |
| 28 | theta -= lr * g |
| 29 | losses.append(loss(theta)) |
| 30 | return theta, losses |
| 31 | |
| 32 | # โโโ SGD with momentum โโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโ |
| 33 | def sgd_momentum(lr=0.01, momentum=0.9, n_steps=500): |
| 34 | theta = np.array([0.0, 0.0]) |
| 35 | v = np.zeros(2) |
| 36 | losses = [] |
| 37 | for _ in range(n_steps): |
| 38 | g = grad_stochastic(theta) |
| 39 | v = momentum * v - lr * g |
| 40 | theta += v |
| 41 | losses.append(loss(theta)) |
| 42 | return theta, losses |
| 43 | |
| 44 | # โโโ Adam โโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโ |
| 45 | def adam(lr=0.1, beta1=0.9, beta2=0.999, eps=1e-8, n_steps=500): |
| 46 | theta = np.array([0.0, 0.0]) |
| 47 | m = np.zeros(2) |
| 48 | v = np.zeros(2) |
| 49 | losses = [] |
| 50 | for t in range(1, n_steps + 1): |
| 51 | g = grad_stochastic(theta) |
| 52 | m = beta1 * m + (1 - beta1) * g |
| 53 | v = beta2 * v + (1 - beta2) * g**2 |
| 54 | m_hat = m / (1 - beta1**t) # bias correction |
| 55 | v_hat = v / (1 - beta2**t) |
| 56 | theta -= lr * m_hat / (np.sqrt(v_hat) + eps) |
| 57 | losses.append(loss(theta)) |
| 58 | return theta, losses |
| 59 | |
| 60 | # โโโ SAM (Sharpness-Aware Minimization) โโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโ |
| 61 | def sam_step(theta, rho=0.05): |
| 62 | """ |
| 63 | SAM two-step update: |
| 64 | 1. Compute gradient at theta |
| 65 | 2. Perturb to worst-case neighborhood: theta_hat = theta + rho * grad / ||grad|| |
| 66 | 3. Compute gradient at theta_hat |
| 67 | 4. Apply this gradient to original theta |
| 68 | """ |
| 69 | g1 = grad_stochastic(theta) |
| 70 | eps = rho * g1 / (np.linalg.norm(g1) + 1e-12) # perturbation |
| 71 | theta_hat = theta + eps # perturbed parameters |
| 72 | g2 = grad_stochastic(theta_hat) # gradient at perturbed point |
| 73 | return g2 # apply this gradient to original theta |
| 74 | |
| 75 | def sam_optimizer(lr=0.01, rho=0.05, n_steps=500): |
| 76 | theta = np.array([0.0, 0.0]) |
| 77 | losses = [] |
| 78 | for _ in range(n_steps): |
| 79 | g = sam_step(theta, rho) |
| 80 | theta -= lr * g |
| 81 | losses.append(loss(theta)) |
| 82 | return theta, losses |
| 83 | |
| 84 | # โโโ Compare โโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโ |
| 85 | results = {} |
| 86 | for name, fn in [("GD", gradient_descent), |
| 87 | ("SGD+Momentum", sgd_momentum), |
| 88 | ("Adam", adam), |
| 89 | ("SAM", sam_optimizer)]: |
| 90 | final_theta, losses = fn() |
| 91 | results[name] = {"final": final_theta, "losses": losses, |
| 92 | "final_loss": losses[-1]} |
| 93 | print(f"{name:<15}: final ฮธ={final_theta.round(4)}, " |
| 94 | f"loss={losses[-1]:.6f} (best: 0.0)") |
| 95 | |
| 96 | # Learning rate warmup + cosine decay schedule |
| 97 | def cosine_warmup_schedule(step: int, warmup_steps: int, |
| 98 | total_steps: int, lr_max: float, |
| 99 | lr_min: float = 1e-5) -> float: |
| 100 | if step < warmup_steps: |
| 101 | return lr_max * step / warmup_steps |
| 102 | progress = (step - warmup_steps) / (total_steps - warmup_steps) |
| 103 | return lr_min + 0.5 * (lr_max - lr_min) * (1 + np.cos(np.pi * progress)) |
| 104 | |
| 105 | # Example schedule |
| 106 | schedule = [cosine_warmup_schedule(t, 100, 1000, 3e-4) for t in range(1000)] |
| 107 | print(f"\nLR schedule: warmup 100 steps โ cosine decay") |
| 108 | print(f" Step 0: {schedule[0]:.2e}") |
| 109 | print(f" Step 50: {schedule[50]:.2e}") |
| 110 | print(f" Step 100: {schedule[100]:.2e} (peak)") |
| 111 | print(f" Step 500: {schedule[500]:.2e}") |
| 112 | print(f" Step 999: {schedule[999]:.2e} (final)") |
Compares GD, SGD+momentum, Adam, and SAM on a deliberately ill-conditioned 2D quadratic (condition number=100). Adam converges fastest due to per-parameter adaptive learning rates that normalize the curvature. SAM's two-step update targets flat minima. The warmup+cosine schedule shows how LR evolves during modern LLM training: linear ramp for stability, then cosine decay for convergence.
Worked Interview Problems
3 problemsProblem
Explain the empirical observation that SGD trained for longer often achieves better test accuracy than Adam, despite Adam's faster convergence.
Solution
Sharp vs flat minima hypothesis: there are many minima with near-zero training loss. Some are 'sharp' (loss increases quickly around them), some are 'flat' (loss changes slowly). Flat minima tend to generalize better because small perturbations to weights (analogous to distribution shift) don't change the loss much.
SGD's implicit bias: stochastic gradient noise effectively prevents the optimizer from settling into narrow, sharp minima. The noise bounces the optimizer out of sharp minima but keeps it in flat ones. This creates an implicit regularization effect toward flat minima.
Adam's behavior: Adam's adaptive learning rate is small for parameters with large gradient history (frequently updated). This allows precise convergence into narrow sharp minima. The reduced effective noise means Adam can find sharper, lower-training-loss solutions that generalize worse.
Empirical evidence (Keskar et al. 2017): explicitly found that large-batch SGD finds sharper minima and generalizes worse than small-batch SGD. This supports the noise-as-regularizer hypothesis for SGD.
Counterarguments: Wilson et al. (2017) showed that Adam with proper tuning achieves similar generalization to SGD on image classification. The 'SGD generalizes better' finding may be more about hyperparameter sensitivity than an inherent property. In practice, many tasks (NLP, LLMs) show no clear winner between Adam and SGD โ the choice depends on architecture and task.
Answer
SGD's mini-batch noise acts as implicit regularizer, pushing toward flat minima that generalize better. Adam's per-parameter adaptation allows precise convergence to sharp minima with lower training loss but potentially worse generalization. Effect is task-dependent โ LLMs are trained with AdamW, image classifiers sometimes benefit from SGD+momentum.
Common Mistakes
Confusing Adam with AdamW. Adam with weight_decay reduces the effective weight decay for parameters with high gradient variance (because the adaptive LR is small for these). AdamW fixes this by decoupling weight decay from the gradient update. All modern LLM training uses AdamW โ saying 'we use Adam' for an LLM training question is a red flag.
Not knowing the bias correction in Adam. In the first step, without bias correction, mฬ = m / (1 - ฮฒโยน) = gโ / (1-0.9) = 10 ร gโ. With bias correction: mฬ = m / (1-0.9^1) = gโ โ the correct estimate. Without bias correction, Adam makes enormous updates in early training.
Saying SGD always generalizes better than Adam. This is a common oversimplification. For NLP/LLM tasks, Adam consistently wins. For vision tasks, SGD sometimes matches or beats Adam. The right answer: 'It's task-dependent. For transformers I'd start with AdamW.'
Not knowing that SAM requires 2ร compute per step. SAM's two-step process (compute gradient at ฮธ, perturb to ฮธ+ฮต, compute gradient at ฮธ+ฮต, apply to ฮธ) requires two forward+backward passes per update. This doubles compute but improves generalization by ~1-3%. This tradeoff should always be mentioned when discussing SAM.
Confusing the warmup length heuristic. Warmup length is NOT a fixed number โ it depends on total training steps, model size, and LR. Rule: warmup โ 1-5% of total steps, or until Adam's moment estimates have stabilized (roughly ฮฒโ^{-1/(1-ฮฒโ)} = 1000 steps for ฮฒโ=0.999). Large models with aggressive LR need longer warmup.