Objective: Implement plain gradient descent and gradient descent with momentum as small, reusable optimizer functions, then compare their convergence speed and stability on a deliberately awkward loss surface. All in runnable, in-browser Python.
Task 1: Define an Awkward Loss Surface
We'll use the same elongated-bowl idea from the lesson, steep along one axis and shallow along the other, since it's exactly the shape where plain gradient descent struggles and momentum helps.
import numpy as np
def loss(w):
return w[0] ** 2 + 10 * w[1] ** 2
def grad(w):
return np.array([2 * w[0], 20 * w[1]])
w_start = np.array([4.0, 1.0])
print("Starting loss:", loss(w_start))
print("Starting gradient:", grad(w_start))
Task 2: Implement Plain Gradient Descent
import numpy as np
def loss(w):
return w[0] ** 2 + 10 * w[1] ** 2
def grad(w):
return np.array([2 * w[0], 20 * w[1]])
def plain_gd(w_start, lr, n_steps):
w = w_start.copy()
history = [loss(w)]
for _ in range(n_steps):
w = w - lr * grad(w)
history.append(loss(w))
return w, history
w_final, history = plain_gd(np.array([4.0, 1.0]), lr=0.04, n_steps=50)
print("Final w:", np.round(w_final, 5))
print("Final loss:", round(history[-1], 6))
print("Loss after 5 steps:", [round(h, 3) for h in history[:6]])
Task 3: Implement Gradient Descent with Momentum
import numpy as np
def loss(w):
return w[0] ** 2 + 10 * w[1] ** 2
def grad(w):
return np.array([2 * w[0], 20 * w[1]])
def momentum_gd(w_start, lr, beta, n_steps):
w = w_start.copy()
v = np.zeros_like(w)
history = [loss(w)]
for _ in range(n_steps):
g = grad(w)
v = beta * v + (1 - beta) * g
w = w - lr * v
history.append(loss(w))
return w, history
w_final, history = momentum_gd(np.array([4.0, 1.0]), lr=0.04, beta=0.9, n_steps=50)
print("Final w:", np.round(w_final, 5))
print("Final loss:", round(history[-1], 6))
print("Loss after 5 steps:", [round(h, 3) for h in history[:6]])
Compare the "loss after 5 steps" printed by each implementation. With identical learning rates, momentum should be making faster initial progress by accumulating velocity in the consistent direction.
Task 4: Compare Convergence Side by Side
import numpy as np
import matplotlib.pyplot as plt
def loss(w):
return w[0] ** 2 + 10 * w[1] ** 2
def grad(w):
return np.array([2 * w[0], 20 * w[1]])
def plain_gd(w_start, lr, n_steps):
w = w_start.copy()
history = [loss(w)]
for _ in range(n_steps):
w = w - lr * grad(w)
history.append(loss(w))
return history
def momentum_gd(w_start, lr, beta, n_steps):
w = w_start.copy()
v = np.zeros_like(w)
history = [loss(w)]
for _ in range(n_steps):
g = grad(w)
v = beta * v + (1 - beta) * g
w = w - lr * v
history.append(loss(w))
return history
w_start = np.array([4.0, 1.0])
history_plain = plain_gd(w_start, lr=0.04, n_steps=60)
history_momentum = momentum_gd(w_start, lr=0.04, beta=0.9, n_steps=60)
plt.plot(history_plain, label="Plain GD")
plt.plot(history_momentum, label="GD + momentum")
plt.yscale("log")
plt.xlabel("Step"); plt.ylabel("Loss (log scale)")
plt.legend(); plt.title("Convergence: plain GD vs. momentum")
show_plot()
On the log-scale plot, momentum should reach a low loss in noticeably fewer steps. Try pushing
lr higher for plain GD in Task 2's code: at some point it will start oscillating or diverging
outright on the steep axis, while momentum (because it averages gradients) tolerates a somewhat higher
effective learning rate before showing the same instability. That tradeoff, between step size and stability,
is exactly what the next lesson on learning rate schedules addresses directly.