Implementation exercises for optimization — implementing gradient descent variants manually and empirically comparing optimizer and learning rate schedule behavior.
🟢 Problem 1: Implement plain gradient descent on a simple 2D loss function
Task: Given the function \(f(x, y) = x^2 + 10y^2\) (an elongated bowl shape), implement gradient descent to find its minimum, and track the path taken.
def f(x, y): return x**2 + 10 * y**2
def grad_f(x, y): return np.array([2*x, 20*y])
def gradient_descent(start, lr, steps):
point = np.array(start, dtype=float)
path = [point.copy()]
for _ in range(steps):
point -= lr * grad_f(*point)
path.append(point.copy())
return np.array(path)
path = gradient_descent(start=[5.0, 5.0], lr=0.05, steps=50)
print(f"Final point: {path[-1]}") # should approach [0, 0]
print(f"Final loss: {f(*path[-1]):.6f}")Hint if stuck: Try different learning rates — a rate too high for the steep \(y\) direction will cause oscillation or divergence, even if it works fine for the shallower \(x\) direction. This directly illustrates why a single learning rate can be problematic for elongated (poorly conditioned) loss surfaces.
🟡 Problem 2: Implement momentum from scratch and compare against plain gradient descent
Task: Add momentum to your gradient descent implementation and compare convergence speed against plain gradient descent on the same function from Problem 1.
def gradient_descent_momentum(start, lr, momentum, steps):
point = np.array(start, dtype=float)
velocity = np.zeros_like(point)
path = [point.copy()]
for _ in range(steps):
velocity = momentum * velocity - lr * grad_f(*point)
point += velocity
path.append(point.copy())
return np.array(path)
path_plain = gradient_descent(start=[5.0, 5.0], lr=0.05, steps=50)
path_momentum = gradient_descent_momentum(start=[5.0, 5.0], lr=0.05, momentum=0.9, steps=50)
print(f"Plain GD loss after 50 steps: {f(*path_plain[-1]):.6f}")
print(f"Momentum GD loss after 50 steps: {f(*path_momentum[-1]):.6f}")
# Momentum should reach a lower loss in the same number of steps🔴 Problem 3: Implement a simplified version of Adam from scratch
Task: Implement Adam's core update rule (first and second moment estimates with bias correction) and verify it converges on the same test function.
\[ m_t = \beta_1 m_{t-1} + (1-\beta_1)g_t, \quad v_t = \beta_2 v_{t-1} + (1-\beta_2)g_t^2 \] \[ \hat{m}_t = \frac{m_t}{1-\beta_1^t}, \quad \hat{v}_t = \frac{v_t}{1-\beta_2^t}, \quad \theta_t = \theta_{t-1} - \frac{\eta}{\sqrt{\hat{v}_t}+\epsilon}\hat{m}_t \]def adam_optimizer(start, lr, steps, beta1=0.9, beta2=0.999, eps=1e-8):
point = np.array(start, dtype=float)
m = np.zeros_like(point)
v = np.zeros_like(point)
path = [point.copy()]
for t in range(1, steps + 1):
g = grad_f(*point)
m = beta1 * m + (1 - beta1) * g
v = beta2 * v + (1 - beta2) * (g ** 2)
m_hat = m / (1 - beta1 ** t) # bias correction -- important in early steps
v_hat = v / (1 - beta2 ** t)
point -= lr * m_hat / (np.sqrt(v_hat) + eps)
path.append(point.copy())
return np.array(path)
path_adam = adam_optimizer(start=[5.0, 5.0], lr=0.5, steps=50)
print(f"Adam loss after 50 steps: {f(*path_adam[-1]):.6f}")Hint if stuck: Bias correction matters most in the first few steps, when \(m\) and \(v\) are still close to their zero initialization — without it, early updates would be artificially small. Try removing the bias correction terms and observe how the first few steps behave differently.
🟡 Problem 4: Implement and visualize a cosine annealing learning rate schedule
Task: Implement the cosine annealing formula and plot the resulting learning rate curve over training steps.
def cosine_annealing_lr(step, total_steps, lr_max=0.01, lr_min=0.0001):
return lr_min + 0.5 * (lr_max - lr_min) * (1 + np.cos(np.pi * step / total_steps))
steps = np.arange(0, 100)
lrs = [cosine_annealing_lr(s, total_steps=100) for s in steps]
plt.plot(steps, lrs)
plt.xlabel('Training step')
plt.ylabel('Learning rate')
plt.title('Cosine Annealing Schedule')
plt.show()