Sharpness-Aware Minimization（SAM）は、ニューラルネットワークの訓練における最適化手法であり、損失ランドスケープ内の平坦な極小値を明示的に探すことを目的としている。2021年にPierre Foretらによって導入されたSAMは、標準的な勾配降下法の更新を修正し、損失が急激に変化する領域をペナルティ化することで、未見データに対するモデルの汎化を向上させる。確率的勾配降下法（SGD）やAdamといった従来の最適化手法が訓練損失のみを最小化するのに対し、SAMは各パラメータ点の近傍も考慮し、摂動に対して頑健な解を見つける。
SAMの核となるアイデアは、ニューラルネットワークが訓練データに同等に適合する多くの局所極小値を持つが、そのうちのいくつかは他よりも良好に汎化するという観察に由来する。平坦な極小値、すなわち損失曲面がパラメータ変化に対して比較的鈍感な領域は、経験的に良好な汎化と関連付けられる。SAMはこれをmin-max最適化問題として定式化する。各ステップで、まず現在のパラメータの周囲の小さな球内で損失を最大化する摂動を計算し、次にその摂動点での勾配を用いてパラメータを更新する。これにより、最適化器は単一の点だけでなく、近傍全体で一様に低い損失を持つ盆地へと移動するよう促される。
SAMは、画像分類、言語モデリング、強化学習など、さまざまなアーキテクチャとタスクにわたって精度を向上させることが示されている。特に、データ拡張や重み減衰などの他の正則化手法と組み合わせた場合に効果的である。この手法は、計算オーバーヘッドが最小限であり、順伝播・逆伝播のコストがおよそ2倍になるが、多くの場合、有意な精度向上をもたらす。導入以来、適応的SAMやlookahead SAMなど、計算負荷をさらに軽減するいくつかの変種が提案されている。
動機：シャープな極小値と平坦な極小値
深層学習の経験的成功は、大規模モデルが高次元の損失ランドスケープで良好な解を見つける能力に起因することが多い。しかし、すべての極小値が同等であるわけではない。研究により、SGDは平坦な極小値に収束する傾向があり、これらは小さなパラメータ摂動に対してより頑健で、しばしば良好な汎化に対応することが示されている。対照的に、シャープな極小値は、極小値から離れるにつれて損失が急激に増加する狭い盆地であり、これらは攻撃的な最適化から生じることが多く、過学習につながる可能性がある。
平坦な極小値に関する初期の理論的研究は、1990年代のニューラルネットワーク汎化の研究に遡るが、厳密な枠組みは欠如していた。SAMは、明示的な正則化項を必要とせずにシャープネスを直接ペナルティ化する最適化目的を提供する。イプシロン半径内の最悪ケースの摂動を計算することで、SAMはその近傍で損失ランドスケープが一様に低くなることを効果的に促し、平坦性を促進する。
形式的定義とアルゴリズム
\(w\) をモデルパラメータ、\(L(w)\) を訓練損失とする。SAMは以下の問題を解く：
\[ \min_w \max_{\|\epsilon\| \le \rho} L(w + \epsilon) \]
ここで、\(\rho\) は摂動半径を制御するハイパーパラメータである。実際には、内側の最大化は勾配上昇の1ステップで近似される：\(\hat{\epsilon} = \rho \frac{\nabla L(w)}{\|\nabla L(w)\|}\)。更新規則は次のようになる：
\[ w_{t+1} = w_t - \eta \nabla L(w_t + \hat{\epsilon}_t) \]
これには、摂動を見つけるための勾配計算とパラメータ更新のための勾配計算の、ステップごとに2回の勾配計算が必要である。この手法はベースの最適化器に依存せず、SGD、Adam、または任意の勾配ベースのアルゴリズムと組み合わせて使用できる。実際には、SAMは重み減衰や運動量と組み合わせられることが多く、さらなる改善が得られる。
ハイパーパラメータとチューニング
主要なハイパーパラメータは摂動半径 \(\rho\) であり、考慮される近傍のサイズを制御する。\(\rho\) が大きいほど平坦な極小値が促進されるが、収束が遅くなる可能性がある。正規化されたデータセットでは、典型的な値は0.01から0.1の範囲である。学習率と重み減衰は \(\rho\) と併せてチューニングする必要がある。いくつかの研究では、SAMはより大きなバッチサイズと学習率スケジュールの恩恵を受けることが示唆されている。適応的SAM（ASAM）などの適応的変種は、パラメータのスケールに基づいて半径を調整し、安定性を向上させる。
変種と拡張
元の論文以来、いくつかの改善が提案されている。適応的SAM（ASAM）は、摂動をパラメータの大きさで正規化してスケール不変性を達成し、CIFAR-10やImageNetなどのタスクでの性能を向上させる。Lookahead SAMは、Lookahead最適化器とSAMを組み合わせて分散を低減する。他の変種には、確率的重み平均（SWA）を用いたSharpness-Aware Minimizationや、ラベル平滑化を用いたSAMが含まれる。これらの手法は、計算コストの削減または汎化のさらなる向上を目的としている。
さまざまな分野での応用
SAMは、幅広いMachine learningタスクに成功裏に適用されている。コンピュータビジョンでは、CIFAR-10やImageNetなどのデータセットでResidual Network (ResNet)アーキテクチャの画像分類精度を向上させる。自然言語処理では、SAMはLarge language modelなどのトランスフォーマーのファインチューニングに役立ち、特にデータが乏しい場合に有効である。また、医用画像、創薬、強化学習でも使用されている。Generative AIでは、SAMは安定した画像生成のためのモデル訓練を支援する。その汎用性により、PyTorchやTensorFlowなどの人気ライブラリに実装され、学術研究と産業の両方で貴重なツールとなっている。
他の最適化手法との関係
SAMはAdam (Optimizer)やSGDの代替ではなく、損失ランドスケープを修正するアドオンである。これは、訓練を安定化させる勾配クリッピングなどの手法と比較されることが多いが、SAMは収束だけでなく汎化に対処する。また、Batch Normalizationと類似点があり、両者は異なるメカニズムを通じて損失曲面の幾何学に暗黙的に影響を与える。SAMはData AugmentationやDropoutと組み合わせてモデルをさらに正則化できる。平坦な極小値の概念は、BAIR (Berkeley AI Research)コミュニティによる損失ランドスケープ可視化やエントロピーベースの正則化に関する研究とも関連している。
理論的洞察
研究は、平坦な極小値がなぜ良好に汎化するかについて理論的正当性を提供している。過剰パラメータ化されたモデルでは、関数クラスの複雑性はしばしば極小値のシャープネスによって測定される。平坦な極小値は通常、過学習しない低複雑性の解に対応する。SAMのmin-max定式化は、モデルを最悪ケースの摂動に対して頑健にする敵対的訓練の一形態として解釈できる。いくつかの研究はSAMをベイズ推論と関連付け、パラメータ上の事後分布を近似することを示唆している。しかし、完全な理論的理解は未解決のままであり、活発な研究が続いている。
計算コストとトレードオフ
SAMの主な欠点は、更新ごとに2回の順伝播・逆伝播が必要なため、計算コストが2倍になることである。これは、非常に大規模なモデルやリアルタイムアプリケーションでは prohibitive となる可能性がある。しかし、精度の向上がオーバーヘッドを正当化することが多く、特に訓練時間よりも最終的なモデル品質が重要である場合に有効である。いくつかの変種は、摂動の更新頻度を減らす（例：数ステップごと）か、二次近似を使用することでコストを償却しようと試みている。標準と適応的変種の選択は、特定のユースケースに依存する。
産業での採用
主要なAI研究組織は、さまざまなアプリケーションでSAMを採用している。Google DeepMindは、強化学習と大規模訓練でSAMを探求している。OpenAIは、モデルの頑健性向上の文脈で平坦な極小値に言及している。AppleとSamsung Electronicsは、効率的な訓練とファインチューニングのためにオンデバイス機械学習でSAMを適用している。Amazon Web ServicesとGoogle Cloudは、クラウドベースの訓練用にSAM実装を含むライブラリを提供している。この手法は人気のある深層学習フレームワークにも統合されており、Kaggleコンペティションで広く使用されている。
他の正則化手法との比較
SAMは、重み減衰、ドロップアウト、Batch Normalizationなどの従来の正則化子と補完的である。これらの手法は個々のパラメータや活性化に作用するのに対し、SAMは損失ランドスケープ全体に作用する。SAMは、損失を増加させるように設計された摂動を用いる敵対的訓練の一形態と見なすことができる。これは、勾配の大きさのみを制限する勾配クリッピングとは異なる。SAMは、訓練サンプルの順序に焦点を当てるCurriculum Learningとも異なる。計算コストの点では、SAMは単純な正則化よりも高価であるが、精度の面でより良いリターンをもたらすことが多い。
制限と考慮事項
その利点にもかかわらず、SAMには制限がある。ステップごとの追加の順伝播・逆伝播により、訓練時間がおよそ2倍になり、大規模なモデルやデータセットでは prohibitive となる可能性がある。\(\rho\) の選択は常に簡単ではなく、タスクごとに慎重なチューニングが必要な場合がある。非常にノイズの多いラベルなどの一部の設定では、SAMは改善をもたらさず、むしろ有害である可能性もある。さらに、モデルがすでに十分に正則化されている場合、SAMの効果は顕著ではない。研究者は、摂動計算にデータのサブセットを使用するか、2つの勾配ステップを償却するなど、コストを削減する方法を探求している。
深層学習実践への影響
SAMは、深層学習実践者のツールボックスにおける標準的なツールとなっている。ドメインシフトが懸念される場合、事前訓練済みモデルのファインチューニングのデフォルト選択となることが多い。ビジョンベンチマークでの多くの最先端の結果は、訓練パイプラインのコンポーネントとしてSAMに依存している。産業では、SAMはApple、Intel、Nokia Bell Labsなどの企業によってさまざまなアプリケーションで採用されており、AWS TrainiumやGoogle CloudなどのクラウドAIプラットフォームに統合されている。この手法のシンプルさと有効性により、研究と本番環境の両方で人気のある選択肢となっている。
制限と課題
その成功にもかかわらず、SAMには制限がある。訓練時間の2倍化は、Natural language processingで使用されるTransformer (architecture)などの非常に大規模なモデルでは prohibitive となる可能性がある。摂動半径は慎重にチューニングする必要があり、最適な値はデータセットやアーキテクチャによって異なる場合がある。一部の設定では、特にベースモデルがすでに十分に正則化されている場合、SAMの標準的な手法に対する改善はわずかである。さらに、SAMの理論的保証はまだ不完全であり、時系列やグラフデータなどの他のモダリティでの性能はあまり探求されていない。進行中の研究はこれらの問題に対処し、より効率的な近似と適応的スキームを提案している。
将来の方向性
SAMの成功は、幾何学的観点から汎化を理解することへの関心を喚起している。研究者は、情報理論、PAC-Bayes、頑健性との関連を探求している。効率的なSAM変種は、Amazon Web ServicesやMicrosoft Azureなどのクラウド環境など、計算リソースが限られているオンデバイス訓練用に開発されている。SAMとTransformer (architecture)アーキテクチャおよびDeep learningフレームワークとの統合は、活発な研究分野である。モデルが大きくなるにつれて、追加データなしで汎化を改善するSAMの能力はますます価値が高まっている。
結論
Sharpness-Aware Minimizationは、良好な汎化の鍵である平坦な極小値を明示的にターゲットにすることで、ニューラルネットワークの最適化における重要な進歩を表している。そのシンプルでありながら強力なアイデア、すなわち最悪ケースの摂動に対して損失を最小化することは、多くの分野で採用され、精度と頑健性の向上に効果的であることが証明されている。計算コストにもかかわらず、SAMとその変種は現在、実践者のキットにおける標準的なツールであり、従来の最適化手法で達成可能なものを超えてモデル性能を押し上げる信頼性の高い方法を提供している。深層学習が進化し続けるにつれて、SAMのさらなる改良がさらに大きな利益をもたらす可能性がある。