PythiloLesson 13 of 15 · progress kept in this browser

Learning rate and curvature

Near a minimum, a loss looks like a bowl: f(x) = 0.5 x @ A @ x, with A symmetric. Its gradient is A @ x, and the matrix A, the Hessian, holds the curvature. Its eigenvalues are the curvatures along its eigenvectors.

One gradient step is x - lr A x. Along an eigenvector with eigenvalue lam, that multiplies the coordinate by 1 - lr lam. When |1 - lr lam| is below 1 the coordinate shrinks; above 1 it grows every step and training diverges. So the learning rate must stay below 2 / lam_max.

import numpy as np

A = np.diag([1.0, 20.0])
x0 = np.array([1.0, 1.0])
for lr in (0.05, 0.11):
    x = x0.copy()
    for _ in range(50):
        x = x - lr * (A @ x)
    print(lr, np.abs(x).round(4))

The steepest direction caps the step size, while the flattest decides how many steps you need. Their ratio, the condition number, is why some losses train slowly at any learning rate.

Your task

Write max_stable_lr(A) returning 2 / lam_max for a symmetric matrix as a float (use np.linalg.eigvalsh). Write converges(A, lr) returning whether gradient descent on 0.5 x @ A @ x shrinks every direction, and descend(A, x0, lr, steps) returning x after that many steps.

The first run loads NumPy into this browser.