Michael Brenndoerfer | Data & AI, Private Equity, Technology

Michael Brenndoerfer

Fine-tuning Learning Rates › Learning Rate Warmup for Fine-tuning

Michael BrenndoerferNovember 27, 20259 min read

Language AI HandbookAll chapters

Part 34 of 67Fine-tuning Fundamentals

Chapter 252Fine-tuning Learning Rates › Learning Rate Warmup for Fine-tuning

Reading section 5 of 12 · Learning Rate Warmup for Fine-tuning · 9 minSection 5 of 12 · 9 min

Section 5 of 12 · 9 min read

Fine-tuning Learning Rates › Learning Rate Warmup for Fine-tuningOne section · Sections in this chapter5 / 12 ↓

Sections in this chapter12

Reading tools Definitions & inline links

Choose your expertise level to adjust how many terms are explained. Beginners see more tooltips, experts see fewer to maintain reading flow. Hover over underlined terms for instant definitions.

Article links

Make inline references clickable

Warmup refers to starting training with a very small learning rate and gradually increasing it to the target value over a specified number of steps. This technique is especially important for fine-tuning. While warmup was originally motivated by considerations specific to the Adam optimizer and very deep networks, its benefits during fine-tuning arise from slightly different concerns: the unique vulnerability of pre-trained representations during the earliest phase of adaptation.

Think of warmup as the first few minutes of a workout. Jumping straight into maximum-intensity exercise risks injury precisely because the body is not yet prepared for peak exertion. Similarly, starting fine-tuning at the full target learning rate risks "injuring" the pre-trained representations before the optimizer has had a chance to build reliable gradient estimates. Warmup gives a controlled ramp-up period where the optimizer calibrates itself before making full-size updates.

The warmup phase serves several distinct but complementary purposes that interact to stabilize the beginning of fine-tuning. Understanding each of these purposes helps you reason about how long warmup should last and what can go wrong when it is omitted. In practice, models fine-tuned without warmup often show erratic early loss behavior, sometimes followed by eventual recovery, but frequently exhibiting worse final performance than runs that included warmup from the start.

Why Warmup Matters for Fine-tuningLink Copied

At the start of fine-tuning, several factors make large learning rates dangerous:

Task head initialization: The newly added classification head is randomly initialized while the pre-trained backbone is well-optimized. This creates an asymmetry in the network: the backbone produces real representations, but the randomly initialized head turns these into in effect random predictions. The gradients flowing back from this random head can be large and poorly directed. Large initial updates from the random head can destabilize the entire network, propagating noise deep into the pre-trained layers. Warmup allows the task head to "catch up" to the backbone before the learning rate becomes large enough to cause damage.
Gradient variance: Early mini-batches may not be representative of the full dataset, leading to noisy gradient estimates. The first few batches might happen to contain unusual examples or reflect sampling artifacts rather than the true data distribution. Small learning rates during warmup prevent over-reacting to this noise, giving the optimizer time to aggregate gradient information from more representative samples.
Adam statistics: As we discussed in the Adam optimizer chapter, Adam maintains running estimates of gradient moments, specifically the first moment (mean) and second moment (variance) of the gradients. These estimates are unreliable early in training because they are initialized to zero and require many updates to become accurate. The bias correction terms in Adam help, but they do not fully compensate for the lack of historical information. Warmup gives these moment estimates time to stabilize before making large parameter updates that rely on their accuracy.

The interaction between these three factors is multiplicative in the worst case. A poorly calibrated task head generating large, noisy gradients, evaluated against a dataset subset that is not representative, feeding into an Adam optimizer with unreliable moment estimates, could easily result in a catastrophic early update that erases key pre-trained representations. Warmup breaks this chain by so that no single factor has time to fully develop before the update magnitude is brought under control.

Linear warmupLink Copied

The most common warmup schedule increases the learning rate linearly from 0 to the target value over the warmup period. This approach is simple to implement, easy to understand, and effective in practice.

We want to compute a learning rate that starts at zero on the first step and reaches the target value exactly at the end of the warmup period. At step $t$ during warmup, the learning rate is:

\eta_t = \frac{t}{T_{\text{warmup}}} \cdot \eta_{\text{target}}

where:

$\eta_t$ : the learning rate at step $t$ (what we compute and apply to the optimizer)
$t$ : the current training step, starting at 0
$T_{\text{warmup}}$ : the total number of warmup steps (defining the duration of the ramp-up phase)
$\eta_{\text{target}}$ : the target learning rate to reach at the end of warmup (the peak rate)

Why does this formula make sense? The ratio $\frac{t}{T_{\text{warmup}}}$ is the fraction of the warmup phase completed, which grows linearly from 0 to 1 as $t$ increases from 0 to $T_{\text{warmup}}$ . Multiplying this fraction by $\eta_{\text{target}}$ scales the learning rate linearly from 0 up to the target value. The linear relationship means that the learning rate increases by a constant amount at each step. This gives a predictable and smooth ramp-up that is easy to reason about and debug.

For fine-tuning, warmup typically spans 5-10% of total training steps. For a fine-tuning run of 10,000 steps, you might use 500-1,000 warmup steps. This duration allows Adam's moment estimates to stabilize and the task head to begin learning patterns without wasting training time in an overly cautious exploration period.

In[10]:

Code

def linear_warmup_schedule(step, warmup_steps, target_lr):
    """
    Calculate learning rate during linear warmup.

    Args:
        step: Current training step
        warmup_steps: Number of warmup steps
        target_lr: Target learning rate after warmup

    Returns:
        Learning rate for this step
    """
    if step < warmup_steps:
        return target_lr * (step / warmup_steps)
    return target_lr


def warmup_with_decay_schedule(step, warmup_steps, total_steps, target_lr):
    """
    Linear warmup followed by linear decay to zero.

    This is the most common schedule for fine-tuning BERT-style models.
    """
    if step < warmup_steps:
        # Linear warmup
        return target_lr * (step / warmup_steps)
    else:
        # Linear decay
        decay_steps = total_steps - warmup_steps
        remaining_steps = total_steps - step
        return target_lr * (remaining_steps / decay_steps)

Let's visualize the warmup schedule:

In[11]:

Code

import matplotlib.pyplot as plt
import numpy as np

use_book_style()
plt.rcParams["figure.figsize"] = (6.0, 4.0)

total_steps = 10000
warmup_steps = 1000
target_lr = 2e-5

steps = np.arange(total_steps)
lrs = [linear_warmup_schedule(s, warmup_steps, target_lr) for s in steps]

fig, ax = plt.subplots()
ax.plot(steps, [lr * 1e5 for lr in lrs], linewidth=1)
ax.axvline(
    x=warmup_steps,
    color=theme_color("red"),
    linestyle="--",
    label="End of warmup",
)
ax.fill_between(
    steps[: warmup_steps + 1],
    0,
    [lr * 1e5 for lr in lrs[: warmup_steps + 1]],
    alpha=0.3,
    color=theme_color("orange"),
    label="Warmup phase",
)
ax.set_xlabel("Training Step")
ax.set_ylabel(r"Learning Rate ($\times 10^{-5}$)")
ax.set_title("Linear Warmup Schedule")
ax.legend()
ax.set_xlim(0, total_steps)
ax.set_ylim(0, target_lr * 1.1 * 1e5)
polish_axes(ax)

This schedule prevents the optimizer from making dangerously large updates during the initial training phase, when gradient estimates are most volatile. The shaded warmup region is the protective phase where the model's pre-trained knowledge is most vulnerable to disruption.

Out[13]:

Visualization

Line plot showing three learning rate warmup curves for 5%, 10%, and 15% warmup ratios over the first 3000 training steps. — Comparison of warmup duration effects over the first 3,000 training steps. Shorter warmup (5%, 500 steps) reaches the peak learning rate faster but gives less stabilization time. Longer warmup (15%, 1,500 steps) gives more protection but delays the effective learning phase. The 10% default (1,000 steps) balances both concerns for most fine-tuning tasks.

Gradual Unfreezing as WarmupLink Copied

A more aggressive warmup strategy combines warmup with gradual unfreezing. Instead of training all layers from the start with small learning rates, you begin by training only the task head, then progressively unfreeze layers from top to bottom as training proceeds:

Epoch 1: Train only the task head
Epoch 2: Unfreeze top transformer layers
Epoch 3: Unfreeze middle layers
Epoch 4+: Train all layers

This approach, introduced in ULMFiT, gives an extreme form of "warmup" where lower layers literally receive zero gradient until later in training. The philosophy is that by the time lower layers are unfrozen, the task head and upper layers have already adapted substantially. This gives more stable and informative gradients. The gradients that flow into the lower layers at the point of unfreezing are no longer the confused, high-variance signals of a randomly initialized network, but the more structured, task-relevant signals of a partially adapted one.

While effective, gradual unfreezing requires more epochs and careful tuning of the unfreezing schedule. It has fallen somewhat out of favor for transformer fine-tuning because LLRD reaches similar goals with less complexity and in fewer training epochs. However, it remains a useful technique when maximum caution is warranted, particularly in low-data regimes where even a brief period of high-variance gradient updates could cause irreversible damage to the pre-trained representations. If you find that standard LLRD with warmup still results in unstable early training, gradual unfreezing is worth adding to your toolkit.