[論文レビュー] Investigating the Benefits of Projection Head for Representation Learning

Yihao Xue, Eric Gan+3 more|arXiv (Cornell University)|Mar 18, 2024

Educational Games and GamificationPsychology被引用数 3

ひとこと要約

本論文は、対照的学習における表現学習を向上させるプロジェクションヘッドの理由を理論的に分析し、深層部の特徴の専門化を軽減し、分布シフト下でも頑健性を向上させることを示している。特徴の重み付けの暗黙的層別化により、プロジェクション前の表現がより一般化可能であることが明らかになった。また、学習可能なプロジェクションヘッドと同等の性能を示す固定再重み付けヘッドの提案により、より解釈可能な代替手段を提供している。

ABSTRACT

An effective technique for obtaining high-quality representations is adding a projection head on top of the encoder during training, then discarding it and using the pre-projection representations. Despite its proven practical effectiveness, the reason behind the success of this technique is poorly understood. The pre-projection representations are not directly optimized by the loss function, raising the question: what makes them better? In this work, we provide a rigorous theoretical answer to this question. We start by examining linear models trained with self-supervised contrastive loss. We reveal that the implicit bias of training algorithms leads to layer-wise progressive feature weighting, where features become increasingly unequal as we go deeper into the layers. Consequently, lower layers tend to have more normalized and less specialized representations. We theoretically characterize scenarios where such representations are more beneficial, highlighting the intricate interplay between data augmentation and input features. Additionally, we demonstrate that introducing non-linearity into the network allows lower layers to learn features that are completely absent in higher layers. Finally, we show how this mechanism improves the robustness in supervised contrastive learning and supervised learning. We empirically validate our results through various experiments on CIFAR-10/100, UrbanCars and shifted versions of ImageNet. We also introduce a potential alternative to projection head, which offers a more interpretable and controllable design.

研究の動機と目的

• 対照的自己教師付き学習におけるプロジェクションヘッドの経験的成功の背後にある理論的メカニズムを理解すること。
• 直接最適化されていないにもかかわらず、プロジェクション前の表現がより優れた下流性能を達成する理由を調査すること。
• データ拡張とネットワークの非線形性が特徴分布および一般化に与える影響を分析すること。
• 結果を教師あり対照的学習および教師あり学習に拡張し、頑健性の向上とクラス崩壊の低減を示すこと。
• 標準的なプロジェクションヘッドのより解釈可能な代替手段として、固定再重み付けヘッドの提案と検証を行うこと。

提案手法

• 対照的損失で訓練された線形モデルの理論的分析により、最適化バイアスに起因する層別進行的特徴重み付けが明らかになった。
• プロジェクション前の表現がより有益となる条件の導出、特にデータ拡張によって有用な特徴が歪められる場合に特に有効である。
• 非線形ネットワークへの分析の拡張により、低層が高層やプロジェクションヘッド出力に存在しない特徴を学習することが示された。
• CIFAR-10/100、UrbanCars、およびシフトされたImageNetバージョンを用いた実験的検証により、プロジェクション前・後の表現を比較した。
• 学習可能なプロジェクションヘッドの代替手段としての固定再重み付けヘッドの導入と評価を行い、同等の性能を達成した。
• 半実験的データセット（MNISTをCIFAR-10に適用）と、OODベンチマークを用いたImageNetでの微調整を用いて、頑健性をテストした。

Figure 1: Weights of features in a two-layer fully connected ReLU network trained with CL. Left : With all features having equal strength, those that are more disrupted by augmentation have smaller/zero weights. Right : With augmentation treating all features equally. features with the largest/small

実験結果

リサーチクエスチョン

• RQ1 直接最適化されていないにもかかわらず、なぜプロジェクションヘッド前の表現がその後の表現よりも一般化性が高いか？
• RQ2 最適化の暗黙的バイアスが、対照的学習における層ごとの進行的特徴重み付けをどのように引き起こすか？
• RQ3 データ拡張が有用な特徴を損なう場合や、特徴が強すぎる・弱すぎる場合に、プロジェクションヘッドが頑健性をどのように向上させるか？
• RQ4 低層がプロジェクションヘッド出力に存在しない特徴を学習できるか？ その影響は一般化にどのように現れるか？
• RQ5 固定再重み付けヘッドは学習可能なプロジェクションヘッドと同等の性能を達成できるか？ また、より解釈可能な代替手段として機能できるか？

主な発見

• データ拡張によって有用な特徴が歪められる場合、または特徴が強すぎる・弱すぎる場合、プロジェクション前の表現はプロジェクション後の表現よりも頑健で、より一般化性が高くなる。
• 最適化の暗黙的バイアスにより、層別進行的特徴重み付けが生じ、深層部が自己教師付き学習の目的により専門化するようになり、低層部はより正規化され、特徴の専門化が抑えられる。
• 非線形活性化は特徴の専門化を強化し、低層がプロジェクションヘッド出力に存在しない完全に新しい特徴を学習可能にし、一般化性の向上に寄与する。
• シフトされたImageNetベンチマークでは、プロジェクション前の表現がプロジェクション後や元のモデルよりも優れた分布外精度を達成し、背景バイアスに対する頑健性が向上していることが示された。
• 固定再重み付けヘッドは標準的なプロジェクションヘッドと同等の性能を達成しており、特徴再重み付けが本質的なメカニズムであるという理論的知見を裏付け、より解釈可能な設計の代替手段である可能性を示している。
• 教師あり対照的学習および教師あり学習において、プロジェクション前の表現はクラス崩壊を低減し、特にUrbanCarsおよびImageNetにおいて分布シフトへの少データ適応性を向上させた。

Figure 2: Left: Weights of features in a two-layer fully connected ReLU network trained using SCL. The subclass feature is not represented post-projection but is represented pre-projection. Right: As a result, the four subclasses are only separable in pre-projection representations.