[論文レビュー] Sharpness-Aware Training for Free

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

Advanced Neural Network Applications被引用数 15

ひとこと要約

本稿では、現在の重みと過去の重みにおけるモデル出力のKLダイバージェンスに基づくトラジェクトリーロスを最小化することで、ほぼゼロの計算オーバーヘッドでシャープネスに配慮した学習を実現する新規手法、Sharpness-Aware Training for Free (SAF) を提案する。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.

研究の動機と目的

• シャープネスに配慮した最小化（SAM）の高い計算コストに対処すること。SAMは1ステップあたり2回の逆伝播を必要とし、学習時間を2倍に増加させる。
• 標準的な最適化に比べてほぼ追加計算を要しないシャープネスに配慮した学習手法の開発。
• SGDと同等の速度で動作しながら、SAMと同等またはそれ以上の一般化性能を維持または向上させること。
• メモリと計算リソースが限られる大規模データセット（例：ImageNet-1kやそれ以上）において、効率的なシャープネスに配慮した学習を可能にすること。
• ImageNet-21K や JFT-300M などの極めて大規模なデータセットへの適用を想定し、メモリ効率の良い変種、MESA の導入。

提案手法

• 現在の重みと過去の重みを使用したニューラルネットワークの出力分布間のKLダイバージェンスに基づく新規なトラジェクトリーロスを提案。
• このトラジェクトリーロスを、SAMが用いるシャープネス測定の代理として用い、学習軌道に沿った損失の変化率を捉える。
• 標準的な最適化に統合し、主な学習損失とともにトラジェクトリーロスを最小化することで、急激な損失の低下（シャープな低下）をペナルティ化。
• 訓練中に追加のフォワードパスを実行せずに、過去のステップでのモデル出力を記録することで、トラジェクトリーロスを計算。
• 重みの指数移動平均（EMA）を用いて過去のモデル出力を生成する、メモリ効率の良い変種、MESA を導入。
• 2段階の訓練プロセスを採用：まず標準的なSGDでモデルを学習し、次に以前のステップの出力を保持してトラジェクトリーロスを最小化。

実験結果

リサーチクエスチョン

• RQ1 標準的なSGDと比較して、ほぼ追加計算コストゼロでシャープネスに配慮した学習を実現できるか？
• RQ2 出力分布のズレに基づくトラジェクトリーベースのロスが、SAMで用いられるシャープネス測定を効果的に近似できるか？
• RQ3 提案手法SAFが、SGDと同等の速度を維持しながらSAMの一般化性能に匹敵または上回るか？
• RQ4 画像分類データセット（例：ImageNet-1k や ImageNet-21k）において、精度とメモリ効率の両面でスケーリング可能か？
• RQ5 ImageNet-21K や JFT-300M などの極めて大規模な学習タスクに応用可能な、メモリ効率の良い変種（例：MESA）を設計できるか？

主な発見

• ResNet-50を用いたImageNet-1kでは、SAFが96.37%のトップ1精度を達成し、SAMの96.50%と同等の性能を示したが、SGDの97.6%の学習速度で動作した。
• ResNet-101では、SAFが96.93%のトップ1精度を達成し、SAMと同等の性能を示したが、SGDの99.2%の速度で学習した。
• Wide-ResNet-28-10では、SAFが97.08%のトップ1精度を達成し、SAMの97.07%をわずかに上回ったが、ほぼ同等の学習速度で動作した。
• MESAはPyramidNet-110で97.46%のトップ1精度を達成し、SAMおよびSAFを上回ったが、SGDの84.3%の学習速度で動作した。
• SAFはSAMの100%に比べてわずか0.7%の追加計算コストに留まり、実際にはほぼゼロのオーバーヘッドを実現した。
• KLダイバージェンスに基づくトラジェクトリーロスは、シャープネスのダイナミクスを効果的に捉え、追加の逆伝播なしに平坦な最小値に収束を可能にした。