[论文解读] Towards Efficient and Scalable Sharpness-Aware Minimization

Yong Liu, Siqi Mai+3 more|arXiv (Cornell University)|Mar 5, 2022

Advanced Neural Network Applications被引用 5

一句话总结

本文提出 LookSAM，一种计算效率更高的 Sharpness-Aware Minimization (SAM) 变体，通过周期性地计算内部梯度上升步骤，并在迭代间重用增强尖锐度的方向。该方法在保持 SAM 级别泛化性能的同时，实现了类似一阶优化器的效率，使 Vision Transformers 在 64k 批大小下训练时间不足一小时——创下 ViT 训练的新速度纪录。

ABSTRACT

Recently, Sharpness-Aware Minimization (SAM), which connects the geometry of the loss landscape and generalization, has demonstrated significant performance boosts on training large-scale models such as vision transformers. However, the update rule of SAM requires two sequential (non-parallelizable) gradient computations at each step, which can double the computational overhead. In this paper, we propose a novel algorithm LookSAM - that only periodically calculates the inner gradient ascent, to significantly reduce the additional training cost of SAM. The empirical results illustrate that LookSAM achieves similar accuracy gains to SAM while being tremendously faster - it enjoys comparable computational complexity with first-order optimizers such as SGD or Adam. To further evaluate the performance and scalability of LookSAM, we incorporate a layer-wise modification and perform experiments in the large-batch training scenario, which is more prone to converge to sharp local minima. We are the first to successfully scale up the batch size when training Vision Transformers (ViTs). With a 64k batch size, we are able to train ViTs from scratch in minutes while maintaining competitive performance.

研究动机与目标

• 降低 SAM 的计算开销，因其每步需两次连续的梯度计算，导致训练时间翻倍。
• 在显著提升大规模模型（如 Vision Transformers）训练效率的同时，保持 SAM 的泛化优势。
• 将 SAM 扩展至大批次训练，该训练范式易受尖锐极小值和层间非均匀不稳定性的影响。
• 开发 LookSAM 的逐层自适应版本（Look-LayerSAM），实现在极高批次大小下的稳定且高效训练。
• 通过实现比之前 SOTA 设置快 8 倍的训练速度，为 Vision Transformers 建立新的训练速度纪录。

提出的方法

• LookSAM 在每次步骤中不再重新计算内部梯度上升，而是重用前一迭代中获得的尖锐度促进方向。
• 该方法将 SAM 的更新分解为与标准 SGD 方向平行和正交的分量，仅在稳定时重用正交分量。
• 重用频率超参数 $k$ 控制内部梯度上升的重新计算频率，以在速度与精度之间取得平衡。
• Look-LayerSAM 在 LookSAM 的权重扰动中应用逐层缩放，提升了大批次训练下的稳定性。
• 该算法使用梯度重用权重 $\alpha$ 和扰动强度 $\rho$ 来微调泛化与效率之间的权衡。
• 在 ImageNet-1k 上对 ViT-B-16 进行实证验证，并对 $k$、$\alpha$ 和 $\rho$ 进行消融研究。

Figure 1 : Accuracy of SAM-5, SAM and vanilla ViT on ImageNet-1k. SAM-5 indicates the method that calculating SAM gradients every 5 steps.

实验结果

研究问题

• RQ1 SAM 的计算成本能否在不牺牲泛化性能的前提下降低？
• RQ2 周期性重用尖锐度促进方向对模型精度和训练速度有何影响？
• RQ3 LookSAM 能否在保持竞争力性能的同时扩展至极高的批量大小（如 64k）？
• RQ4 逐层扰动缩放如何提升大批次设置下的训练稳定性？
• RQ5 哪些超参数设置（$k$、$\alpha$、$\rho$）能实现速度与精度之间的最优权衡？

主要发现

• 当重用频率 $k=5$ 时，LookSAM 在 ViT-B-16 上达到 79.8% 的验证准确率，与原始 SAM 性能相当。
• 在 $k=5$ 时，LookSAM 的吞吐量从 12,800 提升至 19,051 个样本/秒，显著提升训练速度。
• 当 $k=10$ 时，吞吐量达到 20,480 个样本/秒，准确率仅轻微下降至 78.7%。
• Look-LayerSAM 在 16,384 批大小下达到 78.4% 准确率，在 32,768 批大小下达到 77.1% 准确率，参数设置为 $\alpha=0.7$ 和 $\rho=1.0$，表现出强鲁棒性。
• 该方法使 ViT-B-16 从零开始在 64k 批大小下训练时间不足一小时，准确率超过 75%。
• Look-LayerSAM 相较于 [10] 中的训练设置实现了约 8 倍的速度提升，0.7 小时内完成 ViT-B-16 训练——创下新速度纪录。

Figure 2 : Difference of gradients between every 5 steps for $\boldsymbol{g_{s}}$ , $\boldsymbol{g_{h}}$ , and $\boldsymbol{g_{v}}$ (i.e., $||g_{s}^{t}-g_{s}^{t+k}||$ ) . $\boldsymbol{g_{v}}$ that leads to a smoother region changes much slower than $\boldsymbol{g_{s}}$ and $\boldsymbol{g_{h}}$ .