Learning Rate Warmup

Learning Objectives

By the end of this section, you will:

1. Understand why warmup prevents early training failure
2. Compare linear and exponential warmup strategies
3. Choose appropriate warmup duration for your model
4. Implement warmup schedulers in PyTorch
5. Diagnose warmup-related training issues

Why This Matters: Large learning rates at the start of training can cause gradient explosion, NaN losses, or convergence to poor local minima. Warmup starts with a tiny learning rate and gradually increases it, giving the optimizer time to estimate gradient statistics before taking large steps.

Why Warmup is Necessary

Several factors make early training particularly fragile.

The Cold Start Problem

At initialization, Adam's moment estimates are unreliable:

• First moment (m): Initialized to 0, takes ~10 steps to stabilize
• Second moment (v): Initialized to 0, takes ~100 steps to stabilize
• Bias correction: Helps but doesn't fully compensate

Large learning rates with inaccurate moment estimates cause erratic updates.

Gradient Magnitude at Initialization

When Warmup is Most Important

Warmup Strategies

There are several ways to increase the learning rate during warmup.

Linear Warmup

The most common approach—linearly increase from 0 to target:

Where:

• $\eta_t$: Learning rate at step t
• $\eta_{\text{target}}$: Target learning rate (e.g., 1e-3)
• $T_{\text{warmup}}$: Warmup duration (steps or epochs)

1Linear Warmup Schedule:
2
3LR
4  η ─┤                    ────────────
5    │                  ╱
6    │                ╱
7    │              ╱
8    │            ╱
9    │          ╱
10    │        ╱
11    │      ╱
12    │    ╱
13  0 ─┼──╱───────────────────────────────
14    └──┬────────┬────────────────────
15       0    T_warmup            Epochs
16
17Linear: η(t) = η_target × t / T_warmup

Exponential Warmup

Start very small and grow exponentially:

Exponential warmup spends more time at low learning rates, which can be beneficial for very sensitive models.

Gradual Warmup (Recommended)

Start from a small non-zero value:

With $\eta_{\text{start}} = \eta_{\text{target}} / 10$, this provides a gentler start than pure linear warmup.

Comparison

Warmup Duration Selection

How long should warmup last?

Rules of Thumb

Batch-Based vs. Epoch-Based

Warmup can be specified in steps (batches) or epochs:

AMNL Warmup Configuration

Implementation

Our research implementation uses a simple but effective warmup approach that integrates cleanly with the training loop.

AMNL Research Implementation

Epoch 0: factor=0.1, Epoch 5: factor=0.55, Epoch 10: factor=1.0

0.1 + 0.9 * (5/10) = 0.1 + 0.45 = 0.55

1def get_lr_warmup_factor(epoch: int, warmup_epochs: int = 10) -> float:
2    """
3    Learning rate warmup (from V5 - it worked!)
4
5    Linearly ramp up learning rate from 0.1x to 1.0x over warmup_epochs.
6    """
7    if epoch < warmup_epochs:
8        # Linear warmup: 0.1 → 1.0 over warmup_epochs
9        return 0.1 + 0.9 * (epoch / warmup_epochs)
10    return 1.0

Integration with Training Loop

1# Training loop with warmup
2for epoch in range(epochs):
3    # Apply learning rate warmup
4    warmup_factor = get_lr_warmup_factor(epoch, warmup_epochs=10)
5    if epoch < 10:
6        for param_group in optimizer.param_groups:
7            param_group['lr'] = learning_rate * warmup_factor
8
9    # Training phase
10    model.train()
11    for batch_idx, (sequences, targets) in enumerate(train_loader):
12        # ... training code ...
13        pass
14
15    # ReduceLROnPlateau scheduler step (only after warmup)
16    if epoch >= 10:
17        scheduler.step(rmse_last)

Combined Schedule Visualization

1# Visualize the combined warmup + plateau schedule
2import matplotlib.pyplot as plt
3
4epochs = range(100)
5lrs = []
6base_lr = 1e-3
7warmup_epochs = 10
8
9# Simulate warmup phase
10for e in epochs:
11    if e < warmup_epochs:
12        factor = 0.1 + 0.9 * (e / warmup_epochs)
13        lrs.append(base_lr * factor)
14    else:
15        # After warmup, assume some plateau reductions
16        if e < 50:
17            lrs.append(base_lr)
18        elif e < 70:
19            lrs.append(base_lr * 0.5)  # First reduction
20        else:
21            lrs.append(base_lr * 0.25)  # Second reduction
22
23plt.figure(figsize=(10, 5))
24plt.plot(epochs, lrs, 'b-', linewidth=2)
25plt.axvline(x=10, color='red', linestyle='--', alpha=0.5, label='Warmup ends')
26plt.xlabel('Epoch')
27plt.ylabel('Learning Rate')
28plt.title('Warmup + ReduceLROnPlateau Schedule')
29plt.legend()
30plt.grid(True, alpha=0.3)
31plt.show()

Summary

In this section, we covered learning rate warmup:

1. Why needed: Adam's moments need time to stabilize
2. Linear warmup: $\eta_t = \eta \cdot t / T$(recommended)
3. Duration: 5-10% of total epochs
4. Combined scheduler: Warmup + cosine decay
5. Step-level updates: Call scheduler.step() after each batch

Looking Ahead: After warmup, we need a strategy for the main training phase. The next section covers cosine annealing with warm restarts—a schedule that can escape local minima through periodic learning rate increases.

With warmup understood, we explore cosine annealing schedules.