Nesterov Momentum

Minimise a function with Nesterov accelerated gradient and return the final point.

Classical momentum computes the gradient at the current point, then moves along the velocity. Nesterov momentum looks ahead first: it computes the gradient at the point the velocity is about to carry it to. This lets the optimiser slow down before it overshoots.

Start with x = x0 and a zero velocity, then repeat steps times:

lookahead = x - lr * beta * v
v         = beta * v + grad_fn(lookahead)
x         = x - lr * v

Function signature

def nesterov_minimize(grad_fn, x0, lr, beta, steps):

Example

Input

grad_fn = lambda x: 2 * x        # gradient of f(x) = x ** 2
nesterov_minimize(grad_fn, np.array([1.0]), lr=0.1, beta=0.9, steps=2)

Output

array([0.496])

Reasoning

Step 1: the velocity is zero, so the lookahead is 1.0. The gradient there is 2.0, the velocity becomes 2.0, and x = 1.0 - 0.2 = 0.8.

Step 2: the lookahead is 0.8 - 0.1 * 0.9 * 2.0 = 0.62. The gradient there is 1.24, the velocity becomes 0.9 * 2.0 + 1.24 = 3.04, and x = 0.8 - 0.304 = 0.496. Classical momentum would use the gradient at 0.8 instead and give 0.46.

Constraints and notes

• x0 is a 1D array. grad_fn takes and returns an array of the same shape.
• Call grad_fn exactly once per step, at the lookahead point.
• Do not modify x0.