Mastering Learning Rate Schedules

Introduction

In our previous lesson, we established a principled approach to initializing network weights, ensuring that signal variance remains stable during both the forward and backward passes. This was a critical step in preventing the vanishing and exploding gradient problems from the outset. We now possess a robust toolkit: advanced optimizers like Adam, regularization techniques like Dropout, and intelligent initialization strategies.

However, a crucial piece of the optimization puzzle remains—the learning rate, $\alpha$. Until now, we have treated it as a static hyperparameter. While a well-chosen fixed learning rate can work, it is often a compromise. In this lesson, we will explore a more dynamic approach. Our objective is to implement a learning rate schedule and explain its purpose. We will see how strategically varying the learning rate during training can lead to faster convergence and better final models, moving us from a fixed setting to an adaptive optimization strategy.

The Learning Rate's "Goldilocks" Dilemma

We know that a learning rate that is too high can cause the optimizer to overshoot the minimum and diverge, while one that is too small results in painfully slow convergence. But let's think about this not as a static choice, but in the context of the training process over time.

Here's a question for you: What are the ideal characteristics of a learning rate at the beginning of training, when the randomly initialized weights place us far from an optimal solution, versus at the end of training, when we are approaching a minimum in the loss landscape?

Think about the trade-off between the speed of descent and the precision of convergence.

At the start of training, the weights are far from optimal, and the gradient is typically large and points roughly in the right direction. A larger learning rate is beneficial here, as it allows the model to take big steps and make rapid progress toward a good region of the loss landscape.

However, as training progresses and the model approaches a local or global minimum, the loss landscape becomes flatter and more complex. In this phase, a large learning rate becomes a liability. It can cause the optimizer to constantly overshoot the minimum, leading to oscillations and preventing it from settling into a precise solution. Here, a smaller learning rate is needed to take more careful, smaller steps, allowing the model to finely tune its weights and converge into the basin of the minimum.

This inherent tension—the need for speed early on and precision later—is the core motivation for learning rate scheduling.

The following video provides an excellent visual intuition for this problem.

Learning Rate Schedules: A Principled Approach

A learning rate schedule (or scheduler) is a pre-defined strategy for adjusting the learning rate during training. Instead of a fixed value, the learning rate becomes a function of the current epoch or training step. The overarching goal is to harness the benefits of a large learning rate in the beginning and a small learning rate towards the end.

By systematically decreasing the learning rate, we are effectively damping the oscillations caused by noisy mini-batch gradients, especially as we get closer to a minimum. This helps stabilize the final stages of training.

Common Scheduling Policies

There are many ways to decrease the learning rate over time. We'll focus on a few of the most common and effective policies.

Step Decay

This is arguably the most intuitive scheduling policy. The learning rate is held constant for a set number of epochs and then "steps down," or is reduced by a multiplicative factor. This process is repeated at pre-defined epochs. For example, one might start with $\alpha =0.1$, run for 30 epochs, decay to $\alpha =0.01$, run for another 30 epochs, decay to $\alpha =0.001$, and so on.

The intuition is to let the optimizer make significant progress and converge as much as possible with the current learning rate. Once the loss plateaus, we decrease $\alpha$ to allow for more refined adjustments and settle into a better minimum.

The image below perfectly illustrates this process.

This method is also known as a piecewise constant schedule or, in frameworks like PyTorch, a MultiStepLR scheduler.

Exponential Decay

Instead of discrete steps, we can also decay the learning rate smoothly and continuously. Exponential decay follows the formula:

where ${\alpha }_t$ is the learning rate at epoch $t$, ${\alpha }_0$ is the initial learning rate, and $\gamma$ is a decay rate less than 1 (e.g., 0.95). After each epoch, the learning rate is multiplied by $\gamma$, resulting in an exponential decrease over time.

Another common formulation, which you will often see in deep learning frameworks, is:

where $d$ is a decay hyperparameter. This is a form of inverse decay.

These methods are simple to implement and only require tuning one or two extra hyperparameters ($\gamma$ or $d$).

Other Important Policies

While Step and Exponential decay are foundational, two other concepts are prevalent in modern deep learning research:

1. Cosine Decay: This popular schedule decreases the learning rate following the shape of a cosine curve. It starts high, decreases slowly, then accelerates its decay towards the middle of training, and finally slows down again to a very small value at the end. It's empirically very effective, especially in computer vision.

2. Warmup: Counter-intuitively, for very deep or complex models (like Transformers), it can be beneficial to start with a very small learning rate and increase it linearly over the first few epochs—a "warmup" phase—before starting the decay schedule. This prevents the large, unstable gradients at the very beginning of training (due to random initialization) from causing the optimization to diverge.

The "Dive into Deep Learning" resource provides a comprehensive look at these policies. You can explore this for more depth.

Implementation from Scratch

Let's now translate this theory into practice. The learning outcome for this lesson is to implement a scheduler. We will implement a StepDecay scheduler and integrate it into a conceptual training loop.

First, let's define our scheduler. It can be a simple function that takes the current epoch and the initial learning rate, and returns the updated learning rate based on a pre-defined schedule.

def step_decay_schedule(initial_lr: float, current_epoch: int, decay_factor: float, decay_epochs: int):
    """
    Implements a step decay learning rate schedule.

    Args:
        initial_lr: The learning rate at the beginning of training.
        current_epoch: The current epoch number (starting from 0).
        decay_factor: The factor by which to decay the learning rate (e.g., 0.1).
        decay_epochs: The number of epochs after which to decay.

    Returns:
        The learning rate for the current epoch.
    """
    return initial_lr * (decay_factor ** (current_epoch // decay_epochs))

Now, let's see how this function plugs into a standard training loop. At the beginning of each epoch, we calculate the new learning rate and update our optimizer with it.

import torch




# --- Hyperparameters ---
INITIAL_LR = 0.1
NUM_EPOCHS = 90
DECAY_FACTOR = 0.1
DECAY_EPOCHS = 30




# --- Dummy Model and Optimizer ---
model = torch.nn.Linear(10, 2) # A simple example model
optimizer = torch.optim.SGD(model.parameters(), lr=INITIAL_LR)




# --- Training Loop ---
for epoch in range(NUM_EPOCHS):



    # 1. Calculate the new learning rate for this epoch
    new_lr = step_decay_schedule(INITIAL_LR, epoch, DECAY_FACTOR, DECAY_EPOCHS)
    



    # 2. Update the optimizer's learning rate
    # This is the key step for integrating the scheduler
    for param_group in optimizer.param_groups:
        param_group['lr'] = new_lr

    print(f"Epoch {epoch+1}/{NUM_EPOCHS}, Current LR: {optimizer.param_groups[0]['lr']:.4f}")




    # --- Your training steps would go here ---
    # model.train()
    # for data, labels in train_loader:
    #     optimizer.zero_grad()
    #     outputs = model(data)
    #     loss = loss_fn(outputs, labels)
    #     loss.backward()
    #     optimizer.step()
    # -----------------------------------------




# Example Output:
# Epoch 1/90, Current LR: 0.1000
# ...
# Epoch 30/90, Current LR: 0.1000
# Epoch 31/90, Current LR: 0.0100
# ...
# Epoch 60/90, Current LR: 0.0100
# Epoch 61/90, Current LR: 0.0010
# ...

This implementation directly fulfills our learning objective. You can see how simple it is to control the learning rate dynamically. While deep learning libraries like PyTorch and TensorFlow have powerful built-in lr_scheduler modules that handle this automatically, understanding how to build it from scratch reveals that there is no "magic" involved—it is simply a programmatic way of adjusting a hyperparameter at each step of an iterative process.

Conclusion

Today, we've addressed one of the final core components of the deep learning training process: the dynamic adjustment of the learning rate. By moving from a static value to a scheduled one, we can design an optimization process that is both fast and precise.

Key Takeaways:

• Purpose: Learning rate schedules aim to resolve the conflict between needing a large learning rate for fast initial progress and a small learning rate for stable final convergence.
• Mechanism: By systematically decreasing $\alpha$ over time, schedulers help the optimizer take large steps when far from a minimum and smaller, more precise steps when close to one, damping oscillations and improving final model accuracy.
• Common Policies:
  • Step Decay: Reduces the learning rate by a fixed factor at specific epoch intervals. It's simple and effective.
  • Exponential Decay: Smoothly decreases the learning rate at every step.
  • Cosine Decay & Warmup: More advanced but powerful techniques common in state-of-the-art research.
• Implementation: A scheduler can be implemented as a simple function that is called at the start of each epoch to update the optimizer's learning rate parameter.

Preview of the Next Lesson:

We have now assembled a complete, general-purpose toolkit for training deep neural networks. We understand the optimization process, regularization, initialization, and learning rate control. With these fundamentals secure, we will now shift our focus to specialized architectures.

Our journey into computer graphics and state-of-the-art AI begins in earnest in the next module with Convolutional Neural Networks (CNNs). We will discover how the convolution operation, a cornerstone of classical image processing, was adapted to create networks that are exceptionally powerful for visual data. This will be our first major step toward understanding the models that power modern computer vision and, eventually, neural rendering.