SGD with Momentum

Write one update step of stochastic gradient descent with momentum.

Plain gradient descent moves each parameter by the current gradient only. With momentum, the optimiser keeps a velocity: a running sum of past gradients that decays by a factor beta at each step. The parameters then move along the velocity, not along the raw gradient. This smooths out noisy gradients and speeds up progress along directions where the gradient keeps the same sign.

Use the same form as PyTorch:

v_new      = beta * v + grad
params_new = params - lr * v_new

Return the new parameters and the new velocity. Do not change the input arrays.

Function signature

def momentum_update(params, grads, velocity, lr, beta=0.9):

Example

Input

params   = np.array([1.0, -2.0])
grads    = np.array([0.5, -1.0])
velocity = np.array([0.1, 0.2])
momentum_update(params, grads, velocity, lr=0.1, beta=0.9)

Output

(array([ 0.941, -1.918]), array([ 0.59, -0.82]))

Reasoning

The velocity becomes 0.9 * [0.1, 0.2] + [0.5, -1.0] = [0.59, -0.82]. The parameters then move by 0.1 times that velocity: 1.0 - 0.059 = 0.941 and -2.0 + 0.082 = -1.918.

Constraints and notes

• params, grads and velocity have the same shape, which can be any shape.
• With beta = 0 the step is plain gradient descent.
• Return new arrays. Do not modify the inputs in place.