作者|周展鹏,上海交通大学博士生

【ICLR 2025预讲会】系列内容

ICLR 2025预讲会系列文章来源于 DAMO 开发者矩阵与 AI Time 联合举办 ICLR 2025预讲会整理成稿,旨在帮助大家率先了解深度学习领域的最新研究方向和成果。本文为文章作者的观点/研究数据,仅供参考,不代表本账号的观点和研究内容。

摘要

本文深入探究了深度学习优化算法SAM(Sharpness-Aware Minimization)在训练后期的优化动力学及隐式偏差。通过实验发现,即使仅在训练的最后几个周期使用SAM,其仍可找到比 SGD 更平坦的解,且与全程使用SAM效果相当,凸显了训练后期动力学对模型泛化能力的重要性。

 

理论分析揭示了 SAM 后期优化的两阶段过程:先从较尖锐的解逃离但保持在当前山谷,再快速收敛到更平坦的解。

图片

理解深度神经网络的良好泛化(generalization)能力是现代深度学习的核心问题之一。最近的研究表明神经网络的泛化能力和损失图景(loss landscape)的平坦度(flatness),或者陡峭度(sharpness)有关。很多研究尝试设计新的梯度更新算法来显式/隐式地调节最终收敛到的解(minima/solution)的sharpness。其中,Foret et al. (2021) 提出了Sharpness-Aware Minimization(SAM)。SAM在CV和NLP等领域都能显著提升模型的泛化能力和鲁棒性(robustness)。

A diagram of a person's hand

AI-generated content may be incorrect.

尽管 SAM 在实际应用中取得了不错的成绩,但是仍然不能完全理解 SAM 的有效性。有研究指出Foret et al. (2021)原文中根据 PAC-Bayes 理论给出的 SAM泛化界(Bound)并不能充分解释 SAM 的有效性。很多研究尝试对 SAM 的动力学进行渐进分析,然而和实际仍有差距。理解 SAM 有效性的背后机理仍然是一个开放问题。

最近的研究指出,基于梯度更新的优化算法的成功可以归结于其隐式偏差(implicit bias):倾向于寻找具有良好性质的minima。一个比较经典的例子是:Stochastic Gradient Descent(SGD)相比于full-batch GD倾向于选择更加平坦(flat)的minima。这是由 SGD 中 mini-batch 所带来的 noise 决定的。

类似地,尽管 SAM 的设计灵感来自于 landscape 的flatness/sharpness,其具体实现(见公式2-3)并没有显式地优化solution的flatness/sharpness。理解SAM的隐式偏差对于我们理解 SAM 的有效性是非常重要的。

原始 SAM 的优化目标:

图片

实际 SAM 的优化算法:

图片

图片

A graph of a number of people

AI-generated content may be incorrect.

作者发现了 SAM 优化动力学(training dynamics)中一个有趣的现象:即使在训练的最后几个 epoch 使用SAM,SAM一样可以找到平坦的解。具体而言,使用 SGD 训练模型并且在最后几个 epoch 切换到 SAM 继续训练。更进一步地,他们发现从 sharpness 的角度,switching method和全程 SAM 也具有相似的表现(如图三所示)。

因此,将发现总结为:SAM在网络训练后期也能找到比 SGD 更加平坦的解。相比于前人研究中“训练早期的 dynamics 更加重要”的观点,作者发现强调了训练后期 dynamics 对泛化能力的重要性。

A black text on a white background

AI-generated content may be incorrect.

图片

A close up of a math problem

AI-generated content may be incorrect.

图片

图片

理解上述发现对于作者解释 SAM 的隐式偏差非常重要。接下来,又从理论角度构建了训练后期 SAM 的动力学图景(picture)。具体而言,将 SAM 在后期的优化分成了两个阶段(Phase):在第一个阶段,SAM会从相对 sharp 的解中逃离(例如,前期使用 SGD 所找到的相对 sharp 的解),但仍然停留在当前的山谷(valley)中;第二个阶段,SAM会以指数快的速度收敛到一个更加 flat 的解(见图1b的示意图)。

在表1中,把一个动力学图景拆分成 4 个claim,并且每一个 claim 对应一个定理。一个 Picture 解释了 SAM 在后期选择更加 flat 的 minima 的原因,为我们解释 SAM 的有效性提供了新的视角。

值得一提的是,作者还探索了从 SAM 切换到 SGD 的训练方法。实验发现,仅仅在训练早期使用 SAM 对模型最终的泛化能力的影响甚微。因此,便猜想训练后期使用的优化算法会对模型的最终性能产生更加重要的影响。

基于此,作者将发现从 SAM 推广到了Adversarial Training(AT)。类似地,仅在最后几个 epoch 使用 AT 训练模型,一样可以达到和全程使用 AT 训练模型类似的对抗鲁棒的模型。

Logo

DAMO开发者矩阵,由阿里巴巴达摩院和中国互联网协会联合发起,致力于探讨最前沿的技术趋势与应用成果,搭建高质量的交流与分享平台,推动技术创新与产业应用链接,围绕“人工智能与新型计算”构建开放共享的开发者生态。

更多推荐