A learning rate warmup can be applied to various learning rate schedules, such as a cosine schedule, to stabilize training and improve initial convergence. In deep learning frameworks, this is often configured using a warmup steps parameter. For example, a cosine scheduler can be configured to linearly increase the learning rate for the first 5 steps before applying standard cosine decay. Plotting the learning rate schedule over epochs visually demonstrates this initial linear increase followed by the cooling-down (decay) period.

```python
num_epochs = 20
scheduler = CosineScheduler(num_epochs, warmup_steps=5, base_lr=0.3, final_lr=0.01)
d2l.plot(torch.arange(num_epochs), [scheduler(t) for t in range(num_epochs)])
```

Claude

Google

To address the dilemma of choosing an initial learning rate that is either too small (causing slow progress) or too large (causing divergence), a simple strategy called optimizer warmup is used. During a warmup period, the learning rate gradually increases—typically linearly—from a small value to its initial maximum, after which it cools down until the end of the optimization process.

Optimizer Warmup

A cosine learning rate schedule, proposed by Loshchilov and Hutter (2016), dynamically adjusts the learning rate by following the shape of a cosine curve. It relies on the observation that the learning rate should not decrease too drastically at the beginning of training, and that the solution should be refined at the end using a very small learning rate. For learning rates in the range $$t \in [0, T]$$, this results in a schedule with the functional form:

$$\eta_t = \eta_T + \frac{\eta_0 - \eta_T}{2} \left(1 + \cos(\pi t/T)\right)$$

Here, $$\eta_0$$ is the initial learning rate and $$\eta_T$$ is the target rate at the maximum update step $$T$$. For steps $$t > T$$, the learning rate is simply pinned to $$\eta_T$$ without increasing it again.

Cosine Learning Rate Schedule

Dive into Deep Learning

The fundamental intuition behind using an optimizer warmup phase is that random parameter initialization, especially in advanced or very deep neural networks, often leads to unstable optimization and significant early divergence. A warmup period mitigates this by starting with a small learning rate, which effectively limits the amount of parameter divergence in the parts of the network that take the most time to make initial progress. Once the parameters have stabilized, the learning rate can be safely increased to avoid slow training.

Intuition Behind Optimizer Warmup

Research has shown that applying a warmup phase during optimization limits the amount of parameter divergence in very deep neural networks. Because the network weights are randomly initialized, the parts of the network that require the most time to make progress are highly susceptible to significant divergence early in training. Gradually increasing the learning rate during a warmup period mitigates this instability, leading to better initial convergence.

Effect of Warmup on Parameter Divergence

Learning Rate Warmup Schedule Example

To empirically observe the benefits of a learning rate warmup, a neural network can be trained using an optimizer configured with a warmup scheduler. The training metrics typically show that the network converges better initially—especially during the warmup epochs—compared to training without it. This improved early performance stabilizes the optimization process for advanced networks.

```python
net = net_fn()
trainer = torch.optim.SGD(net.parameters(), lr=0.3)
train(net, train_iter, test_iter, num_epochs, loss, trainer, device, scheduler)
```

Learning Rate Warmup Training Example

A cosine learning rate scheduler can be implemented from scratch as a custom Python class. The class calculates the decayed learning rate based on the current step according to a cosine curve. It optionally includes a warmup phase where the learning rate increases linearly. Once past the warmup steps but within the maximum update steps, the learning rate is determined by the formula $$\eta_t = \eta_T + \frac{\eta_0 - \eta_T}{2} \left(1 + \cos\left(\frac{\pi (t - t_{\text{warmup}})}{T_{\text{max\_steps}}}\right)\right)$$. The following code demonstrates this implementation and plots the resulting schedule:

```python
class CosineScheduler:
    def __init__(self, max_update, base_lr=0.01, final_lr=0,
               warmup_steps=0, warmup_begin_lr=0):
        self.base_lr_orig = base_lr
        self.max_update = max_update
        self.final_lr = final_lr
        self.warmup_steps = warmup_steps
        self.warmup_begin_lr = warmup_begin_lr
        self.max_steps = self.max_update - self.warmup_steps

    def get_warmup_lr(self, epoch):
        increase = (self.base_lr_orig - self.warmup_begin_lr) \
                       * float(epoch) / float(self.warmup_steps)
        return self.warmup_begin_lr + increase

    def __call__(self, epoch):
        if epoch < self.warmup_steps:
            return self.get_warmup_lr(epoch)
        if epoch <= self.max_update:
            self.base_lr = self.final_lr + (
                self.base_lr_orig - self.final_lr) * (1 + math.cos(
                math.pi * (epoch - self.warmup_steps) / self.max_steps)) / 2
        return self.base_lr

scheduler = CosineScheduler(max_update=20, base_lr=0.3, final_lr=0.01)
d2l.plot(torch.arange(num_epochs), [scheduler(t) for t in range(num_epochs)])
```

Cosine Learning Rate Scheduler Implementation

To empirically demonstrate the impact of a cosine learning rate schedule, a neural network can be trained using an optimizer configured with this scheduler. The schedule dictates that the learning rate decreases gently at the beginning, drops more steeply, and then refines the solution at the end with a very small step size. While this approach can lead to improved training results and final accuracy—particularly in computer vision tasks—such improvements are heuristic and not strictly guaranteed across all models and datasets.

```python
net = net_fn()
trainer = torch.optim.SGD(net.parameters(), lr=0.3)
train(net, train_iter, test_iter, num_epochs, loss, trainer, device,
      scheduler)
```

Learn Before

Related