Nesterov

前段时间被人问到了 Nesterov，我以为我很了解，结果发现我自己也讲不明白 Nesterov 的 “Ilya version” 是怎么来的。遂花了点时间想了一下，感觉纯是闹麻了。下面就让我来批判一下。

Nesterov 的假设

Nesterov 的思想很简单，我用 $\theta + \eta M$ 处的梯度来更新 $\theta$，保证了下次更新方向是一个考虑了动量的综合，避免因为动量“冲过头”。

在实践上，我们不必真的在 $\theta + \eta M$ 处计算梯度，有一个 Ilya 提出的“近似”版本：

标准的 nesterov 更新过程：

$$\left\{ \begin{aligned} \hat{M}_{t} &= \beta_1 M_{t-1} + (1-\beta_1) G(\theta_t + \eta M_{t-1}) \\ V_{t} &= \beta_2 V_{t-1} + (1-\beta_2) G^2(\theta_t)\\ \theta_{t+1} &= \theta_t + \eta \frac{\hat{M}_{t}}{\sqrt{V_{t}+\epsilon}} \end{aligned} \right.$$

其中，二阶矩还是得用的 $\theta_t$ 这一点的？直观上考虑到它作为正则化的作用。

通过公式变换：

$$\begin{aligned} \theta_{t+1} &= \theta_t + \eta \frac{M_{t+1}}{\sqrt{V_{t+1}+\epsilon}} = \theta_t + \eta \frac{\beta_1 M_t + (1-\beta_1) G(\theta_t + \eta M_t)}{\sqrt{V_{t+1}+\epsilon}} \\ &= \theta_t + \eta \frac{\beta_1 M_t + (1-\beta_1)G(\theta_t)} {\sqrt{V_{t+1}+\epsilon}} + \eta \frac{1-\beta_1}{\sqrt{V_{t+1}+\epsilon}} [G(\theta_t + \eta M_t) - G(\theta_t)] \end{aligned}$$

前面两项就是标准的梯度下降。这里我们需要计算两次梯度，我们都知道，这是不可能的。

因此实际的 Nesterov 过程采用的是一种近似，更新公式是：

$$\left\{ \begin{aligned} M_{t} &= \beta_1 M_{t-1} + (1-\beta_1) G(\theta_t) \\ \hat{M}_{t} &= \beta_1 M_{t} + (1-\beta_1) G(\theta_t)\\ V_{t} &= \beta_2 V_{t-1} + (1-\beta_2) G^2(\theta_t)\\ \theta_{t+1} &= \theta_t + \eta \frac{\hat{M}_{t}}{\sqrt{V_{t}+\epsilon}} \end{aligned} \right.$$

也就是说，在动量上应用了两次梯度的滑动平均。这有道理吗？我们可以考虑它相比标准 Adam 的增量：

$$\begin{aligned} \Delta &= [(\beta_1 M_t +(1-\beta_1)G(\theta_t)) - M_t] \\ &= (1-\beta_1)(G(\theta_t)-M_t) \end{aligned}$$

对比一下理想公式的增量，我们就能知道这里的假设是

$$[G(\theta_t + \eta M_t) - G(\theta_t)] \approx G(\theta_t)-M_t$$

这个合理吗？其实很不直观，只能尝试强行解释。对左式泰勒：

$$\eta H(\theta_t)M_t \approx G(\theta_t)-M_t$$

先不管有的没的可不可逆，强行推：

$$M_t \approx (\eta H(\theta_t)+\beta_1 I)^{-1} G(\theta_t)$$

而我们的更新量就是：

$$\eta M_t \approx (H(\theta_t) + \frac{\beta_1}{\eta}I)^{-1} G(\theta_t)$$

这是一项阻尼强度为 $\frac{\beta_1}{\eta}$ 的 hessian 曲率矫正。所以这个假设的核心是，动量可以视为一种对真实梯度的白化。

首先第一点，从理论的角度来说，这个假设有点强过头了，它不仅要求方向，实际上连大小也限制了。看起来，它直接要求了动量就是牛顿法更新量，这个我只能说哈哈了。

其次第二点，从实践的角度来说，一般 $H(\theta_t)$ 的元素的均值可能是在 $1e^{-5}$ 左右？相比于 $\beta_1 = 0.9, \eta = 1e^{-6}$ 这种水平的设置，阻尼太大了！！！ 可以说是完全抹平了曲率的影响。

所以只有当上述这个很强的假设成立时，近似才成立。

实践

我没仔细看 NAdam 之类的论文，在我自己心里反正是很难评了。我在一些别的比较 solid 的文章里，比如说 《Fantastic Pretraining Optimizers and Where to Find Them》 也见过一些 NAdam 的结果，好像还行。我自己试过在 muon 上加上 Nesterov（用的是 Dion 的代码），可以说是毫无作用。

一个简单的修正方法是，再加一个超参 $k$ 上去，令 $M_t \approx k \beta_1 (\eta H(\theta_t)+\beta_1 I)^{-1} G(\theta_t)$，这样起码解放了大小的限制。回推上去，等价于 $\hat{M}_{t} = \beta_1 M_{t} + (1-\beta_1 + k\beta_1) G(\theta_t)$。

这里仍然有很多问题：

• $k$ 是否是一个可以设置的超参，在整个训练过程中近似不变？

• $k$ 作为一个 scala 是否合理？还是说它必须是一个矩阵，哪怕是一个对角阵？

如果把 $k$ 当作一个scala的话，那其实等价于调节 $\beta_1$，也就是 $\frac{\beta_1}{1-\beta_1} = \frac{\beta_1’}{1-\beta_1’ + k \beta_1’}$罢了。可以通过观测 $k$，提出一个 Nesterov 解耦的 $\beta_{11} , \beta_{12}$，其实怎么说都行，但感觉非常无聊。

不太明白 Ilya 当时怎么想的。也许读者有什么更深入的见解，欢迎来信和我分享，不吝赐教。