[论文解读] A Loss Curvature Perspective on Training Instability in Deep Learning

Justin Gilmer, Behrooz Ghorbani+7 more|arXiv (Cornell University)|Oct 8, 2021

Domain Adaptation and Few-Shot Learning参考文献 34被引用 6

一句话总结

本文通过损失曲率的视角研究深度学习中的训练不稳定性，表明成功的模型架构和训练启发式方法（如权重初始化、归一化、梯度裁剪和初始学习率热身）通过将损失Hessian的最大特征值（$ olambda_1$）保持在$2/\eta$以下来稳定训练，其中$\eta$为学习率。其核心贡献在于提供了一个统一的视角：这些看似不同的方法均通过引导优化过程早期进入更平坦、更稳定的区域，来缓解损失曲率的不良条件这一根本性失败模式。

ABSTRACT

In this work, we study the evolution of the loss Hessian across many classification tasks in order to understand the effect the curvature of the loss has on the training dynamics. Whereas prior work has focused on how different learning rates affect the loss Hessian observed during training, we also analyze the effects of model initialization, architectural choices, and common training heuristics such as gradient clipping and learning rate warmup. Our results demonstrate that successful model and hyperparameter choices allow the early optimization trajectory to either avoid -- or navigate out of -- regions of high curvature and into flatter regions that tolerate a higher learning rate. Our results suggest a unifying perspective on how disparate mitigation strategies for training instability ultimately address the same underlying failure mode of neural network optimization, namely poor conditioning. Inspired by the conditioning perspective, we show that learning rate warmup can improve training stability just as much as batch normalization, layer normalization, MetaInit, GradInit, and Fixup initialization.

研究动机与目标

• 理解尽管在架构上有所进步，为何深度神经网络的训练仍会出现不稳定性。
• 研究损失曲率（特别是损失Hessian的最大特征值$\lambda_1$）如何影响优化稳定性。
• 考察模型初始化、归一化、学习率热身和梯度裁剪在训练过程中控制$\lambda_1$的作用。
• 将多种不同的训练启发式方法统一于同一机制之下：通过减少训练初期的曲率来提升稳定性。

提出的方法

• 通过SGD+动量在多个分类任务中，实证测量损失Hessian最大特征值（$\lambda_1$）在训练过程中的演化。
• 系统性地改变学习率、热身时长、初始化方法（如MetaInit、Fixup、GradInit）以及架构选择（如归一化层的位置）。
• 在WideResNet和ResNet-50等大规模模型上进行实验，追踪$\lambda_1$随时间的动力学变化。
• 比较学习率热身与归一化、初始化和梯度裁剪在降低$\lambda_1$方面的效果。
• 分析$\lambda_1$与理论稳定性边界$2/\eta$之间的关系，以验证稳定性条件。
• 将分析扩展至Adam优化的Transformer模型，表明尽管采用自适应优化，$\lambda_1$与学习率之间仍存在一致的适应关系。

实验结果

研究问题

• RQ1 损失Hessian的最大特征值（$\lambda_1$）在训练过程中如何演化？它是否始终低于理论稳定性阈值$2/\eta$？
• RQ2 不同初始化策略在训练初期在多大程度上降低了$\lambda_1$？这如何使更高学习率的稳定训练成为可能？
• RQ3 学习率热身是否通过逐步降低$\lambda_1$来稳定训练？与架构归一化或更优初始化相比，其效果如何？
• RQ4 学习率热身是否足以实现与MetaInit或Fixup等先进初始化方法相当的性能？
• RQ5 高损失曲率如何影响大批次训练的扩展性？能否通过降低曲率来提升数据并行效率？

主要发现

• 使用SGD+动量训练时，仅当$\lambda_1 \lesssim 2/\eta$时训练才稳定；该条件在多种模型和任务中均得到实证验证。
• 合适的初始化策略（如MetaInit、GradInit、Fixup）可在训练初期降低$\lambda_1$，从而实现更高学习率下的稳定训练。
• 学习率热身可逐步降低训练初期的$\lambda_1$，其效果与归一化和先进初始化方法相当。
• 学习率热身无需复杂初始化方案，即可实现与MetaInit、Fixup和GradInit相当的性能。
• 高损失曲率会降低大批次训练的扩展效率，而能降低曲率的干预措施可显著提升数据并行效率。
• $\lambda_1$对学习率的适应性在不同模型和优化器（包括Adam）中保持一致，表明存在一种与经典优化理论相悖的根本性优化动态。