Introduction

In deep learning, setting the right learning rate is crucial for model convergence and performance. A fixed learning rate may not always be optimal throughout training, which is where learning rate schedules come in. Learning rate schedules adjust the learning rate over time to improve training stability and efficiency.

What is a Learning Rate Schedule?

A learning rate schedule defines how the learning rate changes over the course of training. This adjustment helps models converge faster, avoid oscillations, and sometimes escape local minima.

Types of Learning Rate Schedules

1. Step Decay

Step decay reduces the learning rate by a fixed factor after a specific number of epochs.

Formula

If the initial learning rate is represented as alpha and the decay factor as gamma, the learning rate after each step is updated by: alpha = alpha * gamma

Benefits

• Simple to Implement: Easy to tune.
• Stable Convergence: Often leads to stable training.

2. Exponential Decay

In exponential decay, the learning rate decreases exponentially over time.

Formula

alpha_t = alpha_0 * e^(-k * t)

where alpha_0 is the initial learning rate, k is the decay rate, and t is the epoch or step.

Benefits

• Smooth Decay: Reduces learning rate gradually.
• Less Oscillation: Helps in steady convergence.

3. Cosine Annealing

Cosine annealing gradually reduces the learning rate following a cosine function, allowing it to "warm restart" periodically.

Formula

alpha_t = alpha_min + 0.5 * (alpha_max - alpha_min) * (1 + cos((t * pi) / T))
where T is the period and t is the current step.

Benefits

• Effective for Cyclic Training: Often used with cyclical learning rates.
• Encourages Exploration: Helps the model explore new solutions.

4. Cyclical Learning Rate (CLR)

In cyclical learning rates, the learning rate oscillates between a lower and upper bound, restarting every few steps.

Benefits

• Boosts Generalization: Helps escape saddle points.
• Adaptable: Effective for noisy, non-convex landscapes.

5. Warmup + Decay

The warmup phase starts with a low learning rate, gradually increasing it to the desired value. This is often followed by one of the decay strategies.

Benefits

• Improved Stability: Reduces large updates at the beginning.
• Ideal for Large Models: Often used for transformer models.

Choosing the Right Schedule

The best learning rate schedule depends on the model, dataset, and training environment. Here are some general recommendations:

• Simple Models: Step decay is usually sufficient.
• Deep Networks: Cosine annealing or CLR can improve performance.
• Transformers: Warmup with exponential decay or cosine annealing is commonly used.

Example Code (PyTorch)

Here's a sample implementation for a cosine annealing schedule with warmup.

import torch.optim as optim

# Assuming `model` and `initial_lr` are defined
optimizer = optim.Adam(model.parameters(), lr=initial_lr)
scheduler = optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=100)

for epoch in range(num_epochs):
    train_one_epoch(model, optimizer)  # Custom function for training
    scheduler.step()

Conclusion

Learning rate schedules are key for maximizing the efficiency and stability of training deep learning models. By adapting the learning rate, models can achieve faster convergence and better generalization.