Optimization and math foundations for ML research

Why it matters

Many research-engineering failures are optimization failures in disguise. An idea "does not work" because the learning rate was tuned for the baseline only, the warmup was missing, or the comparison stopped before one run had decayed its learning rate. Research interviewers therefore ask you to reason from first principles: why a step size diverges, what batch size does to gradient noise, why AdamW is not Adam with L2. You do not need convergence proofs for deep networks (few exist), but you need the convex intuitions that explain most of what you observe, and you need to know which claims in the literature are empirical regularities rather than theorems.

Core ideas

The quadratic picture. Near a minimum a smooth loss looks like a quadratic with Hessian H. Gradient descent multiplies the error along each eigenvector by (1 − η·λᵢ). It converges only if every factor has magnitude below 1, so η must stay below 2/λ_max, while the flattest direction (λ_min) sets the speed. The condition number κ = λ_max/λ_min is why plain descent crawls along narrow valleys. Momentum accumulates velocity along consistently signed directions and damps oscillation across the valley; on quadratics, well-tuned momentum needs on the order of √κ rather than κ iterations.

Convexity as a lens. For convex, L-smooth functions a step of 1/L guarantees descent and every local minimum is global. Networks are non-convex, yet the same local reasoning explains step-size limits, why normalization layers help (better conditioning), and why a learning rate that is stable early can become unstable when curvature grows during training.

Stochastic gradients. A minibatch gradient is an unbiased estimate whose variance falls roughly as 1/B. Doubling B halves the noise but also halves the number of updates per epoch. Up to a problem-dependent critical batch size, larger batches with proportionally larger learning rates train about as well per example; beyond it, returns diminish. The linear scaling rule with warmup is an empirical recipe for that regime. Noise also acts as a regularizer; the claim that small batches find flatter minima that generalize better is debated, partly because flatness depends on how the network is parameterized.

Adaptive methods. Adam keeps running estimates of each coordinate's first and second gradient moments, with bias correction for the early steps, and divides by the root of the second moment, making steps roughly invariant to gradient scale. An L2 term added to the loss is then rescaled by that denominator, so heavily updated weights are regularized less; AdamW applies decay directly to the weights instead.

Schedules. Warmup is usually explained by unreliable early moment estimates and high initial curvature. Decay (cosine, linear, step, inverse square root) lowers the noise floor late in training, which is why loss often drops sharply when the learning rate falls. Consequence for experiments: compare runs at the end of their own schedules, not at an intermediate step of a longer schedule.

Scaling laws (awareness). Empirical fits show loss falling roughly as a power law in parameters, data and compute over several orders of magnitude. Compute-optimal analyses ask how to split a fixed budget between model size and tokens, and different studies reached different answers partly because of schedule and fitting choices. Treat an exponent as a measurement with uncertainty, valid over the range where it was fitted.

Common misconceptions

• "Adam removes the need to tune the learning rate." It normalizes per-coordinate scale, but the global step size and warmup still dominate stability and final loss.
• "L2 regularization and weight decay are the same thing." They coincide for plain SGD, not for adaptive optimizers, which is the motivation for AdamW.
• "A bigger batch is always faster." Past the critical batch size each extra example adds little; wall-clock gains then come only from better hardware utilization.
• "Non-convex means convex theory is useless." Local quadratic models explain step-size limits, conditioning and the benefits of momentum and normalization remarkably well.
• "The curve plateaued, so training converged." A plateau often reflects the noise floor of the current learning rate; decaying it frequently reveals further progress.
• "Scaling laws predict any model's loss." They are fits for one data distribution, architecture family and recipe; extrapolation is a hypothesis to test.

How to approach it

Rebuild the quadratic picture first: run gradient descent and momentum on an ill-conditioned 2-D quadratic and watch divergence just above 2/λ_max. Then train a small network with SGD and AdamW under two schedules across several seeds while logging gradient norms. At graduate level, interviewers check that you can explain the update rules and learning-rate tradeoffs; at mid and senior level, expect "your run diverged after a few thousand steps, what do you check?" and "how would you compare two optimizers fairly?", which links to experimental controls.