SGD with Momentum

The problem

Plain SGD (stochastic gradient descent) moves every weight against the gradient of the current mini-batch. That gradient is noisy, so the path zigzags. Momentum smooths it: it keeps a running direction, the velocity, built from all the recent gradients, and moves the weights along that instead.

You will implement one step of SGD with momentum, in the form used by PyTorch's torch.optim.SGD.

One step, in two lines

First update the velocity, then move the weights with the new velocity:

where:

• $\theta$ (theta) is the parameter array param, and $g$ its gradient grad;
• $v$ is the velocity velocity, with the same shape as param, and all zeros before the first step;
• $\eta$ (eta) is the learning rate lr;
• $\mu$ (mu) is momentum, usually $0.9$: the share of the old velocity that is kept;
• $\leftarrow$ means "replace with".

In words: the new velocity is 90% of the old one plus the current gradient, and the parameters move a step of size lr against it.

Example. One parameter, lr = 0.1, momentum = 0.9, starting at param = 1.0 with velocity = 0.0. Three steps:

step	grad	new velocity	new param
1	0.5	$0.9 \cdot 0 + 0.5 = 0.5$	$1.0 - 0.1 \cdot 0.5 = 0.95$
2	0.5	$0.9 \cdot 0.5 + 0.5 = 0.95$	$0.95 - 0.1 \cdot 0.95 = 0.855$
3	$-0.2$	$0.9 \cdot 0.95 - 0.2 = 0.655$	$0.855 - 0.1 \cdot 0.655 \approx 0.790$

In step 2 the same gradient moves the parameter almost twice as far as in step 1: the velocity has built up. In step 3 the gradient changes sign, but the velocity is still positive, so the parameter keeps going down, only more slowly: one noisy gradient does not reverse the direction.

With momentum = 0 the velocity is just the current gradient, and this is plain SGD.

What to implement

sgd_momentum_step(param, grad, velocity, lr, momentum=0.9): return the pair (new_param, new_velocity).

Requirements

• Do not modify the input arrays: return new ones. (Real optimizers update in place for speed; here returning values makes each step easy to test.)
• momentum=0 must be plain SGD.
• Works for parameters of any shape.

Hint

The learning rate multiplies the velocity, not the gradient. The form $v \leftarrow \mu v + \eta g$, then $\theta \leftarrow \theta - v$, is the one of Sutskever et al. (2013): a fine method, but not the one asked for here. The theory page explains when the two differ.