针对边缘稳定性下 Adam 优化器的杆流模型

# 位于稳定性边缘的 Adam 的杆流模型

来源：https://arxiv.org/html/2605.06821

###### 摘要

Cohen 等人 \[2022 (https://arxiv.org/html/2605.06821#bib.bib10)\] 观察到，诸如 Adam 等自适应梯度方法在**稳定性边缘**（edge of stability, EoS）运行。虽然针对稳定性边缘处梯度下降的连续时间建模已有大量工作，但将这些模型扩展到动量方法方面仍显不足。在梯度下降场景中，Regis 和 Chewi \[2026 (https://arxiv.org/html/2605.06821#bib.bib1)\] 引入了**杆流**（rod flow），该方法将连续迭代视为一个扩展的一维对象——即“杆”。在这里，我们通过在工作于参数和第一矩 $(w,m)$ 的联合相空间，并将第二矩 $\nu$ 视为平滑辅助变量的方式，将杆流扩展到 Adam。我们还为重球动量（heavy ball momentum）、Nesterov 动量以及 RMSProp、Adam 和 NAdam 的标量版本和分量版本开发了杆流模型。对于所有八种优化器，我们在代表性的机器学习架构上对杆流进行了实证评估，结果表明，在稳定性边缘区域内，杆流追踪离散迭代的准确性显著高于相应的稳定流。

## 1 引言

神经网络通过基于梯度的优化器最小化损失函数来进行训练。Cohen 等人 \[2021 (https://arxiv.org/html/2605.06821#bib.bib2)\] 观察到，全批量梯度下降在**稳定性边缘**（EoS）运行：Hessian 矩阵的最大特征值，称为**尖锐度**（sharpness），首先上升（称为**渐进尖锐化**阶段），然后悬浮在稳定性阈值 $2/\eta$ 附近，其中 $\eta$ 是学习率。Cohen 等人 \[2022 (https://arxiv.org/html/2605.06821#bib.bib10)\] 将这一图景扩展到动量方法和自适应梯度方法，表明每种优化器都表现出其自身的稳定性边缘。相关量——**预条件尖锐度**（preconditioned sharpness）——并非悬浮在 $2/\eta$，而是悬浮在依赖于优化器的超参数相关阈值处（表 2 (https://arxiv.org/html/2605.06821#A5.T2)）。

在实践中，机器学习中占主导地位的优化器是 Adam \[Kingma and Ba, 2015 (https://arxiv.org/html/2605.06821#bib.bib26)\]，它与梯度下降有两个不同之处。首先，它是一种动量方法：它不是用当前点的梯度更新参数，而是维护过去梯度的指数移动平均 $m$，并基于此动量进行更新。其次，它是一种自适应方法：在更新参数时，Adam 将动量的每个分量除以 $\sqrt{\nu}$，其中 $\nu$ 是梯度平方的指数移动平均。其效果是用每个参数梯度的典型幅度来归一化其更新。

连续时间模型长期以来作为分析离散时间优化器的宝贵理论工具 \[Liet al., 2017 (https://arxiv.org/html/2605.06821#bib.bib20), Barrett and Dherin, 2021 (https://arxiv.org/html/2605.06821#bib.bib8), Shiet al., 2021 (https://arxiv.org/html/2605.06821#bib.bib18)\]，阐明了它们的隐式偏差和渐近行为。对于现代深度学习的基石 Adam 来说，这样的模型尤其有价值。然而，开发一个在稳定性边缘仍然有效的此类模型并非易事。在 EoS 处，离散迭代沿 Hessian 的最尖锐方向振荡。但是，朴素的连续时间极限——称为“稳定流”——不允许存在振荡。虽然关于稳定性边缘处梯度下降的连续时间理论丰富且不断增长 \[Aroraet al., 2022 (https://arxiv.org/html/2605.06821#bib.bib5), Damianet al., 2023 (https://arxiv.org/html/2605.06821#bib.bib4), Cohenet al., 2025 (https://arxiv.org/html/2605.06821#bib.bib3), Regis and Chewi, 2026 (https://arxiv.org/html/2605.06821#bib.bib1)\]，但针对 Adam 及其密切相关方法（重球动量、Nesterov 动量、NAdam）的相应理论仍发展不足。现有的 Adam 连续时间分析要么采用构造上消除振荡动力学的消失步长极限 \[Belotto da Silva and Gazeau, 2020 (https://arxiv.org/html/2605.06821#bib.bib39), Barakat and Bianchi, 2021 (https://arxiv.org/html/2605.06821#bib.bib38)\]，要么研究未能触及振荡 EoS 区域的有限步长效应 \[Malladiet al., 2022 (https://arxiv.org/html/2605.06821#bib.bib33), Cattaneoet al., 2024 (https://arxiv.org/html/2605.06821#bib.bib31)\]。实践者实际使用的区域——在稳定性边缘运行的有限步长 Adam——目前缺乏有原则的连续时间模型。

### 1.1 我们的方法

最近，Regis 和 Chewi \[2026 (https://arxiv.org/html/2605.06821#bib.bib1)\] 引入了**杆流**：一种稳定性边缘处梯度下降的连续时间模型。杆流追踪平滑演变的中心 $\bar{w}$（即连续迭代的平均值）和范围张量 $\Sigma$（即连续迭代之间半位移的外积）。物理图景很直观：与其跟随沿着 Hessian 尖锐方向来回振荡的点，不如跟随一个扩展的一维对象——“杆”，其长度和方向编码了振荡信息。通过扩展杆流框架以适用于动量和自适应梯度方法，我们开发了稳定性边缘处 Adam 的连续时间模型。

将杆流扩展到 Adam 需要两项修改。首先，杆必须从权重空间 $w \in \mathbb{R}^d$ 提升到位相空间 $z=(w,m) \in \mathbb{R}^{2d}$。对于在稳定性边缘运行的动量方法，第一矩 $m$ 与 $w$ 同步振荡，因此杆的中心和半差必须在两个坐标上定义。此扩展处理了重球和 Nesterov 动量。其次，预条件器必须作为平滑辅助变量进行追踪。RMSProp 和 Adam 中的第二矩 $\nu$ 是*平方*梯度的指数移动平均；平方操作消除了符号翻转，因此即使 $w$ 在稳定性边缘振荡，$\nu$ 也会平滑变化。因此，它自身不需要杆结构。因此，Adam 被建模为由 $\nu$ 定义的自适应度量下的相空间杆流。

### 1.2 贡献

- 重球和 Nesterov 动量的杆流推导，将原始的权重空间模型扩展到相空间（第 4 节 (https://arxiv.org/html/2605.06821#S4)）。
- 预条件器的平滑辅助处理，得出了标量和分量 RMSProp 的杆流（第 5 节 (https://arxiv.org/html/2605.06821#S5)）。
- Adam 和 NAdam 的杆流，结合了相空间和预条件器扩展（第 6 节 (https://arxiv.org/html/2605.06821#S6)）。
- 对 Adam 的实证评估，表明杆流准确追踪离散迭代，并显示出预条件尖锐度在理论预测阈值处的自稳定化（第 7 节 (https://arxiv.org/html/2605.06821#S7)）。

## 2 相关工作

**稳定性边缘。** 对于全批量梯度下降，Cohen 等人 \[2021 (https://arxiv.org/html/2605.06821#bib.bib2)\] 记录了跨架构的渐进尖锐化和稳定性边缘，尖锐度上升至 $2/\eta$ 然后悬浮。理论解释包括自稳定化 \[Damianet al., 2023 (https://arxiv.org/html/2605.06821#bib.bib4)\]、两步分析 \[Ahnet al., 2023 (https://arxiv.org/html/2605.06821#bib.bib6), Chen and Bruna, 2023 (https://arxiv.org/html/2605.06821#bib.bib9)\]、不稳定区域的收敛性 \[Ahnet al., 2022 (https://arxiv.org/html/2605.06821#bib.bib15), Aroraet al., 2022 (https://arxiv.org/html/2605.06821#bib.bib5), Zhuet al., 2023 (https://arxiv.org/html/2605.06821#bib.bib14)\]、二次回归中的渐进尖锐化 \[Agarwalaet al., 2023 (https://arxiv.org/html/2605.06821#bib.bib7)\]、弹射机制 \[Lewkowyczet al., 2020 (https://arxiv.org/html/2605.06821#bib.bib11), Ghoshet al., 2023 (https://arxiv.org/html/2605.06821#bib.bib34)\] 及其动量诱导放大 \[Phunyaphibarnet al., 2024 (https://arxiv.org/html/2605.06821#bib.bib47)\]、损失景观分析 \[Maet al., 2022 (https://arxiv.org/html/2605.06821#bib.bib21)\]、分岔理论 \[Song and Yun, 2023 (https://arxiv.org/html/2605.06821#bib.bib19)\]、混沌 GD 动力学 \[Kong and Tao, 2020 (https://arxiv.org/html/2605.06821#bib.bib36), Chenet al., 2024 (https://arxiv.org/html/2605.06821#bib.bib22)\]、正则性诱导隐式偏差 \[Wanget al., 2025 (https://arxiv.org/html/2605.06821#bib.bib30)\] 以及 EoS 处的 NTK 演化 \[Jianget al., 2025 (https://arxiv.org/html/2605.06821#bib.bib23)\]。Cohen 等人 \[2022 (https://arxiv.org/html/2605.06821#bib.bib10)\] 实证识别了自适应梯度方法的稳定性边缘。

**Adam。** 原始 Adam 方法由 Kingma and Ba \[2015 (https://arxiv.org/html/2605.06821#bib.bib26)\] 提出；变体包括 NAdam \[Dozat, 2016 (https://arxiv.org/html/2605.06821#bib.bib27)\]、AMSGrad \[Reddiet al., 2018 (https://arxiv.org/html/2605.06821#bib.bib28)\] 和 AdamW \[Loshchilov and Hutter, 2019 (https://arxiv.org/html/2605.06821#bib.bib29)\]。我们以其标准 EMA 形式分析 Adam 和 NAdam。Adam 的连续时间分析均在小子步或极小值附近区域运行，而非贯穿 EoS：基础 ODE 极限 \[Barakat and Bianchi, 2021 (https://arxiv.org/html/2605.06821#bib.bib38), Dereichet al., 2025 (https://arxiv.org/html/2605.06821#bib.bib42)\]、通用自适应 ODE 框架 \[Belotto da Silva and Gazeau, 2020 (https://arxiv.org/html/2605.06821#bib.bib39)\]、带缩放规则的 SDE 近似 \[Malladiet al., 2022 (https://arxiv.org/html/2605.06821#bib.bib33)\]、IMEX 时间步进 \[Bhattacharjeeet al., 2024 (https://arxiv.org/html/2605.06821#bib.bib43)\] 和积分微分视角 \[Heredia, 2024 (https://arxiv.org/html/2605.06821#bib.bib44)\]、控制论框架 \[Chakrabarti and Chopra, 2024 (https://arxiv.org/html/2605.06821#bib.bib45)\]、超参数稳定区域的稳定性 \[Gould and Tanaka, 2024 (https://arxiv.org/html/2605.06821#bib.bib32)\]、小学习率处的反向误差分析 \[Cattaneoet al., 2024 (https://arxiv.org/html/2605.06821#bib.bib31)\]、极小值附近慢 SDE 尖锐度降低 \[Liet al., 2025 (https://arxiv.org/html/2605.06821#bib.bib40)\]、同质网络上的 Adam 流隐式偏差 \[Wanget al., 2021 (https://arxiv.org/html/2605.06821#bib.bib41)\] 以及 EoS 诱导损失尖峰的机制解释 \[Baiet al., 2025 (https://arxiv.org/html/2605.06821#bib.bib46)\]。这些都没有捕捉到 EoS 期间有限步长的 Adam。

**连续时间模型。** 梯度流是梯度下降的经典连续极限，但在稳定性边缘失败，因为它无法捕捉振荡。Roscaet al. \[2023 (https://arxiv.org/html/2605.06821#bib.bib12)\] 研究了梯度流模型中的不稳定性。Shiet al. \[2021 (https://arxiv.org/html/2605.06821#bib.bib18)\] 为动量方法推导了高分辨率 ODE，而 Liet al. \[2017 (https://arxiv.org/html/2605.06821#bib.bib20)\] 为 SGD 开发了随机修正方程。来自反向误差分析的修正方程在数值分析中是标准的 \[Wilkinson, 1963 (https://arxiv.org/html/2605.06821#bib.bib24), 1965 (https://arxiv.org/html/2605.06821#bib.bib25), Haireret al., 2006 (https://arxiv.org/html/2605.06821#bib.bib16)\]；Barrett and Dherin \[2021 (https://arxiv.org/html/2605.06821#bib.bib8)\] 和 Smithet al. \[2021 (https://arxiv.org/html/2605.06821#bib.bib13)\] 使用此框架来理解梯度下降中的隐式正则化，Di Giovacchinoet al. \[2024 (https://arxiv.org/html/2605.06821#bib.bib17)\] 提供了方法论验证。最相关的工作是 Central Flow \[Cohenet al., 2025 (https://arxiv.org/html/2605.06821#bib.bib3)\]，它为稳定性边缘处的梯度下降、标量 RMSProp 和分量 RMSProp 推导了连续时间模型。本工作的直接前身是 Regis and Chewi \[2026 (https://arxiv.org/html/2605.06821#bib.bib1)\]，它为 EoS 处的梯度下降引入了杆流。我们将杆流框架扩展到相空间和自适应方法。

## 3 梯度下降的杆流

我们简要回顾 Regis and Chewi \[2026 (https://arxiv.org/html/2605.06821#bib.bib1)\] 中引入的梯度下降的杆流。完整的重新推导见附录 B (https://arxiv.org/html/2605.06821#A2)。考虑学习率为 $\eta$ 的梯度下降：
$$ w_{t+1} = w_t - \eta \nabla L(w_t). \tag{1} $$
定义连续迭代的中心和半差：
$$ \bar{w}_t = \frac{1}{2}(w_{t+1} + w_t), \tag{2} $$
$$ \delta_t = \frac{1}{2}(w_{t+1} - w_t). \tag{3} $$
杆流框架利用这样一个事实：即使单个迭代 $w_t$ 在稳定性边缘快速振荡，中心 $\bar{w}_t$ 和范围张量 $\delta_t \otimes \delta_t$ 也平滑变化——这使得它们适合连续时间 ODE 近似。为简洁起见，我们写 $L_{\pm} \coloneqq L(\bar{w} \pm \delta)$，其中 $\nabla L_{\pm}$ 和 $\nabla^2 L_{\pm}$ 表示相应的梯度和 Hessian。$\bar{w}$ 和 $\delta_t \otimes \delta_t$ 的离散差分方程给出如下：
$$ \bar{w}_{t+1} - \bar{w}_t = -\frac{\eta}{2}(\nabla L_+ + \nabla L_-), \tag{4} $$
$$ \delta_{t+1} \otimes \delta_{t+1} - \delta_t \otimes \delta_t = \frac{\eta^2}{4}(\nabla L_+ \otimes \nabla L_+ + \nabla L_- \otimes \nabla L_-) - 2\delta_t \otimes \delta_t. \tag{5} $$
令 $\Sigma(t)$ 表示范围张量 $\delta_t \otimes \delta_t$ 的连续时间模拟量。在连续时间中，$\delta$ 被识别为 $\Sigma$ 的主特征向量，乘以主特征值的平方根。将方程 (4) 和 (5) 提升为 ODE 得到梯度下降的杆流 ODE：
$$ \frac{d\bar{w}}{dt} = -\eta \bar{g}, \tag{6} $$
$$ \frac{d\Sigma}{dt} = \frac{\eta^2}{4}(\nabla L_+ \otimes \nabla L_+ + \nabla L_- \otimes \nabla L_-) - 2\Sigma. \tag{7} $$
其中 $\bar{g} = (\nabla L_+ + \nabla L_-)/2$ 是杆两端梯度的平均值。注意，$\bar{w}_t$ 和 $\delta_t \otimes \delta_t$ 的离散差分方程是*精确*的。杆流 ODE 中唯一的近似是用连续时间导数替换离散差分。传统上，这种插值需要反向误差分析。Regis and Chewi \[2026 (https://arxiv.org/html/2605.06821#bib.bib1)\] 中的原始推导在方程 (6) 中包含 $O(\eta^2)$ 的反向误差分析修正，为了简单起见，我们在全文中省略了它，因为它不会定性影响动力学。有关反向误差分析的进一步讨论，请参见附录 F (https://arxiv.org/html/2605.06821#A6)。

## 4 动量方法的杆流

参见标题
参见标题
**图 1：相空间杆流。** 左：相空间杆的图示。端点代表连续 p