Optimizer race

Four optimizers start from the same point on the same surface. Tap the map to move the start.

The Rosenbrock function: easy to find the valley, slow to follow its gently curving floor to the minimum.

Tap to move the start

Gradient descent
Momentum
RMSProp
Adam

Challenges

0 of 4 complete

On Banana valley, race until any racer gets its loss below 0.001. (not yet)
Scale the step sizes to 3x or more on Banana valley, then race until a racer gets its loss below 0.0001. (not yet)
On Saddle point, race until a racer slides off the ridge to a loss below -0.25. Watch which ones lead. (not yet)
On Many valleys, race until two racers come to rest in different valleys. (not yet)

Try this

Race on the Banana valley with Fast forward, then scale the step sizes to 4x.
On the Saddle point, count how many steps each optimizer needs to slide off the ridge.

Four optimizers in NumPyPython with NumPy

import numpy as np

def loss(p):
    # The banana valley (a Rosenbrock function), lowest at (1, 1).
    x, y = p
    return (1 - x) ** 2 + 10 * (y - x * x) ** 2

def grad(p):
    x, y = p
    return np.array([-2 * (1 - x) - 40 * x * (y - x * x), 20 * (y - x * x)])

def race(name, lr, max_steps=1500):
    p = np.array([-1.6, 2.4])  # the demo's start
    v = np.zeros(2)            # velocity (momentum) or mean gradient (Adam)
    s = np.zeros(2)            # running mean of squared gradients
    for t in range(1, max_steps + 1):
        g = grad(p)
        if name == "gradient descent":
            p = p - lr * g
        elif name == "momentum":
            v = 0.9 * v + g
            p = p - lr * v
        elif name == "RMSProp":
            s = 0.99 * s + (1 - 0.99) * g ** 2
            p = p - lr * g / (np.sqrt(s) + 1e-8)
        elif name == "Adam":
            v = 0.9 * v + (1 - 0.9) * g
            s = 0.999 * s + (1 - 0.999) * g ** 2
            v_hat = v / (1 - 0.9 ** t)    # bias correction for the
            s_hat = s / (1 - 0.999 ** t)  # first few steps
            p = p - lr * v_hat / (np.sqrt(s_hat) + 1e-8)
        if loss(p) < 1e-3:
            return t
    return None

# Each optimizer with its base step size from the demo.
for name, lr in [("gradient descent", 0.01), ("momentum", 0.002), ("RMSProp", 0.02), ("Adam", 0.05)]:
    print(f"{name}: loss below 0.001 after {race(name, lr)} steps")

gradient descent: loss below 0.001 after 746 steps
momentum: loss below 0.001 after 341 steps
RMSProp: loss below 0.001 after 382 steps
Adam: loss below 0.001 after 210 steps