[論文レビュー] Learning and Leveraging World Models in Visual Representation Learning

Quentin Garrido, Mahmoud Assran+4 more|arXiv (Cornell University)|Mar 1, 2024

Advanced Image and Video Retrieval Techniques被引用数 7

ひとこと要約

Image World Models (IWM) は Joint Embedding Predictive Architecture に基づいて、変換の影響を予測する再利用可能な潜在世界モデルを学習することを紹介する。予測子のファインチューニングと、対照学習型とマスクド画像モデリング型アプローチを結ぶ、制御可能な表現抽象化を示す。

ABSTRACT

Joint-Embedding Predictive Architecture (JEPA) has emerged as a promising self-supervised approach that learns by leveraging a world model. While previously limited to predicting missing parts of an input, we explore how to generalize the JEPA prediction task to a broader set of corruptions. We introduce Image World Models, an approach that goes beyond masked image modeling and learns to predict the effect of global photometric transformations in latent space. We study the recipe of learning performant IWMs and show that it relies on three key aspects: conditioning, prediction difficulty, and capacity. Additionally, we show that the predictive world model learned by IWM can be adapted through finetuning to solve diverse tasks; a fine-tuned IWM world model matches or surpasses the performance of previous self-supervised methods. Finally, we show that learning with an IWM allows one to control the abstraction level of the learned representations, learning invariant representations such as contrastive methods, or equivariant representations such as masked image modelling.

研究の動機と目的

• JEPA フレームワーク内で再利用可能な Image World Model (IWM) を学習する方法を検討する。
• IWM を成功させる鍵となる要因を特定する：条件付け、変換の複雑さ、予測子の能力。
• 下流の識別タスクのファインチューニングプロトコルとマルチタスク効率を実証する。
• 世界モデルの容量が表現抽象化（不変性 vs 同時変換性）に与える影響を示す。

提案手法

• JEPA を拡張して、変換された表現を予測する潜在空間の世界モデルを学習する。
• ソース x とターゲット y を選択したオーグメンテーション下で訓練する；予測子 p_phi は z_x から z_y を一致させ、L2 ロスを用いる。
• 予測子を変換情報で条件付ける（系列条件付けまたは特徴条件付けを介して、デフォルトは特徴条件付け）。
• Augmented targets の銀行に対して Mean Reciprocal Rank (MRR) で世界モデルの品質を評価する。
• predictor ファインチューニング（エンコーダーのファインチューニングと比較）とマルチタスクファインチューニングを通じた下流転送を評価する。 不変モードと同時変換モードを研究する。
• 世界モデルの容量と変換強度がパフォーマンスと表現抽象化に与える影響を特徴づける。

実験結果

リサーチクエスチョン

• RQ1 学習された潜在世界モデルを再利用して、下流の識別的ビジョンタスクを改善できるか。
• RQ2 predictor conditioning、変換の複雑さ、モデル容量が Image World Models の品質と有用性にどう影響するか。
• RQ3 世界モデルにおける不変性と同時変換性が下流の性能と表現抽象化レベルにどう影響するか。
• RQ4 IWMs は複数のビジョンタスクにまたがる効率的なマルチタスクファインチューニングを可能にするか。

主な発見

• 予測子を変換情報で条件付けることは不可欠であり、条件付けがない場合の MRR は 0 であるのに対し、系列条件付けまたは特徴条件付けは高い MRR（約 0.8）をもたらす。
• より強力で複雑な変換と予測子の深さの増加は、世界モデルの忠実度を高める（MRR が高くなる）。
• 同時変換性を持つ IWMs は予測子ファインチューニングをより効果的に行え、ランダム予測子に対して顕著な利得があり、多くの設定でエンコーダーのファインチューニング効率と同等以上を達成できる。
• 不変 IWMs は線形評価でより良い傾向を示す一方、同時変換性 IWMs は予測子のファインチューニング時およびマルチタスクの場面で優れる。
• IWMs による predictor ファインチューニングはエンコーダーのファインチューニングよりパラメータ効率が高く、マルチタスクファインチューニングはタスク間の効率向上を示す。
• IWMs は対照的な類似（不変）と MIM 的（同時変換性）な表現抽象化のスペクトラムを提供し、制御可能なトレードオフを実現する。