Lecture 11: Optimization and Learning Rates

Lecture notes on optimization methods, learning rate schedules, momentum, and adaptive optimizers in deep learning.

1. Introduction to Optimization

Training a deep neural network involves finding model parameters that minimize a loss function. Optimization algorithms update these parameters iteratively using gradients computed through backpropagation.

The basic gradient descent update is

\theta_{t+1} = \theta_t - \eta_t \nabla_\theta L(\theta_t)

where:

\theta_t
represents the model parameters at iteration t.
\eta_t
is the learning rate.
\nabla_\theta L(\theta_t)
is the gradient of the loss function.

The learning rate determines how far the parameters move in the direction opposite to the gradient.

1.1 Challenges in Deep Learning Optimization

Deep neural networks typically have non-convex loss functions. This makes optimization challenging because the loss surface may contain saddle points, flat regions, and regions with very different curvatures.

A learning rate that is too large can cause unstable updates or oscillations. A learning rate that is too small can make training unnecessarily slow.

Choosing an appropriate optimization algorithm and learning rate schedule is therefore important for efficient training.

2. Batch Size and Learning Rate

In deep learning, the batch size determines how many training examples are used to estimate the gradient before updating the model parameters. The choice of batch size affects the stability, computational efficiency, and behavior of the optimization process.

2.1 Full-Batch, Stochastic, and Mini-Batch Gradient Descent

There are three common approaches to computing gradients:

Full-batch gradient descent: Uses the entire training dataset to compute each gradient update. It provides an accurate gradient estimate but can be computationally expensive.
Stochastic gradient descent (SGD): Uses one training example per update. Its gradients are noisy, but updates are computationally inexpensive.
Mini-batch gradient descent: Uses a subset of the training data for each update. This is the most common approach in modern deep learning because it balances computational efficiency and gradient stability.

For a mini-batch of size B, the gradient estimate is

g_t = \frac{1}{B}\sum_{i=1}^{B} \nabla_\theta \ell_i(\theta_t)

where \ell_i is the loss for training example i in the mini-batch.

The parameter update becomes

\theta_{t+1} = \theta_t - \eta_t g_t

2.2 How Batch Size Affects Optimization

Smaller batches generally produce noisier gradient estimates. This noise can help the optimizer explore the loss surface, but it can also cause fluctuations during training.

Larger batches generally produce more stable gradient estimates and can take advantage of parallel computation on GPUs. However, they require more memory and may provide fewer parameter updates per epoch.

2.3 Relationship Between Batch Size and Learning Rate

Batch size and learning rate are closely related. When the batch size increases, the gradient estimate often becomes less noisy, which may allow a larger learning rate.

However, increasing the learning rate proportionally with batch size is not always optimal. The appropriate learning rate depends on the model architecture, optimizer, dataset, and training conditions.

In practice, batch size and learning rate should be tuned together rather than independently.

3. Learning Rate Schedules

The learning rate controls the size of each parameter update during optimization. Instead of using a constant learning rate throughout training, we can adjust it over time using a learning rate schedule.

A relatively large learning rate can help the model make faster progress early in training, while a smaller learning rate can help stabilize training and refine the parameters later.

3.1 Exponential Decay

Exponential decay gradually reduces the learning rate by multiplying it by a constant factor:

\eta_t = \eta_0 \gamma^t

where \eta_0 is the initial learning rate and 0 < \gamma < 1 is the decay factor.

This schedule produces a smooth decrease in the learning rate over time.

3.2 Step Decay

Step decay reduces the learning rate at specific intervals rather than continuously.

For example, the learning rate can be halved every K epochs:

\eta_t = \eta_0 \left(\frac{1}{2}\right)^{\lfloor t/K \rfloor}

This allows the optimizer to make larger updates initially and smaller updates later.

3.3 Inverse Time Decay

Inverse time decay decreases the learning rate according to

\eta_t = \frac{\eta_0}{1 + kt}

where k > 0 controls how quickly the learning rate decreases.

Unlike exponential decay, inverse time decay decreases more slowly during later stages of training.

3.4 Cyclical Learning Rates

Cyclical learning rate schedules periodically increase and decrease the learning rate between predefined bounds.

The purpose is to allow the optimizer to explore different regions of the loss surface instead of continuously decreasing the step size.

A larger learning rate can encourage exploration, while a smaller learning rate allows more careful parameter updates.

3.5 Cosine Annealing

Cosine annealing decreases the learning rate following a cosine-shaped curve:

\eta_t = \eta_{\min} + \frac{1}{2}(\eta_{\max}-\eta_{\min})\left(1+\cos\left(\frac{\pi t}{T}\right)\right)

where T is the duration of the schedule, and \eta_{\max} and \eta_{\min} are the maximum and minimum learning rates.

This schedule smoothly decreases the learning rate toward its minimum value.

3.6 Learning Rate Warmup

Learning rate warmup starts training with a small learning rate and gradually increases it to the target value.

A simple linear warmup schedule is

\eta_t = \eta_{\max}\frac{t}{T_w}

for 0 \leq t \leq T_w, where T_w is the number of warmup steps. Warmup can help prevent unstable parameter updates at the beginning of training, particularly when using large batch sizes.

After warmup, another learning rate schedule, such as cosine annealing, is often applied.

3.7 Early Stopping

Early stopping is another practical strategy for controlling training.

Instead of choosing a fixed number of training epochs, we monitor performance on a validation dataset and stop training when validation performance no longer improves.

Early stopping can reduce unnecessary computation and help prevent overfitting. Unlike learning rate decay, it determines when training should stop rather than changing the step size.

4. Momentum and Stochastic Gradient Descent

4.1 Limitations of Standard SGD

Stochastic Gradient Descent (SGD) updates model parameters in the opposite direction of the gradient:

\theta_{t+1} = \theta_t - \eta \nabla L(\theta_t)

Although SGD is simple and widely used, it can be inefficient when the loss surface has different curvatures in different directions.

For example, consider a narrow valley in the loss landscape. The gradient may point toward the opposite side of the valley rather than along the valley toward the minimum. As a result, SGD can oscillate from one side to the other, slowing down convergence.

4.2 Momentum

Momentum improves SGD by incorporating information from previous gradient updates. Instead of relying only on the current gradient, it maintains a running average of past gradients.

A common formulation is:

v_t = \mu v_{t-1} + g_t

\theta_{t+1} = \theta_t - \eta v_t

where:

v_t
is the accumulated gradient direction.
\mu
is the momentum coefficient, typically between 0.9 and 0.99.
g_t = \nabla L(\theta_t)
is the current gradient.
\eta
is the learning rate.

Momentum reduces oscillations by combining gradient information across iterations. In a narrow valley, gradient components pointing across the valley tend to cancel, while components pointing along the valley accumulate.

This allows the optimizer to move more efficiently toward the minimum.

4.3 Momentum in PyTorch

PyTorch supports momentum through the SGD optimizer:

import torch

optimizer = torch.optim.SGD(
    model.parameters(),
    lr=0.01,
    momentum=0.9
)

The momentum parameter controls how strongly previous gradient updates influence the current update.

PyTorch also provides a dampening parameter, which modifies the contribution of new gradients to the momentum buffer.

5. Adaptive Optimization Methods

5.1 Why Adaptive Learning Rates?

Standard SGD uses the same learning rate for all model parameters. However, different parameters may have gradients of different magnitudes.

Adaptive optimization methods address this problem by adjusting the effective learning rate for each parameter based on its gradient history.

These methods can improve training efficiency and reduce the need for manually tuning a single learning rate.

5.2 AdaGrad

AdaGrad adapts the learning rate using the accumulated squared gradients.

Let g_t be the gradient at iteration t. The accumulated squared gradient is

G_t = G_{t-1} + g_t^2

The parameter update is

\theta_{t+1} = \theta_t - \frac{\eta}{\sqrt{G_t}+\epsilon}g_t

All operations involving G_t are performed element-wise.

AdaGrad gives smaller effective learning rates to parameters with large accumulated gradients.

However, because G_t continuously increases, the effective learning rate may become too small during long training runs.

5.3 RMSProp

RMSProp addresses AdaGrad’s continuously decreasing learning rates by using an exponential moving average of squared gradients.

s_t = \beta s_{t-1} + (1-\beta)g_t^2

The parameter update is

\theta_{t+1} = \theta_t - \frac{\eta}{\sqrt{s_t}+\epsilon}g_t

Here, \beta controls how much information from previous gradients is retained.

Unlike AdaGrad, RMSProp gives more weight to recent gradients rather than accumulating all past squared gradients equally.

5.4 Adam Optimizer

Adam stands for Adaptive Moment Estimation. It combines ideas from momentum and adaptive learning rates.

Adam maintains two moving averages:

The first moment estimates the average gradient.
The second moment estimates the average squared gradient.

m_t = \beta_1 m_{t-1} + (1-\beta_1)g_t

v_t = \beta_2 v_{t-1} + (1-\beta_2)g_t^2

Because both moving averages are initialized at zero, Adam applies bias correction:

\hat{m}_t = \frac{m_t}{1-\beta_1^t}

\hat{v}_t = \frac{v_t}{1-\beta_2^t}

The parameter update is

\theta_{t+1} = \theta_t - \eta \frac{\hat{m}_t}{\sqrt{\hat{v}_t}+\epsilon}

Common default values are:

\beta_1 = 0.9
\beta_2 = 0.999
\epsilon = 10^{-8}

Adam uses momentum-like gradient averaging while also scaling updates based on recent squared gradients.

5.5 SGD with Momentum vs. Adam

SGD with momentum uses accumulated gradient information to reduce oscillations and accelerate optimization.

Adam additionally adapts the effective learning rate for each parameter using estimates of the first and second moments.

Adam often makes rapid progress early in training and can work well with limited learning-rate tuning. However, SGD with momentum can sometimes achieve better generalization, depending on the model and task.

Neither optimizer is universally better.

5.6 Adam in PyTorch

PyTorch provides Adam through torch.optim.Adam.

import torch

optimizer = torch.optim.Adam(
    model.parameters(),
    lr=0.001,
    betas=(0.9, 0.999),
    eps=1e-8
)

The lr parameter specifies the base learning rate, while betas controls the exponential moving averages of the gradient and squared gradient.

Adam is widely used in deep learning because it combines adaptive parameter updates with momentum-like behavior.