[Paper] 稳定性边缘的泛化

Source: arXiv - 2604.19740v1

概述

现代深度学习从业者注意到，使用 large learning rates——将优化器推向“edge of stability”——进行训练，往往能够获得出乎意料的优秀测试性能。本文将这种混沌的训练 regime 重新表述为一个 random dynamical system，该系统会收敛到低维分形吸引子上，并推导出一种基于 sharpness dimension 的新型泛化界限，解释了为何这种不稳定性可能是有益的。

关键贡献

• 随机动力系统视角： 表明随机优化器（SGD、Adam 等）表现为随机动力系统，其长期状态形成分形吸引子，而不是单一点。
• 锐度维度： 引入一种基于吸引子的 Lyapunov（分形）维度 的新颖复杂度度量，捕获完整的 Hessian 谱而不仅仅是其迹或谱范数。
• 泛化界限： 证明测试误差随锐度维度尺度增长，将混沌动力学直接关联到泛化性能。
• 实证验证： 在多层感知器和 Transformer 模型上演示该理论，复现 “edge‑of‑stability” 学习曲线，并阐明 grokking 现象。
• 实用诊断工具： 提供从训练日志估计锐度维度的工具，使开发者能够监控模型何时进入有益的混沌 regime。

方法论

1. 将优化器建模为随机动力系统 (RDS)。
   • 每一次 SGD/Adam 更新被视为随机映射 (x_{t+1}=f_{\theta_t}(x_t)+\xi_t)，其中 (\xi_t) 捕获梯度噪声。
2. 分析长期吸引子。
   • 使用 Lyapunov 理论的概念，作者证明在大学习率下，RDS 不会收敛到固定点，而是收敛到一个 分形吸引子，其内在（Lyapunov）维度为 (d_L)。
3. 定义锐度维度。
   • 他们从 完整的 Hessian 谱 以及其主子矩阵的行列式计算 (d_L)，得到一个标量，反映在混沌 regime 下损失曲面的 “尖锐” 或 “平坦” 程度。
4. 推导泛化界。
   • 通过将 PAC‑Bayesian 论证扩展到分形吸引子，证明期望测试损失被一个与 (\sqrt{d_L / n}) 成正比的项所界定（其中 (n) 为训练样本数）。
5. 实验流程。
   • 在 MNIST/Fashion‑MNIST 上训练 MLP，在语言建模任务上训练 transformer，进行学习率的全范围搜索。
   • 使用 Hessian 的随机 Lanczos 二次求积估计锐度维度，并将其与验证准确率和 grokking 曲线进行比较。

Results & Findings

• 仅当学习率超过临界阈值时，会出现分形吸引子，这与先前实证工作中识别的“稳定边缘”相吻合。
• Sharpness dimension 与最终测试性能呈强相关（Pearson ≈ 0.85），在所有实验中优于传统的 Sharpness 指标（Hessian 的迹、谱范数）。
• 在 grokking 实验中，Sharpness dimension 在测试准确率突跃前急剧下降，暗示模型动力学转向了一个更低维的吸引子，从而编码出更稳健的解。

实际影响

• 学习率调节: 与其把大学习率视为风险，开发者可以有意识地将其推入混沌区间，并监控锐度维度，以确保仍然位于“好”边缘的一侧。
• 训练诊断: 锐度维度可以在训练过程中实时估计（例如每几百步一次），利用廉价的 Hessian‑vector 产品，提前警示优化器是否正漂移到对泛化有害的过度混沌区域。
• 有限数据下的模型选择: 由于界随 (\sqrt{d_L/n}) 缩放，天然倾向于落在低维吸引子上的模型（例如某些 Transformer 架构）在训练数据稀缺时可能更为合适。
• 理解 grokking: 该理论为模型为何会在训练多个 epoch 后才突然泛化提供了具体解释——这对应于吸引子的 维度塌缩。实践者可以通过安排学习率衰减，在期望的时机触发塌缩。
• 正则化的替代方案: 传统的权重衰减或批归一化旨在平坦化损失景观；锐度维度暗示 受控混沌 可以作为一种替代正则化手段，或许能够降低对强显式惩罚的需求。

局限性与未来工作

• Hessian 估计开销： 对于非常大的模型，精确计算完整的 Hessian 谱仍然代价高昂；当前方法依赖可能带噪声的随机近似。
• 平稳噪声假设： RDS 分析假设梯度噪声是 i.i.d. 的，这在高度非平稳的数据流或课程学习中可能不成立。
• 架构范围： 实验主要聚焦于相对较小的 MLP 和 transformer‑style 语言模型；将理论扩展到卷积网络、图神经网络或强化学习代理仍是一个未解问题。
• 理论紧致性： 推导的界限虽有洞见，但相较于经验差距仍显宽松；细化常数并探索更紧的分形维度基界限是一个有前景的方向。

要点: 通过将大学习率训练视为随机动力系统并引入锐度维度，本文为开发者提供了一种新颖且理论扎实的工具，以利用边缘稳定性 regime 实现更好的泛化。

作者

• Mario Tuci
• Caner Korkmaz
• Umut Şimşekli
• Tolga Birdal

论文信息

• arXiv ID: 2604.19740v1
• 分类: cs.LG, cs.AI, cs.CV, stat.ML
• 出版日期: 2026年4月21日
• PDF: Download PDF