The Role of Learning Rate Scheduling in Training a Neural Network

What is the role of learning rate scheduling in training a neural network

If there’s one hyperparameter that can single-handedly make or break a training run, it’s the learning rate. Set it wrong and your model either crawls toward convergence at a snail’s pace or blows up into NaN losses within a dozen steps. But here’s the thing I learned the hard way: a single fixed learning rate for an entire training run is almost never optimal. What you actually want is a learning rate that changes intelligently over time — that’s the whole idea behind learning rate scheduling.

Why a Fixed Learning Rate Isn’t Enough

Recall the basic gradient descent update:

$$\theta_{t+1} = \theta_t – \eta \nabla_\theta J(\theta_t)$$

The learning rate $\eta$ controls how big a step you take in the direction of the negative gradient. Early in training, when parameters are far from any good solution, you generally want larger steps to make fast progress. But as training progresses and parameters approach a good minimum, large steps become a liability — they cause the optimizer to overshoot, oscillate around the minimum, or even bounce out of a good region entirely.

This creates a fundamental tension: no single fixed value of $\eta$ is ideal for the entire training process. A learning rate schedule resolves this tension by varying $\eta$ over time according to some predefined (or adaptive) rule.

The General Idea

Formally, instead of using a constant $\eta$, we define $\eta$ as a function of the current training step or epoch:

$$\eta_t = f(t)$$

$$\theta_{t+1} = \theta_t – \eta_t \nabla_\theta J(\theta_t)$$

The function $f(t)$ is the “schedule.” Different schedules encode different beliefs about how the learning rate should evolve — decay smoothly, decay in sudden steps, oscillate periodically, or ramp up first before decaying.

Common Learning Rate Schedules

1. Step Decay

The learning rate is multiplied by a decay factor $\gamma$ every fixed number of epochs $s$:

$$\eta_t = \eta_0 \cdot \gamma^{\lfloor t / s \rfloor}$$

For example, starting at $\eta_0 = 0.1$, decaying by $\gamma = 0.1$ every 30 epochs, gives you $0.1 \to 0.01 \to 0.001$ at epochs 30 and 60.

2. Exponential Decay

A smoother alternative to step decay, where the learning rate decays continuously:

$$\eta_t = \eta_0 \cdot e^{-kt}$$

where $k$ is a decay rate constant.

3. Cosine Annealing

Popularized in modern deep learning, this schedule smoothly decreases the learning rate following a cosine curve from $\eta_0$ down to a minimum value $\eta_{min}$ over $T$ total steps:

$$\eta_t = \eta_{min} + \frac{1}{2}(\eta_0 – \eta_{min})\left(1 + \cos\left(\frac{t\pi}{T}\right)\right)$$

Cosine annealing has the nice property of decaying slowly at first, then quickly through the middle, then slowly again near the end — often outperforming step decay in practice.

4. Warmup

Especially important for transformer-based architectures, warmup starts the learning rate near zero and linearly ramps it up over the first few hundred or thousand steps, before switching to a decay schedule:

$$\eta_t = \eta_0 \cdot \frac{t}{T_{warmup}}, \quad t \le T_{warmup}$$

Warmup helps stabilize early training, particularly for architectures sensitive to poor initialization or large early gradients (like transformers with layer normalization).

5. Cyclical Learning Rates and Warm Restarts

Instead of monotonically decreasing, cyclical schedules oscillate the learning rate between a lower and upper bound repeatedly throughout training. Stochastic Gradient Descent with Warm Restarts (SGDR) periodically resets the learning rate to a high value and lets cosine annealing decay it again, effectively giving the optimizer repeated “fresh starts” that can help it escape local minima and explore more of the loss landscape.

6. Reduce on Plateau

Rather than following a fixed formula, this adaptive approach monitors a validation metric and reduces the learning rate (e.g., by a factor of 10) whenever that metric stops improving for a specified number of epochs (“patience”). This is a pragmatic, data-driven schedule that reacts to actual training dynamics rather than a predetermined curve.

Comparison Table

ScheduleBehaviorBest suited for
Step decaySudden drops at fixed intervalsClassic CNN training (e.g., ResNet on ImageNet)
Exponential decaySmooth continuous decayGeneral-purpose training
Cosine annealingSmooth curve, slow-fast-slow decayModern CV and NLP models
WarmupRamps up then decaysTransformers, large batch training
Cyclical / warm restartsRepeated up-down cyclesEscaping local minima, ensembling effects
Reduce on plateauReactive to validation metricWhen you don’t want to hand-pick a schedule shape

Visualizing a Typical Schedule (Warmup + Cosine Decay)

flowchart LR
    A[Step 0: LR near zero] --> B[Warmup phase: LR increases linearly]
    B --> C[Peak LR reached]
    C --> D[Cosine decay phase: LR decreases smoothly]
    D --> E[Final steps: LR approaches minimum]

Implementing Schedules in Code

Step decay in PyTorch:

import torch

model = torch.nn.Linear(10, 1)
optimizer = torch.optim.SGD(model.parameters(), lr=0.1, momentum=0.9)
scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=30, gamma=0.1)

for epoch in range(90):
    train_one_epoch(model, optimizer)
    scheduler.step()

Cosine annealing in PyTorch:

scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=100, eta_min=1e-6)

Warmup followed by cosine decay (a common pattern for transformers):

from torch.optim.lr_scheduler import LambdaLR
import math

def warmup_cosine_schedule(step, warmup_steps, total_steps):
    if step < warmup_steps:
        return step / max(1, warmup_steps)
    progress = (step - warmup_steps) / max(1, total_steps - warmup_steps)
    return 0.5 * (1 + math.cos(math.pi * progress))

scheduler = LambdaLR(optimizer, lr_lambda=lambda step: warmup_cosine_schedule(step, 1000, 100000))

Reduce on plateau in PyTorch:

scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(optimizer, mode='min', factor=0.1, patience=5)

for epoch in range(epochs):
    train_loss = train_one_epoch(model, optimizer)
    val_loss = validate(model)
    scheduler.step(val_loss)

In Keras:

lr_schedule = tf.keras.optimizers.schedules.CosineDecay(
    initial_learning_rate=0.01, decay_steps=10000
)
optimizer = tf.keras.optimizers.Adam(learning_rate=lr_schedule)

The Relationship Between Learning Rate and Batch Size

A subtlety that trips up a lot of practitioners: learning rate schedules don’t operate in isolation from other hyperparameters, especially batch size. Since gradient noise decreases as batch size increases (as discussed in the SGD article), larger batches generally tolerate — and often benefit from — larger learning rates. The commonly cited “linear scaling rule” suggests that if you multiply your batch size by $k$, you should roughly multiply your peak learning rate by $k$ as well, while extending your warmup period proportionally to keep early training stable. Ignoring this relationship is one of the most common reasons a “known good” learning rate schedule suddenly performs poorly after someone changes batch size for a different hardware setup.

The One-Cycle Policy

A related but distinct approach to learning rate scheduling, proposed by Leslie Smith, is the one-cycle policy. Rather than simply warming up and then decaying, one-cycle training ramps the learning rate up to a relatively high maximum over the first part of training, then decays it — often below its starting value — over the remainder, while simultaneously moving momentum in the opposite direction (lower momentum when learning rate is high, higher momentum when learning rate is low). This combination has been shown to enable what’s sometimes called “super-convergence” — reaching good accuracy in a fraction of the epochs that a more conservative schedule would require. It illustrates a broader point: learning rate and momentum scheduling are often best thought of together rather than in isolation.

Diagnosing Whether Your Schedule Is Working

A few signals to watch for when evaluating whether your chosen learning rate schedule is appropriate for your training run:

  • Loss spikes right after a scheduled drop: this can indicate the drop was too aggressive, or that momentum/optimizer state wasn’t well-adapted to the new, smaller learning rate.
  • Validation loss plateaus long before the schedule finishes decaying: your schedule may be decaying too slowly, wasting training time at an unnecessarily high learning rate.
  • Training loss doesn’t move at all in the first several hundred steps: warmup may be too long, or the very earliest learning rate may be too close to zero for meaningful progress.
  • Loss is jagged and noisy throughout training, never really smoothing out: the learning rate might be too high throughout the entire schedule, not just at a specific phase — consider lowering the peak value rather than just adjusting the decay shape.

Learning Rate Range Test

Before committing to a specific schedule, a practical technique called the learning rate range test (also introduced by Leslie Smith) can help you find a reasonable range of learning rates to schedule between. The idea: run a short training pass where the learning rate increases linearly or exponentially from a very small value to a very large one over a few hundred or thousand steps, while recording the loss at each step. Plotting loss against learning rate typically reveals a clear pattern — loss decreases as the learning rate increases from too-small values, reaches a minimum somewhere in a reasonable middle range, and then increases sharply once the learning rate becomes too large. The learning rate just before that sharp increase is often a strong candidate for your schedule’s peak value.

# Conceptual sketch of a learning rate range test
lrs = np.geomspace(1e-6, 1.0, num=500)
losses = []
for lr in lrs:
    set_optimizer_lr(optimizer, lr)
    loss = train_one_step(model, optimizer, batch)
    losses.append(loss)
# Plot losses against lrs on a log-x-axis to find the ideal range

Why Scheduling Matters: The Intuition

Think of learning rate scheduling as similar to how a person searching for the lowest point in a foggy valley might behave. At first, with the fog thick and visibility low, you take big confident strides in the direction that seems downhill — you’re not worried about overshooting because you’re still far from the target. As you get closer and start to sense you’re near the bottom, you naturally slow down and take smaller, more careful steps to avoid overshooting past the true lowest point. Learning rate scheduling encodes exactly this behavior into the optimization process.

Advantages of Learning Rate Scheduling

  1. Faster early progress — larger learning rates early in training accelerate initial convergence.
  2. Better final convergence — smaller learning rates later in training allow fine, precise adjustments near a good minimum.
  3. Improved stability — warmup schedules in particular prevent early divergence caused by poor initialization or large initial gradients.
  4. Better generalization — some schedules (cyclical, warm restarts) have been empirically shown to lead to flatter minima, which often generalize better to unseen data.
  5. Reduced need for manual re-tuning — reactive schedules like reduce-on-plateau adapt automatically to the actual training dynamics rather than requiring you to guess the right decay curve in advance.

Disadvantages and Limitations

  1. More hyperparameters — a schedule introduces its own parameters (decay rate, step size, warmup length, patience) that themselves need tuning.
  2. Schedule choice can be non-obvious — different architectures and datasets often respond best to different schedule shapes, and there’s no universal “best” schedule.
  3. Reactive schedules can lag — reduce-on-plateau only reacts after a metric has already stopped improving, which can waste several epochs before adjusting.
  4. Interacts with other hyperparameters — batch size, optimizer choice, and weight decay all interact with the learning rate schedule, making isolated tuning difficult.

Real-World Use Cases

  • Image classification (ResNet on ImageNet): Classic step decay schedules (dropping the learning rate by 10x at fixed epochs) remain a benchmark approach.
  • Transformer training (BERT, GPT): Linear or cosine warmup followed by linear or cosine decay is now close to universal.
  • Fine-tuning pretrained models: Small, often constant or slowly-decaying learning rates prevent catastrophic forgetting of pretrained weights.
  • Super-convergence / one-cycle policy: Leslie Smith’s one-cycle learning rate policy — ramp up quickly, then decay — has enabled dramatically faster training in some computer vision tasks.

Best Practices

  • Always pair a learning rate schedule with a reasonable starting learning rate; a schedule cannot fix a badly chosen $\eta_0$.
  • Use warmup for transformer architectures and any model prone to early-training instability.
  • Try cosine annealing as a strong general-purpose default when you’re unsure which schedule to pick.
  • Use reduce-on-plateau when you want a low-maintenance, data-driven approach without manually specifying a decay curve.
  • Log and plot your learning rate over time alongside your loss curve — visualizing both together makes it much easier to diagnose training issues.
  • Re-tune your schedule’s hyperparameters whenever you significantly change batch size, since the two interact strongly (larger batches often pair well with larger peak learning rates).

Frequently Asked Questions

Do I need a learning rate schedule if I’m already using Adam? Almost always, yes. Adam’s adaptive per-parameter scaling addresses relative differences between parameters, but it doesn’t automatically know when the overall training process has entered a phase where smaller global steps would help fine-tune toward a good minimum. Even with Adam, a schedule — especially warmup followed by decay for larger models — is standard practice rather than an optional extra.

How long should warmup last? This depends heavily on model size and batch size, but as a rough starting point, warmup periods commonly span anywhere from a few hundred steps for smaller models to several thousand steps for very large transformer models trained with large batch sizes. If you observe instability (loss spikes or divergence) in the first portion of training, extending warmup is often the first thing worth trying before reducing the peak learning rate itself.

What’s the difference between decaying by epoch versus by step? Decaying by epoch (as in classic step decay) ties the schedule to how many full passes through the dataset have completed, which can behave inconsistently if you change dataset size or batch size. Decaying by step ties the schedule directly to the number of optimizer updates performed, which tends to transfer more consistently across different dataset sizes and batch configurations — this is why most modern large-scale training recipes (particularly for transformers) define schedules in terms of training steps rather than epochs.

Can I combine multiple schedule types? Yes, and this is actually the norm rather than the exception for large-scale training — warmup combined with cosine decay is a very common combination, and reduce-on-plateau can even be layered on top of a base schedule as a safety mechanism for detecting unexpected stalls. Just be cautious about overly complex compound schedules, since they add more hyperparameters that need to be validated together rather than in isolation.

Does the choice of schedule matter more than the choice of optimizer? For many practical problems, yes — a poorly scheduled Adam run can perform worse than a well-scheduled SGD run, and vice versa. Learning rate scheduling and optimizer choice interact strongly enough that it’s usually more productive to tune them together (via a small grid or a learning rate range test) rather than fixing one and only tuning the other.

A Quick Mental Model to Keep

Think of the learning rate schedule as answering a single, ongoing question throughout training: “given how far along we are, how confident should we be in taking a big step right now?” Early on, when you’re far from a good solution and there’s little risk in taking big strides, the answer is “very confident.” Later on, once you’re near a good minimum and want precision rather than speed, the answer shifts to “much less confident.” Every schedule discussed in this article is just a different mathematical way of encoding that shifting confidence over time.

Summary

Learning rate scheduling exists to resolve a simple but important tension: the optimal learning rate is different at the beginning of training than it is near the end. By varying the learning rate over time — whether through predetermined curves like step decay and cosine annealing, or reactive strategies like reduce-on-plateau — you get faster initial convergence, more stable training, and better final results than a single fixed learning rate could ever achieve. Modern deep learning, especially in transformer-based architectures, has come to rely on carefully designed schedules (typically warmup followed by decay) as a near-mandatory part of the training recipe rather than an optional extra.

References and Further Reading

  • Smith, L. N. (2017). “Cyclical Learning Rates for Training Neural Networks.” https://arxiv.org/abs/1506.01186
  • Loshchilov, I., & Hutter, F. (2017). “SGDR: Stochastic Gradient Descent with Warm Restarts.” ICLR. https://arxiv.org/abs/1608.03983
  • Goyal, P., et al. (2017). “Accurate, Large Minibatch SGD: Training ImageNet in 1 Hour.” (Introduces linear warmup) https://arxiv.org/abs/1706.02677
  • PyTorch learning rate scheduler documentation: https://pytorch.org/docs/stable/optim.html#how-to-adjust-learning-rate
  • TensorFlow learning rate schedules documentation: https://www.tensorflow.org/api_docs/python/tf/keras/optimizers/schedules
Total
1
Shares

Leave a Reply

Previous Post
Explain the concept of early stopping in neural network training.

Early Stopping in Neural Network Training: A Practical Guide

Next Post
What is the Adam optimizer and how does it work

The Adam Optimizer: How It Works and Why It’s the Default Choice

Related Posts