Adam Optimizer Step

hardFREECalculus & OptimizationOptimizationCV: Mathematical Foundations

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

Inputs: two parameters $\theta_0 = [1.0, 2.0]$ with gradients $g_1 = [0.1, 0.2]$ ; both moment estimates start at zero and this is the first step, $t = 1$ . All maths is element-wise.
First moment: $m_1 = 0.9 \cdot 0 + 0.1 \cdot g_1 = [0.01, 0.02]$ , still pulled toward its zero start.
Second moment: $v_1 = 0.999 \cdot 0 + 0.001 \cdot g_1^2 = 0.001 \cdot [0.01, 0.04] = [0.00001, 0.00004]$ .
Bias correction undoes that pull toward zero: $\hat{m}_1 = m_1 / (1 - 0.9) = [0.1, 0.2]$ and $\hat{v}_1 = v_1 / (1 - 0.999) = [0.01, 0.04]$ . At $t = 1$ they recover exactly $g_1$ and $g_1^2$ .
Update: $\sqrt{\hat{v}_1} = [0.1, 0.2]$ , so the ratio $\hat{m}_1 / (\sqrt{\hat{v}_1} + \epsilon) \approx [1, 1]$ and $\theta_1 = [1, 2] - 0.001 \cdot [1, 1] = [0.999, 1.999]$ . The $\epsilon$ only shifts the 10th decimal, lost when rounding to 6 places.
Notice: both parameters moved by exactly $\alpha$ even though one gradient is twice the other. The step size ignores the gradient's raw scale.

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

Learn this in

Neural NetworksML Study Plan

More from CV: Mathematical Foundations

Power Iteration for Dominant Eigenvectorhard SVD Image Compressionhard Batch Affine Transform with PyTorchmedium Covariance Matrix from Image Patcheshard Gram-Schmidt Orthonormalizationmedium

Concept primer

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:

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)

Compute:

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

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:

Update first moment (biased):

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

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

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

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}

Compute the normalized step:

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

Update parameters:

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

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.

solution.py

Python 3.13

Test results