TLDR: This research paper analyzes the “cooldown” phase of the Warmup-Stable-Decay (WSD) learning rate scheduler in transformer training. It reveals that the cooldown shape dictates a bias-variance trade-off, with “sqrt” and “lowered linear 0.7” shapes being optimal. The study also finds that tuning AdamW hyperparameters, especially beta2, during cooldown is as crucial as shape selection, and empirically supports the “river valley” loss landscape concept.
In the world of artificial intelligence, especially when training large language models like Transformers, how you manage the learning rate—the speed at which a model learns—is incredibly important. One popular method is the Warmup-Stable-Decay (WSD) learning rate scheduler, which has three distinct phases: an initial “warmup” where the learning rate gradually increases, a “stable” phase where it maintains a constant rate, and a final “decay” or “cooldown” phase where it reduces rates to zero.
While the WSD scheduler has shown great promise, often matching or even outperforming other common methods like cosine decay, the final “cooldown” phase has remained a bit of a mystery. This is surprising because this phase is where a significant drop in the model’s error (perplexity) often occurs, leading to better performance.
A recent research paper, “Training Dynamics of the Cooldown Stage in Warmup-Stable-Decay Learning Rate Scheduler,” by Aleksandr Dremov, Alexander Hägele, Atli Kosson, and Martin Jaggi from EPFL, Switzerland, dives deep into this often-overlooked cooldown stage. You can find the full paper here: Research Paper.
The authors provide a comprehensive analysis, focusing exclusively on the cooldown phase. Their key finding is that the shape of this cooldown phase—how the learning rate decreases over time—is crucial. Different shapes reveal a fundamental “bias-variance trade-off” in the resulting models. Think of it this way: a more aggressive cooldown (higher learning rates) allows the model to “explore” the learning landscape more, potentially finding better solutions but also leading to more varied results across different training runs (higher variance). Conversely, a less aggressive cooldown (lower learning rates) encourages “exploitation” of the current learning area, leading to more consistent but potentially less optimal solutions (lower variance, higher bias).
The study found that shapes like “sqrt” and “lowered linear 0.7” consistently performed best because they strike an optimal balance between this exploration and exploitation. These shapes allow the model to refine its learning effectively without getting stuck or becoming too unstable.
Beyond the cooldown shape, the researchers also investigated the impact of tuning the AdamW optimizer’s hyperparameters, particularly the beta parameters. They discovered that adjusting these parameters during the cooldown phase can lead to performance improvements comparable to those achieved by selecting the best cooldown shape. Notably, higher values of the beta2 parameter consistently led to better results, emphasizing that even subtle optimizer settings can have a significant impact.
From a visual perspective, the paper supports the “river valley” analogy for the loss landscape during training. This analogy suggests that during the stable phase, the model descends along a broad “river,” and during the cooldown, it descends directly into the “river itself” or a narrow “basin.” The visualizations provided in the paper empirically support this concept, showing how optimization transitions from a broader exploration to a more focused descent.
The research also touched upon model averaging, a technique where multiple models trained with different data portions are combined. They found that while averaging models from high-variance, low-bias cooldown shapes (like “mirror cosine”) can improve performance, a single, longer training run still generally outperforms averaged models.
Also Read:
- Decoding LLM Learning: How Uncertainty and Position Drive Reinforcement Learning Success
- Enhancing Trustworthiness in Language Models: A Deep Dive into Calibration and Label Smoothing
In conclusion, this work offers valuable practical recommendations for configuring the WSD scheduler. It highlights that optimizing the cooldown phase, including its shape and the AdamW optimizer settings, is just as important as traditional hyperparameter tuning for achieving peak performance in transformer training.
