小批量噪声通过主导子空间波动降低锐度

# 小批量噪声通过主导子空间波动降低尖锐度
来源：https://arxiv.org/html/2607.23012
\\hldauthor\\Name

Junho So\\Emailjhso@ajou\.ac\.kr \\NameDongwook Shin\\Emaildws@ajou\.ac\.kr \\addr韩国亚洲大学数学系

###### 摘要

在 SGD 训练过程中，梯度通常与损失函数 Hessian 矩阵的前 \(k\) 个特征向量所张成的**主导子空间**高度对齐。虽然这自然意味着损失减少主要发生在这个空间内，但先前的研究表明，在此主导子空间内的更新在减少损失方面并没有取得有意义的进展。在这项工作中，我们认为，主导子空间更应被理解为解释小批量 SGD 尖锐度动态的关键子空间，而非损失减少的主要空间。为了解释主导子空间在降低前 \(k\) 尖锐度中的作用，我们展示了在主导方向上对梯度波动进行平均如何产生一个尖锐度校正项，并推导出由小批量噪声在主导方向上诱导的尖锐度校正项。实验结果表明，将推导出的校正项添加到 GD 中，会使 GD 的尖锐度演化更接近 SGD。

## 1 引言

理解深度神经网络的训练动态是机器学习的主要课题之一。在随机梯度下降（SGD）训练中，人们观察到损失函数的 Hessian 矩阵通常具有少量较大的离群特征值 (sagun2017eigenvalues; DBLP:journals/corr/SagunEGDB17; pmlr-v97-ghorbani19b; pmlr-v97-papyan19a; JMLR:v21:20-933)，并且训练梯度已知与由前 \(k\) 个 Hessian 特征向量张成的**主导子空间**高度对齐 (gurari2018gradientdescenthappenstiny; pmlr-v97-ghorbani19b)。这表明 SGD 的有效训练动态可能是低维的，并且自然导致人们期望损失减少主要发生在这个空间内。

然而，先前的工作 (song2025does) 报道了与这种解读相矛盾的结果。具体来说，当将 SGD 更新投影到主导子空间上时，训练无法实现有意义的损失减少。相反，当将更新投影到**体子空间**（即主导子空间的正交补空间）上时，损失减少与标准 SGD 类似。那么，主导子空间对训练毫无贡献吗？

在本文中，我们认为主导子空间的主要作用在于降低损失景观的前 \(k\) 尖锐度。具体来说，我们观察到，尽管主导投影更新对损失减少影响甚微，但它显著降低了尖锐度。然后，通过涉及沿不同子空间（主导和随机）的受控扰动的实验，我们表明只有主导方向上的扰动会降低尖锐度。为了解释这种效应，我们首先展示对主导方向上的梯度波动进行平均会产生一个确定性的尖锐度校正项。然后，我们推导出由主导方向上的小批量噪声诱导的确定性校正项。最后，我们凭经验证明，将该校正项添加到 GD 会产生与 SGD 相似的尖锐度动态。

## 2 设置与动机

参见说明

参见说明

参见说明

图 1：主导分量不减少损失，但降低尖锐度。(a) 训练梯度与主导子空间对齐。(b) Dom-SGD 无法减少损失，而 Bulk-SGD 继续训练。(c) Dom-SGD 降低 \(S_k\)，而 Bulk-SGD 则保持比 SGD 更高的 \(S_k\)。

#### 设置。

设 \(L: \mathbb{R}^d \to \mathbb{R}\) 为损失函数。对于一个小批量 \(B\) 及其小批量损失 \(L_B\)，将小批量噪声定义为 \(\xi_B(\theta) := \nabla L_B(\theta) - \nabla L(\theta)\)，其中 \(\mathbb{E}_B[\xi_B(\theta) \mid \theta] = 0\)。我们将小批量 SGD 写为 \(\theta_{t+1} = \theta_t - \eta (\nabla L(\theta_t) + \xi_{B_t}(\theta_t))\)，将全批量 GD 写为 \(\theta_{t+1} = \theta_t - \eta \nabla L(\theta_t)\)。设 \((\lambda_i(\theta), e_i(\theta))\) 为 \(H(\theta) = \nabla^2 L(\theta)\) 的特征对，按 \(\lambda_1(\theta) \ge \cdots \ge \lambda_d(\theta)\) 排序。我们将**主导子空间**定义为由前 \(k\) 个特征向量张成的空间，\(E_{\mathrm{dom}}(\theta) = \mathrm{span}\{e_1(\theta), \ldots, e_k(\theta)\}\)，并将其正交补空间 \(E_{\mathrm{bulk}}(\theta) = E_{\mathrm{dom}}(\theta)^\perp\) 定义为**体子空间**。我们还定义到这些空间的投影为主导投影 \(P_{\mathrm{dom}}(\theta)\) 和体投影 \(P_{\mathrm{bulk}}(\theta) = I - P_{\mathrm{dom}}(\theta)\)。将前 \(k\) 尖锐度定义为 \(S_k(\theta) = \sum_{j=1}^k \lambda_j(\theta)\)。我们将关注点限制在稳定学习率区间，即在训练过程中保持 \(\lambda_1(\theta) < 2 / \eta\)。

### 2.1 动机与主要观察

让我们定义主导对齐度量 \(\chi_k(\theta) = \|P_{\mathrm{dom}}(\theta) \nabla L(\theta)\| / \|\nabla L(\theta)\|\)，它量化了沿 SGD 轨迹位于主导子空间中的训练损失梯度的相对大小。如图 1 (a) 所示，从训练早期开始，SGD 训练梯度就高度集中在主导子空间中，并且在后期阶段这种趋势变得更强 (gurari2018gradientdescenthappenstiny; pmlr-v97-ghorbani19b)。这一观察表明 SGD 动态与主导子空间密切相关。

参见说明

图 2：河谷直觉。

然而，如图 1 (b) 所示，将 SGD 更新投影到主导子空间上的 Dom-SGD 并不会显著降低训练损失。相反，仅保留体分量的 Bulk-SGD 与 SGD 一样有效地降低了损失 (song2025does)。先前的工作 (song2025does; wen2025understanding) 使用**河谷直觉**解释了这一现象。在这种观点下，如图 2 所示，主导方向对应于高曲率的谷壁，其上的更新不会降低损失，而体方向则沿着谷底，损失减少主要发生在那里。

尽管如此，先前的工作报告称，仅使用 Bulk-SGD 训练的模型无法完全恢复 SGD 所取得的最终准确率提升 (zakarin2025acceleratingneuralnetworktraining)。这表明，尽管主导子空间不主要负责损失减少，但它仍然可能在训练中发挥有意义的作用。其贡献可能在于优化动态的某些方面，而这些方面不能仅由损失减少直接体现。

在这项工作中，我们认为主导子空间在降低损失景观的前 \(k\) 尖锐度方面起着重要作用。如图 1 (c) 所示，Bulk-SGD 的损失曲线几乎与 SGD 相同，但其前 \(k\) 尖锐度变得比 SGD 更高。相比之下，Dom-SGD 在保持高损失的同时，显著降低了前 \(k\) 尖锐度，这表明主导子空间与尖锐度降低相关，而非与损失降低相关。在下一节中，我们将进行实验，考察沿主导方向的运动是否确实是降低前 \(k\) 尖锐度的因素。

## 3 主导子空间波动降低尖锐度

参见说明

参见说明

参见说明

图 3：主导扰动降低前 \(k\) 尖锐度。我们比较了 GD、带有主导零均值扰动的 GD（GD+Dom）以及带有随机零均值扰动的 GD（GD+Rand）。两个扰动运行在秩和噪声尺度上匹配（\(k=10, \rho=0.1\)）。(a) 主导扰动诱导强主导对齐。(b) 只有主导扰动降低 \(S_k\)。(c) 两个扰动运行几乎不改变 GD 损失曲线。所有实验均在一个 MLP 上运行。

先前的观察表明，主导子空间与尖锐度降低相关，而非与损失降低相关。然而，Dom-SGD 和 Bulk-SGD 之间的简单比较并不能揭示这种尖锐度降低是源于主导子空间中的**梯度分量**，还是源于该子空间内的**随机运动**。

为了分离这些效应，我们进行了一项受控的噪声注入实验，该实验保持平均更新等于全批量 GD，同时仅在不同子空间添加零均值扰动。具体来说，我们比较了受扰动的 GD 动态 \(\theta_{t+1} = \theta_t - \eta (\nabla L(\theta_t) + \zeta_t)\)，其中 \(\zeta_t\) 从不同子空间中采样。对于主导扰动，我们设置 \(\zeta_t \sim \mathcal{N}(0, \rho^2 P_{\mathrm{dom}}(\theta_t))\)，对于随机扰动，我们设置 \(\zeta_t \sim \mathcal{N}(0, \rho^2 P_{\mathrm{rand}})\)，其中 \(P_{\mathrm{rand}}\) 是在初始化时固定的一个随机 \(k\) 维正交投影。

首先，如图 3 (c) 所示，GD、带有随机扰动的 GD（GD+Rand）和带有主导扰动的 GD（GD+Dom）遵循几乎相同的训练损失曲线，表明这些扰动不会显著改变损失降低的速率。然而，如图 3 (a) 和 3 (b) 所示，它们的轨迹显示出显著的几何差异：GD 和 GD+Rand 保持较低的 \(\chi_k\) 和较高的 \(S_k\)，而 GD+Dom 则将 \(\chi_k\) 驱动到接近 1，并保持比 GD 或 GD+Rand 更低的前 \(k\) 尖锐度 \(S_k\)。

这些结果展示了两点。首先，正如先前的观察所表明的，与主导子空间对齐并不意味着损失减少沿着主导方向发生。其次，更重要的是，仅靠主导子空间内的随机运动实际上就能影响尖锐度动态，同时几乎不改变损失曲线。综合来看，这些结果表明，主导子空间的作用不在于提供损失减少的方向，而在于改变与主导方向相关的**局部几何**，尤其是前 \(k\) 尖锐度。在下一节中，我们将展示这种效应源于零均值位移的协方差与 Hessian 的局部变化之间的相互作用。

## 4 尖锐度校正项

在本节中，我们表明，当参数在主导方向围绕参考点波动时，对梯度进行平均会留下一个**尖锐度校正项**。为了说明这一点，考虑一个来自参考点 \(\theta_c\) 的位移 \(\delta \in E_{\mathrm{dom}}(\theta_c)\)，且 \(\mathbb{E}[\delta] = 0\)，并将邻近参数写为 \(\theta_c + \delta\)。使用泰勒展开在 \(\theta_c\) 附近展开 \(\nabla L(\theta_c + \delta)\)，我们得到：

\[
\nabla L(\theta_c + \delta) = \nabla L(\theta_c) + H(\theta_c)\delta + \frac{1}{2}\nabla(\delta^\top H \delta)(\theta_c) + O(\|\delta\|^3). \tag{1}
\]

式 (1) 中的第一项是参考点处的梯度，第二项是位移的线性项，第三项是二阶项，捕捉 Hessian 的局部变化。对 \(\delta\) 进行平均得到：

\[
\mathbb{E}_\delta[\nabla L(\theta_c + \delta)] = \nabla L(\theta_c) + \underbrace{H(\theta_c)\mathbb{E}[\delta]}_{=0} + \frac{1}{2}\nabla\operatorname{Tr}\!\left(H(\theta_c)C\right) + O(\mathbb{E}\|\delta\|^3), \qquad C := \mathbb{E}[\delta\delta^\top]. \tag{2}
\]

这里，\(C\) 在求导时视为固定。如式 (2) 所示，平均梯度是参考点处的梯度与附加项 \(\frac{1}{2}\nabla\operatorname{Tr}\!\left(H(\theta_c)C\right)\) 的和。该项由位移协方差 \(C\) 决定，其作用是降低由 \(\operatorname{Tr}(HC)\) 表示的协方差加权曲率。当 \(C\) 支撑在主导子空间内时，它会降低与前 \(k\) 尖锐度相关的加权曲率，因此我们将其称为**尖锐度校正项**。

#### 小批量诱导的位移协方差。

最后，我们推导出由局部递归中小批量噪声诱导的位移协方差，并将其代入尖锐度校正项。这表明小批量噪声的尖锐度效应可以表示为一个确定性的校正项。

为了定义小批量噪声在参考点 \(\theta_c\) 周围产生的局部位移，我们比较了小批量 SGD 更新与全批量 GD 更新。考虑一个局部状态 \(\theta_c + \delta_s\)，其中 \(\delta_s \in E_{\mathrm{dom}}(\theta_c)\)，并设 \(P_c := P_{\mathrm{dom}}(\theta_c)\)。下一个位移 \(\delta_{s+1}\) 定义为

\[
\delta_{s+1} := P_c(\theta_{\mathrm{mb}}^+ - \theta_{\mathrm{gd}}^+), \qquad \theta_{\mathrm{mb}}^+ := \theta_c + \delta_s - \eta \nabla L_{B_s}(\theta_c + \delta_s), \qquad \theta_{\mathrm{gd}}^+ := \theta_c - \eta \nabla L(\theta_c).
\]

然后，我们在 \(\theta_c\) 附近线性化在 \(\theta_c + \delta_s\) 处的小批量梯度：

\[
\nabla L_{B_s}(\theta_c + \delta_s) \approx \nabla L(\theta_c) + H_c \delta_s + \xi_{B_s}(\theta_c), \qquad H_c := H(\theta_c).
\]

将其代入 \(\delta_{s+1}\) 的定义中，得到

\[
\delta_{s+1} = P_c \bigl( \delta_s - \eta H_c \delta_s - \eta \xi_{B_s}(\theta_c) \bigr) = A_c \delta_s - \eta P_c \xi_{B_s}(\theta_c), \qquad A_c := P_c (I - \eta H_c) P_c. \tag{3}
\]

对式 (3) 取二阶矩得到

\[
C_{s+1} = A_c C_s A_c^\top + \eta^2 \Sigma_{\mathrm{dom}}(\theta_c), \quad (C_s := \mathbb{E}[\delta_s \delta_s^\top]), \quad (\Sigma_{\mathrm{dom}}(\theta_c) := P_c \operatorname{Cov}(\xi_B(\theta_c)) P_c).
\]

这里，\(\Sigma_{\mathrm{dom}}\) 表示投影到主导子空间上的小批量噪声协方差。