2024阿里巴巴数学竞赛决赛(应用与计算赛道第2题)

今天写一下第二题我的个人解答，这道题源自NIPS2022的一篇论文《The alignment property of SGD noise and how it helps select flat minima: A stability analysis》，问题的背景是与梯度下降(GD)相比，随机梯度下降(SGD)在过参数化的神经网络训练过程当中可以收敛到一个更加平缓(flat)的极值点，而且平缓程度与神经网络的参数量无关。

T2. 假设 $ F(x; w) $ 是一个输出标量的深度神经网络，其中 $ x $ 是输入，$ w $ 表示权重。假设 $ F $ 关于 $ w $ 连续可微，并且对于训练数据 $\{ x_j, y_j \}_{j=1}^n $ 过参数化，即存在 $ w^\ast $ 使得对所有 $ j $ 满足 $ F(x_j; w^\ast) = y_j $。为了研究训练神经网络时 $ w^* $ 的局部优化行为，我们考虑线性化神经网络 $ \widetilde{F}(x; w) = F(x; w^\ast) + (w - w^\ast)^T \nabla F(x; w^\ast) $，其损失函数为

$$\text{Loss}(w) := \frac{1}{2m} \sum_{j=1}^m (y_j - \widetilde{F}(x_j; w))^2$$

令 $ s $ 表示学习率，梯度下降法为

$$w_{i+1} = w_i - s \nabla \text{Loss}(w_i)$$

而随机梯度下降法为

$$w_{i+1} = w_i - s (\nabla \text{Loss}(w_i) + \epsilon_i)$$

其中噪声项满足 $ \mathbb{E} \epsilon_i = 0 $ 和 $ \mathbb{E} \epsilon_i \epsilon_i^\top = M(w_i)/b $，$ b $ 是 mini-batch 的大小。假设协方差矩阵 $ \Sigma $ 为

$$\Sigma = \frac{1}{m} \sum_{j=1}^m \nabla F(x_j; w^\ast) \nabla F(x_j; w^\ast)^T$$

在以下意义上对齐：

$$\frac{\text{Tr}(M(w) \Sigma)}{2 \text{Loss}(w) \| \Sigma \|_F^2} \geq \delta$$

对于 $ s > 0 $ 和所有 $ w $ 成立。这里 $| \cdot |_F$ 表示 Frobenius 范数。

(1) 对于梯度下降，证明如果 $ \Sigma $ 的谱范数满足

$$\| \Sigma \|_2 \leq \frac{2}{s}$$

则梯度下降是局部稳定的(即对于所有 $ i $，$\text{Loss}(w_i)$ 是有界的)。(注意，这蕴含了一个依赖维度的 $\Vert \Sigma \Vert_F \leq \frac{2\sqrt{d}}{s} $，其中 $ d $ 是 $w$ 的维度)。

(2) 对于随机梯度下降，如果 $\mathbb{E}\text{Loss}(w_i)$ 对于所有 $ i $ 都有界，则以下独立于维度的不等式必须成立

$$\| \Sigma \|_F \leq \frac{\sqrt{b \delta}}{s}$$

(1) $Proof.$ 由过参数化的条件有， $F(x_j; w^\ast) = y_j$，那么

$$\begin{aligned} Loss(w_i)&=\frac{1}{2m} \sum_{j=1}^m(F(x_j; w^\ast) + (w_i - w^\ast)^\top \nabla F(x_j; w^\ast)-y_j)^2 \\\\ &= \frac{1}{2m} \sum_{j=1}^m((w_i - w^\ast)^\top \nabla F(x_j; w^\ast))^2 \end{aligned}$$

对 $Loss(w_i)$ 求梯度可得

$$\begin{aligned} \nabla Loss(w_i)&=\frac{1}{m} \sum_{j=1}^m\nabla F(x_j; w^\ast)\nabla F(x_j; w^\ast)^\top (w_i-w^\ast) \\\\ &=\Sigma (w_i-w^\ast) \\\\ \end{aligned}$$

由梯度下降法的更新格式，有

$$\begin{aligned} w_{i+1}-w^\ast&=w_i - s\nabla Loss(w_i)-w^\ast \\\\ &=(I-s\Sigma)(w_i-w^\ast) \\\\ \end{aligned}$$

对上式两边同时取 $2$ 范数，因为 $I-s\Sigma$ 为对称矩阵，因此一定存在矩阵 $P$ 使得 $I-s\Sigma$ 可以正交对角化，故

$$\begin{aligned} ||I-s\Sigma||_2&=||P(I-s\Sigma)P^\top||_2 \\\\ &=||I-s\Lambda||_2 \\\\ &=\rho (I-s\Lambda) \end{aligned}$$

由 $\Sigma$ 为半正定矩阵，以及 $\rho (\Sigma) \leq \frac{2}{s}$，可得 $I-s\Lambda$ 对角线上元素范围为 $[-1,1]$，即 $\rho (I-s\Lambda) \leq 1$

所以

$$\begin{aligned} ||w_{i+1}-w^\ast||_2\leq ||w_{i}-w^\ast||_2 \leq ... \leq ||w_{0}-w^\ast||_2 \end{aligned}$$

由

$$\begin{aligned} Loss(w_i) = \frac{1}{2} (w_i - w^\ast)^\top \Sigma (w_i - w^\ast) \end{aligned}$$

可得

$$\begin{aligned} |Loss(w_i)| &\leq \frac{1}{2} ||(w_i - w^\ast)||_2^2 ||\Sigma||_F \\\\ &\leq \frac{\sqrt{d}||(w_i - w^\ast)||_2^2}{s} \end{aligned}$$

故梯度下降法是局部稳定的.

第二小问的证明过程参考论文 《The alignment property of SGD noise and how it helps select flat minima: A stability analysis》中的 Theorem 3.3 的证明
(2) $Proof.$ 记 $\theta_i = w_i - w^\ast$，那么有

$$\begin{aligned} Loss(w_{i}) &=\frac{1}{2} \theta_i^\top \Sigma \theta_i \\\\ \nabla Loss(w_{i}) &=\Sigma \theta_i \end{aligned}$$

计算$\mathbb{E}Loss(w_{i+1})$，并由 $\mathbb{E}\epsilon_i = 0$ 和 $\mathbb{E} \epsilon_i \epsilon_i^\top = M(w_i)/b$，有

$$\begin{aligned} \mathbb{E}Loss(w_{i+1}) &=\mathbb{E}\left[\frac{1}{2} (w_{i+1}-w^\ast)^\top \Sigma(w_{i+1}-w^\ast)\right]\\\\ &= \mathbb{E}\left[\frac{1}{2} (\theta_i - s (\nabla \text{Loss}(w_i) + \epsilon_i))^\top \Sigma (\theta_i - s (\nabla \text{Loss}(w_i) + \epsilon_i))\right] \\\\ &= \mathbb{E}\left[\frac{1}{2} \theta_i^\top \Sigma \theta_i -s\theta_i^\top \Sigma^2 \theta_i - s\theta_i^\top \Sigma \epsilon_i + \frac{1}{2} s^2 (\Sigma \theta_i + \epsilon_i)^\top \Sigma (\Sigma \theta_i + \epsilon_i) \right] \\\\ &= \mathbb{E}\left[\frac{1}{2} \theta_i^\top \Sigma \theta_i -s\theta_i^\top \Sigma^2 \theta_i + \frac{1}{2} s^2 (\Sigma \theta_i + \epsilon_i)^\top \Sigma (\Sigma \theta_i + \epsilon_i) \right] \\\\ &= \mathbb{E}\left[\frac{1}{2} \theta_i^\top \Sigma \theta_i -s\theta_i^\top \Sigma^2 \theta_i \right]+\frac{1}{2} s^2 \mathbb{E}\left[(\Sigma \theta_i + \epsilon_i)^\top \Sigma (\Sigma \theta_i + \epsilon_i)\right] \\\\ &= \mathbb{E}\left[\frac{1}{2} \theta_i^\top \Sigma \theta_i -s\theta_i^\top \Sigma^2 \theta_i \right]+\frac{1}{2} s^2 \mathbb{E}\left[\theta_i^\top \Sigma^3 \theta_i + \epsilon_i^\top \Sigma \epsilon_i + \epsilon_i^\top \Sigma^2 \theta_i \right] \\\\ &= \mathbb{E}\left[\frac{1}{2} \theta_i^\top \Sigma \theta_i -s\theta_i^\top \Sigma^2 \theta_i \right]+\frac{1}{2} s^2 \mathbb{E}\left[\theta_i^\top \Sigma^3 \theta_i \right] + \frac{1}{2} s^2 \mathbb{E}\left[\epsilon_i^\top \Sigma \epsilon_i \right] \\\\ \end{aligned}$$

考虑期望 $ \mathbb{E}[x^T A x] $，其中$\mathbb{E}x = 0$ 和 $\mathbb{E} x x^\top = M(w_i)/b$

$$\begin{aligned} \mathbb{E}[x^T A x] &= \mathbb{E}\left[\sum_{i=1}^n \sum_{j=1}^n x_i A_{ij} x_j\right] \\\\ &= \sum_{i=1}^n \sum_{j=1}^n A_{ij} \mathbb{E}[x_i x_j] \end{aligned}$$

对于 $ i \neq j $，因为 $ x $ 的均值为 0， $ E[x_i x_j] = 0 $(这里假设各个分量是相互独立的)，故

$$\mathbb{E}[x^T A x] = \sum_{i=1}^n A_{ii} \mathbb{E}[x_i^2] = \text{Tr}(A \mathbb{E}x x^\top)$$

即

$$\mathbb{E}\left[\epsilon_i^\top \Sigma \epsilon_i \right] = \text{Tr}(\Sigma \mathbb{E}\epsilon_i \epsilon_i^\top) = \frac{1}{b} \text{Tr}(M(w_i) \Sigma )$$

记 $r(\theta) = 1 - 2s \frac{\theta^T \Sigma^2 \theta}{\theta^T \Sigma \theta} + s^2 \frac{\theta^T \Sigma^3 \theta}{\theta^T \Sigma \theta}$，那么

$$\mathbb{E}Loss(w_{i+1}) = \mathbb{E}\left[r(\theta_i)Loss(w_i) \right] + \frac{s^2}{2b} \text{Tr}(M(w_i) \Sigma )$$

下证 $r(\theta) \geq 0$，令

$$u = \frac{\Sigma^{1/2} \theta}{ || \Sigma^{1/2} \theta||_2 }$$

并对 $\Sigma$ 进行谱分解

$$\Sigma = \sum_j \lambda_j e_j e_j^T$$

假设 $u$ 在 ${e_i}$ 这组标准正交基下的表示如下

$$u = \sum_j a_j e_j$$

且由 $||u||_2 = 1$ 得 $\sum_j a_j^2 = 1$，那么

$$\begin{aligned} r(\theta) &= 1-2s u^\top \Sigma u + s^2 u^T \Sigma^2 u \\\\ &= 1 + \sum_j (-2s \lambda_j - s^2 \lambda_j^2) a_j^2 \\\\ &\geq 1 + \inf_{\lambda \geq 0} (-2s \lambda - s^2 \lambda^2)=0 \end{aligned}$$

由 $r(\theta_i)Loss(w_i) \geq 0$ 和 $\frac{\text{Tr}(M(w) \Sigma)}{2 \text{Loss}(w) | \Sigma |_F^2} \geq \delta$，有

$$\begin{aligned} \mathbb{E}Loss(w_{i+1}) &= \mathbb{E}\left[r(\theta_i)Loss(w_i) \right] + \frac{s^2}{2b} \text{Tr}(M(w_i) \Sigma ) \\\\ & \geq \frac{s^2}{2b} \text{Tr}(M(w_i) \Sigma ) \\\\ & \geq \frac{s^2}{b} \delta Loss(w_i) ||\Sigma||_F^2 \end{aligned}$$

因为 $\mathbb{E}Loss(w_{i+1})$ 是有界的，所以

$$\frac{s^2}{b} \delta ||\Sigma||_F^2 \leq 1$$

否则，令 $q = \frac{s^2}{b} \delta ||\Sigma||_F^2$

$$\begin{aligned} \mathbb{E}Loss(w_{i+1}) & \geq \frac{s^2}{b} \delta Loss(w_i) ||\Sigma||_F^2 \\\\ &= q^i Loss(w_1) \end{aligned}$$

$\mathbb{E}Loss(w_{i+1})$ 呈指数增长，故 $| \Sigma |_F \leq \frac{\sqrt{b \delta}}{s}$，证毕.