Edge of Stochastic Stability: Revisiting the Edge of Stability for SGD¶

会议: NeurIPS 2026
arXiv: 2412.20553
代码: https://github.com/arseniqum/edge-of-stochastic-stability
领域: 优化算法
关键词: 随机稳定性边缘, Batch Sharpness, SGD动力学, 渐进锐化, 弹射机制

一句话总结¶

揭示了小批量随机梯度下降（SGD）在均方误差神经网络训练中并不以全批次海森矩阵最大特征值 $\lambda_{\max}$ 为稳定性边界，而是由小批量锐度（Batch Sharpness，即小批量海森矩阵沿其随机梯度方向的期望瑞利商）在 $2/\eta$ 处饱和所主导的“随机稳定性边缘”（EoSS）机制，阐明了小批次选择平坦极小值的内生动力学根源。

研究背景与动机¶

深度学习优化中全批次梯度下降（GD）的稳定机制在过去几年取得了突破性进展。Cohen 等人发现使用恒定学习率 $\eta$ 的全批次 GD 会进入“稳定性边缘”（Edge of Stability, EoS）区域：损失函数全批次海森矩阵的最大特征值 $\lambda_{\max}$ 持续攀升并最终悬停在经典二次稳定边界 $2/\eta$ 附近振荡，尽管损失不再单调递减，网络仍能持续有效优化。然而在实际工业界和学术界的深度学习训练中，绝大多数任务都依赖小批量随机梯度下降（SGD），且研究人员长期观察到小批量 SGD 倾向于收敛到全批次曲率更平坦、泛化更优的极小值点，且其全批次 $\lambda_{\max}$ 往往显著低于 $2/\eta$。

现有理论将 EoS 的核心特征归结为损失函数的振荡，但这在 SGD 场景下直接失效。SGD 由于小批量梯度噪声的存在，几乎从训练伊始就伴随着无处不在的损失随机游走与非单调波动，导致“损失振荡”无法作为系统接近不稳定边界的诊断指标；同时，以全批次海森矩阵或者基于全批次梯度的指标来评估 SGD 稳定边界不仅缺乏紧确性，也无法定量解释为何学习率 $\eta$ 与批大小 $b$ 会系统性地改变解的平坦度。

本文切入的角度在于将视角从小样本噪声掩盖下的表观损失曲线，转向基于局部二次逼近的“单向不稳定性判定标准”（Certificates of Instability）：考察若冻结局部几何并在二次模型上运行算法，系统是否会发散并发生弹射（Catapult）。核心 idea：SGD 的稳定性不由全局全批次曲率决定，而是由小批量锐度（Batch Sharpness，即小批量海森矩阵沿对应随机梯度方向的期望瑞利商）在 $2/\eta$ 处的饱和所支配，该机制构成了随机稳定性边缘（EoSS），并以内生副产物的方式将全批次最大特征值抑制在远低于 $2/\eta$ 的水平。

方法详解¶

整体框架¶

本文构建了针对随机优化算法的单向不稳定性认证框架。SGD 每一步更新并非在单一固定的全局损失曲面上行进，而是接触独立采样的小批量损失子曲面 $L_B(\theta)$。每个更新步具有独立的梯度方向 $\nabla L_B(\theta)$ 与局部曲率 $\mathcal{H}(L_B)$。因此，决定动力学稳定性的核心统计量必须是小批量与方向感知的。

整个理论与分析框架遵循自底向上的递进流程：首先在局部二次逼近模型中确立不稳定性充要等价性（违反判定准则 $\Leftrightarrow$ 发生弹射逃离置信紧集 $\Leftrightarrow$ 出现任意大相对损失尖峰）；接着提出并证明 Batch Sharpness 作为一阶梯级不稳定性准则的数学有效性（Theorem 1）；随后从理论上解耦噪声驱动型振荡（Type-1，由梯度噪声互协方差主导）与曲率驱动型振荡（Type-2，由 EoSS 主导）；最后揭示全批次曲率 $\lambda_{\max}$ 在 EoSS 达成后的副产物演化与批大小缩放规律。

%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
    A["SGD 参数迭代点与小批量采样<br/>θ_t, B ~ P_b, η"] --> B["小批量方向曲率与 Batch Sharpness<br/>计算期望瑞利商 S_B(θ)"]
    B --> C["不稳定性认证与二次发散检验<br/>Batch Sharpness ≥ (2+ϵ)/η"]
    C -->|大于 2/η 触发乘性发散| D["梯度模长爆炸与动力学弹射<br/>E[||∇L_B(θ_t+1)||^2] 指数增长"]
    C -->|稳定在 2/η 边缘| E["EoSS 随机稳定性边缘自组织<br/>Batch Sharpness 饱和在 2/η"]
    D -->|损失剧增跳出局部二次域| F["重新触发渐进锐化与平坦域恢复"]
    E --> G["全批次曲率抑制与批大小平坦偏差<br/>λ_max < 2/η，小批次导致更大抑制"]

关键设计¶

1. 不稳定性判定准则与动力学弹射等价性：确立严格的单向发散认证机制 针对随机优化中传统下降引理和一致光滑性假设的失效，本文避开构造全空间紧确收敛条件的死胡同，转而定义单向不稳定性准则（Definition 1）：对于开放参数置信邻域 $U$，若标量统计量 $f(\theta_0) > c$，则 SGD 轨迹在有限时间内必将离开 $U$ 的任意紧子集。在参数点 $\theta_t$ 的局部二次逼近模型 $\widetilde{L}_i(\theta) = \frac{1}{2}(\theta - x_i)^\top \mathcal{H}_i (\theta - x_i)$ 下，本文证明了三大现象在数学上严格等价：破坏有效不稳定性判据、发生动力学弹射（Catapult，即脱离局部二次信任域）、以及在二次模型上观测到任意大相对振幅的损失尖峰（$\widetilde{L}(\theta_s) \geq \alpha \widetilde{L}(\theta_t)$）。这一等价性为经验诊断算法是否处于 EoSS 提供了可操作的超参数微扰检验法：若某一几何量在训练中处于紧确饱和状态，则调大步长 $\eta$ 或缩小批大小 $b$ 必将在局部诱发剧烈的弹射。

2. 小批量锐度（Batch Sharpness）：方向感知的小批量曲率统计量 为刻画 SGD 实际面对的离散随机曲面几何，本文定义小批量锐度（Batch Sharpness）为小批量海森矩阵沿其对应小批量梯度方向的期望瑞利商： $$ \text{Batch Sharpness}(\theta) := \mathbb{E}{B \sim \mathcal{P}b} \left[ \frac{\nabla L_B(\theta)^\top \mathcal{H}(L_B) \nabla L_B(\theta)}{|\nabla L_B(\theta)|^2} \right] $$ 该定义将期望置于分式外部，真实反映了优化器单步更新实际感受到的曲率分布。本文在定理 1（Theorem 1）中严格证明：当系统处于二次模型逼近域且存在 $\epsilon > 0$ 使得 $\text{Batch Sharpness}(\theta_t) \geq \frac{2+\epsilon}{\eta}$ 时，以单步小批量梯度模长二次型 $V_B(\theta_t) = \|\nabla L_B(\theta_t)\|^2$ 为李雅普诺夫函数，条件期望满足乘性发散： $$ \mathbb{E}\left[\frac{|\nabla L_B(\theta(\theta_t))^2 \geq (1+\epsilon)^2 $$ 该结论仅依赖海森半正定性质、柯西-施瓦茨不等式与简森不等式，表明一旦小批量锐度突破 })|^2}{|\nabla L_B(\theta_t)|^2} \,\middle|\, \theta_t\right] \geq (1 - \eta \cdot \text{Batch Sharpness$2/\eta$，小批量梯度模长将沿轨迹发生几何级数爆炸，迫使轨迹发生弹射。

3. 双重振荡解耦（Type-1 vs Type-2）：区分梯度噪声与曲率不稳定性 针对经典 SGD 损失曲线上无休止震荡的问题，本文将振荡分解为两类本质不同的动力学机制：Type-1 为稳态梯度噪声驱动型振荡，Type-2 为曲率驱动的 EoSS 不稳定振荡。通过二阶泰勒展开，本文定义了梯度-噪声相互作用指标（GNI）： $$ \text{GNI}(\theta) := \frac{\mathbb{E}_{B \sim \mathcal{P}_b}\left[\nabla L_B(\theta)^\top \mathcal{H}(L) \nabla L_B(\theta)\right]}{|\nabla L(\theta)|^2} $$ 在命题 1（Proposition 1）与命题 3 中，作者严格推导了强凸极小值附近稳态离散李雅普诺夫方程解的协方差结构：由于极小值附近全批次梯度均值衰减为零，而噪声方差常驻，系统在平坦区域也会迅速演化至 $\text{GNI} \approx 2/\eta [1 + \mathcal{O}(\eta)]$ 的平衡态。这意味着 GNI 达到 $2/\eta$ 仅仅是进入噪声游走（Type-1）的伴生现象，微扰超参数并不会引起发散或重塑曲率；而只有当系统经历渐进锐化并使 Batch Sharpness 触碰 $2/\eta$ 时，才会进入真正的曲率饱和与自稳定状态（Type-2）。

4. 全批次特征值压制与临界批大小：解释小批次平坦泛化偏差 在 EoSS 框架下，全批次海森矩阵的最大特征值 $\lambda_{\max}$ 的稳定不再是主导因素，而纯粹是 Batch Sharpness 饱和后的内生副产物。单样本海森矩阵由模型梯度外积构成，样本间的梯度方向存在固有未对齐性（余弦相似度通常仅约为 0.1）。当批大小 $b$ 较小时，小批量海森矩阵由少量样本叠加而成，其主特征向量与小批量梯度的对齐程度更高，导致小批量锐度与全局 $\lambda_{\max}$ 之间形成显著的正向几何间隙（Gap）。由于饱和机制强制将小批量锐度锁定在 $2/\eta$，反向迫使全局 $\lambda_{\max}$ 只能停留在远低于 $2/\eta$ 的平台。此外，曲率平台随批大小的依赖曲线存在明显的“临界批大小”（Critical Batch Size, $b_c$）：当 $b < b_c$ 时，$\lambda_{\max}$ 随批大小缩小急剧下降；当 $b > b_c$ 时曲线逐渐平缓，动力学平滑过渡至全批次 GD 机制。

损失函数 / 训练策略¶

论文主要在视觉分类基准上进行系统验证，采用均方误差损失函数（MSE）与恒定学习率 SGD（无动量）： $$ L(\theta) = \frac{1}{2N} \sum_{i=1}^N |f_\theta(x_i) - y_i|_2^2 $$ 训练策略覆盖不同网络架构（2层多层感知机 MLP、5层卷积网络 CNN、ResNet-14、ResNet-20），在 CIFAR-10（8k/32k/全量子集）及 SVHN 数据集上进行系统性消融。在超参数介入实验中，在训练不同阶段分别实施步长倍增/减半、批大小翻倍/减半操作，用于诊断动力学是否处于真正的稳定性边缘。

实验关键数据¶

主实验¶

论文全面测量了不同架构、不同步长与批大小下，Batch Sharpness 与全批次 $\lambda_{\max}$ 的终态稳定数值。下表汇总了在 CIFAR-10（8k子集）上训练收敛时各曲率统计量与理论边界 $2/\eta$ 的定量对比：

网络架构	步长 $\eta$	批大小 $b$	理论阈值 $2/\eta$	Batch Sharpness 终态	全批次 $\lambda_{\max}$ 终态	稳定性状态
MLP (512-512)	0.010	16	200.0	198.4 ± 3.2	82.5 ± 4.1	严格处于 EoSS ($\lambda_{\max}$ 被抑制)
MLP (512-512)	0.020	32	100.0	99.1 ± 1.8	64.2 ± 3.0	严格处于 EoSS ($\lambda_{\max}$ 被抑制)
MLP (512-512)	0.030	64	66.7	66.2 ± 1.5	51.8 ± 2.4	严格处于 EoSS ($\lambda_{\max}$ 被抑制)
5-Layer CNN	0.030	8	66.7	67.1 ± 2.0	28.4 ± 1.9	严格处于 EoSS (小批次强平坦抑制)
5-Layer CNN	0.030	32	66.7	66.5 ± 1.6	42.1 ± 2.2	严格处于 EoSS (曲率随批大小上升)
5-Layer CNN	0.050	16	40.0	40.3 ± 1.1	24.6 ± 1.5	严格处于 EoSS
ResNet-14	0.020	32	100.0	98.7 ± 2.3	58.9 ± 3.5	严格处于 EoSS
ResNet-20	0.020	16	100.0	100.8 ± 2.7	45.3 ± 2.8	严格处于 EoSS

消融实验¶

为严格论证“SGD 并非简单的全局 GD 叠加梯度高斯噪声”，论文在批大小 16、步长 $\eta=0.01$（$2/\eta=200$）的条件下对比了真实 SGD 与多种噪声注入代理及 SDE 连续时间模拟的动力学终态曲率表现：

动力学模型	模拟机制配置	Batch Sharpness 终态	全批次 $\lambda_{\max}$ 终态	是否重现 EoSS 抑制现象	说明
真实 mini-batch SGD	真实样本重采样 ($b=16$)	~200.0	~85.0	是	原始基准，$\lambda_{\max}$ 远低于 $2/\eta$
各向异性高斯重采样 GD	匹配单样本海森一阶与二阶矩	~200.0	~86.2	是	唯一成功保留海森矩阵高阶波动信息的噪声模型
对角噪声注入 (GLD diag)	仅注入 SGD 经验协方差对角高斯	~128.0	~199.5	否	丧失海森扰动，曲率退化至全批次 EoS 边界
各向同性噪声 (GLD dynamic)	注入标量匹配平均方差高斯噪声	~115.0	~201.2	否	忽略参数相关性，$\lambda_{\max}$ 直接撞上 $2/\eta$
连续时间欧拉-马山 SDE	步长 0.0005 离散化维纳过程	~110.0	~200.8	否	忽略小批量离散几何，无法解释平坦极小值

关键发现¶

Batch Sharpness 的绝对紧确性：跨越不同神经网络架构（MLP、CNN、ResNet）、学习率及数据集（CIFAR-10、SVHN），Batch Sharpness 均展现出极为一致的“渐进锐化 $\rightarrow$ 饱和在 $2/\eta$ 边界”的动力学过程，与批大小完全解耦。
超参数扰动的因果检验：在 Batch Sharpness 尚未达到 $2/\eta$ 之前加大步长或减小批大小，模型不发生发散；而在其饱和后施加扰动，会立即激发出数十倍的弹射尖峰（Catapult），随后由于自稳定机制降落到更平坦的曲率盆地重新开始渐进锐化。
连续扩散与 SDE 模型的局限：单纯在梯度上叠加高斯噪声或使用经典朗之万/SDE 动力学建模，无法保留小批量海森矩阵的离散统计涨落，其稳态全部收敛至全批次 EoS 边界（$\lambda_{\max} \approx 2/\eta$），揭示了 SDE 难以准确预测 SGD 隐式偏差的深层物理原因。

亮点与洞察¶

将 EoS 理论推广至现代小批量实际场景：首次回答了 Cohen 等人在 2021 年提出的核心遗留问题，给出了量化 SGD 动力学不稳定相变点的显式曲率指标，将离散非凸优化理论推进到无需依靠全批次假定。
给出小批次选择平坦极小值的机理解释：建立了“小批次 $\rightarrow$ 样本梯度与单批次海森矩阵对齐性提高 $\rightarrow$ Batch Sharpness 与全局 $\lambda_{\max}$ 差距拉大 $\rightarrow$ 饱和在 $2/\eta$ 倒逼全局 $\lambda_{\max}$ 保持低位”的完整因果链条，终结了关于小批次为何产生平坦解的经验争论。
指出经典下降引理在深度学习分析中的系统性失效：证明了神经网络训练中 SGD 并非在平滑函数的经典收敛区工作，而是常态化地处于局部二次发散与非单调自组织的平衡边缘。

局限与展望¶

任务与损失函数范围局限：目前全部实证实验均基于均方误差损失（MSE）及图像分类任务，交叉熵（Cross-Entropy）损失因其在无限远处平坦流形上的发散漂移特性，尚未给出系统的 EoSS 定量刻画。
未覆盖自适应与动量优化器：实验聚焦于恒定步长无动量 SGD，对于工业界广泛采用的 AdamW、动量 SGD 等自适应算法，其方向感知小批量稳定准则的形式仍有待拓展。
自稳定（Self-Stabilization）显式解析机制留白：本文证实了二次模型下的不稳定性判定，但神经网络高阶项（如三阶导数与梯度的对齐）如何具体在随机设定下产生耗散并拉回发散轨迹，仍缺乏闭式解析理论。

评分¶

新颖性: ⭐⭐⭐⭐⭐ 成功将 EoS 理论推进到小批量 SGD，提出了极具洞察力的 Batch Sharpness 统计量。
实验充分度: ⭐⭐⭐⭐⭐ 涵盖了跨架构、跨数据集、高精度海森谱跟踪、超参数因果干预与多种噪声代理消融。
写作质量: ⭐⭐⭐⭐⭐ 数学证明严谨，概念界定（Type-1 vs Type-2、单向不稳定性认证）清晰，论证链条完整。
价值: ⭐⭐⭐⭐⭐ 为深度学习优化理论、泛化隐式正则化以及离散时间优化算法提供了里程碑式的理论认知。

动力学模型	模拟机制配置	Batch Sharpness 终态	全批次 \(\lambda_{\max}\) 终态	是否重现 EoSS 抑制现象	说明
真实 mini-batch SGD	真实样本重采样 (\(b=16\))	~200.0	~85.0	是	原始基准，\(\lambda_{\max}\) 远低于 \(2/\eta\)
各向异性高斯重采样 GD	匹配单样本海森一阶与二阶矩	~200.0	~86.2	是	唯一成功保留海森矩阵高阶波动信息的噪声模型
对角噪声注入 (GLD diag)	仅注入 SGD 经验协方差对角高斯	~128.0	~199.5	否	丧失海森扰动，曲率退化至全批次 EoS 边界
各向同性噪声 (GLD dynamic)	注入标量匹配平均方差高斯噪声	~115.0	~201.2	否	忽略参数相关性，\(\lambda_{\max}\) 直接撞上 \(2/\eta\)
连续时间欧拉-马山 SDE	步长 0.0005 离散化维纳过程	~110.0	~200.8	否	忽略小批量离散几何，无法解释平坦极小值