SGD with Momentum - PyTorch Basics

SGD with Momentum

What you'll build: sgd_momentum_step(...), one update of stochastic gradient descent with momentum, in PyTorch's formulation.

Background

Plain gradient descent zig-zags across narrow valleys and crawls along flat directions. Momentum keeps a velocity: a decaying sum of past gradients. Components that agree from step to step add up and accelerate, components that flip sign cancel out. With momentum = 0.9, a constant gradient eventually produces steps 10 times larger than plain SGD (1 / (1 - 0.9)).

Algorithm

Element-wise, at each step:

1. Velocity: v = momentum * v + g
2. Update: params = params - lr * v

This is exactly what torch.optim.SGD(params, lr, momentum) does (it calls v the momentum_buffer). Some textbooks write the equivalent form v = momentum * v - lr * g; params = params + v; it matches only while lr is constant, and the tests use PyTorch's form.

Comparison

Optimizer	Update rule
Vanilla SGD	p = p - lr * g
SGD + momentum	v = μ * v + g; p = p - lr * v

Inputs

• params, grads, velocity: float arrays of the same shape (velocity starts at zeros).
• lr: learning rate (default 0.01); momentum: coefficient μ in [0, 1) (default 0.9).

Expected Return Format

Return a dictionary with these keys, each the same shape as params:

• 'params': updated parameters
• 'velocity': updated velocity

Constraints

• Return new arrays rather than modifying the inputs in place.
• PyTorch track: use tensor ops; don't use torch.optim.SGD (the hidden tests compare against it).

Follow-up

• Nesterov momentum evaluates the gradient at the look-ahead point. How does PyTorch's nesterov=True change the update?
• Why does a large momentum often need a smaller learning rate?

Input: sgd_momentum_step(np.array([1.0]), np.array([1.0]), np.zeros(1), lr=0.1, momentum=0.9)
Output: {'params': array([0.9]), 'velocity': array([1.])}

Input: second step with the same gradient, velocity=[1.0]
Output: {'params': array([0.71]), 'velocity': array([1.9])}

import numpy as np

def sgd_momentum_step(params, grads, velocity, lr=0.01, momentum=0.9):
    """
    Perform one step of SGD with momentum.

    Args:
        params: Current parameters (np.ndarray)
        grads: Current gradients (np.ndarray, same shape as params)
        velocity: Current velocity (np.ndarray, same shape as params)
        lr: Learning rate
        momentum: Momentum coefficient

    Returns:
        dict with 'params' and 'velocity' (updated, same shape as params)
    """
    # Your code here
    pass