Warmup-Stable-Decay Learning Rates and the River Valley Loss Landscape

Wen, Li, Wang, Hall, Liang, Ma · 2024 · "Understanding Warmup-Stable-Decay Learning Rates: A River Valley Loss Landscape Perspective"
arXiv 2410.05192 Interactive page

Schedules compared Illustrative shapes

Pick a schedule. Token axis marks the 50 / 100 / 200B budgets from the paper's setup; curve shapes are sketches.

Why the loss looks broken Illustrative

Observed loss = river progress + hill (bounce) term. Decay removes the bounce and reveals progress already made.

The canyon, simulated Illustrative 2-D toy

Noisy gradient descent on L = 1 − a·x + ½·h·y². x runs along the river, y across the walls. Noise acts only across the walls. Same noise sequence in every run.

Loss over time

Orange: total loss. Green dashed: river component only.

Trade-off readout

Final river progress vs mean bounce loss (last 50 steps).

How the theory is built

Click a stage. Stages 1–4 are proved for SGD under strong simplifications. All drawings are schematic.

Evidence

How much is actually shown? Editorial reading, not paper scores

Hover a cell. Warmer = stronger support for that claim from that kind of evidence, as discussed in the episode.
Hover a cell for notes.

SGD → Adam

What would settle it

References