Optimisers

An optimiser updates a model's weights and biases to reduce its loss. It uses the gradients calculated during backpropagation.

The basic idea is to move each parameter in a direction that reduces the loss. The optimiser determines how to calculate each update. Its Learning Rate controls the update scale.

Main choices

Optimiser	Information used	Effect on the update
SGD	Current gradient	Uses the gradient and a learning rate
SGD with momentum	Current gradient and past update direction	Can maintain progress through small changes in direction
RMSprop	Current gradient and moving average of squared gradients	Adjusts the scale of each parameter's update
Adam	Moving averages of gradients and squared gradients	Combines a stored direction with an adaptive update scale

Momentum stores information from previous gradients. Adam, or Adaptive Moment Estimation, also corrects the initial bias in its moving averages. See the Adam paper.

The batch determines which examples supply the gradient.

Choosing an optimiser:

• Use validation results to compare choices under the same data split
• Tune the learning rate for each choice
• Use a learning rate schedule when the rate must change during training
• Treat training speed and final validation quality as separate measurements

Code

PyTorch

import torch
from torch import nn
 
X = torch.tensor([[1., 0.], [0., 1.]])
y = torch.tensor([0, 1])
model = nn.Linear(2, 2)
criterion = nn.CrossEntropyLoss()
optimizer = torch.optim.Adam(model.parameters(), lr=0.001)
 
optimizer.zero_grad()
logits = model(X)
loss = criterion(logits, y)
loss.backward()
optimizer.step()

Read the data flow:

• X contains 2 samples with 2 features each, so its shape is (2, 2)
• y contains the correct class index for each sample
• Linear(2, 2) produces 2 class scores per sample
• model.parameters() gives Adam access to the model's trainable weights and biases
• zero_grad() clears gradients from the previous update
• loss.backward() calculates gradients and stores them in each parameter's .grad
• optimizer.step() uses these gradients to update the parameters

To use another update rule, replace the optimizer assignment before training:

optimizer = torch.optim.SGD(model.parameters(), lr=0.01, momentum=0.9)
# Alternative
optimizer = torch.optim.RMSprop(model.parameters(), lr=0.001)

Use one assignment for each training run. Creating the optimiser outside the batch loop preserves its stored history. See the PyTorch optimiser API.

TensorFlow and Keras

import numpy as np
import tensorflow as tf
 
X = np.array([[1., 0.], [0., 1.]], dtype="float32")
y = np.array([0, 1], dtype="int32")
model = tf.keras.Sequential([
    tf.keras.Input(shape=(2,)),
    tf.keras.layers.Dense(2),
])
optimizer = tf.keras.optimizers.Adam(learning_rate=0.001)
model.compile(
    optimizer=optimizer,
    loss=tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True),
)
model.fit(X, y, batch_size=2, epochs=1, verbose=0)

Read the corresponding steps:

• Dense(2) produces the same 2 class scores per sample
• from_logits=True tells the loss to accept these raw scores
• compile() connects the optimiser and loss to the model
• fit() performs the forward pass, gradient calculation, and parameter update
• Replace Adam(...) with SGD(learning_rate=0.01, momentum=0.9) or RMSprop(learning_rate=0.001) to change the rule