为什么 Nesterov 动量有效：谱分析视角
考虑带 Polyak 和 Nesterov 动量的 Lion-$\mathcal{K}$ 算法： \(\begin{align} M_{t+1} &= \beta_{\text{p}} M_t - (1 - \beta_{\text{p}}) \nabla F(X_t),\\ \tilde{M}_{t+1} &= \beta_{\text{n}} M_{t+1} - (1 - \beta_{\text{n}}) \nabla F(X_t),\\ X_{t+1} &= X_t + \eta(\nabla \mathcal{K}(\tilde{M}_{t+1}) - \lambda X_t), \end{align}\) 其中 $\eta > 0$ 是学习率，$\beta_{\text{n}}, \beta_{\text{p}} \in [0,1]$ 分别是 Nesterov 和 Polyak 动量系数，$\lambda \geq 0$ 是权重衰减系数，$\mathcal{K}: \mathbb{X} \rightarrow \mathbb{R}$ 是一个凸函数，$\nabla \mathcal{K}$ 表示其梯度（如果不可微则为次梯度）。
为简单起见，我们分析确定性情况，不考虑权重衰减，并设 $\nabla \mathcal{K}$ 为恒等映射：$\nabla \mathcal{K} = I$ 和 $\lambda = 0$。这简化为对 $F$ 做带 Nesterov 动量的梯度下降。
当选择 $\nabla \mathcal{K}=\text{msign}$ 时，这就恢复为带 Nesterov 动量的 Muon 优化器。需要指出的是，在强凸情形下，目前没有工作证明带 Polyak 动量的梯度下降可以实现全局加速收敛，而 Nesterov 动量则有全局加速收敛保证。
经过代数处理，我们得到： \(X_{t+1} = X_t + \beta_{\text{p}} \underbrace{(X_t - X_{t-1})}_{\text{动量}} - \eta \beta_{\text{p}} (1 - \beta_{\text{n}}) \underbrace{(\nabla F(X_t) - \nabla F(X_{t-1}))}_{\text{梯度校正}} - \eta (1 - \beta_{\text{p}}) \nabla F(X_t).\) 推导过程详见附录。这个递归式用 $X_t$、$X_{t-1}$、$\nabla F(X_t)$ 和 $\nabla F(X_{t-1})$ 来表示 $X_{t+1}$，没有显式的 $M$ 或 $\tilde{M}$ 变量。适合通过谱方法分析收敛性。
该递归式有两个关键组成部分：动量项 $X_t - X_{t-1}$ 和梯度校正项 $\nabla F(X_t) - \nabla F(X_{t-1})$。
考虑如下近似：
\[\frac{\nabla F(X_t) - \nabla F(X_{t-1})}{\eta} \approx \nabla^2 F(X_t) \dot{X}_t,\]这说明校正项捕捉了沿方向 $\dot{X}_t$ 的二阶曲率信息 $\nabla^2 F(X_t)$。这提供了类似海森矩阵的信息，而无需额外计算。
设置 $\beta_{\text{n}} = 1$ 会消除校正项（因为 $1 - \beta_{\text{n}} = 0$），退化为纯 Polyak 动量法。
Nesterov 的另一种形式，即隐式速度格式，在一个前瞻点评估梯度：$\nabla F(X_t + \beta (X_t - X_{t-1})/\eta)$。对其泰勒展开可得：
\[\nabla F(X_t + \beta (X_t - X_{t-1})/\eta) \approx \nabla F(X_t) + \beta \nabla^2 F(X_t) \dot{X}_t.\]实际上也是在估计方向曲率信息。如这篇论文所示，两种形式在收敛速度上是等价的。
我们将对二次目标函数 $F(x) = \frac{1}{2} x^\top A x$ 进行谱分析，其中 $A \succeq 0$ 是对称正定矩阵，其特征值在 $[\mu, L]$ 区间内（$0 < \mu \leq \cdots \leq L$，条件数 $\kappa = L/\mu$）。最小值点为 $x_\star = 0$ 且 $F(x_\star) = 0$，梯度为 $\nabla F(x) = A x$。
将递归式简写为：
\[X_{t+1} = X_t + \alpha (X_t - X_{t-1}) - \beta (\nabla F(X_t) - \nabla F(X_{t-1})) - \gamma \nabla F(X_t),\]其中 $\alpha = \beta_{\text{p}}$，$\beta = \eta \beta_{\text{p}} (1 - \beta_{\text{n}})$，$\gamma = \eta (1 - \beta_{\text{p}})$。
引入辅助变量 $V_t = X_t - X_{t-1}$ 构成一个线性系统：
\(\begin{pmatrix} I & \mathbf{0}_{n \times n} \\ -\alpha I + \beta \nabla F & I \end{pmatrix} \begin{pmatrix} X_{t+1} \\ V_{t+1} \end{pmatrix} = \begin{pmatrix} I - \gamma \nabla F & I \\ -\alpha I + \beta \nabla F & \mathbf{0}_{n \times n} \end{pmatrix} \begin{pmatrix} X_t \\ V_t \end{pmatrix}.\) 对于二次函数情况，$\nabla F = A$，所以： \(\begin{pmatrix} I & \mathbf{0}_{n \times n} \\ -\alpha I + \beta A & I \end{pmatrix} \begin{pmatrix} X_{t+1} \\ V_{t+1} \end{pmatrix} = \begin{pmatrix} I - \gamma A & I \\ -\alpha I + \beta A & \mathbf{0}_{n \times n} \end{pmatrix} \begin{pmatrix} X_t \\ V_t \end{pmatrix}.\)
对左侧矩阵求逆，得到迭代矩阵：
\[\begin{pmatrix} X_{t+1} \\ V_{t+1} \end{pmatrix} = \begin{pmatrix} I & \mathbf{0}_{n \times n} \\ \alpha I - \beta A & I \end{pmatrix} \begin{pmatrix} I - \gamma A & I \\ -\alpha I + \beta A & \mathbf{0}_{n \times n} \end{pmatrix} \begin{pmatrix} X_t \\ V_t \end{pmatrix} = \begin{pmatrix} I - \gamma A & I \\ \gamma A (\beta A - \alpha I) & \alpha I - \beta A \end{pmatrix} \begin{pmatrix} X_t \\ V_t \end{pmatrix}. \tag{$\dagger$}\]由于 $A = Q \Lambda Q^\top$（$Q$ 是正交矩阵，$\Lambda$ 是对角矩阵），于是
\[\begin{pmatrix} Q^\top X_{t+1} \\ Q^\top V_{t+1} \end{pmatrix} = \underbrace{ \begin{pmatrix} I - \gamma \Lambda & I \\ \gamma \Lambda (\beta \Lambda - \alpha I) & \alpha I - \beta \Lambda \end{pmatrix} }_{M(\alpha, \beta, \gamma)} \begin{pmatrix} Q^\top X_t \\ Q^\top V_t \end{pmatrix}. \tag{$\dagger$}\]谱半径 $\rho(\alpha, \beta, \gamma) = |M(\alpha, \beta, \gamma)|_2$ 控制着收敛速率：$\rho < 1$ 越小，收敛越快。$M$ 的特征值由 $2 \times 2$ 的块 $G(\lambda)$ 的特征值构成。总的谱半径是这些块谱半径的最大值。这个 $2 \times 2$ 的块定义为（即通过对称的行列变换把 $M$ 重组为块对角矩阵，这个过程不改变特征值）：
\[G(\lambda) = \begin{pmatrix} 1 - \gamma \lambda & 1 \\ \gamma (\beta \lambda - \alpha) \lambda & \alpha - \beta \lambda \end{pmatrix}.\]其特征多项式为：
\[\det(rI - G(\lambda)) = r^2 - [1 + \alpha - (\gamma + \beta) \lambda] r + (\alpha - \beta \lambda) = 0,\]判别式为：
\[\Delta(\lambda) = [1 + \alpha - (\beta + \gamma) \lambda]^2 - 4 (\alpha - \beta \lambda).\]$\Delta(\lambda)$（作为 $\lambda$ 的二次函数）的首项系数为正，因此最大值在端点 $\lambda = \mu$ 或 $L$ 处取得。当我们选择 $0\leq \alpha,\beta\leq 1$ 和 $\gamma=1/L$ 时，通常有 $\Delta(L) < 0$。
设置 $\beta_{\text{p}} = 0$ 和 $\beta_{\text{n}} = 1$，于是 $\alpha = \beta = 0$ 且 $\gamma = \eta$。其谱半径为：
\[\rho(0, 0, \eta) = \|I - \gamma A\|_2 = \max\{|1 - \gamma \mu|, |1 - \gamma L|\}.\]选取标准步长 $\eta = 1/L$ ，则 $\rho = 1 - \mu/L=1-1/\kappa$。
设置 $\beta_{\text{n}} = 1$，于是 $\beta = 0$, $\alpha = \beta_{\text{p}}$, $\gamma = \eta (1 - \beta_{\text{p}})=1/L$。矩阵简化为：
\[G_{\text{HB}} = \begin{pmatrix} I - \gamma A & I \\ -\alpha \gamma A & \alpha I \end{pmatrix}.\]在特征基中，每个 $2 \times 2$ 的块是：
\[G_{\text{HB}}(\lambda) = \begin{pmatrix} 1 - \gamma \lambda & 1 \\ -\alpha \gamma \lambda & \alpha \end{pmatrix}, \quad \lambda \in [\mu, L].\]谱半径为：
\[\rho_{\text{HB}} = \max_{\lambda \in [\mu, L]} \rho(G_{\text{HB}}(\lambda)).\]特征值是特征多项式的解：
\[\det\begin{pmatrix} 1 - \gamma \lambda - r & 1 \\ -\alpha \gamma \lambda & \alpha - r \end{pmatrix} = r^2 - (1 + \alpha - \gamma \lambda) r + \alpha = 0.\]如果对所有 $\lambda$，$\Delta(\lambda)=(1 + \alpha - \gamma \lambda)^2 - 4\alpha \leq 0$ 均成立，这对应于 $\Delta(\mu)\leq 0\Leftrightarrow 1 - \sqrt{\mu/L} \leq \sqrt{\alpha}$，那么特征值为复数，其模为 $\sqrt{\alpha}$。这说明该方法收敛速度的最快可达 \(\|G_{\text{HB}}(\lambda)\|=\sqrt{\alpha}\approx1 - 1/\sqrt{\kappa}，\) 这比梯度下降的 $1 - 1/\kappa$ 要好。
现在引入 Nesterov 动量：$\alpha = \beta_{\text{p}}$, $\beta = \eta \beta_{\text{p}} (1 - \beta_{\text{n}})$, $\gamma = \eta (1 - \beta_{\text{p}})$。确保 $\Delta(\mu) \leq 0$ 需要：
\[1 - \sqrt{\frac{\mu}{L}} \leq \sqrt{\alpha - \beta \mu}.\]如果这对所有 $\lambda \in [\mu, L]$ 都成立，那么在整个区间上 $\Delta(\lambda) \leq 0$。此时，$G(\lambda)$ 的特征值是模为 $\sqrt{\alpha - \beta \lambda}$ 的复数。因此，在这个方向的收敛速度最快可达：
\[\|G(\lambda)\|=\sqrt{\alpha - \beta \lambda}\leq \sqrt{\alpha - \beta \mu}\leq 1 - \sqrt{\frac{\mu}{L}},\]当 $\lambda>\mu$ 且 $\sqrt{\alpha - \beta \mu}\approx 1 - \sqrt{\mu/L}=1-1/\sqrt{\kappa}$ 时，此方向的收敛速度比重球法统一的 $\sqrt{\alpha} \approx 1 - 1/\sqrt{\kappa}$ 更快。
这种谱分析的视角表明，Polyak 动量将谱半径从 $1 - 1/\kappa$（梯度下降）统一减小到 $1 - 1/\sqrt{\kappa}$。Nesterov 更进一步，通过梯度校正使得不同方向的收敛速度依赖于 $\lambda$，且在 $\lambda > \mu$ 时的理论最优比 Polyak 动量更好。这解释了 Nesterov 在理论上的优势，尤其是在处理条件数较大的病态问题时，其感知二阶曲率的能力更强，而每次迭代的成本只比标准动量多一次矩阵加法。同时，也启示了我们，参数选择的重要性，如果不能保证都是复特征值，则动量法的收敛速度未必会比梯度法更好。
让我们逐步简化系统，以消除辅助变量 $M_t$ 和 $\tilde{M}_t$，从而得到一个仅关于 $X_t$ 的递归式。
第一步：应用简化。 当 $\nabla \mathcal{K} = I$ 和 $\lambda = 0$ 时，$X_{t+1}$ 的更新变为：
\[X_{t+1} = X_t + \eta \tilde{M}_{t+1}.\]第二步：用 $M_{t+1}$ 表示 $\tilde{M}_{t+1}$。 从第二个方程：
\[\tilde{M}_{t+1} = \beta_{\text{n}} M_{t+1} - (1 - \beta_{\text{n}}) \nabla F(X_t).\]第三步：代入 $X_{t+1}$ 方程。 我们得到：
\[X_{t+1} = X_t + \eta \left[ \beta_{\text{n}} M_{t+1} - (1 - \beta_{\text{n}}) \nabla F(X_t) \right].\]解出 $M_{t+1}$：
\[M_{t+1} = \frac{1}{\beta_{\text{n}} \eta} (X_{t+1} - X_t) + \frac{1 - \beta_{\text{n}}}{\beta_{\text{n}}} \nabla F(X_t).\]第四步：类似地，从上一步表示 $M_t$。 同理可得：
\[M_t = \frac{1}{\beta_{\text{n}} \eta} (X_t - X_{t-1}) + \frac{1 - \beta_{\text{n}}}{\beta_{\text{n}}} \nabla F(X_{t-1}).\]第五步：代入第一个方程。 代入 $M_{t+1}$ 和 $M_t$ 的表达式：
\[\frac{1}{\beta_{\text{n}} \eta} (X_{t+1} - X_t) + \frac{1 - \beta_{\text{n}}}{\beta_{\text{n}}} \nabla F(X_t) = \beta_{\text{p}} \left[ \frac{1}{\beta_{\text{n}} \eta} (X_t - X_{t-1}) + \frac{1 - \beta_{\text{n}}}{\beta_{\text{n}}} \nabla F(X_{t-1}) \right] - (1 - \beta_{\text{p}}) \nabla F(X_t).\]第六步：重新整理，得到只含 $X$ 的递归式。 经过代数处理，我们得到：
\[\begin{align} X_{t+1} &= X_t + \beta_{\text{p}} (X_t - X_{t-1}) - \eta \left[ (1 - \beta_{\text{p}}) \beta_{\text{n}} + (1 - \beta_{\text{n}}) \right] \nabla F(X_t) + \eta \beta_{\text{p}} (1 - \beta_{\text{n}}) \nabla F(X_{t-1}) \\ &= X_t + \beta_{\text{p}} \underbrace{(X_t - X_{t-1})}_{\text{动量}} - \eta \beta_{\text{p}} (1 - \beta_{\text{n}}) \underbrace{(\nabla F(X_t) - \nabla F(X_{t-1}))}_{\text{梯度校正}} - \eta (1 - \beta_{\text{p}}) \nabla F(X_t). \end{align}\]