[論文レビュー] The Mechanism of Prediction Head in Non-contrastive Self-supervised Learning

Zixin Wen, Yuanzhi Li|arXiv (Cornell University)|May 12, 2022

Stochastic Gradient Optimization Techniques被引用数 6

ひとこと要約

この論文は、非対照的自己教師あり学習における次元的崩壊を防ぐトレーナブルな予測ヘッドのメカニズムを解明している。実験的および理論的分析を通じて、置換効果と加速効果という2つの主要な効果を特定した。予測ヘッドにより、強い特徴が弱い特徴の代わりに使用可能になり、学習が促進されることで、損失関数の最適解に存在する自明な崩壊解を避けることができる。

ABSTRACT

Recently the surprising discovery of the Bootstrap Your Own Latent (BYOL) method by Grill et al. shows the negative term in contrastive loss can be removed if we add the so-called prediction head to the network. This initiated the research of non-contrastive self-supervised learning. It is mysterious why even when there exist trivial collapsed global optimal solutions, neural networks trained by (stochastic) gradient descent can still learn competitive representations. This phenomenon is a typical example of implicit bias in deep learning and remains little understood. In this work, we present our empirical and theoretical discoveries on non-contrastive self-supervised learning. Empirically, we find that when the prediction head is initialized as an identity matrix with only its off-diagonal entries being trainable, the network can learn competitive representations even though the trivial optima still exist in the training objective. Theoretically, we present a framework to understand the behavior of the trainable, but identity-initialized prediction head. Under a simple setting, we characterized the substitution effect and acceleration effect of the prediction head. The substitution effect happens when learning the stronger features in some neurons can substitute for learning these features in other neurons through updating the prediction head. And the acceleration effect happens when the substituted features can accelerate the learning of other weaker features to prevent them from being ignored. These two effects enable the neural networks to learn all the features rather than focus only on learning the stronger features, which is likely the cause of the dimensional collapse phenomenon. To the best of our knowledge, this is also the first end-to-end optimization guarantee for non-contrastive methods using nonlinear neural networks with a trainable prediction head and normalization.

研究の動機と目的

• 確率的勾配降下法で学習するニューラルネットワークが、自明な崩壊解が全局最適であるにもかかわらず、非対照的自己教師あり学習でなぜ次元的崩壊を回避するのかを理解すること。
• 予測ヘッドが対称性を破り、多様で崩壊しない表現の学習を可能にする役割を調査すること。
• 非対照的対照的学習に類似した訓練における予測ヘッドがもたらす暗黙のバイアスを説明する理論的枠組みを提供すること。
• 特に非対角成分の上昇と下降の動きを特徴づける、アイデンティティ初期化された予測ヘッドを有する状況下での特徴学習のダイナミクスを同定すること。

提案手法

• 予測ヘッドを恒等行列として初期化し、非対角成分のみを更新するという設定で、トレーニングダイナミクスを実験的に調査する。
• 学習の強い特徴の獲得、置換効果、加速効果、収束の4段階に分けたトレーニングプロセスの理論的モデル化を行う。
• 置換効果を定義・分析する。強いニューロンの学習が、予測ヘッドの更新を通じて他のニューロンの学習を代替可能にする。
• 加速効果を定義・分析する。代替された特徴が、弱い特徴の学習を促進し、最適化過程で無視されないようにする。
• 非線形ニューラルネットワークにトレーナブルな予測ヘッドとバッチ正則化を導入した簡略化された線形モデルを用いて最適化の保証を導出する。
• テンソルパワー法の境界を活用し、学習率や初期化条件の変化に伴う勾配更新と収束の分析を行う。

実験結果

リサーチクエスチョン

• RQ1 確率的勾配降下法で学習するニューラルネットワークが、自明な崩壊解が全局最適であるにもかかわらず、非対照的自己教師あり学習でなぜ次元的崩壊を回避するのか？
• RQ2 予測ヘッドが恒等写像を学習してしまう可能性がある中で、なぜ多様な特徴の学習が可能になるのか？
• RQ3 予測ヘッドの非対角成分はトレーニング中にどのような役割を果たし、なぜ一貫した上昇と下降のトレースを示すのか？
• RQ4 予測ヘッドの置換効果と加速効果がどのように組み合わさって特徴の忘却を防ぎ、完全な表現学習を可能にするのか？
• RQ5 トレーナブルな予測ヘッドと正則化を備えた非対照的自己教師あり学習に対して、エンド・ツー・エンドの最適化保証を提供できるか？

主な発見

• 予測ヘッドを恒等行列として初期化し、非対角成分のみを学習対象とする場合でも、ネットワークは競争力のある表現を学習する。これは、ヘッドの構造そのものが、単なるパラメータとは別に重要であることを示している。
• トレーニング中に予測ヘッドの非対角成分に一貫した上昇と下降のトレースが観察され、動的な特徴の置換が起こっている可能性を示唆している。
• 置換効果とは、あるニューロンの強い特徴が、予測ヘッドの更新を通じて他のニューロンの学習を代替可能になる現象である。
• 加速効果とは、代替された特徴が弱い特徴の学習を促進し、最適化過程で無視されないようにする現象である。
• 置換効果と加速効果が併存することで、ネットワークは最も強い特徴にのみ集中せず、すべての特徴を学習可能となり、これが次元的崩壊の回避に寄与している。
• 本論文は、非対照的自己教師あり学習において、トレーナブルな予測ヘッドと正則化を備えた非線形ニューラルネットワークに対する、初めてのエンド・ツー・エンドの最適化保証を提供している。