warmup-stable-decay schedule
In plain wordsA plan for how fast a model learns during training: speed up briefly, hold steady for most of the run, then slow down sharply at the end — which lets labs stop or extend training without starting over.
A learning rate schedulelearning rate schedule, n.: A plan for changing how big a model's learning steps are during training — usually starting cautiously, then taking big steps, then smaller and smaller ones as it settles in. with three phases: a short linear warmuplearning rate warmup, n.: Starting training with tiny learning steps and gradually increasing them over the first stretch of training, so the model does not lurch wildly before it has found its footing. to the peak rate; a long stable phase at a constant rate; and a decay or cooldown phase, often the last 10–20% of steps, in which the rate falls quickly toward zero along a linear, square-root or exponential curve. Loss improves slowly in the stable phase and drops sharply during decay. Unlike cosine annealingcosine annealing, n.: A popular plan for shrinking a model's step size during training: steps start large, shrink slowly, then faster, then level off near zero, following the shape of half a wave., which fixes the total length in advance, WSD lets a run continue indefinitely from a stable-phase checkpointcheckpoint, n.: A saved snapshot of a model partway through or at the end of training, like a save point in a video game: training can resume from it after a crash, and finished snapshots are what people download and run. and branch off decays of several lengths. That makes continued pretrainingcontinued pretraining, n.: Taking an already-trained AI model and resuming its broad, read-everything training on new material — a specialty such as medicine or law, a new language, or fresher data — rather than starting from scratch., mid-trainingmid-training, n.: A middle stage in building a language model, between its broad initial training and its tuning into an assistant: it keeps learning the same way, but from a smaller, carefully chosen diet of math, code, reasoning or long documents. data changes and scaling-law sweeps much cheaper, since one long run yields models for many token budgets. Final quality matches or slightly exceeds cosine when the cooldown is long enough. One explanation pictures a 'river valley' loss landscape: the high rate moves along the valley while bouncing between its walls, and the decay settles onto the floor.
With WSD we can decay from any stable-phase checkpoint, so the 1T- and 2T-token models share one run.
Origin
Named by Shengding Hu and colleagues in their technical report on the MiniCPM models (April 2024). Schedules with no fixed end point that finish with a short cooldown were used earlier, for example a reciprocal-square-root decay with linear cooldown in Google's vision-transformer scaling work (Xiaohua Zhai and colleagues, 2022); Alexander Hägele and colleagues (May 2024) showed that constant-plus-cooldown schedules support scaling laws without fixed training lengths.
Usage note
In the formula, f decreases from 1 to about 0. The decay phase is where annealingannealing, n.: The last stretch of training, when the model's learning is slowed to a stop — often while it is fed a small helping of the best data — so it settles into its final, sharpest version. on high-quality data usually happens, so 'cooldown', 'decay' and 'annealing' are often used together.
Related
- learning rate schedulelearning rate schedule, n.: A plan for changing how big a model's learning steps are during training — usually starting cautiously, then taking big steps, then smaller and smaller ones as it settles in.
- cosine annealingcosine annealing, n.: A popular plan for shrinking a model's step size during training: steps start large, shrink slowly, then faster, then level off near zero, following the shape of half a wave.
- annealingannealing, n.: The last stretch of training, when the model's learning is slowed to a stop — often while it is fed a small helping of the best data — so it settles into its final, sharpest version.
- learning rate warmuplearning rate warmup, n.: Starting training with tiny learning steps and gradually increasing them over the first stretch of training, so the model does not lurch wildly before it has found its footing.
- continued pretrainingcontinued pretraining, n.: Taking an already-trained AI model and resuming its broad, read-everything training on new material — a specialty such as medicine or law, a new language, or fresher data — rather than starting from scratch.
- IsoFLOP analysisIsoFLOP analysis, n.: A way of finding the best model size for a given amount of computing: train several models of different sizes at exactly the same computing cost and see which one ends up best.
- cosine annealingcosine annealing, n.: A popular plan for shrinking a model's step size during training: steps start large, shrink slowly, then faster, then level off near zero, following the shape of half a wave.
Sources
- Hu, S., Tu, Y., Han, X., He, C., Cui, G. et al. (2024). MiniCPM: Unveiling the Potential of Small Language Models with Scalable Training Strategies. arXiv:2404.06395. arxiv.org/abs/2404.06395
- Hägele, A., Bakouch, E., Kosson, A., Ben Allal, L., Von Werra, L. & Jaggi, M. (2024). Scaling Laws and Compute-Optimal Training Beyond Fixed Training Durations. arXiv:2405.18392. arxiv.org/abs/2405.18392
Mentioned in
annealingannealing, n.: The last stretch of training, when the model's learning is slowed to a stop — often while it is fed a small helping of the best data — so it settles into its final, sharpest version., continued pretrainingcontinued pretraining, n.: Taking an already-trained AI model and resuming its broad, read-everything training on new material — a specialty such as medicine or law, a new language, or fresher data — rather than starting from scratch., cosine annealingcosine annealing, n.: A popular plan for shrinking a model's step size during training: steps start large, shrink slowly, then faster, then level off near zero, following the shape of half a wave., learning rate schedulelearning rate schedule, n.: A plan for changing how big a model's learning steps are during training — usually starting cautiously, then taking big steps, then smaller and smaller ones as it settles in., loss curveloss curve, n.: A graph of a model's error as training goes on. A healthy curve drops steeply and then keeps sloping gently downward; bumps, jumps or flat stretches can warn engineers that something is wrong., mid-trainingmid-training, n.: A middle stage in building a language model, between its broad initial training and its tuning into an assistant: it keeps learning the same way, but from a smaller, carefully chosen diet of math, code, reasoning or long documents., IsoFLOP analysisIsoFLOP analysis, n.: A way of finding the best model size for a given amount of computing: train several models of different sizes at exactly the same computing cost and see which one ends up best., schedule-free optimizationschedule-free optimization, n.: A way of training that needs no plan for slowing learning down over time, so training can be stopped or extended at any point and still produce a well-settled model.