Medium Primitives

Implement Momentum Update

Momentum keeps a running buffer of past gradients. Each step retains a fraction of the old buffer and adds the current gradient. That updated buffer determines the weight change.

This exercise uses the following convention:

new_velocity = momentum * velocity + gradients

new_weights = weights - lr * new_velocity

For weights [1, 2], gradients [0.5, -0.5], old velocity [1, 0], lr=0.1, and momentum=0.9, the new velocity is [1.4, -0.5] and the new weights are [0.86, 2.05]. The weights use the new velocity, not the buffer from the previous step.

Your task

Implement momentum_step(weights, gradients, velocity, lr, momentum). The three floating-point tensors have the same shape. lr is a nonnegative float and momentum is a float in [0, 1).

Return a Python dictionary with exactly these fields:

"new_weights": the updated weight tensor.
"new_velocity": the updated momentum-buffer tensor.

Both tensors keep the input shape. Do not return a tuple, and do not modify the input tensors. This is one update using supplied gradients; you do not calculate gradients or maintain hidden optimizer state.

The buffer stores a weighted sum of gradients, not a learning-rate-scaled parameter change. Do not multiply the current gradient by 1 - momentum, and do not add weight decay or a Nesterov correction.

With momentum=0, the new velocity equals the current gradient. With a zero current gradient, a nonzero old velocity can still move the weights. With lr=0, the weight values stay unchanged but the velocity still updates.

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 Momentum Update

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