[論文レビュー] Switch EMA: A Free Lunch for Better Flatness and Sharpness

Siyuan Li, Zicheng Liu+10 more|arXiv (Cornell University)|Feb 14, 2024

Soil Mechanics and Vehicle DynamicsEngineering被引用数 3

ひとこと要約

本論文は、各訓練エポックの終了時に主モデルに動的に指数移動平均（EMA）パラメータを切り替える、シンプルで効果的な手法であるSwitch EMA（SEMA）を提案する。速やかに収束する鋭い更新と、一般化性能を最適化した平坦なEMA状態を交互に切り替えることで、追加の計算コストなしに、視覚および言語タスクにおいて優れたモデル一般化性能と高速な収束を達成する。これは深層学習の訓練に即座に適用可能な「無料の昼食」を提供する。

ABSTRACT

Exponential Moving Average (EMA) is a widely used weight averaging (WA) regularization to learn flat optima for better generalizations without extra cost in deep neural network (DNN) optimization. Despite achieving better flatness, existing WA methods might fall into worse final performances or require extra test-time computations. This work unveils the full potential of EMA with a single line of modification, i.e., switching the EMA parameters to the original model after each epoch, dubbed as Switch EMA (SEMA). From both theoretical and empirical aspects, we demonstrate that SEMA can help DNNs to reach generalization optima that better trade-off between flatness and sharpness. To verify the effectiveness of SEMA, we conduct comparison experiments with discriminative, generative, and regression tasks on vision and language datasets, including image classification, self-supervised learning, object detection and segmentation, image generation, video prediction, attribute regression, and language modeling. Comprehensive results with popular optimizers and networks show that SEMA is a free lunch for DNN training by improving performances and boosting convergence speeds.

研究の動機と目的

• 既存の重み平均化（WA）手法が平坦性を向上させるが、最終的な性能を向上させないか、追加の推論コストを要するという限界を解消すること。
• 鋭さ（高速な収束）と平坦性（一般化）の利点を統合した、単一で効率的な訓練戦略を確立すること。
• 各エポック終了後にEMAパラメータを主モデルに切り替えるという最小限の変更を提案し、最適化ダイナミクスと最終的なモデル性能の両方を向上させること。
• SEMAが追加の計算コストや推論オーバーヘッドなしに、多様なタスクとアーキテクチャにおいて精度と収束速度を向上させることを実証すること。

提案手法

• SEMAは動的スイッチング機構を導入する：各エポックの終了時に、EMA平均化された重みが主モデルのパラメータにコピーされる。
• これにより、二段階の訓練ループが形成される。主モデルは鋭い損失領域を素早く探索し、EMAモデルはより平坦で一般化に優れた経路を追跡する。
• スイッチングプロセスにより、各新しいエポックが平坦で一般化に最適化された状態から開始され、より良い極小値への高速収束が可能になる。
• この方法は完全にプラグアンドプレイであり、既存のEMA実装にわずか1行のコード変更を加えるだけで可能である。
• SEMAは訓練中のみに作用するため、SWAやアンサンブルベース手法で一般的なテスト時の推論オーバーヘッドを回避する。
• 速やかな鋭い更新と安定した平坦なEMA状態の相補的強みを活用して、最適化ダイナミクスのバランスを図る。

実験結果

リサーチクエスチョン

• RQ1 EMAと主モデルの重み間での単純なスイッチング機構が、深層ニューラルネットワーク最適化における平坦性と鋭さの両方を向上させられるか？
• RQ2 各エポック終了後にEMAパラメータを主モデルに切り替えることで、標準的なEMAやベースライン訓練と比較して、より優れた一般化性能と高速な収束が達成できるか？
• RQ3 SEMAが追加の計算コストや推論オーバーヘッドなしに、性能向上を達成できるか？
• RQ4 多様なタスクにおいて、SWA、SAM、EMAといった既存手法と比較して、SEMAは最適化ダイナミクスと最終的なモデル精度の両面で優位性を示せるか？

主な発見

• SEMAは、画像分類、物体検出、自己教師付き学習、動画予測、言語モデリングを含む12の多様なベンチマークで一貫してテスト精度を向上させる。
• DeiT-Sを用いたImageNet-1Kでは、SEMAが84.2％のトップ-1精度を達成し、ベースライン（83.8％）とEMA（84.0％）を上回り、収束が速い。
• ResNet-50 Cascade Mask R-CNNを用いたCOCO物体検出では、SEMAがベースラインより0.7％、EMAより0.5％のAP向上を達成する。
• MoCo.V3を用いたCIFAR-100では、SEMAが線形プローブ精度を88.9％まで向上させ、EMA（88.3％）とベースライン（87.6％）を上回る。
• AgeDBにおける年齢推定回帰では、SEMAがEMAより0.4年、ベースラインより0.6年のMAE低減を達成する。
• SEMAは評価されたすべてのタスクで収束速度を加速させ、訓練曲線がEMAやベースラインモデルと比較して早期にかつ急激に性能向上を示す。