Adam Optimizer From Scratch in NumPy, With Bias Correction

Starting from parameters theta, apply one Adam step for each gradient vector in grads, in order, with step counter t = 1, 2, …:

m ← β₁m + (1 − β₁)g, v ← β₂v + (1 − β₂)g², m̂ = m/(1 − β₁^t), v̂ = v/(1 − β₂^t), θ ← θ − lr·m̂/(√v̂ + eps)

with m and v starting at zero. Return the final parameters.

Examples

Input theta = [1], grads = [[0.5]], lr = 0.1, beta1 = 0.9, beta2 = 0.999, eps = 1.000e-8

Output [0.9]

The first step moves by almost exactly lr, whatever the gradient's size.

Input theta = [1, -2], grads = [[100, -0.001]], lr = 0.01, beta1 = 0.9, beta2 = 0.999, eps = 1.000e-8

Output [0.99, -1.99]

Gradients of 100 and 0.001 produce steps of the same size.

Input theta = [0], grads = [[1], [-1], [1], [-1]], lr = 0.1, beta1 = 0.9, beta2 = 0.999, eps = 1.000e-8

Output [-0.1231]

An oscillating gradient: the first moment averages it away.

Constraints

1 ≤ len(theta) ≤ 1000
1 ≤ len(grads) ≤ 500
0 < β₁, β₂ < 1
Checked to a relative tolerance of 10⁻⁹

Hints

Hint 1

Keep m and v as arrays the shape of theta, both starting at zero, and a step counter starting at 1.

Hint 2

The bias correction divides by 1 − βᵗ, where t is the step number, not the index from zero.

Hint 3

eps is added after the square root, outside it.

Where candidates go wrong

Starting t at 0: the first correction divides by 1 − β⁰ = 0.
Putting eps inside the square root. PyTorch and the original paper add it outside.
Squaring m instead of g in the second moment.
Applying the bias correction to m and v themselves rather than to copies, so the error compounds.

The idea

m is an exponential moving average of the gradient and v of its elementwise square. Both start at zero, so after t steps their expectations are (1 − β^t) times the quantities they estimate. Dividing by 1 − β^t removes that factor. Without it, the second moment is tiny in the first steps (β₂ = 0.999 means 1 − β₂ = 0.001), and the steps are enormous.

At t = 1, m̂ = g and v̂ = g², so the step is lr·g/(|g| + eps): the learning rate times the sign of the gradient. That scale invariance is Adam's main property and the reason one learning rate works across layers whose gradients differ by orders of magnitude.

Time O(T · P) for T steps and P parameters · Space O(P)

Reference solution

import numpy as np

def adam(theta, grads, lr, beta1, beta2, eps):
    theta = np.array(theta, dtype=float)
    m = np.zeros_like(theta)
    v = np.zeros_like(theta)
    for t, g in enumerate(grads, start=1):
        g = np.asarray(g, dtype=float)
        m = beta1 * m + (1 - beta1) * g            # first moment
        v = beta2 * v + (1 - beta2) * g * g        # second moment
        m_hat = m / (1 - beta1 ** t)               # both averages start at zero: correct for it
        v_hat = v / (1 - beta2 ** t)
        theta -= lr * m_hat / (np.sqrt(v_hat) + eps)
    return theta

Console⌘/Ctrl + Enter runs

Python runs in your browser with NumPy. The first run downloads the runtime, which takes a few seconds.

Why interviewers ask it

Adam trains almost every network in use, and few people who use it can write it down. The interviewer is looking for the bias correction and an explanation of why it is needed, more than for the update itself.

The theory behind it is in the lesson Backpropagation from scratch.

Keep going

Easy10 min

The Adam optimiser, with bias correction