稳定边缘选择性塑造数据分布上的学习

# Edge of Stability 选择性塑造数据分布上的学习

**来源：** <https://arxiv.org/html/2606.04212>

Shauna Kwag\* MIT kwags@mit\.edu  
Anakha Ganesh\* MIT anakhag@mit\.edu  
Tomaso Poggio MIT tp@ai\.mit\.edu  
Pierfrancesco Beneventano MIT pierb@mit\.edu

###### 摘要

现有的关于稳定性边缘（Edge of Stability, EoS）的分析将其视为优化的全局属性。我们证明它也具有选择性：稳定性约束重新分配了对训练分布各子集的学习，增强了对某些组的进展，同时抑制了对其他组的进展。通过使用从相同训练状态进入或退出 EoS 状态的分支干预，我们因果性地展示了这种权衡，并确定了某个组受益的两个必要条件。首先，该组的聚合梯度必须与 Hessian 矩阵的顶部特征向量对齐。我们通过一种受控扰动来隔离这一机制，该扰动保持距离但随机化方向，从而破坏对齐并消除优势。其次，该组必须在训练过程中保持非零的梯度幅度。在交叉熵损失下，梯度饱和使置信度高的组与学习解耦，将优势转移到输出异常值（output-outliers），其梯度持续存在。综合来看，这些结果表明 EoS 不仅作为稳定性边界，还作为一种机制，管理着数据分布上学习的分配。

††footnotetext:\* 同等贡献。

## 1 引言

深度神经网络对优化器和超参数表现出强烈的敏感性。与经典凸优化中这些选择不影响所达到的最小值不同，训练选择如学习率、批量大小和优化器会影响找到的解 [41 (https://arxiv.org/html/2606.04212#bib.bib10),23 (https://arxiv.org/html/2606.04212#bib.bib20),18 (https://arxiv.org/html/2606.04212#bib.bib8)]。理解这种隐式偏置背后的机制是深度学习理论的一个关键目标。一个结构性解释来自稳定性边缘（Edge of Stability, EoS）文献：在全批次和大批次梯度下降下，Hessian 矩阵的顶部特征值会自我稳定在接近稳定性阈值处，该阈值取决于优化器和超参数 [40 (https://arxiv.org/html/2606.04212#bib.bib36),19 (https://arxiv.org/html/2606.04212#bib.bib37),20 (https://arxiv.org/html/2606.04212#bib.bib38),9 (https://arxiv.org/html/2606.04212#bib.bib12),8 (https://arxiv.org/html/2606.04212#bib.bib39),4 (https://arxiv.org/html/2606.04212#bib.bib29),3 (https://arxiv.org/html/2606.04212#bib.bib32),16 (https://arxiv.org/html/2606.04212#bib.bib40),22 (https://arxiv.org/html/2606.04212#bib.bib1),35 (https://arxiv.org/html/2606.04212#bib.bib2)]。在这个阈值处，优化器在离散时间稳定性的边界上运行，限制了训练过程中损失景观中哪些区域仍然可访问。

虽然 EoS 现象已被充分证实，但对其后果的了解却少得多。特别是，目前尚不清楚在稳定性阈值附近操作是否提供了任何功能上的益处，或者这些稳定性约束如何影响数据分布上的优化。先前的工作主要通过参数空间中的曲率和优化轨迹来表征 EoS，留下了这些动态如何影响训练过程中哪些示例被学习的疑问。我们问：EoS 的发生实际上是否有实际后果？如果有，训练分布的哪些子集从 EoS 中受益？哪些没有？是什么决定了这种分配？我们对第一个问题给出了肯定的答案。然后，为了研究这种分配，我们从 \(P_{X,Y}\) 的几何结构出发定义了四个原型组，这些组在输入典型性、标签一致性和边界接近度上有所不同，独立于模型或损失。

**信息：** 我们发现 EoS 诱导了一种 **选择性学习** 机制：稳定性约束不均衡地分配优化努力，增强了那些梯度持续与顶部 Hessian 特征方向对齐的子集，同时抑制了其他子集。这与经典直觉相反：虽然下降引理将最陡方向视为限制学习的稳定性边界，但与该方向的对齐恰恰决定了学习如何在整个分布上分配。我们的发现与一个长期存在的争论联系起来，即低曲率是否改善泛化 [15 (https://arxiv.org/html/2606.04212#bib.bib6),23 (https://arxiv.org/html/2606.04212#bib.bib20),14 (https://arxiv.org/html/2606.04212#bib.bib22),11 (https://arxiv.org/html/2606.04212#bib.bib26)]。EoS 是尖锐度被主动约束在稳定性阈值附近的机制，因此是测试低曲率实际赋予什么的自然设置。我们的结果表明答案并非全局性的：功能益处取决于哪个子集主导顶部 Hessian 特征方向，并随着训练分布的几何组成而变化。从这个角度看，平坦性是由数据几何决定的方向性属性，而非解的标量属性。更广泛地说，我们的结果向着连接两个 largely separate 的研究方向迈出了一步：参数空间中的隐式正则化和数据分布上的归纳偏差。

我们的论文做出了以下贡献：

- • **EoS 是选择性的，而非全局性的**（第 3. 节 (https://arxiv.org/html/2606.04212#S3)）。通过使用从共享训练轨迹进入或退出 EoS 的分支干预，我们发现稳定性约束并非均匀的瓶颈：它选择性地有益于训练分布的某些子集，同时抑制其他子集。这种权衡在不同架构和优化器上定性可复现（附录 C (https://arxiv.org/html/2606.04212#A3)）。
- • **选择性由对齐 × 持续性决定**（第 4. 节 (https://arxiv.org/html/2606.04212#S4)）。EoS 的优势被那些聚合梯度既与顶部 Hessian 特征向量对齐又在整个训练过程中保持非零的子集所捕获。两个受控的反事实分别隔离了每个因素：随机方向位移消除了对齐，而交叉熵饱和消除了持续性。在这两种情况下，EoS 优势消失或转移到保留缺失因素的子集。这两个因素都来自于将自稳定化框架 [10 (https://arxiv.org/html/2606.04212#bib.bib13)] 从全局损失扩展到每个子集损失（第 2.2 节 (https://arxiv.org/html/2606.04212#S2.SS2)；完整推导见附录 F (https://arxiv.org/html/2606.04212#A6)）。
- • **几何形状改变了受益者**（第 5. 节 (https://arxiv.org/html/2606.04212#S5)）。在我们基于 CIFAR-10 的 MLP 实验中（先前 EoS 工作的标准设置 [9 (https://arxiv.org/html/2606.04212#bib.bib12)]），从 EoS 受益的子集是那些在输入空间中距离类质心最远的子集。改变训练分布的几何组成会连续地改变哪个子集受益。初步证据表明，由此产生的泛化行为也随之改变：当几何上遥远的示例位于决策边界附近时，鲁棒性得到改善；当它们远离训练分布时，分布外泛化得到改善。

![引用图注](略) 图 1：原型概念分类。数据样本根据输入空间中相对于特定类别簇质心（\(\mu_0,\mu_1\)）的几何接近度进行分类。

![引用图注](略) 图 2：CIFAR-10 原型组在输入空间中的可视化。每个类别（汽车 vs. 卡车）展示了三个代表性样本。

## 2 预备知识

### 2.1 稳定性边缘

深度网络中的梯度下降通常在稳定性边缘（EoS）运行，这是一种训练过程保持在稳定与非稳定更新边界附近而不发生发散的状态 [9 (https://arxiv.org/html/2606.04212#bib.bib12)]。考虑全批次梯度下降 \(\theta_{t+1} = \theta_t - \eta \nabla L(\theta_t)\)，其 Hessian 矩阵为 \(H_t = \nabla^2 L(\theta_t)\)，特征对为 \((\lambda_i, v_i)\)。沿着 \(v_i\) 的扰动每一步被乘以 \(1 - \eta \lambda_i\)，因此离散时间稳定性要求 \(\eta \lambda_i < 2\)。EoS 是顶部特征值 \(\lambda_1\)（尖锐度）满足边界条件的机制：\(\eta \lambda_1 \approx 2\)。此时，沿着相应特征方向 \(\mathbf{v}_1\) 的乘子接近 \(-1\)，产生符号交替的振荡，这些振荡在训练过程中线性上不稳定但仍保持有界 [9 (https://arxiv.org/html/2606.04212#bib.bib12)]。

### 2.2 自稳定化

Damian 等人 [10 (https://arxiv.org/html/2606.04212#bib.bib13)] 的自稳定化机制解释了 EoS 处损失的非单调下降。一旦 \(\lambda_1\) 超过 \(2/\eta\)，梯度下降沿着 \(\mathbf{v}_1\) 发展出快速的周期-2 振荡。高阶项将此振荡转换为一个缓慢的校正漂移，朝向较低的尖锐度，从而使循环平均后的动态保持在活跃边界 \(\{\lambda_1 \approx 2/\eta\}\) 附近。在这种循环平均描述下，EoS 漂移与普通梯度下降不同，多了一个额外的尖锐度减少分量：\(-(\alpha/\beta) \nabla \lambda_1\)，其中 \(\alpha = -\langle \nabla L, \nabla \lambda_1 \rangle\)，\(\beta = \|\nabla \lambda_1\|^2\)。对于子集损失 \(\ell_k\)，这个额外漂移贡献了 \(-(\alpha/\beta) \langle \nabla \ell_k, \nabla \lambda_1 \rangle\)。我们利用这个已建立的框架来启发我们的实验设计。选择器 \(\langle \nabla \ell_k, \nabla \lambda_1 \rangle\) 依赖于方向因子（\(\nabla \ell_k\) 与 \(\nabla \lambda_1\) 之间的对齐）和幅度因子（\(\|\nabla \ell_k\|\) 在训练中的持续性）。这两个因素在第 4 节 (https://arxiv.org/html/2606.04212#S4) 中通过互补的受控干预来研究。经验上，我们通过测量每个子集数据耦合到不稳定模式的程度，即与顶部 Hessian 特征方向 \(\mathbf{v}_1\) 的对齐来跟踪方向因子。附录 F (https://arxiv.org/html/2606.04212#A6) 推导了每个子集损失的分解，并表明在单模式近似下，测量的 \(\cos^2 \theta_k\) 作为理论 EoS 选择器 \(Q_k = \langle \nabla \ell_k, \nabla \lambda_1 \rangle\) 的经验代理。这启发了我们的实验：分支干预测试了留在 EoS 上的效果，随机方向位移隔离了对齐，而交叉熵饱和隔离了梯度持续性。

### 2.3 原型分类

#### 定义。
我们将训练分布划分为四个组，由 \(P_{X,Y}\) 的联合几何定义，独立于任何训练模型（图 2 (https://arxiv.org/html/2606.04212#S1.F2)）。
*   **内部点**是靠近类质心 \(\mu_c\) 且具有正确标签的高密度点。
*   **边界点**是位于类间边界附近的分布内样本，通过其局部邻域中的高标签歧义性识别。
*   **输入异常值**是几何上非典型的输入，在输入空间中远离 \(\mu_c\)，但保留正确标签。
*   **输出异常值**是被分配了错误标签的高密度输入。

每个组的代表性示例如图 2 (https://arxiv.org/html/2606.04212#S1.F2) 所示。内部点和边界点分别通过排序质心距离和 \(k\)-NN 标签歧义性从现有训练数据中识别，而输入异常值和输出异常值则通过合成构造来隔离输入空间非典型性和标签不一致性的影响。

#### 构造。
我们在一个二分类 CIFAR-10 任务（汽车 vs. 卡车，\(n=10,000\)）上实例化该分类 [26 (https://arxiv.org/html/2606.04212#bib.bib4)]。内部点候选是每个类中质心距离 \(\|x_i - \mu_c\|\) 最小的 \(M=3m\) 个点；边界点候选是每个类中 \(k\)-NN 标签组成（\(k=50\)）最接近均匀分布的 \(m\) 个点。从内部点候选池中，我们为每个类采样三个大小为 \(m=25\) 的不相交子集：第一个保留原始输入和标签（内部点）；第二个被分配翻转标签 \(1-c\)（输出异常值）；第三个通过将每个示例从相反类别的质心推开以 \(x_i \pm \alpha v_{\mathrm{diff}}\) 外推（若 \(y_i=1\) 则取 \(+\)，若 \(y_i=0\) 则取 \(-\)），其中 \(v_{\mathrm{diff}} = \mu_1 - \mu_0\) 是未经归一化的质心差（输入异常值）。我们设置 \(\alpha=3\)，使得输入异常值自然成为距离其类质心最远的组；生成的像素值可能超出有效输入范围，但为了保持位移幅度，我们不进行裁剪。边界点直接从歧义池中抽取。最终训练集包含 \(n=10,000\) 个样本，其中包括 200 个原型标记点（每组 50 个，每类 25 个），并在训练过程中进行跟踪。

### 2.4 度量

#### 每组损失 \(\ell_k\)。
对于每个原型组 \(k \in \{\text{inlier, boundary, input-outlier, output-outlier}\}\)，其索引集为 \(P_k\)，我们定义 \(\ell_k = \frac{1}{|P_k|} \sum_{i \in P_k} \ell(f(x_i), y_i)\) 为组 \(k\) 上样本的平均损失。跟踪 \(\ell_k\) 随训练的变化揭示了各组的 learn 顺序，以及哪些组在 EoS 下受到稳定性约束的差异性影响。

#### 方向耦合。
令 \(\nabla \ell_k\) 表示限制在原型组 \(k\) 上的损失梯度。我们通过下式衡量组 \(k\) 与 EoS 约束模式的方向耦合：
\[
\cos^2\theta_k = \frac{(\nabla \ell_k \cdot v_1)^2}{\|\nabla \ell_k\|^2} \in [0,1],
\]
其中 \(v_1\) 是顶部 Hessian 特征向量（假设为单位范数）。当 \(\cos^2\theta_k \approx 1\) 时，组梯度几乎与 \(v_1\) 对齐；当 \(\cos^2\theta_k \approx 0\) 时，则几乎正交。该量衡量一个组的梯度与 EoS 动力学约束方向的对齐强度。

与学习的联系来自自稳定化。在 EoS 下，振荡-稳定化循环主要沿 \(v_1\) 产生净参数移动 [10 (https://arxiv.org/html/2606.04212#bib.bib13)]。因此，损失主要对那些梯度与 \(v_1\) 对齐的组下降，而与 \(v_1\) 正交的组进展有限（图 3 (https://arxiv.org/html/2606.04212#S2.F3)）。

![引用图注](略) 图 3：EoS 下的方向耦合。优化器沿 \(v_1\) 振荡（红色锯齿形）。当一个组的梯度 \(\nabla \ell_k\) 与 \(v_1\) 对齐（左）时，自稳定化会减少该组的损失。当 \(\nabla \ell_k\) 与 \(v_1\) 正交（右）时，该组与振荡解耦，其损失未受益。

#### 曲率影响。
虽然 \(\cos^2\theta_k\) 衡量方向，但它不捕获梯度幅度。我们报告平方投影 \((\nabla \ell_k \cdot v_1)^2 = \|\nabla \ell_k\|^2 \cdot \cos^2\theta_k\)，