Universal Dynamics of Warmup Stable Decay: understanding WSD beyond Transformers

TL;DR

The paper investigates whether Warmup Stable Decay (WSD) benefits extend beyond transformer-based language models by comparing its training dynamics in a 160M-parameter LM and a small CIFAR-10 CNN. It analyzes loss landscapes via river valley interpolations, sharpness evolution under AdamW, and the trajectory directions during stable and cooldown phases. Key findings include convex valley-like loss profiles during the stable phase, a sharp descent during cooldown, and two dominant training directions corresponding to the two phases, alongside evidence of near-quasi-convexity along the optimizer's path. These results suggest that WSD-driven dynamics reflect general high-dimensional optimization geometry rather than architecture-specific phenomena, with implications for continual training and cross-architecture training insights.

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.

Figures (7)

• Figure 1: Typical training dynamics with WSD and Warmup Cosine Annealing schedulers. In both the 160M parameters LM with Pythia 12 layer configuration biderman2023pythia (left) and the 334K parameters CNN with 4 layers (right), we observe that WSD follows the same trend: modest decreases during the stable phase (left of dashed line), followed by sudden and more pronounced gains during the cooldown (right of dashed line). Note that while the former model was trained 3B tokens (way less than 1 epoch), the latter performs 50 epochs.
• Figure 2: Loss evaluated along $\alpha t + (1 - \alpha)t'$, with $\alpha \in [0,1]$, where $t'$ and $t$ denote checkpoints taken at 80% and 100% of the stable phase, or at the start and end of the cooldown phase, on an LM (left) and a CNN (right). Both models display a similar convex valley-shaped profile between the two checkpoints sampled near the end of the stable phase, followed by a monotonic descent over the cooldown period.
• Figure 3: (Left) Estimated sharpness, evaluated on iterates $x_i$ sampled at a fixed rate along the cooldown phase, on an LM (left) and on a CNN (right). In both models, reducing the learning rate leads to sharper regions of the loss surface. (Right) Box plots of $\left\| \nabla^2 \mathcal{L}(\hat{x}) v \right\|$ for directions $v_s$ and $v_d$, derived from multiple training runs for an LM (left) and a CNN (right), where $\hat{x}$ is the checkpoint at the start of the decay. Results indicate that the direction representing the early decay phase $v_d$ aligns more with high-curvature regions.
• Figure 4: (Green) Box plot of $\tau_i = \frac{-\nabla \mathcal{L}(x_i)(x^* - x_i)}{\mathcal{L}(x_i) - \mathcal{L}(x^*)}$, evaluated on a a subset $B$ of iterates $x_i$ sampled regularly during training, with $x^*$ denoting the solution at the end of training. Under weak quasi-convexity, $\tau_i > 0$, which is indeed observed in nearly all cases. (Orange) Box plot of $s_{\cos,i}$ = $\text{cos\_similarity}(-\nabla \mathcal{L}(x_i),\, x_{i+1} - x_i)$, evaluated on $B$. These scores remain positive, suggesting that AdamW updates do not deviate substantially from those of SGD. Results are presented for multiple training runs for an LM (left) and a CNN (right).
• Figure 5: Visualization of the impact, on the considered CNN, of cooldown length on WSD performance, in comparison to Warmup Cosine Annealing. Unlike what is typically observed in transformer models, cooldown dynamics in CNNs can be slower when using the same decay length. For example, under our settings, the learning rate must decay for at least 35% of the total training steps for the WSD performance to match that of Warmup Cosine Annealing.