02 Gradient Descent — Exercises¶

Test your understanding of gradient descent, momentum, and adaptive methods.

Prerequisites. Read theory.md and work through first_principles.ipynb before attempting these.

In [1]:

import random

import matplotlib.pyplot as plt
import numpy as np

%load_ext autoreload
%autoreload 2

SEED = 42
random.seed(SEED)
rng = np.random.default_rng(SEED)

plt.rcParams["figure.dpi"] = 90

import random import matplotlib.pyplot as plt import numpy as np %load_ext autoreload %autoreload 2 SEED = 42 random.seed(SEED) rng = np.random.default_rng(SEED) plt.rcParams["figure.dpi"] = 90

---

Exercise 1 — Hand Calculation: 3 GD Steps on $f(x) = x^2 + 2x + 1$¶

Consider the function $f(x) = x^2 + 2x + 1 = (x + 1)^2$.

• Gradient: $f'(x) = 2x + 2$
• Minimum: $x^\ast = -1$ where $f(x^\ast) = 0$
• Smoothness: $f''(x) = 2$, so $M = 2$

Task: Starting from $x_0 = 3$ with learning rate $\alpha = 0.3$, compute by hand:

Step	$x_t$	$f'(x_t)$	$x_{t+1} = x_t - \alpha f'(x_t)$	$f(x_{t+1})$
0	3.0	?	?	?
1	?	?	?	?
2	?	?	?	?

Questions:

1. Fill in the table above.
2. Is $\alpha = 0.3$ within the convergence bound $\alpha < 2/M$?
3. What would happen if $\alpha = 1.1$ (beyond the bound)?

Expected results:

Step	$x_t$	$f'(x_t)$	$x_{t+1}$	$f(x_{t+1})$
0	3.0	8.0	0.6	2.56
1	0.6	3.2	-0.36	0.4096
2	-0.36	1.28	-0.744	0.065536

In [2]:

# Verify your hand calculations
def f_ex1(x):
    return x**2 + 2 * x + 1

def grad_ex1(x):
    return 2 * x + 2

x = 3.0
alpha = 0.3

# TODO: compute 3 GD steps and fill in the expected values below
expected_x = [3.0]  # x_0
for step in range(3):
    g = grad_ex1(x)
    x_new = x - alpha * g
    print(f"Step {step}: x={x:.4f}, f'(x)={g:.4f}, x_new={x_new:.4f}, f(x_new)={f_ex1(x_new):.6f}")
    expected_x.append(x_new)
    x = x_new

# Deterministic check
assert np.isclose(expected_x[1], 0.6, atol=1e-10), f"x_1 should be 0.6, got {expected_x[1]}"
assert np.isclose(expected_x[2], -0.36, atol=1e-10), f"x_2 should be -0.36, got {expected_x[2]}"
assert np.isclose(expected_x[3], -0.744, atol=1e-10), f"x_3 should be -0.744, got {expected_x[3]}"
print("\nAll hand-calculation checks passed. ✓")

# Question 2: alpha = 0.3 < 2/M = 2/2 = 1.0 → yes, within bound
print(f"\nalpha = {alpha} < 2/M = {2/2} → convergence guaranteed: ✓")

# Verify your hand calculations def f_ex1(x): return x**2 + 2 * x + 1 def grad_ex1(x): return 2 * x + 2 x = 3.0 alpha = 0.3 # TODO: compute 3 GD steps and fill in the expected values below expected_x = [3.0] # x_0 for step in range(3): g = grad_ex1(x) x_new = x - alpha * g print(f"Step {step}: x={x:.4f}, f'(x)={g:.4f}, x_new={x_new:.4f}, f(x_new)={f_ex1(x_new):.6f}") expected_x.append(x_new) x = x_new # Deterministic check assert np.isclose(expected_x[1], 0.6, atol=1e-10), f"x_1 should be 0.6, got {expected_x[1]}" assert np.isclose(expected_x[2], -0.36, atol=1e-10), f"x_2 should be -0.36, got {expected_x[2]}" assert np.isclose(expected_x[3], -0.744, atol=1e-10), f"x_3 should be -0.744, got {expected_x[3]}" print("\nAll hand-calculation checks passed. ✓") # Question 2: alpha = 0.3 < 2/M = 2/2 = 1.0 → yes, within bound print(f"\nalpha = {alpha} < 2/M = {2/2} → convergence guaranteed: ✓")

Step 0: x=3.0000, f'(x)=8.0000, x_new=0.6000, f(x_new)=2.560000
Step 1: x=0.6000, f'(x)=3.2000, x_new=-0.3600, f(x_new)=0.409600
Step 2: x=-0.3600, f'(x)=1.2800, x_new=-0.7440, f(x_new)=0.065536

All hand-calculation checks passed. ✓

alpha = 0.3 < 2/M = 1.0 → convergence guaranteed: ✓

---

Exercise 2 — Coding: Implement Momentum GD with Convergence Check¶

Implement gradient descent with Polyak heavy-ball momentum and verify it converges on a known test problem.

Update rule: $$v_{t+1} = \beta v_t - \alpha \nabla L(\theta_t)$$ $$\theta_{t+1} = \theta_t + v_{t+1}$$

Test problem: $f(x, y) = 3(x - 1)^2 + 10(y + 2)^2$

• Minimum at $(1, -2)$
• $M = 20$, $\mu = 6$, $\kappa = 10/3 \approx 3.33$

Requirements:

1. Implement momentum_gd(grad_fn, x0, lr, beta, max_iter, tol)
2. Return the final point and history of all iterates
3. Converge to within atol=1e-4 of the true minimum $(1, -2)$
4. Compare number of iterations with vanilla GD (momentum should need fewer)

In [3]:

def momentum_gd(grad_fn, x0, lr=0.01, beta=0.9, max_iter=1000, tol=1e-6):
    """Gradient descent with Polyak heavy-ball momentum.

    Args:
        grad_fn: function returning gradient at a point
        x0: initial point (1D array)
        lr: learning rate
        beta: momentum coefficient
        max_iter: maximum iterations
        tol: convergence tolerance on ||x_new - x||

    Returns:
        x: final point
        history: list of iterates
    """
    # TODO: implement momentum GD
    # Hint:
    # 1. Initialize x = x0.copy() and v = zeros_like(x)
    # 2. Loop: v = beta * v - lr * grad_fn(x), then x_new = x + v
    # 3. Stop when ||x_new - x|| < tol
    pass


# Test problem: f(x,y) = 3(x-1)^2 + 10(y+2)^2
def grad_test(w):
    return np.array([6 * (w[0] - 1), 20 * (w[1] + 2)])

def f_test(w):
    return 3 * (w[0] - 1)**2 + 10 * (w[1] + 2)**2

def momentum_gd(grad_fn, x0, lr=0.01, beta=0.9, max_iter=1000, tol=1e-6): """Gradient descent with Polyak heavy-ball momentum. Args: grad_fn: function returning gradient at a point x0: initial point (1D array) lr: learning rate beta: momentum coefficient max_iter: maximum iterations tol: convergence tolerance on ||x_new - x|| Returns: x: final point history: list of iterates """ # TODO: implement momentum GD # Hint: # 1. Initialize x = x0.copy() and v = zeros_like(x) # 2. Loop: v = beta * v - lr * grad_fn(x), then x_new = x + v # 3. Stop when ||x_new - x|| < tol pass # Test problem: f(x,y) = 3(x-1)^2 + 10(y+2)^2 def grad_test(w): return np.array([6 * (w[0] - 1), 20 * (w[1] + 2)]) def f_test(w): return 3 * (w[0] - 1)**2 + 10 * (w[1] + 2)**2

Solution 2¶

In [4]:

def momentum_gd(grad_fn, x0, lr=0.01, beta=0.9, max_iter=1000, tol=1e-6):
    """Gradient descent with Polyak heavy-ball momentum.

    Args:
        grad_fn: function returning gradient at a point
        x0: initial point (1D array)
        lr: learning rate
        beta: momentum coefficient
        max_iter: maximum iterations
        tol: convergence tolerance on ||x_new - x||

    Returns:
        x: final point
        history: list of iterates
    """
    x = np.asarray(x0, dtype=float).copy()
    v = np.zeros_like(x)
    history = [x.copy()]
    for _ in range(max_iter):
        v = beta * v - lr * grad_fn(x)
        x_new = x + v
        history.append(x_new.copy())
        converged = np.linalg.norm(x_new - x) < tol
        x = x_new
        if converged:
            break
    return x, history


# Run momentum GD on the test problem
w0 = np.array([5.0, 3.0])
w_result, hist_result = momentum_gd(grad_test, w0, lr=0.04, beta=0.9, max_iter=500)

# Deterministic check
assert w_result is not None, "momentum_gd returned None — implement it!"
assert np.allclose(w_result, [1.0, -2.0], atol=1e-4), (
    f"Should converge to [1, -2], got [{w_result[0]:.4f}, {w_result[1]:.4f}]"
)
print(f"Momentum GD converged to [{w_result[0]:.6f}, {w_result[1]:.6f}]")
print(f"Iterations: {len(hist_result) - 1}")
print(f"f(result) = {f_test(w_result):.8f}")
print("Convergence check: ✓")

def momentum_gd(grad_fn, x0, lr=0.01, beta=0.9, max_iter=1000, tol=1e-6): """Gradient descent with Polyak heavy-ball momentum. Args: grad_fn: function returning gradient at a point x0: initial point (1D array) lr: learning rate beta: momentum coefficient max_iter: maximum iterations tol: convergence tolerance on ||x_new - x|| Returns: x: final point history: list of iterates """ x = np.asarray(x0, dtype=float).copy() v = np.zeros_like(x) history = [x.copy()] for _ in range(max_iter): v = beta * v - lr * grad_fn(x) x_new = x + v history.append(x_new.copy()) converged = np.linalg.norm(x_new - x) < tol x = x_new if converged: break return x, history # Run momentum GD on the test problem w0 = np.array([5.0, 3.0]) w_result, hist_result = momentum_gd(grad_test, w0, lr=0.04, beta=0.9, max_iter=500) # Deterministic check assert w_result is not None, "momentum_gd returned None — implement it!" assert np.allclose(w_result, [1.0, -2.0], atol=1e-4), ( f"Should converge to [1, -2], got [{w_result[0]:.4f}, {w_result[1]:.4f}]" ) print(f"Momentum GD converged to [{w_result[0]:.6f}, {w_result[1]:.6f}]") print(f"Iterations: {len(hist_result) - 1}") print(f"f(result) = {f_test(w_result):.8f}") print("Convergence check: ✓")

Momentum GD converged to [0.999997, -2.000004]
Iterations: 267
f(result) = 0.00000000
Convergence check: ✓

In [5]:

# Compare with vanilla GD
def vanilla_gd(grad_fn, x0, lr=0.01, max_iter=1000, tol=1e-6):
    x = np.asarray(x0, dtype=float).copy()
    history = [x.copy()]
    for _ in range(max_iter):
        g = grad_fn(x)
        x_new = x - lr * g
        history.append(x_new.copy())
        if np.linalg.norm(x_new - x) < tol:
            break
        x = x_new
    return x_new, history

w_vanilla, hist_vanilla = vanilla_gd(grad_test, w0, lr=0.04, max_iter=500)

# TODO: print and compare iteration counts
# Expected: momentum converges faster than vanilla GD
# print(f"Vanilla GD iterations: {len(hist_vanilla) - 1}")
# print(f"Momentum GD iterations: {len(hist_result) - 1}")

# Compare with vanilla GD def vanilla_gd(grad_fn, x0, lr=0.01, max_iter=1000, tol=1e-6): x = np.asarray(x0, dtype=float).copy() history = [x.copy()] for _ in range(max_iter): g = grad_fn(x) x_new = x - lr * g history.append(x_new.copy()) if np.linalg.norm(x_new - x) < tol: break x = x_new return x_new, history w_vanilla, hist_vanilla = vanilla_gd(grad_test, w0, lr=0.04, max_iter=500) # TODO: print and compare iteration counts # Expected: momentum converges faster than vanilla GD # print(f"Vanilla GD iterations: {len(hist_vanilla) - 1}") # print(f"Momentum GD iterations: {len(hist_result) - 1}")

---

Exercise 3 — Conceptual: Adam and SGD¶

Questions:

3a. Why does Adam combine momentum and adaptive learning rates?¶

Adam maintains two exponential moving averages:

• $m_t$ (first moment ≈ mean gradient) — this is the momentum component
• $s_t$ (second moment ≈ variance of gradient) — this is the adaptive rate component

Explain in your own words:

1. What problem does the momentum ($m_t$) solve? (Hint: oscillation on elongated surfaces)
2. What problem does the adaptive denominator ($\sqrt{s_t}$) solve? (Hint: different scales per parameter)
3. Why is bias correction ($\hat{m}_t = m_t / (1 - \beta_1^t)$) necessary in early iterations?

3b. When might SGD outperform Adam?¶

Consider these scenarios and explain which optimizer you'd prefer:

1. Training a large language model from scratch — does Adam's per-parameter memory overhead matter?
2. Final convergence precision on a well-conditioned convex problem — which gets closer to the true minimum?
3. Generalization on image classification — empirical evidence suggests SGD+momentum sometimes finds flatter minima. Why might that help test accuracy?

(Write your answers as markdown below or discuss in a study group.)

Your answers:

3a.1: ...

3a.2: ...

3a.3: ...

3b.1: ...

3b.2: ...

3b.3: ...

---

Exercise 4 — Coding: Learning Rate Finder¶

Implement a simple learning rate range test: run GD on a problem with exponentially increasing learning rates and plot the loss vs learning rate. The optimal $\alpha$ is typically just before the loss starts increasing.

Problem: $f(x,y) = x^2 + 4y^2$ (same quadratic, $M = 8$, optimal $\alpha = 0.125$)

Tasks:

1. Test $\alpha \in \{10^{-3}, 10^{-2.8}, 10^{-2.6}, \dots, 10^{0}\}$ (about 15 values)
2. Run 20 GD steps for each $\alpha$, record final $f(\theta)$
3. Plot log($\alpha$) vs log($f$) and identify the best $\alpha$
4. Verify that the best $\alpha$ is close to $1/M = 0.125$

In [6]:

def f_lr_test(w):
    return w[0]**2 + 4 * w[1]**2

def grad_lr_test(w):
    return np.array([2 * w[0], 8 * w[1]])

# TODO: implement learning rate range test
# alphas = np.logspace(-3, 0, 15)  # from 0.001 to 1.0
# final_losses = []
# for alpha in alphas:
#     w = np.array([4.0, 3.0])
#     for _ in range(20):
#         ...
#     final_losses.append(f_lr_test(w))

# TODO: plot and find best alpha
# plt.semilogx(alphas, final_losses, 'bo-')
# ...

# TODO: verify best alpha is near 0.125
# best_alpha = alphas[np.argmin(final_losses)]
# print(f"Best alpha found: {best_alpha:.4f}")
# print(f"Theoretical optimal: {1/8:.4f}")

def f_lr_test(w): return w[0]**2 + 4 * w[1]**2 def grad_lr_test(w): return np.array([2 * w[0], 8 * w[1]]) # TODO: implement learning rate range test # alphas = np.logspace(-3, 0, 15) # from 0.001 to 1.0 # final_losses = [] # for alpha in alphas: # w = np.array([4.0, 3.0]) # for _ in range(20): # ... # final_losses.append(f_lr_test(w)) # TODO: plot and find best alpha # plt.semilogx(alphas, final_losses, 'bo-') # ... # TODO: verify best alpha is near 0.125 # best_alpha = alphas[np.argmin(final_losses)] # print(f"Best alpha found: {best_alpha:.4f}") # print(f"Theoretical optimal: {1/8:.4f}")

---

Exercise 5 — Failure Analysis: Ill-Conditioning and Momentum¶

Setup: Consider the quadratic $f(x,y) = x^2 + \kappa \cdot y^2$ for different condition numbers $\kappa \in \{1, 10, 100, 1000\}$.

Tasks:

1. For each $\kappa$, run vanilla GD and momentum GD from $(5, 1)$
2. Use $\alpha = 1/(2\kappa)$ for vanilla GD and $\alpha = 1/(2\kappa)$, $\beta = 0.9$ for momentum
3. Record the number of iterations to reach $f(\theta) < 10^{-6}$
4. Plot iterations vs $\kappa$ for both methods

Questions:

• How does the iteration count scale with $\kappa$ for vanilla GD?
• How does momentum improve this scaling?
• Why does ill-conditioning make GD zig-zag?

In [7]:

# TODO: implement the ill-conditioning experiment
# kappas = [1, 10, 100, 1000]
# iters_gd = []
# iters_mom = []
#
# for kappa in kappas:
#     def grad_kappa(w):
#         return np.array([2 * w[0], 2 * kappa * w[1]])
#
#     def f_kappa(w):
#         return w[0]**2 + kappa * w[1]**2
#
#     lr = 1 / (2 * kappa)
#     ...
#
# TODO: plot iterations vs kappa
# TODO: answer the questions above

# TODO: implement the ill-conditioning experiment # kappas = [1, 10, 100, 1000] # iters_gd = [] # iters_mom = [] # # for kappa in kappas: # def grad_kappa(w): # return np.array([2 * w[0], 2 * kappa * w[1]]) # # def f_kappa(w): # return w[0]**2 + kappa * w[1]**2 # # lr = 1 / (2 * kappa) # ... # # TODO: plot iterations vs kappa # TODO: answer the questions above