Full Fine-tuning › Learning Rate Selection

Choosing the right learning rate requires balancing competing concerns. We want enough change to adapt to the new task, but not so much that we destroy pre-trained knowledge. This balancing act is at the heart of successful fine-tuning. Understanding the available tools helps practitioners address this challenge effectively. The community has developed systematic methods for finding good learning rates without exhaustive search.

The search for a good learning rate is constrained by the nature of the problem. Too low, and the model makes negligible progress within the allotted epochs. Too high, and the pre-trained representations are corrupted early in training, making recovery impossible. The optimal learning rate sits in a relatively narrow band between these failure modes, and the goal of the strategies we discuss is to find that band efficiently.

The learning rate typically follows a schedule that varies it throughout training. The most common schedules for fine-tuning are designed to start with more aggressive learning (once warmup completes) and gradually become more conservative as training progresses. This pattern shows the intuition that early training benefits from larger updates to establish good task-specific representations, while later training benefits from smaller updates to fine-tune details without overshooting optimal values.

Linear decay starts at the target learning rate after warmup and decreases linearly to zero by the end of training. This straightforward approach gives a predictable reduction in learning rate that many practitioners find effective. After the warmup phase, the schedule follows:

$$\text{lr}(t) = \text{lr}_{\text{max}} \times \left(1 - \frac{t - T_{\text{warmup}}}{T_{\text{total}} - T_{\text{warmup}}}\right)$$

where:

• $\text{lr}(t)$: the learning rate at training step $t$ (after warmup)
• $\text{lr}_{\text{max}}$: the maximum learning rate reached after warmup
• $t$: the current training step (where $t \geq T_{\text{warmup}}$)
• $T_{\text{warmup}}$: the number of warmup steps
• $T_{\text{total}}$: the total number of training steps
• The term $\left(1 - \frac{t - T_{\text{warmup}}}{T_{\text{total}} - T_{\text{warmup}}}\right)$: the linear decay factor, starting at 1 and reaching 0 at the end of training
• The numerator $t - T_{\text{warmup}}$: measures how many steps have elapsed since warmup ended
• The denominator $T_{\text{total}} - T_{\text{warmup}}$: measures the total number of training steps after warmup
• Their ratio: is the fraction of post-warmup training completed, growing from 0 to 1

To understand how this formula works, let's examine the decay factor at key points during training. This analysis reveals how the schedule transitions smoothly from full learning rate to zero.

At the start of decay ($t = T_{\text{warmup}}$):

$$\text{decay factor} = 1 - \frac{T_{\text{warmup}} - T_{\text{warmup}}}{T_{\text{total}} - T_{\text{warmup}}} = 1 - \frac{0}{T_{\text{total}} - T_{\text{warmup}}} = 1 - 0 = 1$$

This gives the full learning rate: $\text{lr}(T_{\text{warmup}}) = \text{lr}_{\text{max}} \times 1 = \text{lr}_{\text{max}}$.

At the end of training ($t = T_{\text{total}}$):

$$\text{decay factor} = 1 - \frac{T_{\text{total}} - T_{\text{warmup}}}{T_{\text{total}} - T_{\text{warmup}}} = 1 - 1 = 0$$

This reduces the learning rate to zero: $\text{lr}(T_{\text{total}}) = \text{lr}_{\text{max}} \times 0 = 0$.

The intuition is that as we approach the end of training, we want smaller updates to avoid overshooting the optimal parameters. At intermediate points, the decay factor falls between 0 and 1, creating a smooth linear transition from the maximum learning rate to zero. Starting with larger updates allows faster initial progress, while ending with tiny updates lets fine refinement. This linear decrease gives a gentle transition that helps the model converge smoothly, which makes it the default choice for BERT-style fine-tuning because it works well across many tasks. The simplicity of linear decay also makes it easy to reason about: halfway through training, the learning rate will be at half its maximum value.

Cosine decay decreases the learning rate following a cosine curve with a smoother transition. It decreases slowly at first, more rapidly in the middle, and slowly again at the end. This shape shows the mathematical properties of the cosine function and gives a more gradual transition than linear decay. After warmup, the schedule follows:

$$\text{lr}(t) = \text{lr}_{\text{max}} \times 0.5 \times \left(1 + \cos\left(\pi \times \frac{t - T_{\text{warmup}}}{T_{\text{total}} - T_{\text{warmup}}}\right)\right)$$

where:

• $\text{lr}(t)$: the learning rate at training step $t$ (after warmup)
• $\text{lr}_{\text{max}}$: the maximum learning rate reached after warmup
• $t$: the current training step (where $t \geq T_{\text{warmup}}$)
• $T_{\text{warmup}}$: the number of warmup steps
• $T_{\text{total}}$: the total number of training steps
• $\cos(\cdot)$: the cosine function, which oscillates between -1 and 1, creating the smooth decay curve when combined with the other terms
• The argument $\pi \times \frac{t - T_{\text{warmup}}}{T_{\text{total}} - T_{\text{warmup}}}$: maps training progress linearly from 0 to $\pi$
• The cosine of this argument: decreases from $\cos(0) = 1$ to $\cos(\pi) = -1$ as training progresses
• The addition of 1: shifts the range from $[-1, 1]$ to $[0, 2]$
• The factor 0.5: scales the range from $[0, 2]$ to $[0, 1]$, giving the final decay multiplier

Let's verify the formula at key points to confirm it behaves as expected. Understanding these boundary conditions helps build intuition for how the schedule operates throughout training.

At the start of decay ($t = T_{\text{warmup}}$):

$$\begin{aligned} \text{lr}(T_{\text{warmup}}) &= \text{lr}_{\text{max}} \times 0.5 \times \left(1 + \cos\left(\pi \times \frac{T_{\text{warmup}} - T_{\text{warmup}}}{T_{\text{total}} - T_{\text{warmup}}}\right)\right) && \text{(substitute } t = T_{\text{warmup}} \text{)} \\ &= \text{lr}_{\text{max}} \times 0.5 \times (1 + \cos(0)) && \text{(simplify argument)} \\ &= \text{lr}_{\text{max}} \times 0.5 \times (1 + 1) && \text{(evaluate } \cos(0) = 1 \text{)} \\ &= \text{lr}_{\text{max}} && \text{(compute)} \end{aligned}$$

At the end of training ($t = T_{\text{total}}$):

$$\begin{aligned} \text{lr}(T_{\text{total}}) &= \text{lr}_{\text{max}} \times 0.5 \times \left(1 + \cos\left(\pi \times \frac{T_{\text{total}} - T_{\text{warmup}}}{T_{\text{total}} - T_{\text{warmup}}}\right)\right) && \text{(substitute } t = T_{\text{total}} \text{)} \\ &= \text{lr}_{\text{max}} \times 0.5 \times (1 + \cos(\pi)) && \text{(simplify argument)} \\ &= \text{lr}_{\text{max}} \times 0.5 \times (1 + (-1)) && \text{(evaluate } \cos(\pi) = -1 \text{)} \\ &= 0 && \text{(compute)} \end{aligned}$$

The cosine curve creates a smooth S-shaped transition: it starts with slow decay (letting aggressive early learning), accelerates in the middle (transitioning from exploration to refinement), and then slows again near zero (letting precise convergence). This gentler schedule can be particularly beneficial when fine-tuning requires careful preservation of pre-trained representations. The characteristic shape of the cosine curve means the model spends more time at higher learning rates (good for making progress) and more time at very low learning rates (good for final convergence) compared to linear decay.

Constant keeps the learning rate fixed throughout training. This is simpler but often underperforms scheduled approaches. Without decay, the model continues making relatively large updates even as it approaches convergence, which can cause oscillation around the optimum rather than settling into it. However, constant learning rates can work well for very short training runs or when combined with other regularization techniques. Some modern training recipes for large language models use constant learning rates with an aggressive final decay, which can reach strong results while being simpler to configure.

While default values like $2 \times 10^{-5}$ work well for many BERT fine-tuning scenarios, optimal learning rates vary with the model and task as well as the dataset size. Several strategies can help identify good values, each giving different trade-offs between thoroughness and computational cost.

Grid search evaluates a small set of candidate learning rates, typically varying by factors of 2-3 (e.g., 1e-5, 2e-5, 5e-5). With small fine-tuning datasets, this remains computationally feasible. The approach is systematic and gives clear comparisons across candidates. A grid search over 3-5 candidate learning rates requires 3-5x the compute of a single run, but the improvement from finding the right learning rate can more than compensate for this cost.

Learning rate finder gradually increases the learning rate during a short training run, plotting loss against learning rate. The optimal learning rate is often just before the loss begins to increase sharply. This technique gives quick insights without full training runs. The method exploits the fact that loss should decrease as learning rate increases from very small values, but will eventually start increasing when the learning rate becomes too large. The "sweet spot" is typically just before the loss begins to rise steeply, showing the onset of instability.

Model-specific heuristics use community knowledge. BERT models typically work well with 2e-5 to 3e-5. Larger models like RoBERTa-large often benefit from slightly lower rates. GPT-style decoder models may need different ranges depending on the task. These heuristics encode the collective experience of thousands of practitioners and give reliable starting points. When starting a new fine-tuning project, these community-established defaults are your best first guess.