[論文レビュー] Exploring The Role of Mean Teachers in Self-supervised Masked Auto-Encoders

Youngwan Lee, Jeffrey Willette+3 more|arXiv (Cornell University)|Oct 5, 2022

Domain Adaptation and Few-Shot Learning被引用数 7

ひとこと要約

本稿では、特徴量の類似度に基づいて勾配を調整する条件付きモーメンタム正則化子として機能する指数移動平均（EMA）教師を統合することで、マスクドオートエンコーダー（MAE）を強化する自己教師強化学習手法RC-MAEを提案する。この手法は、MAEや最先端の自己蒸留手法と比較して、収束が速く、メモリ使用量が低減され、下流の視覚タスクにおけるロバスト性と性能が向上する。

ABSTRACT

Masked image modeling (MIM) has become a popular strategy for self-supervised learning~(SSL) of visual representations with Vision Transformers. A representative MIM model, the masked auto-encoder (MAE), randomly masks a subset of image patches and reconstructs the masked patches given the unmasked patches. Concurrently, many recent works in self-supervised learning utilize the student/teacher paradigm which provides the student with an additional target based on the output of a teacher composed of an exponential moving average (EMA) of previous students. Although common, relatively little is known about the dynamics of the interaction between the student and teacher. Through analysis on a simple linear model, we find that the teacher conditionally removes previous gradient directions based on feature similarities which effectively acts as a conditional momentum regularizer. From this analysis, we present a simple SSL method, the Reconstruction-Consistent Masked Auto-Encoder (RC-MAE) by adding an EMA teacher to MAE. We find that RC-MAE converges faster and requires less memory usage than state-of-the-art self-distillation methods during pre-training, which may provide a way to enhance the practicality of prohibitively expensive self-supervised learning of Vision Transformer models. Additionally, we show that RC-MAE achieves more robustness and better performance compared to MAE on downstream tasks such as ImageNet-1K classification, object detection, and instance segmentation.

研究の動機と目的

• EMAベースの蒸留における学生と教師の動的相互作用を理解すること。
• ビジョントランスフォーマーにおけるピクセルレベルのマスクドオートエンコーディング（MIM）において、平均教師の機能的役割を調査すること。
• 既存の自己蒸留手法と比較して、計算コストとメモリコストを低減するより効率的かつ効果的な自己教師型事前学習手法を開発すること。
• 物体検出やインスタンスセグメンテーションを含む分類および密度予測タスクにおける下流性能を向上させること。

提案手法

• 学生の再構築パッチに対する一貫性のターゲットを提供するために、EMA教師を統合した修正版MAE、すなわちRC-MAEを導入する。
• 学生の出力と元のマスク済みパッチを比較する再構築損失 $\mathcal{L}_{\text{r}}$ を使用する。
• 可視パッチにおける学生の予測と教師の出力を一致させる一貫性損失 $\mathcal{L}_{\text{c}}$ を適用する。
• 現在の入力と過去の入力間の特徴量類似度に基づいてモーメンタムを調整するEMアプデートルールを用いて、ゆっくりと更新される教師ネットワークを維持する。
• 線形モデルにおける勾配ダイナミクスを分析し、特徟能力が類似する場合に、教師が過去の勾配方向を除去することで勾配を条件付きで補正していることを示す。
• ViT-BおよびViT-Lバックボーンを用いて、ImageNet-1K、COCO、およびロバストネスベンチマークで本手法を実証的に検証する。

実験結果

リサーチクエスチョン

• RQ1 EMA教師はピクセルレベルのマスクドオートエンコーディングにおける勾配更新にどのように影響を与えるか？
• RQ2 ターゲットを提供する以外に、教師は学生の最適化プロセスにおいてどのような機能的役割を果たしているか？
• RQ3 MAEにEMA教師を単純に統合することで、既存の自己蒸留手法と比較して収束が速く、メモリ使用量が低減するか？
• RQ4 教師の条件付き勾配補正が、視覚タスクにおけるロバスト性と下流性能を向上させるか？
• RQ5 教師の効果は、条件付きモーメンタム正則化の一形態として理論的に説明可能か？

主な発見

• RC-MAEは800エポックでImageNet-1K分類タスクで83.4％のトップ-1正解率を達成し、MAEが1600エポックで達成する性能に相当するため、収束が著しく速いことが示された。
• RC-MAEは、教師に可視パッチのみを入力することで、iBOT、MSN、BootMAEといった最先端の自己蒸留手法と比較して、メモリ使用量と計算コストを低減した。
• ImageNet-C、ImageNet-A、ImageNet-R、ImageNet-Sketchの各ベンチマークにおいて、RC-MAEは一貫したロバストネスの向上を示し、mCEが1.8ポイント低下（49.9 vs. 51.7）し、IN-AおよびIN-Rでは最大1.2ポイントの正解率向上を達成した。
• RC-MAEのアテンションマップは、iBOT や MSN とは異なり、照準パッチの近傍に関連する画像領域に明確に焦点を合わせている。
• 理論的分析により、教師が現在の入力と過去の特徴表現の類似度に基づいて勾配を調整する条件付きモーメンタム正則化子として機能することが明らかになった。
• 線形モデルおよび深層ViTベースアーキテクチャにおける実証的検証により、条件付き勾配補正メカニズムの有効性が確認され、本研究で提案された分析を裏付ける結果が得られた。