Learning Rate Warmup › The Loss Landscape Perspective

To build deeper intuition for warmup, it helps to think about the loss landscape directly. This high-dimensional function maps every possible set of model weights to a loss value. Training moves the weights from a random starting point toward a region of low loss, typically a flat basin that generalizes well to unseen data.

At the start of training, the model's parameters are in a region of the loss landscape that has never been visited before. Random initialization places you somewhere on a roughly random surface. The local gradient at that starting point tells you the direction of steepest descent, but in a high-dimensional, non-convex landscape, following the steepest local gradient blindly is not a reliable strategy. Steepest descent may lead you into a narrow valley with sharp walls (a so-called sharp minimum), or it may take you to a saddle point where the gradient appears to be zero in some directions but the curvature is positive, making it a trap.

Large learning rates at the beginning of training are dangerous precisely because the curvature information at the starting point is unreliable. The loss landscape near a random initialization is not like the landscape near a well-trained point. Near a random initialization, many loss directions are steep and many others are nearly flat, creating a jagged, highly anisotropic surface. An optimizer that takes large steps in this terrain is likely to skip over useful regions and land in areas of high loss.

Research on the connection between learning rate and generalization has revealed an important insight: large learning rates, when applied during the main training phase (after warmup), tend to find flatter minima than small learning rates. Flat minima generalize better because they are robust to small perturbations in the weight space. Sharp minima, by contrast, have a narrow basin of attraction, and even small changes in the weights can push the model into high-loss territory.

This is one of the reasons why warmup is designed to be followed by a relatively high peak learning rate, not a conservative one. The warmup phase gets you to a region where the gradient signal is reliable and the optimizer can make informed decisions. Then the high peak learning rate allows the optimizer to explore broadly, avoiding sharp minima and finding flatter, more generalizable solutions.

The warmup-then-decay pattern divides optimization into distinct phases. During warmup, small steps move the parameters into a region where gradient estimates are more reliable. At peak learning rate, larger steps search for a good basin. During cosine decay, progressively smaller steps refine the parameters within that basin. This sequence matches the empirical observation that models trained with warmup-plus-decay consistently outperform models trained with a fixed learning rate throughout.

There is a deeper mathematical reason why warmup improves optimization on the loss landscape. Adam's second-moment estimate $v_t$ is an approximation to the diagonal of the gradient covariance matrix. This covariance matrix describes how much the gradient varies in each direction in parameter space. Parameters with high gradient variance are uncertain, and Adam compensates by taking smaller steps along those dimensions. Parameters with low gradient variance are reliable, and Adam can take larger steps.

At initialization, the gradient covariance is essentially unknown. The first few gradient samples are dominated by the noise of random initialization, not by the true curvature structure of the loss landscape. As training proceeds through the warmup phase, the optimizer accumulates gradient samples and builds up a picture of the curvature. By the time warmup ends and the learning rate reaches its peak, the optimizer has a reasonably accurate model of which directions in parameter space are risky (high curvature) and which are safe (low curvature). This is exactly the kind of terrain information you want before making large steps.

The warmup duration, viewed through this lens, is the minimum number of gradient samples needed to estimate the curvature structure reliably enough to take large steps safely. This connects back to the earlier derivation: the threshold of $t > 1 / (1 - \beta_2)$ is the number of samples after which the exponential moving average of squared gradients has a coefficient of variation below 10%, meaning it is a reasonably accurate estimate of the true gradient variance in each direction.