THE ADAM STATE RESET

Difficulty: Medium Tags: Optimization, Adam, Optimizer State

Story

A training job resumes from a checkpoint, but its optimizer state was reset while the parameter weights were restored. The first resumed Adam update therefore differs from the update that would have happened if training had continued. Implement one Adam step with an explicitly supplied timestep and moments.

The Math

Given gradient g_t, first moment m_{t-1}, second moment v_{t-1}, and timestep t:

m_t = beta1*m_{t-1} + (1-beta1)*g_t
v_t = beta2*v_{t-1} + (1-beta2)*g_t^2

Bias correction:

m_hat = m_t / (1 - beta1^t)
v_hat = v_t / (1 - beta2^t)

Update:

theta_t = theta_{t-1} - eta * m_hat / (sqrt(v_hat) + eps)

Epsilon goes outside the square root, exactly as written.

Input Format

theta, g, m_prev, v_prev, beta1, beta2, eta, eps, t.

Output Format

(m_t, v_t, theta_t) (8 decimal places).

Constraints

• t >= 1, 0 <= beta1, beta2 < 1, eta > 0, eps > 0, v_prev >= 0.
• Time limit: 1.0 second.

Example

Input

theta=1, g=0.5, m_prev=0, v_prev=0
beta1=0.9, beta2=0.999, eta=0.01, eps=0.00000001, t=1

Output

m_t = 0.05000000
v_t = 0.00025000
theta_t = 0.99000000

Explanation: At t=1, bias correction gives m_hat=0.5 and v_hat=0.25. The update is approximately 0.01*0.5/(0.5+eps), yielding 0.99 to the displayed precision.