Problems
Loading...
1 / 1

Implement Nesterov Momentum (NAG)

Optimization
Easy

Implement one update step of Nesterov Momentum (NAG). Given current parameters, velocity, and gradients, return updated parameters and velocity using the look-ahead gradient approach.

Step 1: Look Ahead Position

wlook=wμvw_{\text{look}} = w - \mu v

Step 2: Update Velocity

vμv+ηg(wlook)v \leftarrow \mu v + \eta g(w_{\text{look}})

Step 3: Update Weights

wwvw \leftarrow w - v

Where: w = parameters, v = velocity, g = gradients, η = learning rate, μ = momentum factor

Function Arguments

  • w: np.ndarray - Current parameters (any shape)
  • v: np.ndarray - Current velocity (same shape as w)
  • grad: np.ndarray - Gradients at look-ahead position (same shape as w)
  • lr: float = 0.01 - Learning rate
  • momentum: float = 0.9 - Momentum factor
Loading visualization...

Examples

Input: w=[1.0, -1.0], v=[0.0, 0.0], grad=[0.5, -0.25], lr=0.1, momentum=0.9

Output: ([0.95, -0.975], [0.05, -0.025])

First step: v becomes lr*grad, then w updated by subtracting new v

Input: w=[1.0, 2.0], v=[0.5, -0.3], grad=[0.1, 0.2], lr=0.1, momentum=0.9

Output: ([0.54, 2.25], [0.46, -0.25])

Momentum carries forward: v = 0.9*v + 0.1*grad, then w -= v

Input: w=[2.0], v=[0.0], grad=[0.0], lr=0.1, momentum=0.9

Output: ([2.0], [0.0])

Zero gradient means no velocity or parameter update

Hint 1

Convert inputs to NumPy arrays first. Update velocity using momentum and learning rate.

Requirements

  • Return tuple (new_w, new_v) with same shapes as inputs
  • Use the exact update formulas above
  • Vectorized implementation only (no Python loops)
  • Handle any array shape (1D, 2D, etc.)

Note: In practice, you would compute the gradient at the look-ahead position w_look = w - μv. For this problem, the grad parameter represents the gradient already computed at that look-ahead position.

Constraints

  • Shapes of w, v, and grad must match
  • 0 ≤ momentum < 1, lr > 0
  • Libraries: NumPy only
Try Similar Problems