Neural Mastery
← Practice
SGD with Momentum Parameter Stepeasy

SGD with Momentum Parameter Step

easy Neural Networks by Hand 15 min

NumPy allowed · plain Python earns 10 bonus points

Gradient descent step 1/4

Loss moves downhill

The update rule changes the parameters on each step.
Open lesson →
01parameters
02prediction
03loss
04gradient
θ = 4.000loss = 16.000∂L/∂θ = 8.000

θ ← θ − 0.14 · ∂L/∂θ

Accelerate stochastic gradient descent and dampen oscillations along ravines using classical momentum.

Task

Implement sgd_momentum_step(param, grad, velocity, lr, momentum): for each index k, update the velocity as v_k = momentum * velocity[k] + lr * grad[k], and then update the parameter as param_k = param[k] - v_k. Return a dictionary {'param': [...], 'velocity': [...]} with each float rounded to 6 decimal places.

Function Signature

sgd_momentum_step(param: list[float], grad: list[float], velocity: list[float], lr: float, momentum: float) -> dict

Examples

Example 1: First step with zero velocity
Input: param = [1, 2] grad = [0.1, -0.2] velocity = [0, 0] lr = 0.1 momentum = 0.9
Output: {'param': [0.99, 2.02], 'velocity': [0.01, -0.02]}
Example 2: Second step with accumulated velocity
Input: param = [0.99, 2.02] grad = [0.1, -0.2] velocity = [0.01, -0.02] lr = 0.1 momentum = 0.9
Output: {'param': [0.971, 2.058], 'velocity': [0.019, -0.038]}
Example 3: Zero gradient inertia
Input: param = [0.5] grad = [0] velocity = [0.1] lr = 0.01 momentum = 0.5
Output: {'param': [0.45], 'velocity': [0.05]}

Constraints

  • len(param) == len(grad) == len(velocity) >= 1
  • 0.0 <= momentum < 1.0
  • lr > 0.0
Python3Saved ✓
def sgd_momentum_step(param: list[float], grad: list[float], velocity: list[float], lr: float, momentum: float) -> dict:
"""Compute one parameter update step using SGD with classical momentum.

Args:
param: Current 1D parameter weights.
grad: Current parameter gradient vector.
velocity: Previous velocity buffer vector.
lr: Learning rate.
momentum: Momentum decay coefficient (e.g. 0.9).

Returns:
Dictionary with updated 'param' and 'velocity' lists, rounded to 6 decimals.
"""
pass

Case 1: First step with zero velocity
Input: param = [1, 2] grad = [0.1, -0.2] velocity = [0, 0] lr = 0.1 momentum = 0.9
Expected: {'param': [0.99, 2.02], 'velocity': [0.01, -0.02]}
Case 2: Second step with accumulated velocity
Input: param = [0.99, 2.02] grad = [0.1, -0.2] velocity = [0.01, -0.02] lr = 0.1 momentum = 0.9
Expected: {'param': [0.971, 2.058], 'velocity': [0.019, -0.038]}
Case 3: Zero gradient inertia
Input: param = [0.5] grad = [0] velocity = [0.1] lr = 0.01 momentum = 0.5
Expected: {'param': [0.45], 'velocity': [0.05]}