[論文レビュー] Towards Efficient and Scalable Sharpness-Aware Minimization

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

Advanced Neural Network Applications被引用数 5

ひとこと要約

この論文は、内側の勾配上昇ステップを定期的に計算し、反復間で鋭さを促進する方向を再利用することで、計算効率の高いSharpness-Aware Minimization (SAM)の変種であるLookSAMを提案する。この手法は、1次最適化子と同等の効率を維持しながらSAMと同等の一般化性能を達成し、64kバッチサイズでVision Transformer (ViT)を1時間未塔でトレーニング可能にし、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の計算コストを削減すること。SAMは1ステップあたり2回の連続的な勾配計算が必要なため、トレーニング時間を2倍に増加させる。
• 大規模モデル（例：Vision Transformers）のトレーニング効率を著しく向上させつつ、SAMの一般化利点を維持すること。
• 鋭い最小値や層間で非一様に現れる不安定性に悩まされやすい大バッチトレーニングに、SAMをスケーリングすること。
• LookSAMの層ごとの適応版（Look-LayerSAM）を開発し、極めて高いバッチサイズでも安定的かつ効率的なトレーニングを可能にすること。
• ViTのトレーニング速度記録を樹立し、先行研究の最良設定比で8倍の高速化を達成すること。

提案手法

• LookSAMは、各ステップで内側の勾配上昇を再計算する代わりに、前回の反復で得られた鋭さを促進する方向を再利用する。
• この手法は、標準のSGD方向に対して平行および直交する成分にSAMの更新を分解し、直交成分を安定している場合に再利用する。
• 再利用頻度を制御するハイパーパrameter $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 速度と正確性の最適なトレードオフを達成するには、どのハイパーパrameter設定（$k$、$\alpha$、$\rho$）が適しているか？

主な発見

• LookSAMは再利用頻度 $k=5$ の場合、ViT-B-16で79.8%の検証正確性を達成し、ネイティブなSAMと同等の性能を示した。
• $k=5$ の場合、LookSAMは1秒あたりのスループットを12,800から19,051サンプルに向上させ、トレーニング速度を顕著に向上させた。
• $k=10$ の場合、スループットは20,480サンプル/秒に達し、正確性は78.7%にわずかに低下した。
• Look-LayerSAMは $\alpha=0.7$ と $\rho=1.0$ の設定で、16,384のバッチサイズで78.4%の正確性、32,768のバッチサイズで77.1%の正確性を達成し、頑健性を示した。
• この手法により、64kバッチサイズでViT-B-16を1時間未塔でスクラッチからトレーニング可能となり、75%以上の正確性を達成した。
• Look-LayerSAMは[10]のトレーニング設定比で約8倍の高速化を達成し、ViT-B-16のトレーニングを0.7時間で完了した—新たな速度記録を樹立した。

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}}$ .