[論文レビュー] On the Pros and Cons of Momentum Encoder in Self-Supervised Visual Representation Learning

Trung Pham, Chaoning Zhang+3 more|arXiv (Cornell University)|Aug 11, 2022

Digital Imaging for Blood Diseases被引用数 5

ひとこと要約

この論文は、自己教師あり視覚表現学習における運動量エンコーダーの役割を調査し、完全なネットワークに運動量を適用するのではなく、プロジェクトャのみに運動量を適用する代替手法を提案する。深層層における勾配不安定性を分析することで、エンコーダ全体に運動量を適用するのではなく、最終的なプロジェクトャにのみ運動量を適用することで、性能の向上を維持しながら、バックボーンでの二重順方向伝搬を回避することで計算コストを削減できることが示された。

ABSTRACT

Exponential Moving Average (EMA or momentum) is widely used in modern self-supervised learning (SSL) approaches, such as MoCo, for enhancing performance. We demonstrate that such momentum can also be plugged into momentum-free SSL frameworks, such as SimCLR, for a performance boost. Despite its wide use as a fundamental component in modern SSL frameworks, the benefit caused by momentum is not well understood. We find that its success can be at least partly attributed to the stability effect. In the first attempt, we analyze how EMA affects each part of the encoder and reveal that the portion near the encoder's input plays an insignificant role while the latter parts have much more influence. By monitoring the gradient of the overall loss with respect to the output of each block in the encoder, we observe that the final layers tend to fluctuate much more than other layers during backpropagation, i.e. less stability. Interestingly, we show that using EMA to the final part of the SSL encoder, i.e. projector, instead of the whole deep network encoder can give comparable or preferable performance. Our proposed projector-only momentum helps maintain the benefit of EMA but avoids the double forward computation.

研究の動機と目的

• 既存の動機を超えて、自己教師あり学習における運動量エンコーダーの実証的利点を理解すること。
• 運動量による性能向上が、特に深層ネットワーク層において安定性に起因するかどうかを調査すること。
• 運動量の計算オーバーヘッドを軽減するため、運動量をエンコーダ全体ではなくプロジェクトャに限定すること。
• プロジェクトャのみの運動量が、顕著に削減された順方向計算量で、同等またはより優れた性能を達成できることを検証すること。

提案手法

• バックボーンの個々のブロックに運動量エンコーダーを分解し、各層における勾配の揺らぎを分析する。
• 各ブロックの出力に対する全体損失の勾配をモニタリングし、不安定性を定量化する。
• 最終的なMLPプロジェクトャのみがEMAを適用するプロジェクトャのみの運動量機構を提案する。バックボーンの重みは共有される。
• 最終層（ブロック4およびプロジェクトャ）にバッチ正規化（BN）を適用し、さらに訓練の安定性を向上させる。
• CIFAR-10およびCIFAR-100で、完全なEMA（例：BYOL）、プロジェクトャのみのEMA、およびベースラインモデルとの間で線形評価精度を比較する。
• 訓練ステップごとのフィルタ重みを可視化し、EMA有無における安定性の違いを示す。

実験結果

リサーチクエスチョン

• RQ1 自己教師あり学習における運動量エンコーダーの性能向上の主なメカニズムは何ですか？
• RQ2 運動量の影響は、エンコーダーのネットワークの異なる段階でどのように変化しますか？
• RQ3 計算コストを削減しながら、運動量の性能向上効果を維持することは可能ですか？
• RQ4 性能を損なわずに、運動量をプロジェクトャに限定することは可能ですか？
• RQ5 バッチ正規化は、プロジェクトャのみの運動量とどのように作用し、訓練の安定性を向上させるのですか？

主な発見

• 運動量による性能向上の主な要因は、エンコーダーの最終層における訓練の安定性の向上に起因している。
• 勾配の揺らぎは、特に最終層（例：ブロック4およびプロジェクトャ）で顕著に高いため、安定性が低く、運動量がその緩和に寄与していることが示された。
• プロジェクトャのみの運動量は、CIFAR-100で88.45％の線形評価精度を達成し、ベースラインのBYOLと同等の性能を維持しながら、バックボーンでの二重順方向伝搬を回避した。
• ブロック4の出力およびプロジェクトャにバッチ正規化を適用することで、プロジェクトャのみのEMAの性能がさらに向上し、88.45％の精度を達成した。これは完全なBYOLベースラインとほぼ同等の結果であった。
• 可視化結果から、運動量がない場合、フィルタは不安定なスパイク（例：約10,000ステップ目）を示す一方、EMAを適用すると一貫して滑らかな学習曲線が得られた。
• 初期エンコーダーステージに運動量を適用しても、効果は最小限にとどまり、最終層が主に利益をもたらすことが確認された。