Plain gradient descent zigzags across a narrow valley: the steep direction caps the learning rate, and the flat one barely moves. Two improvements dominate in practice.

Momentum keeps a running velocity, v = beta v + g, and steps x - lr v. Consistent gradients add up and speed the flat direction; alternating ones cancel.

Adam keeps running averages of the gradient, m, and of its square, s, and steps by lr m / (sqrt(s) + eps), so every coordinate moves about lr regardless of its scale. Both averages start at 0 and are too small at first, so Adam divides them by 1 - beta ** t at step t, which is bias correction.

import numpy as np

g = np.array([0.01, 100.0])
m = 0.1 * g
s = 0.001 * g ** 2
print(m / (np.sqrt(s) + 1e-8))
m_hat = m / (1 - 0.9)
s_hat = s / (1 - 0.999)
print(m_hat / (np.sqrt(s_hat) + 1e-8))

Without the correction the first steps are three times too large. With it, the first step moves every coordinate by lr, whatever its gradient.

Write momentum_step(x, v, g, lr, beta=0.9) returning the new (x, v), and adam_step(x, m, s, g, t, lr, b1=0.9, b2=0.999, eps=1e-8) returning the new (x, m, s) with bias correction, where t counts steps from 1. Use NumPy arrays and do not change the arguments in place.

The explanation opens here once every check passes.