Adam Optimizer Step

Implement a single step of the Adam optimizer, widely used in deep learning.

Adam (Adaptive Moment Estimation) combines momentum and RMSprop, adapting learning rates per-parameter:

Algorithm (single step at time $t$ ):

Compute gradient $g_t$ at current parameters
Update biased first moment: $m_t = \beta_1 m_{t-1} + (1-\beta_1) g_t$
Update biased second moment: $v_t = \beta_2 v_{t-1} + (1-\beta_2) g_t^2$
Bias correction: $\hat{m}_t = \frac{m_t}{1-\beta_1^t}$ , $\hat{v}_t = \frac{v_t}{1-\beta_2^t}$
Update parameters: $\theta_t = \theta_{t-1} - \alpha \frac{\hat{m}_t}{\sqrt{\hat{v}_t} + \epsilon}$

Typical hyperparameters: $\alpha=0.001$ , $\beta_1=0.9$ , $\beta_2=0.999$ , $\epsilon=10^{-8}$

Example:

Input:

params = [1.0, 2.0]
grads = [0.1, 0.2]
m = [0.0, 0.0]
v = [0.0, 0.0]
t = 1
lr = 0.001
beta1 = 0.9
beta2 = 0.999

Output:

{'params': [0.999, 1.999], 'm': [0.01, 0.02], 'v': [0.00001, 0.00004]}

Reasoning:

Step 1: Update first moment (m) $m_1 = 0.9 \times 0 + 0.1 \times [0.1, 0.2] = [0.01, 0.02]$

Step 2: Update second moment (v) $v_1 = 0.999 \times 0 + 0.001 \times [0.01, 0.04] = [0.00001, 0.00004]$

Step 3: Bias correction $\hat{m}_1 = \frac{[0.01, 0.02]}{1-0.9^1} = [0.1, 0.2]$ $\hat{v}_1 = \frac{[0.00001, 0.00004]}{1-0.999^1} = [0.01, 0.04]$

Step 4: Parameter update $\theta_1 = [1, 2] - 0.001 \times \frac{[0.1, 0.2]}{\sqrt{[0.01, 0.04]} + \epsilon}$ $\theta_1 = [1, 2] - 0.001 \times \frac{[0.1, 0.2]}{[0.1, 0.2]}$ $\theta_1 = [1, 2] - 0.001 \times [1, 1] = [0.999, 1.999]$

Constraints:

params: Current parameter values (list)
grads: Gradients at current params (list)
m, v: First and second moment estimates (lists)
t: Current timestep (int >= 1)
Return: Dict with new 'params', 'm', 'v'
Round to 6 decimal places

Example:

Constraints:

Test Results