Learning Rate Warmup › Warmup Variants Beyond Linear

Linear warmup dominates modern practice, but it is worth knowing the other variants that practitioners use and why they exist. Each variant makes a different tradeoff between how fast the learning rate grows in the early steps and how quickly the optimizer transitions to peak-rate training.

Some practitioners use a two-phase approach: start at a fixed small learning rate for a few hundred steps, then ramp linearly to the peak. The logic is that the very first steps are so noisy that even a small linear growth feels premature. By holding constant at, say, $0.01 \times \eta_{\max}$ for the first 100 steps, you give the model a chance to take a handful of consistent small steps before beginning the ramp. In practice, this variant produces results nearly identical to linear warmup, and the extra hyperparameter (the length of the flat phase) is rarely worth the tuning cost.

Polynomial warmup generalizes linear warmup by allowing a power $p$ in the growth function:

$$\eta_t = \eta_{\max} \cdot \left(\frac{t}{T_{\text{warmup}}}\right)^p$$

With $p = 1$ you recover linear warmup. With $p > 1$ (such as $p = 2$ or $p = 3$), the schedule is convex, growing slowly at first and then accelerating. This creates a gentler start than linear warmup at the cost of spending less time at intermediate learning rates. With $p < 1$ (such as $p = 0.5$, the square-root case), the schedule is concave, growing quickly early and then plateauing. This gives the model more time at moderate learning rates during warmup.

For most applications, $p = 1$ (linear) is the right choice. The gains from tuning $p$ are small compared to the cost of searching that hyperparameter space. But polynomial warmup appears in some frameworks (such as the Hugging Face transformers library's get_polynomial_decay_schedule_with_warmup) because it generalizes both the linear and square-root cases.

A less common variant uses the cosine function itself for warmup, analogous to how cosine is used for decay:

$$\eta_t = \eta_{\max} \cdot \frac{1}{2}\left(1 - \cos\!\left(\pi \cdot \frac{t}{T_{\text{warmup}}}\right)\right)$$

This produces an S-shaped growth curve that starts very slowly, accelerates through the middle of warmup, and then slows down as it approaches the peak. The motivation is that this shape mirrors the cosine decay phase, creating a symmetric schedule where both warmup and decay transition smoothly. Some practitioners find this aesthetically pleasing and claim a slight improvement in the earliest training stability, but empirical evidence for this is weak.

A recent and practically important variant is the Warmup-Stable-Decay schedule, which divides training into three phases rather than two. After the warmup ramp, the learning rate stays constant at the peak value (the stable phase) for the majority of training, then decays sharply in the final phase. This pattern emerged from work on continued pretraining and incremental training, where practitioners discovered that a long flat phase followed by a sharp cooldown often outperforms continuous cosine decay. The WSD schedule is increasingly used in modern LLM training because it is more compatible with training runs that are not committed to a fixed total step count from the beginning.