Adam Optimizer

Implement one Adam update step.

Adam keeps two running averages per parameter: m estimates the mean of the gradient and v estimates its uncentred variance.

m = beta1 * m + (1 - beta1) * g
v = beta2 * v + (1 - beta2) * g^2

m_hat = m / (1 - beta1^t)
v_hat = v / (1 - beta2^t)

param = param - lr * m_hat / (sqrt(v_hat) + eps)

t is the 1-based step count, so the first call has t = 1.

The bias correction matters more than it looks. Both averages start at zero, so without dividing by 1 - beta^t the first steps are pulled towards zero and training crawls until the averages warm up. With beta2 = 0.999 that warm up takes thousands of steps.

Dividing by sqrt(v_hat) is what makes Adam scale invariant per parameter: a coordinate with consistently small gradients still gets a full sized step.

Return (new_param, new_m, new_v). Do not modify the arrays you are given.

Function signature

def adam_step(param, grad, m, v, t, lr=0.001, beta1=0.9, beta2=0.999, eps=1e-8):

Example

Input

adam_step(param=np.array([1.0]), grad=np.array([1.0]),
          m=np.zeros(1), v=np.zeros(1), t=1, lr=0.1)

Output

(array([0.9]), array([0.1]), array([0.001]))

Reasoning

After one step m = 0.1 and v = 0.001. Bias correction divides them by 1 - 0.9 and 1 - 0.999, so m_hat = v_hat = 1. The step is lr * 1 / sqrt(1) = 0.1, so the parameter moves from 1.0 to 0.9.

Constraints and notes

• t starts at 1 on the first call.
• eps is added after the square root.
• Return new arrays. Do not change the ones you are given.