Learning Rate Decay › Limitations and Practical Considerations

Learning rate scheduling is powerful but introduces additional hyperparameters that must be set correctly. A poorly chosen schedule can be worse than no schedule at all. If the learning rate decays too quickly, the model gets stuck in a suboptimal region before it has had time to find a good basin. If it decays too slowly, the final convergence is poor and the training may appear to have plateaued.

Each schedule introduces at least one and often several additional hyperparameters beyond the base learning rate. Step decay needs both the step size and the gamma. Cosine annealing needs the minimum learning rate and the total period length. The inverse square root schedule needs the warmup duration. All of these interact with the peak learning rate in non-trivial ways.

In practice, these hyperparameters are often set based on rules of thumb derived from prior experiments rather than being tuned from scratch. This is a practical necessity: running a full hyperparameter sweep over schedule parameters for every new task is computationally infeasible. The standard recipes that have emerged from community experience (like linear warmup to 10% of steps followed by cosine decay) are valuable precisely because they generalize across many settings without requiring explicit tuning.

The interaction between schedules and adaptive optimizers like Adam is not fully understood. Adam's internal adaptation already provides a form of per-parameter learning rate control, and layering a global schedule on top creates a complex combined behavior. Some researchers have argued that with properly tuned Adam, the sensitivity to the global learning rate schedule is lower than with SGD, though in practice schedules still affect performance.

One specific concern is that Adam's second moment estimate $v_t$ accumulates history according to its decay factor $\beta_2$. When the learning rate is decayed and then kept constant for a long time (as in step decay), the effective step size changes both because of the explicit decay and because $v_t$ continues to evolve. The combined effect is not easy to predict analytically. This is part of why cosine and linear schedules, which change the learning rate smoothly and predictably, tend to behave more reliably with Adam than step decay does.

Warmup deserves particular attention in distributed training. When training with large batch sizes across many GPUs, the effective batch size is the product of the per-GPU batch size and the number of GPUs. The linear scaling rule suggests scaling the learning rate proportionally to the batch size, which can push the learning rate very high. Warmup becomes critical in this regime: starting at the full scaled learning rate with a large batch and random initialization can cause catastrophic early instability. Longer warmup periods are used for very large distributed training runs.

The challenge is that the standard warmup durations (4,000 to 10,000 steps) were designed for moderate batch sizes. When batch size is scaled by a factor of 64x for distributed training, the effective number of samples seen per step increases dramatically, and the notion of "a step" in the warmup duration becomes less meaningful. Some large-scale training recipes specify warmup in terms of tokens or samples processed rather than steps, which remains invariant to batch size scaling.

Cycle-based schedules like cosine annealing with warm restarts have an appealing theoretical motivation but require careful setting of the cycle length. A cycle that is too short means the model restarts before it has converged within each basin, wasting each cycle. A cycle that is too long is effectively just cosine annealing without restarts. In practice, the cycle length is often tuned as a hyperparameter or set based on prior experience with similar models.

The restart mechanism can also cause difficulties for checkpoint selection. With a standard monotone schedule, the best checkpoint is usually the final one. With warm restarts, the best checkpoint might be at the end of any cycle. You need to track cycle-end checkpoints explicitly and compare them, rather than simply saving the final model. This adds infrastructure complexity to training pipelines.

The relationship between schedules and generalization is subtle. It is well established empirically that decaying the learning rate to a small value at the end of training improves generalization, not just training loss. The intuition is that a small learning rate at convergence means the model has found a region where the loss surface is locally flat, and flat minima tend to generalize better. But this connection is not a tight theoretical guarantee, and counter-examples exist.

More precisely, the "flat minima generalize better" hypothesis has been the subject of ongoing debate. Some theoretical work argues that the measure of flatness depends on the parameterization and can be an artifact of scale. Others have shown empirically that the correlation between flatness and generalization holds across a wide range of architectures and tasks. In practice, ending training at a small learning rate reliably improves results, even if the theoretical mechanism is not fully settled.