Adam Optimizer Step

Adam is a first-order stochastic optimization algorithm that improves plain gradient descent by keeping exponentially decaying moving averages of both the gradient (a kind of momentum) and the squared gradient (a kind of adaptive per-parameter learning rate, like RMSprop). Conceptually, at each step it uses past gradients to estimate the mean and variance of the gradient for each parameter, then rescales the update so that noisy / high-variance directions get smaller effective steps and consistent directions get larger ones.

Because these moving averages start at zero, their early values are biased toward zero; Adam corrects this with bias correction factors depending on the timestep $t$. After correction, it updates parameters by subtracting a scaled, normalized first moment: the estimated mean gradient divided by the square root of the estimated second moment (plus a small $\epsilon$ to avoid division by zero). This yields stable, per-parameter step sizes that work well in deep learning.

---

1. Background Knowledge

Key concepts you should understand:

• Gradient descent: For parameters $\theta$, learning rate $\alpha$, and gradient $g_t =$\nabla_\theta$L($\theta$_{t-1})$, vanilla SGD does

$$\theta_t = \theta_{t-1} - \alpha g_t.$$

• Momentum: Keeps an exponential moving average of gradients:

$$m_t = \beta_1 m_{t-1} + (1 - \beta_1) g_t,$$

then updates $\theta$ using $m_t$ instead of $g_t$. This smooths noisy gradients and accelerates in consistent directions.

• Adaptive learning rates (RMSprop-style): Keeps an exponential moving average of squared gradients:

$$v_t = \beta_2 v_{t-1} + (1 - \beta_2) g_t^2,$$

and scales the step by $1 / ($\sqrt{v_t}$+$\epsilon$)$ so each parameter gets its own effective step size.

Adam combines both: it keeps first moment $m_t$ and second moment $v_t$, applies bias correction, then does a scaled update.

---

2. Algorithm / Approach Pattern

For this coding task, you are not training a network; you are asked to implement one optimizer step given all necessary inputs.

The general pattern:

1. Inputs:

• Current parameters $\theta_{t-1}$
• Previous first moment $m_{t-1}$
• Previous second moment $v_{t-1}$
• Current gradient $g_t$
• Hyperparameters $\alpha,$\beta$_1,$\beta$_2, \epsilon$
• Current timestep $t$ (1-based)

2. Compute:

• New first moment $m_t$
• New second moment $v_t$
• Bias-corrected moments $\hat{m}_t,$\hat{v}$_t$
• New parameters $\theta_t$

3. Return updated $\theta_t, m_t, v_t$.

All operations are element-wise over the parameter tensors.

---

3. Step-by-Step Strategy

Assuming everything is stored as arrays/tensors of the same shape:

1. Update first moment (biased):

$$m_t = \beta_1 \cdot m_{t-1} + (1 - \beta_1) \cdot g_t$$

2. Update second moment (biased):

• Square the gradient element-wise: $g_t^2$
• Then:

$$v_t = \beta_2 \cdot v_{t-1} + (1 - \beta_2) \cdot g_t^2$$

3. Compute bias-correction denominators:

• You will be given $t$ (or you increment it); compute:

$$\text{bias1\_corr\_denom} = 1 - \beta_1^t$$ $$\text{bias2\_corr\_denom} = 1 - \beta_2^t$$

4. Bias-correct the moments:

$$\hat{m}_t = \frac{m_t}{1 - \beta_1^t}, \quad \hat{v}_t = \frac{v_t}{1 - \beta_2^t}$$

5. Compute the normalized step:

• Denominator: \sqrt{\hat{v}_t} + \epsilon (element-wise sqrt).
• Fraction: $\hat{m}_t / ($\sqrt{$\hat{v}$_t}$+$\epsilon$)$.

6. Update parameters:

$$\theta_t = \theta_{t-1} - \alpha \cdot \frac{\hat{m}_t}{\sqrt{\hat{v}_t} + \epsilon}$$

7. Return:

• New parameters $\theta_t$
• New moments $m_t, v_t$
• (Optionally timestep $t+1$, depending on interface).

Skeleton-style example in pseudocode:

def adam_step(theta, grad, m, v, t, alpha, beta1, beta2, eps):
 m = beta1 * m + (1 - beta1) * grad
 v = beta2 * v + (1 - beta2) * (grad ** 2)

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

 theta = theta - alpha * m_hat / (v_hat**0.5 + eps)
 return theta, m, v

---

4. Common Pitfalls

• Forgetting bias correction: Omitting $\hat{m}_t$ and $\hat{v}_t$ (using $m_t, v_t$ directly) makes early steps too small, diverging from the true Adam algorithm.

• Wrong timestep $t$:

• Make sure $t$ starts at 1 for the first update.

• Do not recompute $\beta_1^t,$\beta$_2^t$ incorrectly (e.g., using $t-1$).

• Element-wise vs scalar ops:

• $g_t^2$, \sqrt{\hat{v}_t}, division, and addition of $\epsilon$ are per-parameter operations.

• Do not accidentally take a global norm or sum.

• Numeric stability:

• Ensure $\epsilon$ is added inside the denominator: sqrt(v_hat) + eps, not sqrt(v_hat + eps) (most Adam definitions use the former).

• Use floating-point types; integer division will break things.

• Shape mismatch:

• theta, grad, m, v must all have the same shape; verify broadcasting doesn’t hide a bug.

---

5. Time & Space Complexity

Let $n$ be the number of parameters:

• Time complexity:

• Each step does a constant number of element-wise operations over $n$ elements: updates for $m$, $v$, bias correction, sqrt, division, and parameter update.

• Overall: $O(n)$ per step.

• Space complexity:

• You store $m$, $v$, $\theta$, and usually $g_t$.

• The optimizer-specific extra memory is from $m$ and $v$: $O(n)$ auxiliary space.