Optimizers#
Optimizers are the algorithms that update a model’s parameters to minimize a function. Building on Gradient Descent, this notebook implements and visualizes more advanced optimizers: Stochastic Gradient Descent (SGD), Momentum, RMSprop, and Adam. Rather than fitting a dataset, we will watch each optimizer navigate a 2D surface, stepping from a shared starting point toward the surface’s minimum. This lets us compare how each optimizer behaves on the exact same problem.
Importing the libraries.
import numpy as np
%matplotlib widget
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D
from celluloid import Camera
import scienceplots
from IPython.display import Image
np.random.seed(0)
plt.style.use(["science", "no-latex"])
Training Dataset#
Instead of a dataset, our optimizers will minimize a mathematical function. We will use the surface:
This is an elongated bowl with its minimum at \((0, 0)\). It is 40 times steeper along the \(y\) axis than along the \(x\) axis, so it forms a long, narrow valley. This is the same shape as the loss surface in the Gradient Descent notebook, where the slope and the intercept had very different scales.
Narrow valleys are common in machine learning, and they are a classic challenge for optimizers. The learning rate has to be small enough not to overshoot across the steep \(y\) direction, which makes progress along the flat \(x\) direction slow. This makes the valley a great way to compare optimizer behavior.
def surface_fn(X, Y):
return 0.05 * X**2 + 2 * Y**2
def generate_function(dims, surface_fn):
x = np.linspace(-10, 10, dims)
y = np.linspace(-3, 3, dims)
X, Y = np.meshgrid(x, y)
Z = surface_fn(X, Y)
return X, Y, Z
dims = 50
X, Y, Z = generate_function(dims, surface_fn)
# Visualize the elongated bowl
fig = plt.figure()
ax = fig.add_subplot(111, projection="3d")
ax.view_init(30, -60)
ax.plot_surface(X, Y, Z, color="grey", alpha=0.4)
ax.set_xlabel("X")
ax.set_ylabel("Y")
ax.set_zlabel("Z")
plt.show()
Optimizer Update Rules#
Every optimizer in this notebook follows the same core idea as gradient descent: repeatedly step the parameters \(\theta\) in the direction that decreases the function the most, which is the negative gradient. A single update step is:
where \(\eta\) is the learning rate and \(\nabla_{\theta} f(\theta_t)\) is the gradient of the function with respect to the parameters at step \(t\). The optimizers below differ in how they use the gradient, such as accumulating past gradients (momentum) or scaling each parameter’s step by a running estimate of its gradient magnitude (adaptive learning rates).
Stochastic Gradient Descent (SGD)#
SGD is the simplest optimizer. It takes a step directly proportional to the negative gradient, scaled by the learning rate:
When training a model, SGD computes the gradient on a random mini-batch of the training data, which is what makes it stochastic. Here we minimize a known function, so the gradient is exact and SGD is plain gradient descent. We keep the name SGD because that is what libraries like PyTorch call this optimizer.
SGD treats every parameter and every direction the same. Because the step depends only on the current gradient and a fixed learning rate, SGD moves very slowly in directions where the gradient is small, such as along the flat floor of a narrow valley.
Momentum#
Momentum keeps a running sum of past gradients, called the velocity \(v_t\), and steps along it:
where \(\mu\) is the momentum coefficient, typically \(0.9\). Gradients that keep pointing the same way, like the gentle slope along the valley floor, add up and speed the optimizer up. Gradients that flip sign, like the steep walls of the valley, cancel out. This is the same update rule as PyTorch’s SGD optimizer with momentum=0.9.
RMSprop#
RMSprop gives each parameter its own adaptive learning rate by keeping an exponentially decaying average of recent squared gradients, \(s_t\):
The square and square root are applied to each parameter separately. Dividing by \(\sqrt{s_{t+1}}\) shrinks the step for parameters with large gradients and enlarges it for parameters with small gradients, so every parameter moves at a similar speed. The small constant \(\epsilon\) prevents division by zero. This adaptivity helps RMSprop make steady progress along the flat valley floor, where SGD is slow.
Adam#
Adam (Adaptive Moment Estimation) combines momentum with RMSprop’s adaptive learning rates. It tracks an exponentially decaying average of the gradients (the first moment \(m_t\)) and of the squared gradients (the second moment \(v_t\)):
Because \(m_t\) and \(v_t\) start at zero, they are biased toward zero early in training, so Adam applies a bias correction:
The parameters are then updated using the corrected moments:
Combining momentum and adaptive scaling makes Adam a robust default optimizer that converges quickly on many problems. However, no optimizer is the fastest on every problem, as we will see below.
Surface Gradient#
Each optimizer needs the gradient of the surface with respect to the current point \(\theta = (x, y)\). For our surface \(z = 0.05x^2 + 2y^2\), the partial derivatives are:
so the gradient is \(\nabla z = (0.1x, 4y)\). At the same distance from the minimum, the gradient along \(y\) is 40 times larger than the gradient along \(x\). This is the signal every optimizer below uses to decide its next step.
Implementing the Optimizers#
Let’s implement each optimizer from scratch as a class with a step method that takes the current parameters and gradient and returns the updated parameters. The stateful optimizers (Momentum, RMSprop, and Adam) keep their running averages between steps.
def surface_gradient(pt):
x, y = pt
return np.array([0.1 * x, 4 * y])
class SGD:
def __init__(self, learning_rate):
self.learning_rate = learning_rate
def step(self, params, gradient):
return params - self.learning_rate * gradient
class Momentum:
def __init__(self, learning_rate, momentum=0.9):
self.learning_rate = learning_rate
self.momentum = momentum
self.velocity = 0
def step(self, params, gradient):
self.velocity = self.momentum * self.velocity + gradient
return params - self.learning_rate * self.velocity
class RMSprop:
def __init__(self, learning_rate, beta=0.99, epsilon=1e-8):
self.learning_rate = learning_rate
self.beta = beta
self.epsilon = epsilon
self.s = 0
def step(self, params, gradient):
self.s = self.beta * self.s + (1 - self.beta) * gradient**2
return params - self.learning_rate * gradient / (np.sqrt(self.s) + self.epsilon)
class Adam:
def __init__(self, learning_rate, beta1=0.9, beta2=0.999, epsilon=1e-8):
self.learning_rate = learning_rate
self.beta1 = beta1
self.beta2 = beta2
self.epsilon = epsilon
self.m = 0
self.v = 0
self.t = 0
def step(self, params, gradient):
self.t += 1
self.m = self.beta1 * self.m + (1 - self.beta1) * gradient
self.v = self.beta2 * self.v + (1 - self.beta2) * gradient**2
m_hat = self.m / (1 - self.beta1**self.t)
v_hat = self.v / (1 - self.beta2**self.t)
return params - self.learning_rate * m_hat / (np.sqrt(v_hat) + self.epsilon)
Graphing functions#
Let’s define helper functions to plot the graphs. create_plots sets up two views: a plot tracking each optimizer’s \(z\) value over the epochs, and a 3D plot of the surface, showing the path each optimizer has taken so far. The contour lines on the floor show the shape of the valley, and the black star marks the minimum at \((0, 0)\). The \(z\) value is plotted on a log scale, because it shrinks by many orders of magnitude as the optimizers approach the minimum. show_epoch controls how often a frame is captured for the animation, sampling frequently early in training and less often later.
def create_plots(X, Y, Z):
fig, ax = plt.subplots(1, 3, figsize=(16 / 9.0 * 4, 4 * 1))
fig.suptitle("Optimizers")
ax[0].set_xlabel("Epoch", fontweight="normal")
ax[0].set_ylabel("Z", fontweight="normal")
ax[0].set_title("Z Value")
ax[0].set_yscale("log")
ax[0].set_ylim((1e-6, 1e2))
ax[1].axis("off")
ax[2].axis("off")
# draw artists in the order they are added, so the paths stay on top of the surface
ax[2] = fig.add_subplot(1, 2, 2, projection="3d", computed_zorder=False)
ax[2].set_xlabel("X")
ax[2].set_ylabel("Y")
ax[2].set_zlabel("Z")
ax[2].set_title("Optimizer Paths")
# look down into the bowl so the paths and the minimum are both visible
ax[2].view_init(60, -70)
ax[2].set_xlim(X.min(), X.max())
ax[2].set_ylim(Y.min(), Y.max())
ax[2].set_zlim(0, Z.max())
camera = Camera(fig)
return ax[0], ax[2], camera
def show_epoch(epoch):
return (
(epoch < 50 and epoch % 2 == 0)
or (epoch < 150 and epoch % 5 == 0)
or epoch % 10 == 0
)
def plot_frame(z_ax, path_ax, X, Y, Z, paths, names, colors):
# paths[i] holds the (x, y, z) positions optimizer i has visited so far
path_ax.plot_surface(X, Y, Z, color="grey", alpha=0.3)
# contour lines projected onto the floor show the shape of the valley
path_ax.contour(X, Y, Z, levels=12, zdir="z", offset=0, colors="gray", linewidths=0.5)
z_legends = []
path_legends = []
for path, color in zip(paths, colors):
path = np.array(path)
epochs = np.arange(1, len(path) + 1)
# clamp tiny values so they stay visible on the log scale
z_values = np.maximum(path[:, 2], 1e-8)
z_legend = z_ax.plot(epochs, z_values, color=color, alpha=0.7)
z_legends.append(z_legend[0])
path_ax.plot(path[:, 0], path[:, 1], path[:, 2], color=color, alpha=0.7, linewidth=1)
path_legend = path_ax.scatter(path[-1, 0], path[-1, 1], path[-1, 2], color=color, s=30)
path_legends.append(path_legend)
# the minimum of the surface at (0, 0)
path_ax.scatter([0], [0], [0], marker="*", color="black", s=120)
z_ax.legend(z_legends, names, loc="upper right")
path_ax.legend(path_legends, names, loc="upper right")
def print_summary(paths, names, threshold=0.01):
for name, path in zip(names, paths):
reached = next((i for i, pt in enumerate(path) if pt[2] < threshold), None)
print(f"{name:>8}: first epoch with z < {threshold}: {reached}, final z: {path[-1][2]:.2e}")
Training the Model#
The fit function runs every optimizer in parallel from the same starting point. In each epoch it evaluates the surface, computes the gradient, and asks each optimizer for its next position, recording the path so the optimizers can be animated side by side.
def fit(optimizers, surface_fn, X, Y, Z, initial_pt, epochs, output_filename, colors):
names = list(optimizers.keys())
points = [initial_pt.copy() for _ in optimizers]
paths = [[] for _ in optimizers]
z_ax, path_ax, camera = create_plots(X, Y, Z)
for idx in range(epochs):
for opt_idx, optimizer in enumerate(optimizers.values()):
pt = points[opt_idx]
paths[opt_idx].append([pt[0], pt[1], surface_fn(pt[0], pt[1])])
gradient = surface_gradient(pt)
points[opt_idx] = optimizer.step(pt, gradient)
if show_epoch(idx):
plot_frame(z_ax, path_ax, X, Y, Z, paths, names, colors)
camera.snap()
animation = camera.animate()
animation.save(output_filename, writer="pillow")
plt.show()
print_summary(paths, names)
Let’s train all four optimizers from the same starting point and compare them. We start at \((-9, 2)\), at one end of the valley and partway up its steep wall. Every optimizer is given the same learning rate, \(\eta = 0.1\), so the only difference is the update rule.
For SGD, a learning rate above \(0.5\) would make it overshoot and diverge along the steep \(y\) axis. At \(\eta = 0.1\), SGD is stable, but each step only shrinks its distance along the flat \(x\) axis by \(1\%\) (\(1 - 2 \cdot 0.05 \cdot 0.1 = 0.99\)).
initial_pt = np.array([-9.0, 2.0])
learning_rate = 0.1
epochs = 300
optimizers = {
"SGD": SGD(learning_rate),
"Momentum": Momentum(learning_rate),
"RMSprop": RMSprop(learning_rate),
"Adam": Adam(learning_rate),
}
colors = ["blue", "orange", "green", "red"]
output_filename = "optimizers.gif"
fit(optimizers, surface_fn, X, Y, Z, initial_pt, epochs, output_filename, colors)
SGD: first epoch with z < 0.01: 299, final z: 9.94e-03
Momentum: first epoch with z < 0.01: 54, final z: 1.80e-13
RMSprop: first epoch with z < 0.01: 60, final z: 6.23e-24
Adam: first epoch with z < 0.01: 144, final z: 9.12e-10
Output GIF#
The animation below shows our from-scratch optimizers descending the valley from the same starting point. The left plot tracks each optimizer’s \(z\) value on a log scale, and the right plot shows the path each optimizer has taken down the surface toward the minimum (black star).
SGD (blue) drops to the valley floor within a few steps, because the gradient along the steep \(y\) axis is large. Along the flat \(x\) axis, it crawls: it needs 299 epochs to reach \(z < 0.01\).
Momentum (orange) zig-zags across the valley, since its velocity carries it past the valley floor. The zig-zags cancel out, while its speed along \(x\) builds up. It reaches \(z < 0.01\) in 54 epochs.
RMSprop (green) scales each axis separately, so it moves along \(x\) almost as fast as along \(y\) and takes the most direct path. It reaches \(z < 0.01\) in 60 epochs.
Adam (red) also scales each axis, but its momentum carries it past the valley floor before it turns. It reaches \(z < 0.01\) in 144 epochs, slower than Momentum and RMSprop but twice as fast as SGD.
On this narrow valley, all three advanced optimizers beat SGD by a wide margin.
Image(filename=output_filename)
PyTorch Implementation#
Machine learning libraries, such as PyTorch, provide highly optimized implementations of these optimizers. We will use PyTorch to run the same four optimizers on the same surface and confirm our from-scratch results.
import torch
import torch.nn as nn
torch.manual_seed(0)
<torch._C.Generator at 0x13d12ba50>
We represent the optimization problem as a small model whose only parameters are the coordinates of a point on the surface. The forward method evaluates the surface function at the current point, which is the value we want to minimize.
class TorchNet(nn.Module):
def __init__(self, initial_pt):
super(TorchNet, self).__init__()
self.pts = torch.nn.Parameter(initial_pt)
def forward(self, surface_fn):
return surface_fn(self.pts[0], self.pts[1])
def get_pts(self):
return self.pts.detach().cpu().numpy()
The training function creates one model and one optimizer per optimizer type, all starting from the same initial point. In each epoch, it evaluates the surface at every optimizer’s current point, backpropagates the gradient, and steps each optimizer.
def torch_fit(
optimizer_types, surface_fn, X, Y, Z, initial_pt, epochs, learning_rate, output_filename, colors
):
names = list(optimizer_types.keys())
torch_models = []
torch_optimizers = []
for make_optimizer in optimizer_types.values():
initial_pt_tensor = torch.tensor(initial_pt, device=device, dtype=torch.float64)
torch_model = TorchNet(initial_pt_tensor).to(device)
torch_optimizer = make_optimizer(torch_model.parameters(), lr=learning_rate)
torch_models.append(torch_model)
torch_optimizers.append(torch_optimizer)
paths = [[] for _ in names]
z_ax, path_ax, camera = create_plots(X, Y, Z)
for idx in range(epochs):
for opt_idx, (model, optimizer) in enumerate(zip(torch_models, torch_optimizers)):
pt = model.get_pts()
loss = model(surface_fn)
paths[opt_idx].append([pt[0], pt[1], loss.item()])
# optimizer
optimizer.zero_grad()
loss.backward()
optimizer.step()
if show_epoch(idx):
plot_frame(z_ax, path_ax, X, Y, Z, paths, names, colors)
camera.snap()
animation = camera.animate()
animation.save(output_filename, writer="pillow")
plt.show()
print_summary(paths, names)
Training the Model with PyTorch#
Let’s run PyTorch’s SGD, SGD with momentum, RMSprop, and Adam from the same starting point and learning rate we used for our from-scratch implementation. PyTorch has no separate Momentum optimizer; momentum is an option of its SGD optimizer.
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
initial_pt = np.array([-9.0, 2.0])
optimizer_types = {
"SGD": torch.optim.SGD,
"Momentum": lambda params, lr: torch.optim.SGD(params, lr=lr, momentum=0.9),
"RMSprop": torch.optim.RMSprop,
"Adam": torch.optim.Adam,
}
colors = ["blue", "orange", "green", "red"]
learning_rate = 0.1
epochs = 300
output_filename_pytorch = "optimizers_pytorch.gif"
torch_fit(
optimizer_types,
surface_fn,
X,
Y,
Z,
initial_pt,
epochs,
learning_rate,
output_filename_pytorch,
colors,
)
SGD: first epoch with z < 0.01: 299, final z: 9.94e-03
Momentum: first epoch with z < 0.01: 54, final z: 1.80e-13
RMSprop: first epoch with z < 0.01: 60, final z: 6.23e-24
Adam: first epoch with z < 0.01: 144, final z: 9.12e-10
Output GIF with PyTorch#
The PyTorch optimizers follow the same paths as our from-scratch implementation. Each optimizer reaches \(z < 0.01\) at the same epoch (see the printed summaries above), confirming our implementations of SGD, Momentum, RMSprop, and Adam.
Image(filename=output_filename_pytorch)
