Nesterov Accelerated gradient is in correct order?

Question

I've implemented Nesterov Accelerated Gradient (NAG) (link, section: "Nesterov Accelerated Gradient"), however not sure it's viable.

$$v_t = \gamma v_{t-1} +\eta\nabla_\theta J(\theta - \gamma v_{t-1})$$ $$\theta = \theta - v_t$$

I attempted the following (consider a minibatch consisting of merely 1 training case):

1. run a forward pass, and get the error
2. Apply the scaled previous momentum, (momentum_tMinOne * coeff) to the weights
3. run a backward pass using the error of each pass.
4. set momentum_tMinOne = currentGradient, but don't apply it to the weights yet
5. repeat forward pass, from 1.

Do I require to remember something like momentum_tMinTwo, or just keeping hold of the previous momentum will suffice?

Am I correct that this way we don't need to keep two matrices for momentums, but only need momentum_tMinOne matrix?

Answer 1

The standard momentum would have these steps:

1. straight away, recompute the new momentum: $$\mu_{t+1} = \mu_{t} \cdot (decayScalar) + (learnRate)\cdot \nabla$$

2. adjust the weights by this new momentum $$\theta_{t+1} := \theta_{t} - \mu_{t+1}$$

---

Nesterov momentum has this:

1. Make a big jump: correct the weights by any $\mu$ we have so far in our posession: $$\theta_{t+1} := \theta_{t} - \mu\cdot (decayScalar)$$

2. Compute the gradient $\nabla$ from the new weights $\theta_{t+1}$

3. Correct these weights by this gradient (right now without any momentum): $$\theta_{t+2} := \theta_{t+1} - (learnRate)\cdot \nabla$$

4. At the very end, re-compute the momentum as follows: $$\mu := \theta_{t+2} - \theta_{t}$$

So, the momentum is updated at the very end. It becomes a vector from "weights before the big jump", pointing towards the "weights after the correction by the fresh gradient".

reference: Geoffrey Hinton Lecture 6C Corsera

---

Re-arranging:

To avoid sticking the gradient computation in the middle of our optimizer function (steps 2 and 3), we can instead re-arrange things as follows:

1. compute gradient for the weights we have thus far.

2. correct such weights by the gradient (right now without any momentum), as follows:

   $$\theta_{t+1}:=\theta_t - (learnRate)\cdot \nabla$$

3. update the momentum: $$\mu := \theta_{t+1} - \theta_{cached}$$ $$\theta_{cached} := \theta_{t+1}$$

4. big jump $$\theta_{t+2} := \theta_{t+1} - \mu\cdot (decayScalar)$$

Notice, this way steps 2,3,4 are all inside of our optimizer. We can compute the gradient, outside of our optimizer (during step 1) making our code much more readable :)

    size_t _numApplyCalled = 0;
      
    //Nesterov Accelerated Gradient.
    //Placed at the end of a backprop, should be followed by a usual forward-propagation
    // https://datascience.stackexchange.com/a/26395/43077
    //
    void apply( float *toChange,  float *newGrad,  size_t count ){

        float learnRate = get(OptimizerVar::LEARN_RATE);//scalar
        float momentumCoeff = get(OptimizerVar::MOMENTUM_1);//scalar

        const bool isFirstEver_apply = _numApplyCalled == 0;

        for (int i=0; i<_arraySize; ++i){
                //correction by gradient alone:
                toChange[i]  -= newGrad[i]*learnRate;

                // determine momentum:
                if (isFirstEver_apply){//nothing was cached yet.
                    _momentumVals[i] = 0.0f;
                }
                else {
                    _momentumVals[i] = toChange[i] - _weightsCached[i];
                }
                //caching, AFTER momentum calc, but BEFORE the jump:
                _weightsCached[i] = toChange[i];
            
                //jump:
                toChange[i] -= _momentumVals[i] * momentumCoeff;
        }//end for

        ++_numApplyCalled;//increments by 1
    }