What do near-optimal learning rate schedules look like?

TL;DR

Naganuma, Agarwala, Kasimbeg & Dahl run a brute-force random search over parameterized LR schedule shape families (Constant, Cosine, Generalized Cosine, Sqrt-decay, Generalized REX, Two-Point Spline, Two-Point Linear, and a fully-free snm family) on three small workloads — synthetic linear regression, CIFAR-10 CNN, and a 8M-param WikiText-103 Transformer — running tens of thousands of training runs to find empirically near-optimal shapes. Two robust findings: (a) once a schedule has warmup + monotonic decay, the base learning rate dominates the shape in predicting final loss; (b) the unconstrained snm family rediscovers warmup + monotonic decay from scratch on deep-learning workloads, even though neither is built in. The linear-regression ground-truth optimal has no warmup and sharp end-decay, so the warmup-and-gradual-decay pattern is specifically a property of deep-learning workloads, not optimization-in-general. AdamW weight decay strongly affects optimal shape; β values affect it less.

Key claims

Across all parameterized schedule families on CIFAR-10 and WikiText-103, the empirically near-optimal shape has a non-trivial linear warmup (5–25% of training) followed by gradual decay; the fully-unconstrained snm family, which can express any shape including flat or non-monotonic, also discovers this pattern via random search [§4.2, Fig. 5].
On linear regression with prescribed uniform covariance and SGD, the closed-form optimal LR schedule has no warmup and ends with sharp decay — qualitatively different from the deep-learning optima — and is reached by an LR very close to the edge-of-stability [§4.1, Fig. 2, Appendix A.4].
Given warmup + decay, the base learning rate is a far stronger predictor of final loss than the choice of shape family; extra hyperparameter budget is better spent fine-tuning base LR than tuning detailed shape parameters [§4.2, Fig. 3].
Flexible families (Two-Point Spline, Sqrt-decay, Generalized Cosine, Two-Point Linear) yield small but measurable wins over the standard fixed-exponent Cosine on CIFAR-10 (median train error 14.6%–14.9% vs 15.4% for Cosine) and on WikiText-103 (smaller median perplexities, but with larger confidence intervals) [§4.2, Fig. 5].
The snm family is the one schedule family the search does not adequately explore: top-snm shapes show high within-family variation and a one-coordinate line-search reveals exploitable directions, while all other families pass the line-search probe [§4.3, Fig. 8].
The search procedure recovers schedules close to the analytically-derived linear-regression optimum (within ~2% of optimal loss when projected into the snm family), validating the random-search-then-reseed-rank protocol [§4.1, Fig. 4].
Among AdamW hyperparameters, weight decay has a strong effect on the optimal LR schedule shape; β₁/β₂ have weaker effects [§4.4].

Method

Each schedule family is parameterized so the shape (a function from [0,1] to [0,1]) is decoupled from the base LR. The search procedure samples ~10⁴–10⁵ shapes per family from prior distributions, scores each shape across a logarithmically-spaced grid of 16 base LRs, and ranks by median final training metric across 4–16 seeds. The top ~50 shapes per family are then re-evaluated with 64 seeds (all pairwise combos of unique inits × data orderings) using DKW confidence intervals to give an unbiased ranking. The shape families are: con (warmup-then-constant), cos-std (standard cosine), cos-gen (cosine with tunable exponent), sqrt (inverse-sqrt decay), rex (REX-parametrized decay), tps/tpl (two-point spline/linear interpolation), and snm (a fully-free two-control-point spline with no warmup-or-decay constraints baked in). Workloads use small models in an optimization-limited regime — 1000–1600 steps, batch sizes 32–512 — chosen so constant-LR runs visibly underperform and shape effects are detectable.

Results

Linear regression (1000 steps, SGD, uniform-spectrum covariance): the analytically optimal schedule has no warmup and a sharp end-decay; flexible families like snm reach within ~2% of optimal loss in the eyeballed selection but slightly worse than analytical optimum after multi-seed re-evaluation [Figs. 2, 4].
CIFAR-10 CNN (1000 steps, AdamW, batch 256): best Cosine median train error 15.4%; best tps/sqrt/cos-gen/tpl reach 14.6%–14.9%; best Constant is much worse, confirming warmup + decay are load-bearing. Optimal warmup fraction is 5–25% of training, with gradual (not sharp) decay [Fig. 5 left].
WikiText-103 Transformer (1600 steps, AdamW, batch 512): same qualitative pattern; flexible families beat Cosine in median perplexity but with larger confidence intervals due to some near-optimal runs becoming unstable [Fig. 5 right]. Shapes are visibly different from Cosine — slower initial decay and a less-sharp tail.
AdamW hyperparameter dependence: weight decay materially shifts the optimal shape; β₁, β₂ shift it less [§4.4].

Why it’s interesting

Cleanly empirical answer to a question every pretraining engineer has wondered: how much does the LR-schedule shape matter beyond “warmup + cosine”? The Naganuma et al. answer — “less than the base LR, but warmup-and-gradual-decay isn’t arbitrary; even a fully-free family rediscovers it” — directly substantiates a piece of folklore that the Hyperparameter scaling laws page has been circling: How to Set the Learning Rate for Large-Scale Pre-training? fits LR*(N, D) under a WSD schedule but takes the WSD shape as given; this paper provides the missing justification for why the warmup-and-decay shape is a sensible inductive bias to fit on top of in the first place. Complements Why Gradients Rapidly Increase Near the End of Training, which explains why end-decay helps mechanistically (weight-decay × LR-decay × Adam interaction), with a complementary empirical “and the shape that minimizes this is workload-dependent, with weight decay as the dominant covariate.” Also contrasts with Seesaw: Accelerating Training by Balancing Learning Rate and Batch Size Scheduling, which argues a different schedule (cosine interleaved with batch-size doubling) is strictly better than vanilla cosine — open whether the Seesaw shape would also be discovered by snm random search if the batch-size axis were exposed to the search.