Nesterov Momentum 工程实现上的trick

Nesterov Momentum 工程实现上的trick

Nesterov Momentum是momentum这种优化方法的一个变种，其参数更新规则这样的：

$$v \leftarrow \alpha v + grad(\theta + \alpha v) \\ \theta \leftarrow \theta - lr \cdot v$$

参数更新规则这样写有一个问题。一般情况下，（以tensorflow为例）optimizationMethod所接受的参数只有计算好的$grad(\theta)$和$\theta$, 那么我们怎么计算$grad(\theta + \alpha v)$的值呢？

查看一下tensorflow的代码，ApplyMomentum 的实现非常简单：

template <typename T>
struct ApplyMomentum<CPUDevice, T> {
  void operator()(const CPUDevice& d, typename TTypes<T>::Flat var,
                  typename TTypes<T>::Flat accum,
                  typename TTypes<T>::ConstScalar lr,
                  typename TTypes<T>::ConstFlat grad,
                  typename TTypes<T>::ConstScalar momentum, bool use_nesterov) {
    accum.device(d) = accum * momentum() + grad;
    if (use_nesterov) {
      var.device(d) -= grad * lr() + accum * momentum() * lr();
    } else {
      var.device(d) -= accum * lr();
    }
  }
};

tensorflow_source_code

但是看起来跟标准的update rule 完全不一样啊。

解释这个问题还是得借用Hinton老爷子的slides。

对于标准的update rule, 一个iteration走过的路径是$0 \rightarrow 1 \rightarrow 2$：$0 \rightarrow 1$ 对应$\alpha v$, $1 \rightarrow 2$ 对应$grad(\theta + \alpha v)$. 但是我们只知道0点的gradient， 不知道1这点的gradient。

既然我们只需要1这点的gradient，那么就把传进来的gradient当成1点的gradient好了。这样，每个iteration的路径就变成了$1 \rightarrow 2 \rightarrow 3$。OK, 假设一个iteration结束了，这时的weight和gradient都是在1点的weight和gradient，而$v$则是0点的$v$。那么下一个iteration，$1 \rightarrow 2$ 就是$grad(\theta)$。要计算$2 \rightarrow 3$, 我们需要知道2点的$v$。怎么算呢？$v_{2} \leftarrow \alpha v_{0} + grad(\theta_{1})$。然后$2 \rightarrow 3$ 就变成了 $\alpha v_{2}$。

因此，我们的update rule变成了：

$$v \leftarrow \alpha v + grad(\theta) \\ \theta \leftarrow \theta - lr \cdot (grad(\theta) + \alpha v)$$

这就和tensorflow的实现一致了。

通过上面的分析，我们知道，如果按照新的update rule计算出的当前的$\theta$和$v$的位置是正确的话（和老的update rule得到结果是一样的），那么下一步新老update rule就是等价的。那么根据数学归纳法，我们只要证明，两种update rule在第一个iteration是等价的，就可以确认两种update rule是等价的了。

在第一个iteration时$v = 0$，因此$0 \rightarrow 1$是一个0向量，那么此时的weight $\theta$既可以理解成0点weight也可以理解成1点的weight。因此两种update rule在初始条件下也是等价。所以，tensorflow的实现是对的。

注：本文参考了这个问题的回答