Universal Dynamics of Warmup Stable Decay: understanding WSD beyond Transformers

AI summary

Overview

Research area: Optimization and training dynamics for deep learning — specifically learning rate scheduling (Warmup Stable Decay, WSD) and the geometry of nonconvex loss landscapes.

Technical level: Intermediate. The paper uses accessible framing but relies on concepts such as Hessian eigenvalue sharpness, principal component analysis (PCA) of optimizer iterates, and weak quasi-convexity.

Scope: A single-sentence scope: the paper empirically compares the WSD training path of Adam/AdamW on a Pythia-like 160M-parameter language model against a 334K-parameter CNN on CIFAR10, showing that the characteristic WSD loss curve and several landscape properties are shared across architectures.

What This Paper Is About

WSD has become popular for training large language models because it reaches performance comparable to warmup cosine annealing and lets training be resumed cheaply before the decay phase. Most studies of WSD have examined only transformer-based language models, which raised the open question of whether WSD's behavior is specific to transformers or is a more general property of the loss landscapes it traverses. This paper tests that question by applying WSD to a small CNN for image classification and comparing the resulting training dynamics, optimizer path, and sharpness behavior with those of a transformer language model.

Key Contributions

1. An architecture-spanning empirical comparison of WSD between a transformer-based language model (about 160M trainable parameters, trained on SlimPajama) and a small CNN (about 334K parameters, trained on CIFAR10).
2. A replication of the "River Valley" loss-landscape visualizations (loss along linear interpolations between checkpoints) on both model types, showing a convex valley-shaped profile near the end of the stable phase and a monotonic descent over the cooldown.
3. Evidence that sharpness (the largest eigenvalue of the loss Hessian) increases along iterates collected while annealing the learning rate, and a directional analysis showing that the early-cooldown PCA direction aligns more with high-curvature Hessian eigenspaces than the late-stable direction.
4. Tests of nearly-convex behavior along the optimizer's trajectory (weak quasi-convexity and the cosine similarity between the negative gradient and the actual update), which hold for both models and partially close the gap between nonconvex AdamW practice and the convex SGD theory previously used to explain WSD.

Main Findings

• Shared macroscopic WSD curve: In both the 160M-parameter LM with Pythia 12-layer configuration (trained on 3B tokens, far less than 1 epoch) and the 334K-parameter 4-layer CNN (trained for 50 epochs), WSD shows modest loss decreases during the stable phase followed by sudden, more pronounced gains during cooldown.
• River Valley generalizes beyond transformers: Loss evaluated along linear interpolations between checkpoints at 80% and 100% of the stable phase shows a convex valley shape, and interpolation between the start and end of cooldown shows monotonic descent, for both the LM and the CNN.
• Sharpness increases during cooldown: Estimated sharpness on iterates sampled at a fixed rate along the cooldown increases as the learning rate is reduced, in both models, consistent with prior work cited by the authors.
• Two distinct movement directions: PCA on iterates sampled uniformly during stable (set S) and decay (set D) phases shows the first component captures at least approximately 40% of the total variance for both models and both phases, indicating each phase is governed by a main yet distinct direction.
• Cooldown direction is higher-curvature: Using sets S' (last 20% of stable) and D' (first 20% of cooldown), the paper finds the inequality ‖∇²ℒ(x̂)v_s‖ < ‖∇²ℒ(x̂)v_d‖ at the decay-start checkpoint x̂ for both models, meaning the early-cooldown direction lies closer to high-curvature regions. The authors interpret this as the smaller decay steps letting the trajectory "see" and follow sharper subspaces previously inaccessible.
• The trajectory is nearly convex: The weak quasi-convexity condition (Def. 2.1 of the cited work) holds in nearly all cases for both models, with τ_i > 0, and the cosine similarity between the negative gradient and the actual update remains positive for both models, suggesting AdamW updates do not deviate substantially from SGD updates.
• A noted caveat: The authors observe that their small CNN does not perfectly fit the data, similar to the transformer, and flag overparametrization as an open factor.

Methodology in Plain English

The authors take two very different models and give them the same training recipe: warm up the learning rate, hold it constant (the stable phase), then linearly decay it (the cooldown). One model is a decoder-only transformer language model with roughly 160M parameters trained on about 3 billion tokens from the cerebras/SlimPajama-627B dataset, with batch size 256 and sequence length 2048, using PlainLM. The other is a small CNN with about 334K parameters trained on CIFAR10 with batch size 128. Both are trained with Adam (AdamW is also referenced in the analysis).

Then the authors probe the training run in several ways. They draw the loss surface by interpolating between pairs of saved checkpoints to see whether it looks like a valley. They measure sharpness as the largest eigenvalue of the loss Hessian along cooldown iterates. They use PCA to extract the dominant parameter-space directions during the stable and decay phases, and they check how those directions interact with the Hessian eigenspaces at the checkpoint where decay begins. Finally, they test two conditions borrowed from optimization theory — weak quasi-convexity and the alignment between the gradient descent direction and the actual update — to see how close the trajectory comes to behaving like convex optimization.

Why This Matters

The results suggest that the unusual WSD loss curve is not an artifact of transformer architecture or of language modeling, but a symptom of geometric properties shared by these different nonconvex loss landscapes. This broadens the theoretical explanation for WSD and connects decades of work on small-scale vision models to modern large-model training practice.

Real-world applications:

• Designing learning rate schedules for large language model training runs, where compute budgets are large and flexibility in choosing the decay start matters.
• Continual or extended pretraining, where training can be restarted from a checkpoint just before decay and the step budget increased without retraining from scratch.
• Transferring scheduling insights across model families, so techniques tuned for transformers can be evaluated on CNNs and other architectures.
• Reducing loss spikes and unlearning behaviors that can occur when rewarming the learning rate after a completed cosine cycle.

Industry relevance: Because the main practical benefit of WSD cited in the paper is computational, these findings matter to organizations that train or keep training large models and want schedule choices that are architecture-agnostic and resumable rather than tied to one model class.

Future Directions

• Investigating the effect of overparametrization on WSD performance, given the authors' observation that their small CNN does not perfectly fit the data, similar to the transformer.
• Determining whether the shared WSD dynamics extend to other architectures, optimizers, and modalities beyond a Pythia-like LM and a small CNN on CIFAR10.
• Exploring the geometry of high-dimensional optimization problems more generally, which the authors explicitly raise as a research question prompted by their consistency findings.
• Reconciling the remaining mismatch between the convex SGD theoretical explanation of WSD and the nonconvex AdamW setting in which the phenomenon is observed.

Target Audience

Researchers and practitioners in machine learning optimization, learning rate scheduling, and large-scale model training who want to know whether WSD's behavior is transformer-specific. It is also useful for readers interested in loss landscape geometry and in empirical tests that bridge optimization theory and deep learning practice. A working familiarity with Eigenvalues, PCA, and basic convex optimization concepts helps, but the paper's framing is accessible enough for graduate students entering the area.

Authors’ abstract

The Warmup Stable Decay (WSD) learning rate scheduler has recently become popular, largely due to its good performance and flexibility when training large language models. It remains an open question whether the remarkable performance of WSD - using a decaying learning rate for only a fraction of training compared to cosine decay - is a phenomenon specific to transformer-based language models that can potentially offer new theoretical insights into their training dynamics. Inspired by the usage of learning rate schedulers as a new lens into understanding landscape geometry (e.g., river valley, connected minima, progressive sharpening), in this work we compare the WSD path of the Adam optimizer on a Pythia-like language model to that of a small CNN trained to classify CIFAR10 images. We observe most training signals, optimizer path features, and sharpness dynamics to be qualitatively similar in such architectures. This consistency points to shared geometric characteristics of the loss landscapes of old and new nonconvex problems, and hints to future research questions around the geometry of high dimensional optimization problems.