1 00:00:01,000 --> 00:00:08,244 [Hal Turing] Alrighty! Thanks for tuning in! Hello AI world! I am your host, Hal Turing, and my co-host is Dr. Ada Shannon. 2 00:00:08,244 --> 00:00:41,867 [Hal Turing] Today we're looking at "Understanding Warmup-Stable-Decay Learning Rates: A River Valley Loss Landscape Perspective." It's by Kaiyue Wen and colleagues Zhiyuan Li, Jason Wang, David Hall, Percy Liang and Tengyu Ma, from Stanford and TTIC, posted to arXiv in October 2024. The puzzle is a loss curve that looks broken. Training loss sits high for most of the run, then falls off a cliff in a short final stretch, and the result still matches or beats the standard schedule. 3 00:00:41,867 --> 00:00:57,378 [Dr. Ada Shannon] And the paper tries to explain that cliff, then uses the explanation to build a schedule that gives you checkpoints for several compute budgets from one run. The stakes are practical. If the explanation is right, that high-looking loss is hidden progress, not stalled training. 4 00:00:57,378 --> 00:01:01,743 [Hal Turing] Let's start with the default. What do most people use, and what's wrong with it? 5 00:01:01,743 --> 00:01:23,523 [Dr. Ada Shannon] Linear warmup, then cosine decay, usually down to about ten percent of the peak rate. The cosine idea comes from the SGDR paper by Loshchilov and Hutter, 2017. The catch is that the curve is stretched to end exactly at a total step count you pick up front. If you later get more compute, the schedule isn't a continuation. You start over. 6 00:01:23,523 --> 00:01:26,588 [Hal Turing] That sounds like an annoyance until you try to fit a scaling law. 7 00:01:26,588 --> 00:01:39,963 [Dr. Ada Shannon] Right. The Chinchilla paper, Hoffmann and colleagues at DeepMind, 2022, found that the schedule has to match the training length to get a fair loss for each budget. So every point on a scaling-law plot needs its own full run. 8 00:01:39,963 --> 00:01:41,356 [Hal Turing] So what's the alternative? 9 00:01:41,356 --> 00:02:05,644 [Dr. Ada Shannon] Warmup-Stable-Decay, or WSD. It was introduced in MiniCPM by Shengding Hu and colleagues, out of OpenBMB, in 2024. After warmup, the learning rate stays constant on a main branch that can run as long as you like. When you want a model for a given budget, you fork a copy and decay the rate quickly. The main branch resumes from the pre-decay checkpoint. One main run gives you many budget-specific models. 10 00:02:05,644 --> 00:02:07,084 [Hal Turing] Was that one team's idea? 11 00:02:07,084 --> 00:02:37,456 [Dr. Ada Shannon] Several groups were circling it. Hägele and colleagues in 2024 wrote Scaling Laws and Compute-Optimal Training Beyond Fixed Training Durations. Ibrahim and colleagues, also 2024, wrote Simple and Scalable Strategies to Continually Pre-train Large Language Models. Tissue and colleagues, 2024, fit a loss-curve scaling law under learning-rate annealing. The Ibrahim paper matters here because it blames the damage from rewarming on the sudden jump in rate. This paper challenges that explanation. 12 00:02:37,456 --> 00:02:41,449 [Hal Turing] Back to the weird curve. Why is the loss high during the stable phase? 13 00:02:41,449 --> 00:03:03,601 [Dr. Ada Shannon] That's the river valley conjecture. Picture a long canyon with steep walls and a gently sloping floor. Across the canyon, curvature is huge. A small move up the wall raises the loss a lot. Along the canyon floor, the river, curvature is tiny but there's a persistent downhill slope. Almost all the long-run progress is along the river. 14 00:03:03,601 --> 00:03:06,527 [Hal Turing] And the optimizer is a ball rolling down this thing with noise. 15 00:03:06,527 --> 00:03:21,806 [Dr. Ada Shannon] A ball taking big noisy steps. Because the steps are big, it can't settle on the floor. It bounces between the walls. The paper splits the loss into two parts. The river component is true progress along the floor. The hill component is the extra loss from bouncing. 16 00:03:21,806 --> 00:03:23,710 [Hal Turing] So what sets how high the bounce goes? 17 00:03:23,710 --> 00:03:37,781 [Dr. Ada Shannon] The step size. For SGD the paper derives a hill term that is linear in the current learning rate. Bounce height scales with step size. Halve the rate and you roughly halve the extra loss, but you also slow your travel along the river. 18 00:03:37,781 --> 00:03:44,979 [Hal Turing] So there's a trade-off. A big rate covers river distance fast but wastes loss on bouncing. A small rate is clean but slow. 19 00:03:44,979 --> 00:03:59,933 [Dr. Ada Shannon] Yes. And the stable phase takes the first option. Loss reads high because of the bounce, yet the position along the river keeps advancing. Then decay shrinks the bounce. The loss drops sharply, and what you see is progress that was already banked. 20 00:03:59,933 --> 00:04:06,481 [Hal Turing] That sounds close to the old convex story. Constant rate gives you a noise ball, and you decay to shrink it. 21 00:04:06,481 --> 00:04:31,326 [Dr. Ada Shannon] It's the same idea with structure added. In that story the noise ball is one size in every direction. Here it lives almost entirely in a few sharp directions, while the progress lives in one flat direction. So a large rate isn't wasteful. It's productive even when the curve looks bad. The picture also echoes the edge-of-stability work from Cohen and colleagues in 2021, where training sits at a sharpness set by the learning rate. 22 00:04:31,326 --> 00:04:35,320 [Hal Turing] Is any of that measured in an actual language model, or is the canyon a drawing? 23 00:04:35,320 --> 00:04:45,537 [Dr. Ada Shannon] Mostly a drawing, and the paper calls it a conjecture about pretraining loss. It isn't a measured property of full-scale LLMs. I'd keep that in mind the whole way through. 24 00:04:45,537 --> 00:04:48,788 [Hal Turing] So why does WSD matter in practice, beyond the pretty story? 25 00:04:48,788 --> 00:05:08,153 [Dr. Ada Shannon] Because the main branch never commits to an endpoint. You can branch off at any budget, evaluate with a decay, and keep going. That makes continual training and scaling-law studies much cheaper. The cost is that you must run a decay, or average weights, before you can judge the model honestly, since the stable-phase loss misleads. 26 00:05:08,153 --> 00:05:17,162 [Hal Turing] Two things I want to hold the paper to. Does the theory survive Adam, which is what everyone actually trains with? And has anyone seen this river in a real LLM? 27 00:05:17,162 --> 00:05:38,803 [Dr. Ada Shannon] Those are the right two questions, and we'll get to them. First we need to see how the theory is built. Here's the chain. With no noise, gradient descent falls into the river and then follows it. Progress along the river is about the learning rate times the number of steps. That's Theorems 3.2 and 3.3. So a bigger step size gives faster progress along the river. 28 00:05:38,803 --> 00:05:44,190 [Hal Turing] What has to be true for that to hold? A river that wanders could throw off the whole picture. 29 00:05:44,190 --> 00:06:08,061 [Dr. Ada Shannon] Two assumptions cover that. The eigengap assumption says the Hessian's smallest eigenvalue, the river direction, is near zero. The second-smallest has to be clearly larger, so every other direction is steeply curved. Slow spinning says the river direction turns gently, with a small constant kappa under 0.01. The gap makes the river well defined. The slow turn means the iterate can keep tracking it. 30 00:06:08,061 --> 00:06:08,943 [Hal Turing] Now add noise. 31 00:06:08,943 --> 00:06:41,126 [Dr. Ada Shannon] With SGD the iterate never settles. It keeps bouncing across the hillsides, and that adds the hill term we talked about: d minus one, times learning rate, times noise variance, over two. Progress along the river still scales with the learning rate, so the loss stays high while the river progress is faster. Then decay shrinks the learning rate. The hill term shrinks with it, and the progress that was hidden shows up as a sharp drop. The paper's decay shape makes one over the learning rate fall linearly in time, because that is optimal on a quadratic. 32 00:06:41,126 --> 00:06:44,562 [Hal Turing] That's clean. What does the stochastic theorem actually assume? 33 00:06:44,562 --> 00:07:08,432 [Dr. Ada Shannon] A lot. The river is a straight line, and the noise is Gaussian with the same variance on every mountain direction and exactly zero along the river. The theorem is proved only under those simplifications. The authors conjecture it generalizes, and they check that on one two-dimensional function, a sine-shaped river. It's also the only place where the river definition is verified directly. 34 00:07:08,432 --> 00:07:12,798 [Hal Turing] So for real models the evidence has to come from somewhere else. What did they try? 35 00:07:12,798 --> 00:07:40,012 [Dr. Ada Shannon] A probe. Take two checkpoints from a GPT-2 124M run on OpenWebText, at 20 and 25 billion tokens. Walk the straight line between their weights and evaluate the loss along it. In the stable phase the loss dips in the middle, like a valley. Between two decay-phase checkpoints it just goes downhill, from about 3.02 toward 2.92. That matches the picture. But it's one slice, through one small model. 36 00:07:40,012 --> 00:07:42,194 [Hal Turing] And why would language have a river at all? 37 00:07:42,194 --> 00:08:10,383 [Dr. Ada Shannon] Their answer is the toy bigram language. Each city has its own distribution over names. One city is almost all Ken. Another splits among Bob, Alex and others. Lemma 4.1 says the sharpness of a city's loss component equals its Gini impurity, a measure of uncertainty. Deterministic cities are flat and form the river. Uncertain cities are sharp and form the hillsides. The toy model reproduces the WSD curve. 38 00:08:10,383 --> 00:08:16,002 [Hal Turing] But a toy built to produce the curve producing the curve isn't much of a test. Did the real-model version hold up? 39 00:08:16,002 --> 00:08:38,619 [Dr. Ada Shannon] Weakly. They fine-tuned GPT-2 on that city-name data. The Spearman correlation between how much decay helped a city and that city's entropy was 0.388. That's moderate. And the deterministic cities actually had slightly higher loss in the decay run than the stable run. That fits the story that the stable phase learns the deterministic parts. It's still a modest fit. 40 00:08:38,619 --> 00:08:40,941 [Hal Turing] Fine. Now the schedule. What is WSD-S? 41 00:08:40,941 --> 00:08:58,123 [Dr. Ada Shannon] Plain WSD decays to get a checkpoint, then rolls back to the pre-decay main branch and discards the decay. If decay also makes river progress, that's wasted. WSD-S keeps the decayed checkpoint and resumes at high learning rate from there. One branch, no rollback. 42 00:08:58,123 --> 00:08:59,238 [Hal Turing] What's the baseline? 43 00:08:59,238 --> 00:09:22,412 [Dr. Ada Shannon] Cosine-Oracle, a separately tuned cosine run for each budget. That's the gold standard, but it costs the sum of the budgets. The setup was LLaMA-architecture models from 0.1B to 1.2B parameters on the Pile, with 4096 context and 4M-token batches. Checkpoints were taken at 50, 100 and 200 billion tokens, with a 10 percent decay. 44 00:09:22,412 --> 00:09:23,433 [Hal Turing] And the result? 45 00:09:23,433 --> 00:09:47,907 [Dr. Ada Shannon] One WSD-S run lands within 0.006 validation loss of the cosine oracles at every size and budget. The oracle route costs 50 plus 100 plus 200, so 350 billion tokens. WSD-S costs 200. Lower loss is better, and 0.006 is small. But they report no seeds or error bars, so we can't say it's outside run-to-run noise. 46 00:09:47,907 --> 00:09:51,483 [Hal Turing] What about the cheap trick of just running cosine cycles back to back? 47 00:09:51,483 --> 00:10:24,130 [Dr. Ada Shannon] That's Cyclic-Cosine, and it loses by about 0.04 at 0.6B. That's several times the worst gap to the oracles. The common explanation is that the loss spike from rewarming the rate is what hurts. But WSD-S also jumps back to a high rate and loses nothing like that. So the authors suspect a long low-rate phase implicitly damages the model for later training. They swapped second-half schedules to test it. In all four combinations, the run that used WSD in the first half did better in the continuation. 48 00:10:24,130 --> 00:10:27,056 [Hal Turing] Is damage the mechanism, or just the label for what they saw? 49 00:10:27,056 --> 00:10:32,350 [Dr. Ada Shannon] A label. The swap shows the effect, not what causes it, and it's at one scale. 50 00:10:32,350 --> 00:10:33,743 [Hal Turing] And against plain WSD? 51 00:10:33,743 --> 00:11:03,093 [Dr. Ada Shannon] WSD-S is slightly better at matched compute. The accounting matters. For checkpoints at 12.5k, 25k and 50k steps, WSD ends up with 12.5k, 26.25k and 53.75k total steps. That's the cost of its decay branches. WSD-S wins on every size except 0.1B at 25k steps. WSD was also given a favorable decay point on the 0.3B model, after a spike. The margin is small, with no variance reported. 52 00:11:03,093 --> 00:11:05,137 [Hal Turing] What would someone actually do with this? 53 00:11:05,137 --> 00:11:22,134 [Dr. Ada Shannon] Decay for about 10 percent of steps. The final loss barely moved between 8 and 12 percent. And don't decay right after a loss spike. That cost 0.01 to 0.02 in their runs. Wait until the loss recovers, or roll back to before the spike and decay from there. 54 00:11:22,134 --> 00:11:31,886 [Hal Turing] Let me take the two critiques I flagged at the start, beginning with Adam. The theory is SGD, and every language model run uses Adam. What survives? 55 00:11:31,886 --> 00:12:09,874 [Dr. Ada Shannon] The shape survives, and the formula doesn't. A bigger step size gives faster river progress and more oscillation, and decay reveals the progress. That matches the empirical WSD curve, and I'd bet on it generalizing. The exact hill term is a different story. It needs isotropic Gaussian noise of constant variance, with exactly zero noise along the river. Real transformer gradient noise is heavy-tailed, which the paper's own related work acknowledges by citing Zhang and colleagues, 2020. It's also anisotropic. And Adam rescales each coordinate, so the step size in a given direction isn't eta. 56 00:12:09,874 --> 00:12:11,546 [Hal Turing] So what replaces the hill term? 57 00:12:11,546 --> 00:12:31,701 [Dr. Ada Shannon] The paper doesn't say. The main text never analyzes Adam. My inference is that the bounce still shrinks with the effective step size, but that's reasoning, not a result. Cohen and colleagues, 2022, showed Adam sits at a preconditioned sharpness threshold. That's a different regime from the one this theory covers. 58 00:12:31,701 --> 00:12:35,323 [Hal Turing] Second critique. Has anyone measured the river in a real model? 59 00:12:35,323 --> 00:13:09,549 [Dr. Ada Shannon] No. The river is defined through the Hessian's smallest eigenvector, but that definition is checked only on the 2-D toy function. No experiment measures a Hessian eigenvalue, an eigengap, or the slow-spinning parameter on a real network. The eigengap assumption wants the second-smallest eigenvalue clearly separated from one near zero. Real networks have a continuum bulk of eigenvalues, so I'd call the assumption likely strained. The paper only requires it near the river, and it suggests a constant might be removable. So strained, not disproved. 60 00:13:09,549 --> 00:13:14,890 [Hal Turing] Then how much of the evidence is just 'loss is high at high learning rate and drops on decay'? 61 00:13:14,890 --> 00:13:42,568 [Dr. Ada Shannon] A lot of it, and many pictures predict that. Noise-scale accounts predict it, and so does plain averaging. Sanyal and colleagues, 2023, showed that weight averaging cuts the same oscillation without any decay phase. The discriminating evidence is thinner. There's the interpolation valley, which is one model and one slice at 124M. There's the bigram toy, which was built to produce the curve. And there's the 0.388 correlation, which is moderate. 62 00:13:42,568 --> 00:13:46,655 [Hal Turing] On the valley, a noise ball around a plain quadratic would also dip in the middle, right? 63 00:13:46,655 --> 00:13:59,844 [Dr. Ada Shannon] It could, and the paper doesn't separate those. The stable-phase losses only span about 3.00 to 3.04, and I don't see replication across seeds or sizes. It's consistent with the river, but it doesn't single it out. 64 00:13:59,844 --> 00:14:01,005 [Hal Turing] What would settle it? 65 00:14:01,005 --> 00:14:24,225 [Dr. Ada Shannon] Three measurements. Track the Hessian spectrum across a WSD run and check for an eigengap. Measure kappa. And test with Adam's actual noise statistics. Also compare WSD-S against constant learning rate plus averaging, since the river picture says those should match. The schedule-free work by Defazio and colleagues, 2024, is the obvious competitor, and it isn't tested here. 66 00:14:24,225 --> 00:14:27,336 [Hal Turing] Now scale. The gaps against the cosine oracles are small. 67 00:14:27,336 --> 00:14:57,801 [Dr. Ada Shannon] Right. Models run from 0.1B to 1.2B, one dataset, final validation loss only, and no variance is reported. The largest gap to the oracles is 0.006. A loss spike followed by a decay costs 0.01 to 0.02, which is bigger than the effect being claimed. So I'd read 'matches the oracle' as 'no visible penalty', not as a measured tie. And the WSD-S margin over plain WSD is small, with one loss, at the smallest model. 68 00:14:57,801 --> 00:15:01,376 [Hal Turing] What about the 0.04 gap against Cyclic-Cosine? That one's large. 69 00:15:01,376 --> 00:15:33,141 [Dr. Ada Shannon] It's the most striking number, and its explanation is a hypothesis. The paper says long low-rate training implicitly damages the model for continuation. But the cosine runs also differ in decay shape. The paper uses a reciprocal-interpolated decay and never tests other shapes inside WSD-S. Hägele and colleagues show shape matters. A cosine-shaped decay inside WSD-S would separate shape from schedule structure. And the damage story isn't derived from the river theory. 70 00:15:33,141 --> 00:15:34,906 [Hal Turing] Is there something here I should trust? 71 00:15:34,906 --> 00:16:03,838 [Dr. Ada Shannon] Yes. At these sizes, one run with cheap branch-off decays gave usable checkpoints for several budgets. That agrees with Hägele's independent results. For scaling-law fitting, it removes per-budget retraining. It also fits the continual pretraining story from Ibrahim and colleagues. An open release could ship one final checkpoint that someone can resume. One caution is that rolling back after loss spikes may be harder at much larger scale. 72 00:16:03,838 --> 00:16:05,138 [Hal Turing] And the blind spots? 73 00:16:05,138 --> 00:16:27,383 [Dr. Ada Shannon] No downstream evals, no data-mix change at the cooldown, which is how many modern recipes work, and nothing on how resuming with Adam's decayed second moments behaves. The paper names its own next steps. Those are a k-dimensional river, curved rivers in the stochastic theory, and tighter constants. I'd add an Adam analysis and Hessian measurements in a real LLM. 74 00:16:27,383 --> 00:16:47,538 [Hal Turing] So here's where I land. The paper gives a coherent picture that predicts the shape of the WSD curve, and a single-run schedule that worked at these sizes. The theory is proved only for SGD under strong simplifications, and the landscape itself is inferred from indirect evidence. Does that sound fair? 75 00:16:47,538 --> 00:16:54,133 [Dr. Ada Shannon] That's fair. Use the schedule, treat the river as a hypothesis, and don't assume the numbers transfer past 1.2B. 76 00:16:54,133 --> 00:16:57,198 [Hal Turing] That's our time. Thanks for listening, and see you next time.