確率的重み平均（SWA）は、深層学習におけるモデル平均化手法であり、ニューラルネットワークの訓練軌跡に沿った複数の点で重みを平均化することで、汎化性能を向上させる。最終的な重みのみを保持する従来の訓練とは異なり、SWAは訓練の後期段階で定期的な間隔で重みを収集し、その算術平均を計算する。この単純な手順は、損失景観においてより平坦な解をもたらすことが多く、これはテスト性能とロバスト性の向上に関連している。SWAは、Pavel Izmailov、Dmitrii Podoprikhin、Timur Garipov、Dmitry Vetrov、Andrew Gordon Wilsonによって、2018年の論文「Averaging Weights Leads to Wider Optima and Better Generalization」で導入された。この手法は、基盤となるモデルアーキテクチャや損失関数の変更を必要とせず、既存の訓練パイプラインへの統合が容易である。
SWAの核となる考え方は、確率的勾配降下法（SGD）またはその変種を用いた訓練中に、モデルパラメータが低損失の領域の周囲で振動するというものである。これらのパラメータを平均化することで、振動を平滑化し、広く平坦な最小値に位置する解を得ることができる。このような平坦な最小値は、鋭い最小値よりも汎化性能が高いと仮定されている。なぜなら、摂動や訓練分布とテスト分布の間のシフトに対して敏感でないためである。SWAは、画像分類、物体検出、言語モデリングなど、さまざまなアーキテクチャとタスクにわたって、精度とキャリブレーションを向上させることが示されている。
メカニズムとアルゴリズム
SWAは2つのフェーズで動作する。最初のフェーズでは、モデルは標準的な学習率スケジュール（例えば、コサインアニーリングやステップ減衰）で訓練され、最小値の近くの領域に到達する。2番目のフェーズでは、訓練は一定または周期的な学習率で継続され、重みは各エポックの終わりまたは固定数の反復後に収集される。SWAの重みは、実行平均として更新される：n_swa = n_swa + 1; w_swa = (w_swa * (n_swa - 1) + w) / n_swa。ここで、wは現在の重み、w_swaは平均化された重みである。最終モデルはw_swaを使用する。
重要な詳細として、バッチ正規化の統計量は平均化後に再計算する必要がある。なぜなら、バッチ正規化層の実行平均と分散は、重みと一緒に平均化されないためである。実際には、SWAの重みが計算された後、訓練データに対するフォワードパスが実行され、バッチ正規化の統計量が更新される。このステップは、ResNetなどのバッチ正規化を使用するモデルにとって重要である。
理論的動機
SWAの成功は、しばしば損失景観の幾何学の観点から説明される。Izmailovらによる研究とその後の研究は、SGDが低損失領域の境界近くの点に収束する傾向がある一方、複数の点を平均化することで解がその領域の中心に向かって移動することを示している。この中心点は通常、より平坦な盆地に位置し、これはより良い汎化性能と関連している。平坦な最小値と汎化の間の関連は、PAC-Bayesバウンドや損失関数の鋭さとの関連を含めて、広く研究されている。
別の観点は、SWAをベイズ推論に関連付ける。特定の仮定の下では、SGDの軌跡は、事後分布からサンプリングするマルコフ連鎖と見なすことができる。この軌跡に沿った平均化は、事後平均を近似し、これは二乗誤差損失に対するベイズ最適予測である。この解釈は、SWAを確率的勾配マルコフ連鎖モンテカルロ法と関連付けるが、SWAはより単純であり、ノイズスケールの注意深い調整を必要としない。
変種と拡張
SWAの性能をさらに向上させるために、いくつかの変種が提案されている。注目すべき拡張の1つは、確率的重み平均ガウス（SWAG）であり、2019年にWesley Maddox、Timur Garipov、Pavel Izmailov、Dmitry Vetrov、Andrew Gordon Wilsonによって導入された。SWAGは、収集された重みの1次および2次モーメントを使用してガウス分布を適合させることで、重みの事後分布を近似する。これにより、不確実性の推定とキャリブレーションの改善が可能になり、ベイズ深層学習に使用できる。
もう1つの変種は、高速幾何学的アンサンブリング（FGE）であり、周期的な学習率スケジュールを使用して、損失景観の異なるモードから重みを収集する。SWAはFGEと組み合わせて、さらに良い性能を達成することができる。さらに、SWAは学習率スケジュール、データ拡張、プルーニングなどの技術と統合され、ロバスト性と効率性を向上させている。
応用
SWAは、幅広い分野で成功裏に適用されている。コンピュータビジョンでは、CIFAR-10、CIFAR-100、ImageNetなどのデータセットで画像分類器の精度を向上させる。例えば、SWAは、アーキテクチャの変更なしに、CIFAR-100上のPreAct ResNet-164のテスト精度を81.15%から82.90%に向上させることが示されている。物体検出では、SWAは訓練の安定化と平均平均精度の向上に役立つ。
自然言語処理では、SWAはTransformerや大規模言語モデルをより効果的に訓練するために使用されている。言語モデリングタスクで過学習を減らし、パープレキシティを改善することが示されている。SWAはまた、生成モデル（GANや拡散モデルなど）にも利益をもたらし、より滑らかな重み更新を提供することで、より安定した訓練と高品質なサンプルを実現する。
標準的な教師あり学習を超えて、SWAは半教師あり学習、転移学習、ドメイン適応にも適用されている。その単純さと汎用性は、機械学習実践者のツールキットにおいて貴重なツールとなっている。
他の技術との比較
SWAは、しばしば他の正則化やアンサンブル手法と比較される。複数の独立して訓練されたモデルからの予測を平均化する従来のアンサンブルとは異なり、SWAは重みを平均化するため、計算コストが低く、単一の訓練実行のみを必要とする。これにより、モデルの保存や推論コストが懸念される場合に特に魅力的である。
ドロップアウトやバッチ正規化などの技術と比較して、SWAは訓練ダイナミクスを変更せず、重みを後処理するだけである。これらの手法と併用することができる。SWAはまた、Adamなどの適応的最適化器とも異なり、これらはパラメータごとの学習率を維持する。SWAは通常、SGDまたはAdamの軌跡の上に適用される。
SWAの1つの制限は、重みの収集を開始するタイミングを決定するスケジュールが必要であることである。収集が早すぎると、平均が最適でない重みに支配される可能性があり、遅すぎると利益が減少する。しかし、実際には、学習率が低い値に減衰した後に収集を開始するという単純な経験則が使用される。
実践的な考慮事項
SWAの実装は簡単である。PyTorchやTensorFlowなどのほとんどの深層学習フレームワークは、SWAのユーティリティや例を提供している。主なステップは次のとおりである：（1）標準的なスケジュールでモデルを訓練する、（2）一定または周期的な学習率で訓練を継続する、（3）間隔で重みを収集する、（4）実行平均を計算する、（5）バッチ正規化の統計量を再計算する。
SWAは、ハイパーパラメータの選択に対してロバストであることが示されている。2番目のフェーズの学習率は、通常、初期スケジュールの最小値または小さな定数に設定される。重みの収集頻度は、各エポックまたは数百回の反復ごとにすることができる。実際には、5から20エポックの間で重みを収集することがしばしば十分である。
影響と評価
SWAの導入は、深層学習コミュニティに大きな影響を与えた。これは、推論中の追加の計算コストなしに汎化性能を向上させる、単純でありながら効果的な方法を提供した。この論文は広く引用され、重み平均化と平坦な最小値に関する多くの後続研究に影響を与えた。SWAは現在、特に学術研究や競技会において、多くの訓練パイプラインで標準的な技術となっている。
SWAはまた、生成モデルや大規模言語モデルの訓練で一般的に使用される指数移動平均（EMA）など、他の平均化手法の開発にも影響を与えた。EMAは、最近の重みにより多くの重みを与える変種であり、その単純さと効果性から実際によく使用される。SWAとEMAは補完的であり、一部のフレームワークではそれらの切り替えが可能である。
制限と将来の方向性
その利点にもかかわらず、SWAには制限がある。損失景観が比較的滑らかであり、重みの平均化が意味を持つことを前提としており、これはすべてのアーキテクチャやタスクに当てはまるとは限らない。例えば、バッチ正規化を使用するモデルは注意深い処理を必要とし、複雑なパラメータ空間を持つモデル（例えば、リカレントネットワーク）はそれほど利益を得られない可能性がある。さらに、SWAは拡張（例えば、SWAG）なしでは不確実性の推定を提供しない。
将来の研究方向には、適応的な平均化スケジュールの開発、SWAとベイズ法の組み合わせ、強化学習やフェデレーテッドラーニングなどの新しい領域へのSWAの適用が含まれる。深層学習が進化し続けるにつれて、SWAのような重み平均化技術は、モデルの性能と信頼性を向上させるための基本的なツールであり続ける。