[论文解读] Sharpness-Aware Training for Free

Jiawei Du, Daquan Zhou+3 more|arXiv (Cornell University)|May 27, 2022

Advanced Neural Network Applications被引用 15

一句话总结

该论文提出了一种名为SAF（Sharpness-Aware Training for Free）的新方法，通过基于当前权重与历史权重之间模型输出的KL散度的轨迹损失，以近乎零的计算开销实现了清晰度感知训练。SAF在ImageNet上的泛化性能与SAM相当，同时训练速度接近标准SGD，因此在大规模模型中具有极高的效率。

ABSTRACT

Modern deep neural networks (DNNs) have achieved state-of-the-art performances but are typically over-parameterized. The over-parameterization may result in undesirably large generalization error in the absence of other customized training strategies. Recently, a line of research under the name of Sharpness-Aware Minimization (SAM) has shown that minimizing a sharpness measure, which reflects the geometry of the loss landscape, can significantly reduce the generalization error. However, SAM-like methods incur a two-fold computational overhead of the given base optimizer (e.g. SGD) for approximating the sharpness measure. In this paper, we propose Sharpness-Aware Training for Free, or SAF, which mitigates the sharp landscape at almost zero additional computational cost over the base optimizer. Intuitively, SAF achieves this by avoiding sudden drops in the loss in the sharp local minima throughout the trajectory of the updates of the weights. Specifically, we suggest a novel trajectory loss, based on the KL-divergence between the outputs of DNNs with the current weights and past weights, as a replacement of the SAM's sharpness measure. This loss captures the rate of change of the training loss along the model's update trajectory. By minimizing it, SAF ensures the convergence to a flat minimum with improved generalization capabilities. Extensive empirical results show that SAF minimizes the sharpness in the same way that SAM does, yielding better results on the ImageNet dataset with essentially the same computational cost as the base optimizer.

研究动机与目标

• 为解决Sharpness-Aware Minimization（SAM）带来的高计算成本问题，后者因每步需两次反向传播而使训练时间翻倍。
• 开发一种清晰度感知训练方法，其额外计算开销几乎可忽略不计，仅略高于标准优化。
• 在与SGD相同的速度下，保持或提升与SAM相当的泛化性能。
• 实现在ImageNet-1k等大规模数据集及其以上规模数据集上的高效清晰度感知训练，其中内存和计算资源至关重要。
• 提出一种内存高效的变体MESA，以支持在ImageNet-21K或JFT-300M等极大规模数据集上的部署。

提出的方法

• 提出一种基于当前权重与历史权重下神经网络输出分布之间KL散度的新型轨迹损失。
• 将该轨迹损失用作SAM中清晰度度量的代理，以捕捉损失沿训练轨迹的变化速率。
• 通过在标准优化中同时最小化主训练损失与轨迹损失，将轨迹损失整合进标准优化过程，从而有效惩罚损失的陡峭或突然下降。
• 通过存储先前步骤的模型输出来计算轨迹损失，而无需在训练期间进行额外前向传播。
• 提出MESA，一种内存高效的变体，利用模型权重的指数移动平均（EMA）来生成历史模型输出。
• 采用两阶段训练流程：首先使用标准SGD训练模型；其次利用早期步骤存储的输出最小化轨迹损失。

实验结果

研究问题

• RQ1 与标准SGD相比，是否可以实现近乎零额外计算开销的清晰度感知训练？
• RQ2 基于输出分布变化的轨迹损失是否能有效近似SAM中使用的清晰度度量？
• RQ3 所提出的SAF方法是否能与SAM的泛化性能相当甚至超越，同时保持SGD的训练速度？
• RQ4 在准确率和内存效率方面，该方法在ImageNet-1k和ImageNet-21k等大规模数据集上的可扩展性如何？
• RQ5 是否可以设计出如MESA这样的内存高效变体，以将SAF的适用性扩展至极大规模学习任务？

主要发现

• 在ResNet-50上，SAF在ImageNet-1k上达到96.37%的top-1准确率，与SAM的96.50%相当，但训练速度达到标准SGD的97.6%。
• 在ResNet-101上，SAF达到96.93%的top-1准确率，与SAM性能一致，同时训练速度达到SGD的99.2%。
• 在Wide-ResNet-28-10上，SAF达到97.08%的top-1准确率，略高于SAM的97.07%，且训练速度几乎完全相同。
• MESA在PyramidNet-110上达到97.46%的top-1准确率，优于SAM和SAF，同时仅消耗SGD 84.3%的训练速度。
• 与SAM的100%相比，SAF仅增加0.7%的计算开销，实际中可视为近乎零开销。
• 基于KL散度的轨迹损失能有效捕捉清晰度动态，并在无需额外反向传播的情况下实现平坦极小值的收敛。