Fine-tuning Learning Rates › Learning Rate Decay

After warmup completes, learning rate decay schedules reduce the rate throughout training. This helps the model converge to a stable minimum rather than oscillating around it. The intuition is that large learning rates are useful early in training when the model needs to make significant adjustments, but become counterproductive later when the model is close to a good solution and only needs fine adjustments.

Think of decay as the difference between making broad brushstrokes and fine detail work when painting. Early in the process, you want big, confident strokes that establish the overall composition. As the painting develops, you switch to finer tools for precise work. A painter who keeps using the broad brush throughout the entire process will never reach the detail that the final piece requires. Similarly, a constant learning rate throughout fine-tuning tends to produce models that have adapted well on average but oscillate around the optimal parameters rather than converging to them precisely.

The choice of decay schedule shapes the final convergence and the entire trajectory of learning. Different schedules allocate the "learning budget" differently across the training run: some front-load it (maintaining higher rates early), others distribute it more uniformly, and others back-load it (maintaining higher rates in the middle of training). These differences can have real effects on final performance, particularly for tasks where the model needs to make large initial adjustments to the fine-tuning task before settling into precise calibration.

The simplest decay schedule reduces the learning rate linearly to zero (or some minimum value) by the end of training. This approach gives a steady, predictable reduction in learning rate that is easy to reason about and implement.

We want to compute a learning rate that starts at the target value immediately after warmup ends and decreases uniformly to zero at the final training step. After warmup ends at step $T_{\text{warmup}}$, the learning rate at step $t$ is:

$$\eta_t = \eta_{\text{target}} \cdot \left(1 - \frac{t - T_{\text{warmup}}}{T_{\text{total}} - T_{\text{warmup}}}\right)$$

where:

• $\eta_t$: the learning rate at step $t$
• $\eta_{\text{target}}$: the learning rate at the start of decay (equal to the peak warmup value)
• $t$: the current training step
• $T_{\text{warmup}}$: the step number where decay begins (end of warmup)
• $T_{\text{total}}$: the total number of training steps

Why does this formula make sense? The fraction $\frac{t - T_{\text{warmup}}}{T_{\text{total}} - T_{\text{warmup}}}$ is the progress through the decay phase, ranging from 0 at the start of decay to 1 at the final step. Subtracting this from 1 yields the remaining proportion of the learning rate. This creates a steady linear decline from $\eta_{\text{target}}$ down to 0. At each step during the decay phase, the learning rate decreases by exactly the same absolute amount. This keeps uniform reduction throughout training.

This is the default schedule for most BERT fine-tuning recipes because it strikes a reasonable balance between simplicity and effectiveness. The linear nature makes it trivial to predict the learning rate at any step, which is helpful for debugging and monitoring training progress.

Cosine decay gives a smoother transition that decays slowly at first, accelerates in the middle, and slows again near the end. This schedule has become increasingly popular because it maintains higher learning rates for longer during the middle of training, giving the optimizer more opportunity to explore the loss surface before settling into a final minimum.

We want to compute a smoothly decaying learning rate using the cosine function's natural properties. After warmup, the learning rate at step $t$ is:

$$\eta_t = \eta_{\text{min}} + \frac{\eta_{\text{target}} - \eta_{\text{min}}}{2} \cdot \left(1 + \cos\left(\pi \cdot \frac{t - T_{\text{warmup}}}{T_{\text{total}} - T_{\text{warmup}}}\right)\right)$$

where:

• $\eta_t$: the learning rate at step $t$
• $\eta_{\text{min}}$: the minimum learning rate to reach at the end of training (often 0)
• $\eta_{\text{target}}$: the target learning rate at the start of decay (peak warmup value)
• $t$: the current training step
• $T_{\text{warmup}}$: the step number where decay begins
• $T_{\text{total}}$: the total number of training steps

Why does this formula produce an S-shaped decay? We can trace the mathematical transformations step by step to see the mechanism clearly:

$$\begin{aligned} p &= \frac{t - T_{\text{warmup}}}{T_{\text{total}} - T_{\text{warmup}}} && \text{(progress fraction: 0 at start, 1 at end)} \\ \theta &= \pi \cdot p && \text{(angle mapping: 0 to } \pi \text{ radians)} \\ c &= \cos(\theta) && \text{(cosine: 1 at start, -1 at end, slow near endpoints)} \\ \eta_t &= \eta_{\text{min}} + \frac{\eta_{\text{target}} - \eta_{\text{min}}}{2}(1 + c) && \text{(scale to desired range)} \end{aligned}$$

The cosine transform is what creates the S-curve behavior. The derivative of $\cos(\theta)$ is $-\sin(\theta)$, which equals zero at $\theta = 0$ and $\theta = \pi$ (the endpoints) but reaches its maximum magnitude at $\theta = \pi/2$ (the midpoint). This means the cosine schedule changes slowly at the beginning and end, and rapidly in the middle, which is exactly the behavior we want: the model maintains higher learning rates early to explore broadly, transitions through a rapid reduction phase in the middle, and then approaches its final value gently without overshooting.

While all three schedules eventually reach zero, they allocate learning rate differently throughout training. Cosine decay maintains higher learning rates for longer in the middle of training, potentially letting the model to escape local minima more effectively than linear decay. Polynomial decay with power 2 drops quickly at first and then more slowly, front-loading the high learning rate period. The choice between these schedules is a trade-off between exploration and convergence speed. Linear decay is the most predictable and easiest to reason about; cosine decay is often slightly better in practice due to its gentler terminal convergence; polynomial decay gives the most control through its exponent parameter.

The choice of decay schedule depends on your training dynamics and the specific characteristics of your task:

• Linear decay: Simple and effective; good default for short fine-tuning runs. The predictable reduction makes it easy to reason about training dynamics and diagnose issues such as learning rate being too high or too low.
• Cosine decay: Often slightly better performance; allows more learning early and late in training. The gentler decay at both ends can help the model both explore better initially and converge more precisely at the end. It is the recommended default when you have sufficient compute for even a modest hyperparameter search.
• Polynomial decay: Useful when you want to front-load learning (higher powers) or extend it (lower powers). The exponent parameter gives direct control over the shape of the decay curve, making this a flexible option for tasks with unusual training dynamics.

Empirically, cosine and linear decay perform similarly for most fine-tuning tasks, with cosine sometimes giving a small but consistent improvement. The difference is often smaller than the effect of other hyperparameters like the base learning rate or batch size, so either choice is reasonable for most applications. When in doubt, cosine decay is the better default because its smooth behavior near both endpoints reduces the risk of numerical instability.