Momentum Update Step

mediumFREEOptimization & TrainingDeep LearningDeep Learning 1

Problem Statement

Update velocity and parameters using Momentum optimization.

Background

Momentum helps accelerate gradients in the right direction and dampens oscillations. The update rules are:

$v = \beta \cdot v + (1 - \beta) \cdot dW$ $W = W - \alpha \cdot v$

Where:

$v$ is the velocity (exponentially weighted average of gradients)
$\beta$ is the momentum coefficient (typically 0.9)
$\alpha$ is the learning rate
$dW$ is the gradient

Your Task

Write a function momentum_step(weight, grad, velocity, beta, learning_rate) that returns a tuple of (updated_weight, updated_velocity).

Output Format

Return a tuple (weight, velocity) with values rounded to 4 decimal places.

Example

Input

weight=0.0, grad=1.0, velocity=0.0, beta=0.9, learning_rate=0.1

Output

(-0.01, 0.1)

Reasoning

Inputs: a scalar parameter weight $W = 0$ , its current gradient grad $dW = 1$ , the running velocity velocity $v = 0$ (a first step), momentum coefficient beta $\beta = 0.9$ and learning rate learning_rate $\alpha = 0.1$ .
Velocity first: the new velocity is a weighted blend of the old velocity and the fresh gradient: $v' = 0.9 \cdot 0 + (1 - 0.9) \cdot 1 = 0.1$ . With $v = 0$ , only the $(1-\beta)$ share of the gradient survives.
Then the weight: the step uses the new velocity, not the old one: $W' = 0 - 0.1 \cdot 0.1 = -0.01$ .
Output: the pair (weight, velocity) = (-0.01, 0.1), each rounded to 4 decimal places. Note the order: weight comes first even though velocity is computed first. Rounding also hides float noise ( $1 - 0.9$ is not exactly $0.1$ ).
Notice: a plain gradient step would move $W$ by $\alpha \cdot dW = 0.1$ ; momentum moves it only $0.01$ here. If the same gradient keeps arriving, $v$ grows toward $dW$ and the steps speed up.

Constraints

-1000 <= weight, grad, velocity <= 1000
0 <= beta < 1
0 < learning_rate <= 1

Learn this in

Linear RegressionML Study Plan Neural NetworksML Study Plan

More from Deep Learning 1

Early Stopping Checkeasy L2 Regularization Penaltyeasy Inverted Dropout Forwardmedium Binary Cross-Entropy Losseasy SGD Parameter Updateeasy

Concept primer

Momentum Optimization: Comprehensive Background

1. Background Knowledge

Momentum is a fundamental optimization technique that accelerates gradient descent by maintaining a velocity vector that accumulates gradients over time. Rather than updating parameters directly with the current gradient, momentum uses an exponentially weighted moving average of past gradients, which provides two key benefits:

Acceleration: Momentum builds up speed in consistent directions, allowing faster convergence toward minima
Noise dampening: Oscillations caused by noisy gradients are reduced, creating smoother optimization trajectories

Core Concept: Exponential Moving Average

The velocity update uses an exponential moving average (EMA), where older gradients decay exponentially. The parameter $\beta$ (typically 0.9) controls this decay:

Higher $\beta$ values (e.g., 0.95) give more weight to historical gradients
Lower $\beta$ values (e.g., 0.5) make the algorithm more responsive to recent gradients

Historical Context

Momentum methods include several variants:

Heavy-ball momentum: Classical approach using a constant momentum coefficient
Nesterov's Accelerated Gradient (NAG): Looks ahead before computing gradients
Quasi-hyperbolic momentum (QHM): Interpolates between current and momentum-based updates

2. Algorithm Approach

The momentum update follows a two-step process:

Step 1: Update Velocity (Exponential Moving Average) $v_t = \beta \cdot v_{t-1} + (1 - \beta) \cdot dW_t$

This combines the previous velocity (scaled by \beta) with the current gradient (scaled by 1-\beta).

Step 2: Update Parameters $W_t = W_{t-1} - \alpha \cdot v_t$

The parameter update uses the accumulated velocity, not the raw gradient.

Why This Works

The velocity acts as a low-pass filter that smooths gradient noise. When gradients consistently point in one direction, velocity accumulates. When gradients oscillate, the averaging effect cancels out noise.

3. Step-by-Step Strategy

Implementation Steps

Compute new velocity using the EMA formula:

Multiply previous velocity by $\beta$
Multiply current gradient by $(1 -$ \beta $)$
Sum these two terms

Update weight using the new velocity:

Multiply velocity by learning rate $\alpha$
Subtract from current weight

Round results to 4 decimal places for output

Walkthrough with Sample Input

Given: weight=0.0, grad=1.0, velocity=0.0, beta=0.9, learning_rate=0.1

Step 1: v_new = 0.9 * 0.0 + (1 - 0.9) * 1.0
 v_new = 0.0 + 0.1 * 1.0 = 0.1

Step 2: w_new = 0.0 - 0.1 * 0.1
 w_new = 0.0 - 0.01 = -0.01

Result: (-0.01, 0.1)

4. Common Pitfalls

Pitfall	Issue	Solution
Order of operations	Computing weight update before velocity update	Always update velocity first, then use new velocity for weight update
Incorrect scaling	Forgetting (1-\beta) factor in velocity update	Ensure gradient is scaled by (1-\beta), not just added directly
Sign errors	Using addition instead of subtraction for weight update	Remember: W = W - \alpha · v (gradient descent, not ascent)
Rounding precision	Floating-point errors accumulating	Round only at the final output, not intermediate steps
Initial velocity	Assuming velocity must start at zero	Velocity can start at any value; zero is conventional but not required

5. Time & Space Complexity

Time Complexity: $O(1)$ — Each momentum step performs a constant number of arithmetic operations regardless of problem size
Space Complexity: $O(1)$ — Only stores scalar values (weight, velocity, gradient) per parameter

Practical Considerations

In practice, momentum is applied to all parameters in a neural network simultaneously, so total complexity scales with the number of parameters. However, each individual update remains constant-time.

Implementation Template

def momentum_step(weight, grad, velocity, beta, learning_rate):
 # Step 1: Update velocity (exponential moving average)
 new_velocity = beta * velocity + (1 - beta) * grad

 # Step 2: Update weight using new velocity
 new_weight = weight - learning_rate * new_velocity

 # Step 3: Round to 4 decimal places
 return (round(new_weight, 4), round(new_velocity, 4))

This implementation directly translates the mathematical formulas into code, maintaining clarity and correctness.

solution.py

Python 3.13

Test results