Demo: Gradient Descent, SGD, and Adam on One Stretched Bowl

The Optimization lecture left us with a bowl that curves \(\kappa = 15\) times harder in one direction than in another, and a single learning rate that cannot serve both.

Here we rebuild that bowl out of data. Each of \(n = 300\) noisy targets \(\mathbf{t}^{(i)} \in \mathbb{R}^2\) contributes its own quadratic loss, and the full loss is their average, so a batch of \(B\) points sees a bowl centered slightly off from the real one. Three ways to walk downhill: full-batch gradient descent, stochastic gradient descent on batches of \(5\), and Adam on the same batches of \(5\).

import numpy as np
import matplotlib.pyplot as plt
import scienceplots

plt.style.use(["science", "no-latex"])
TEAL, CARDINAL, GRAY = "#009090", "#9c1b33", "#c9c9c9"
rng = np.random.default_rng(12)

# a bowl 15 times steeper in one direction than the other, rotated off the axes
kappa = 15.0
theta = np.pi / 6
R = np.array([[np.cos(theta), -np.sin(theta)], [np.sin(theta), np.cos(theta)]])
A = R @ np.diag([1.0, kappa]) @ R.T          # curvatures 1 and 15

n, sigma_t = 300, 0.5
targets = rng.normal(0, sigma_t, size=(n, 2))  # the 300 noisy targets t^(i)

# loss_i(w) = (1/2)(w - t^(i))^T A (w - t^(i)), so the gradient over a set of
# points is A times (w minus that set's average target)
def full_grad(w):
    return A @ (w - targets.mean(0))

def batch_grad(w, B):
    idx = rng.choice(n, size=B, replace=False)
    return A @ (w - targets[idx].mean(0))

true_min = targets.mean(0)
print(f"true minimum at {true_min.round(3)}")

true minimum at [0.005 0.016]

Run all three from the same starting point \(\mathbf{w}^{(0)} = (3.5, 3.0)\), for the same 60 steps.

Adam keeps two running averages, exactly as in the reading: the velocity v (a smoothed gradient) and s (a smoothed squared gradient), each divided by the weight used so far to correct the bias toward zero.

w0 = np.array([3.5, 3.0])
alpha = 1.9 / kappa
n_steps = 60

# full-batch gradient descent
w = w0.copy(); gd_path = [w.copy()]
for _ in range(n_steps):
    w = w - alpha * full_grad(w)
    gd_path.append(w.copy())
gd_path = np.array(gd_path)

# stochastic gradient descent, batch size 5
w = w0.copy(); sgd_path = [w.copy()]
for _ in range(n_steps):
    w = w - alpha * batch_grad(w, 5)
    sgd_path.append(w.copy())
sgd_path = np.array(sgd_path)

# Adam, same tiny batches
w = w0.copy(); v, s = np.zeros(2), np.zeros(2)
beta1, beta2, eps, alpha_adam = 0.9, 0.999, 1e-8, 0.3
adam_path = [w.copy()]
for t in range(1, n_steps + 1):
    g = batch_grad(w, 5)
    v = beta1 * v + (1 - beta1) * g
    s = beta2 * s + (1 - beta2) * g ** 2
    w = w - alpha_adam * (v / (1 - beta1 ** t)) / (np.sqrt(s / (1 - beta2 ** t)) + eps)
    adam_path.append(w.copy())
adam_path = np.array(adam_path)

grid = np.linspace(-4, 4, 200)
G1, G2 = np.meshgrid(grid, grid)
loss = A[0, 0] * G1 ** 2 + 2 * A[0, 1] * G1 * G2 + A[1, 1] * G2 ** 2

fig, ax = plt.subplots(figsize=(6, 6))
ax.contour(G1, G2, loss, levels=15, colors=[GRAY], linewidths=0.8)
ax.plot(*gd_path.T, color=TEAL, linewidth=1.3, marker="o", markersize=2.5,
        label="Gradient descent (full batch)")
ax.plot(*sgd_path.T, color=CARDINAL, linewidth=1.0, marker="o", markersize=2.5,
        alpha=0.85, label="SGD (batch $B=5$)")
ax.plot(*adam_path.T, color="black", linewidth=1.3, marker="o", markersize=2.5,
        label="Adam (batch $B=5$)")
# a white disk clears the converging paths so the minimum's cross stays readable
ax.scatter([0], [0], color="white", marker="o", s=320, edgecolors="none", zorder=6)
ax.scatter([0], [0], color="black", marker="x", s=150, linewidths=2.2, zorder=7)
ax.set_xlim(-4, 4); ax.set_ylim(-4, 4); ax.set_aspect("equal")
ax.set_xlabel("$w_1$"); ax.set_ylabel("$w_2$")
ax.legend(frameon=False, loc="lower left")
plt.show()

The paths show the shape of each method: gradient descent zig-zags across the valley, SGD zig-zags harder because every step aims at a batch’s bowl, Adam does not zig-zag at all.

Distance to the true minimum, step by step, says who actually arrived.

fig, ax = plt.subplots(figsize=(7, 3))
for path, color, label in [(gd_path, TEAL, "Gradient descent"),
                           (sgd_path, CARDINAL, "SGD"),
                           (adam_path, "black", "Adam")]:
    ax.plot(np.linalg.norm(path - true_min, axis=1), color=color, linewidth=1.4, label=label)
    print(f"{label:17s} distance after {n_steps} steps: {np.linalg.norm(path[-1] - true_min):.4f}")
ax.set_yscale("log")
ax.set_xlabel("Step")
ax.set_ylabel("Distance to minimum")
ax.legend(frameon=False)
plt.show()

Gradient descent  distance after 60 steps: 0.0020
SGD               distance after 60 steps: 0.5899
Adam              distance after 60 steps: 0.1770

What if the batches were bigger?

Only the full-batch run keeps descending; the two stochastic runs flatten out at a noise floor. Near the minimum each stochastic step is a random move of length roughly \(\alpha\sigma/\sqrt{B}\), so the floor shrinks as \(\sigma/\sqrt{B}\) does. Change batch_size and rerun. (A single 60-step run is itself a noisy measurement of a noise floor, so expect the drop to be rougher than \(1/\sqrt{B}\) exactly.)

batch_size = 50    # change me!

w = w0.copy(); sgd_path_big = [w.copy()]
for _ in range(n_steps):
    w = w - alpha * batch_grad(w, batch_size)
    sgd_path_big.append(w.copy())
sgd_path_big = np.array(sgd_path_big)

fig, ax = plt.subplots(figsize=(7, 3))
ax.plot(np.linalg.norm(sgd_path - true_min, axis=1), color=CARDINAL, linewidth=1.4,
        label="SGD, batch $B=5$")
ax.plot(np.linalg.norm(sgd_path_big - true_min, axis=1), color=TEAL, linewidth=1.4,
        label=f"SGD, batch $B={batch_size}$")
ax.plot(np.linalg.norm(gd_path - true_min, axis=1), color="black", linewidth=1.2,
        linestyle="--", label="Gradient descent (full batch)")
ax.set_yscale("log")
ax.set_xlabel("Step")
ax.set_ylabel("Distance to minimum")
ax.legend(frameon=False)
plt.show()

Punchline: conditioning and batch noise are two different problems with two different fixes. Momentum and Adam’s per-parameter scaling attack the conditioning, which is why Adam’s path is smooth; only a larger batch or a smaller learning rate lowers the noise floor, which is why Adam still stops short of the minimum.