Optimization: SGD to Adam, and What the Loss Landscape Actually Looks Like

The loss landscape of a neural network is not a bowl with one minimum. It is a high-dimensional surface with flat plateaus, narrow valleys, and saddle points everywhere. The optimizer is what navigates this landscape. SGD, momentum, RMSProp, and Adam each solve a different failure mode of gradient descent. Here is the geometry behind each one.

Gradient descent is simple: compute the gradient, move opposite to it, repeat. The problem is that simple gradient descent fails in almost every interesting neural network loss landscape. Understanding why it fails — and how each optimizer variant fixes it — is more useful than memorising hyperparameter defaults.

What the loss landscape looks like

A neural network with millions of parameters has a loss surface in a space with millions of dimensions. You cannot visualise it, but you can characterise it. Two properties dominate the difficulty of training.

Ill-conditioning: the loss surface has very different curvatures in different directions. In some directions the loss falls rapidly (high curvature), in others it barely changes (low curvature). SGD with a learning rate tuned for the high-curvature directions is too slow in the low-curvature directions. A learning rate tuned for the low-curvature directions causes oscillation in the high-curvature directions.

Saddle points: in high dimensions, a "local minimum" in the traditional sense is rare. Almost every critical point (gradient ≈ 0) is a saddle point — a minimum in some directions and a maximum in others. Gradient descent slows near saddle points because the gradient is small even though the point is not optimal. Plateau regions have the same effect.

SGD: the baseline

w ← w - η * ∇L(w). Simple, well-understood, but sensitive to learning rate and slow on ill-conditioned landscapes. With a large learning rate it overshoots; with a small one it makes negligible progress in low-curvature directions.

Momentum: accumulate velocity

w ← w - v, where v ← β*v + η*∇L(w). Instead of moving in the direction of the current gradient, momentum moves in the direction of the exponentially weighted average of past gradients. In low-curvature directions where gradients are small and consistent, velocity accumulates — the optimizer moves faster. In oscillating dimensions, gradients cancel out — the optimizer moves slower. Momentum implicitly adapts to the local geometry.

SGD + momentum is still widely used for image classification (ResNets, ViTs). Its generalisation properties can exceed Adam because the flat minima it finds are more robust to distribution shift.

RMSProp: per-parameter learning rates

s ← β*s + (1-β)*∇L^2; w ← w - η * ∇L / sqrt(s + ε). RMSProp maintains a running average of squared gradients per parameter. Parameters with large historical gradients get scaled down; parameters with small historical gradients get scaled up. This per-parameter adaptation directly addresses ill-conditioning: the learning rate is automatically adjusted for each dimension's curvature.

Adam: momentum + RMSProp

Adam (Kingma & Ba, 2015) combines both ideas. First moment (like momentum): m ← β1*m + (1-β1)*∇L. Second moment (like RMSProp): v ← β2*v + (1-β2)*∇L^2. Bias-corrected update: w ← w - η * m_hat / (sqrt(v_hat) + ε). The bias correction (dividing by 1-β^t) compensates for the initialisation bias when both moments start at zero.

Adam is robust to the learning rate, adapts per parameter, and converges quickly. It is the default for training Transformers and most modern deep learning architectures.

When SGD beats Adam

Adam can converge to sharper minima than SGD. Sharp minima are sensitive to small input perturbations — the model generalises less well. SGD with momentum finds flatter minima that are more robust. For image classification on ImageNet-scale data, fine-tuned SGD + momentum often beats Adam on final test accuracy by 1-2%. For NLP and Transformers, Adam is usually better because the loss landscape is more ill-conditioned.

Learning rate schedules: the most important hyperparameter

The absolute learning rate matters less than the schedule. Common patterns: cosine annealing (lr decays following a cosine curve, optionally with warm restarts), linear warmup followed by cosine decay (standard for Transformers — warm up for ~4% of training steps, then decay), and one-cycle policy (lr rises to max then falls, often with momentum inversely varying). Warmup is important for Transformers because Adam's second moment estimate is unreliable in the first few steps — a high initial learning rate with an unreliable normaliser causes divergence.

Try on Colab: train a small MLP on CIFAR-10 with SGD (no momentum), SGD + momentum, RMSProp, and Adam. Log the loss curve and final accuracy for each. Then visualise the loss landscape around the final solution for SGD and Adam using random direction projection — the flatness of Adam's minimum vs SGD's is often visible.