Hard Primitives

Implement Adam Optimizer Step

Adam keeps two running statistics for each parameter: m averages the gradients, and v averages their squares. The second statistic measures squared magnitude, not variance about the mean. Together they determine an update scale for each parameter.

Starting these buffers at zero makes their early values smaller. Adam compensates by dividing each updated buffer by a bias-correction factor that depends on the step number.

Your task

Implement adam_step(weights, gradients, m, v, t, lr=0.001, beta1=0.9, beta2=0.999, eps=1e-8) for floating-point tensors of matching shape. The supplied m and v are the uncorrected buffers from before this step. v is nonnegative. t is the step being performed, starting at 1; it is not the number of already completed steps.

Compute these elementwise operations in order:

new_m = beta1 * m + (1 - beta1) * gradients
new_v = beta2 * v + (1 - beta2) * gradients ** 2
m_hat = new_m / (1 - beta1 ** t)
v_hat = new_v / (1 - beta2 ** t)
new_weights = weights - lr * m_hat / (sqrt(v_hat) + eps)

Use nonnegative lr, decay coefficients in [0, 1), and positive eps. Add epsilon after the square root. Do not add weight decay or other optimizer variants. Return a Python dictionary with keys "new_weights", "new_m", and "new_v", each containing a same-shaped tensor. Return the uncorrected updated buffers, not m_hat and v_hat, so they can be passed into the next step. Do not modify the input tensors.

For a first step with weight 1, gradient 0.1, zero buffers, and default parameters, new_m is 0.01 and new_v is 0.00001. Bias correction gives m_hat=0.1 and v_hat=0.01. The new weight is approximately 0.999. On the next call, use the returned buffers and t=2; do not reset them.

The gradients are already supplied. This task performs one arithmetic update, not loss evaluation, gradient calculation, or an optimizer loop.

Understand the ideas behind this exercise. The Problem tab contains the exact implementation requirements.

Use previous gradients to guide the next update

Plain gradient descent uses the current gradient. Momentum also keeps a running update direction:

$$v_t=\beta v_{t-1}+g_t.$$

Here, g_t is the current gradient and beta controls how much of the previous direction remains. The parameter update subtracts the learning rate times v_t. Some definitions multiply the new gradient by (1 - beta) as well. That changes the scale of the updates, so use the convention in the Problem tab.

Adam keeps two running averages

Adam tracks both the gradient and its square, element by element:

$$m_t=\beta_1m_{t-1}+(1-\beta_1)g_t,$$ $$v_t=\beta_2v_{t-1}+(1-\beta_2)g_t^2.$$

The first average helps estimate the update direction. The second tracks the size of recent gradients. It is a mean square, not a variance: Adam does not subtract the mean before squaring.

Starting both averages at zero makes their early values too small. Adam corrects for that initialization using

$$\hat m_t=\frac{m_t}{1-\beta_1^t},\qquad \hat v_t=\frac{v_t}{1-\beta_2^t}.$$

It then subtracts the learning rate times m_hat / (sqrt(v_hat) + epsilon) from the parameters. The first update uses t = 1.

For a first scalar gradient of 2, the corrected averages are 2 and 4. Ignoring epsilon, the update direction is 2 / sqrt(4) = 1. A gradient of 20 would also give direction 1 on that first step. This is why Adam's learning rate does not have the same effect as the learning rate in plain SGD.

Preserve state between steps

Use the newly updated averages when computing the parameter update. Keep those averages and the step counter for the next step; resetting them changes the optimizer.

Follow the specified placement of epsilon. Adding it inside the square root is a different formula. Also remember that a zero current gradient does not always mean no parameter change: stored momentum or weight decay can still cause an update.

Detecting runtime…

Implement Adam Optimizer Step

Your task

Use previous gradients to guide the next update

Adam keeps two running averages

Preserve state between steps

Your code, saved here

Stuck?

PyTorch reference solution