【DL輪読会】Self-Supervised Learning from Images with a Joint-Embedding Predictive Architecture

DEEP LEARNING JP [DL Papers] Self-Supervised Learning from Images with a Joint-Embedding Predictive Architecture 小林 範久 Present Square Co.,Ltd. http://deeplearning.jp/ 1

• http://deeplearning.jp/

書誌情報 タイトル： Self-Supervised Learning from Images with a Joint-Embedding Predictive Architecture https://arxiv.org/abs/2301.08243 著者： Mahmoud Assran, Quentin Duval, Ishan Misra, Piotr Bojanowski, Pascal Vincent, Michael Rabbat, Yann LeCun, Nicolas Ballas 概要： • 画像からの自己教師あり学習のための非生成的アプローチであるI-JEPAを提案。 • I-JEPAは、一つのコンテキストブロックから、同じ画像内の様々なターゲットブロックの表現を予測する。 • (a)十分に大きなスケールでターゲットブロックをサンプリングし、(b)十分に情報量の多い（空間的に分散し た）コンテキストブロックを使用することが重要。 • スケール性があり、線形分類からオブジェクトカウント、深度予測までの幅広いタスクで高い精度を達成。 Copyright (C) Present Square Co., Ltd. All Rights Reserved. 2

• https://arxiv.org/abs/2301.08243

書誌情報 JEPAを画像分野に適用 Copyright (C) Present Square Co., Ltd. All Rights Reserved. 出典：A Path Towards Autonomous Machine Intelligence, Yann LeCun 3

アジェンダ 1. イントロダクション 2. 関連研究 3. 手法 4. 実験 5. 結論 Copyright (C) Present Square Co., Ltd. All Rights Reserved. 4

１. イントロダクション イントロダクション • 自己教師あり学習には一般的に2つのアプローチが存在する：不変性ベースと生成的な方法。 不変性ベースの方法（invariance-based methods） • 画像の同一性を維持することを目指し、一方生成的な方法はマスクされた部分を予測することにより学習を進める。高 い意味レベルの表現を生成できるが、異なるタスクや異なるモダリティ（例えば音声）への一般化が困難である。 生成的な手法（generative methods） • 特にマスクノイズアプローチと呼ばれるもので、ピクセルレベルまたはトークンレベルでの予測を行う。この手法は、画像モダリ ティを超えて容易に一般化できるが、得られる表現の意味レベルが低いという課題がある。 これらの手法の利点を最大限に活用するためには、より複雑な適応メカニズム（例えばエンドツーエンドの微調整）が必要と なり、そのためにはさらなる研究と技術開発が求められる。 Copyright (C) Present Square Co., Ltd. All Rights Reserved. 5

１. イントロダクション イントロダクション • 本研究では、画像変換によるエンコードの過度な事前知識を避け、自己教師あり学習での表現の意味レベルを向上 させる方法を提案。その具体的な手法として「共同埋め込み予測アーキテクチャ（I-JEPA）」を提案。 I-JEPAの基本的なアプローチ • 抽象的な表現空間における欠落情報を予測すること。具体的には、特定のコンテキストブロックが与えられた場合に、 同じ画像内の様々なターゲットブロックの表現を予測する。 • 既存のピクセル/トークン空間での予測による生成的手法と比べ、I-JEPAは抽象的な予測ターゲットを利用する。これ により、不必要なピクセルレベルの詳細を排除し、モデルがより多くの意味的特徴を学習する可能性を提供している。 • マルチブロックマスキング戦略を提案。情報量の多い（空間的に分布した）コンテキストブロックを用いて、画像中の 十分に大きなターゲットブロックを予測することの重要性を強調している。 Copyright (C) Present Square Co., Ltd. All Rights Reserved. 6

２. 関連研究 アーキテクチャの比較 • 互換性のある信号x、yに対して は、同様の埋め込みを出力する よう学習し、一方、互換性のない 入力に対しては異なる埋め込み を生成するように学習する。 Copyright (C) Present Square Co., Ltd. All Rights Reserved. • 互換性のある信号xから、信号yを 直接再構成するように学習する。 • この再構成をより容易にするため、 追加の（潜在的な）変数zを基 にデコーダネットワークを使用する。 • 互換性のある信号xから信号yの埋 め込みを予測するように学習する。 • 予測をより容易にするために、追加 の（潜在的な）変数zを基に予測 ネットワークを使用する。 7

２. 関連研究 関連研究 • 長年にわたり視覚表現学習は、欠損または破損した感覚入力の値を予測することによって探求されてきた。ノイズ 除去オートエンコーダー、コンテキストエンコーダー、画像の着色をノイズ除去タスクとして扱う手法などがある。 • 最近では、ViTを使用して欠損した入力パッチを再構成するマスクされた画像モデリングの文脈で、画像のノイズ除去 のアイデアが再評価されている。MAEは、エンコーダーが可視画像パッチのみを処理することを必要とする効率的なア ーキテクチャを提案している。また、BEiTはトークン化された空間で欠損パッチの値を予測する。 • 他の研究では、SimMIMは古典的なHistogram of Gradients特徴空間に基づく再構成目標を探求し、ピクセ ル空間の再構成よりもいくつかの優位性を示している。 ⇒ JEPAは、ダウンストリームタスクでの広範なファインチューニングを必要としない意味的な表現を学習することを目指している。 Copyright (C) Present Square Co., Ltd. All Rights Reserved. 8

２. 関連研究 関連研究 data2vec 、Context Autoencoders（CAE） • 本研究に最も近い手法である。共に欠損パッチの表現の予測に焦点を当てた手法である。 • data2vecはオンラインターゲットエンコーダを通じて表現を学び、CAEは再構成損失とアライメント制約の和を最適化して 表現を学ぶ。 ⇒これらの手法と比較して、I-JEPAは計算効率の大幅な改善を示し、より意味的な既成表現を学習することができる。 DINO、MSN、iBOT • 結合埋め込みアーキテクチャに基づく手法は、事前学習時に意味的な画像表現を学習するために、手動によるデータ拡 張に依存している。特に、MSNは事前学習時にマスキングを追加的なデータ拡張として使用し、一方iBOTは data2vecスタイルのパッチレベル再構成損失とview-invariance lossを組み合わせている。 • これらの手法は各入力画像の複数のユーザ生成ビューを処理する必要があるため、スケーラビリティが阻害されるという共 通の問題を抱えている。 ⇒ I-JEPAは各画像の単一のビューを処理するだけでよく、I-JEPAで学習したViT-Huge/14は、iBOTで学習したViTSmall/16よりも計算量が少なくて済む利点がある。 Copyright (C) Present Square Co., Ltd. All Rights Reserved. 9

３. 手法 I-JEPA（Image-based Joint-Embedding Predictive Architecture） • 与えられたコンテキストブロックを用いて、同一画像内の複数 （デフォルトは4つ）のターゲットブロックの表現を予測する。 • コンテキストエンコーダ、ターゲットエンコーダ、予測器には、 ViTを使用。MAEに類似しているが、I-JEPA法は非生成 的であり、予測は表現空間内で行われる。 • 予測器は、コンテキストエンコーダの出力を受け取り、位置ト ークン（色で示される）を条件にして、特定の位置のターゲ ットブロックの表現を予測する。ターゲットの表現はターゲット エンコーダの出力に対応し、その重みはコンテキストエンコーダ の重みの指数移動平均によって更新。 Copyright (C) Present Square Co., Ltd. All Rights Reserved. 10

３. 手法 JEPA と I-JEPA JEPA Copyright (C) Present Square Co., Ltd. All Rights Reserved. I-JEPA 11

３. 手法 Context • I-JEPAは、(0.85, 1.0)の範囲のランダムなスケールと単位ア スペクト比を用いて、画像から1つのブロック 𝑥 をサンプリングする。 これにより、コンテキストを得るための初期情報が生成される。 • コンテキストブロックから重複する領域を除去することで、自明で ない予測タスクが保証される。 𝑥 𝑠𝑥 • その後、マスクされたコンテキストブロック 𝑥 はコンテキストエンコ ーダ 𝑓𝜃 により処理され、対応するパッチレベル表現 𝑠𝑥 = {𝑠𝑥𝑗 }𝑗∈𝐵𝑖 が得られる。 Copyright (C) Present Square Co., Ltd. All Rights Reserved. 𝑠𝑦 12

３. 手法 Targets • I-JEPAは、ターゲットを画像ブロックの表現として対応さ せる。入力画像はN個の非重複パッチのシーケンスに変 換され、それらはターゲット・エンコーダ𝑓𝜃ഥ によってパッチレ ベル表現𝑠𝑦 に変換される。 • この変換により、各パッチは自身の一意の表現𝑠𝑦𝑘 を得る。 • ターゲットの表現𝑠𝑦 から 𝑀 個のブロック（重複可能）を ランダムにサンプリングし、損失のターゲットを確立する。 通常、𝑀は4と設定され、アスペクト比（0.75〜1.5） とスケール（0.15〜0.2）の範囲でブロックをサンプリン グする。特徴は、ターゲット・ブロックが入力をマスクするの ではなく、ターゲット・エンコーダの出力をマスクすることで 生成されること。 𝑦 Copyright (C) Present Square Co., Ltd. All Rights Reserved. 𝑠𝑦 13

３. 手法 コンテキストとターゲット・マスキング戦略の例 • 与えられた画像に対して、特定のスケールとアスペク ト比の範囲内でターゲットブロックをランダムに4つサン プリングし、さらに異なるスケール範囲でコンテキストブ ロックをサンプリングする。 • この戦略により、ターゲットブロックは意味的な情報を 持ち、一方、コンテキストブロックは豊富な情報を提 供し、効率的に処理可能な疎な表現を保証する。 Copyright (C) Present Square Co., Ltd. All Rights Reserved. 14

３. 手法 Prediction • コンテキスト・エンコーダの出力 𝑠𝑥 から𝑀個のターゲット・ブロック表現 𝑠𝑦 (1), ... , 𝑠𝑦 (𝑀) を予測する。これは予測器 𝑔Φ (・, ・) を用いて行われ、各パッチのマスク トークンとコンテキスト・エンコーダ 𝑠𝑥 の出力を入力として、パッチレベルの予測を生 成する。 • マスクトークンは、位置埋め込みが追加された共有学習可能なベクトルによってパラ メータ化される。これにより、ターゲットブロックの予測は𝑀回行われ、それぞれの予測 は対応するマスク・トークンを条件として行われる。これにより、ターゲットブロックの予 測𝑠Ƹ𝑦 (1), … , 𝑠Ƹ𝑦 (𝑀)が得られる。 Loss 𝑠𝑦Ƹ (1) 𝑠Ƹ𝑦 (2) 𝑠Ƹ𝑦 (3) 𝑠𝑦 (1) 𝑠𝑦 (2) 𝑠𝑦 (3) • 損失は、予測されたパッチレベル表現 sො 𝑦 𝑖 とターゲットパッチレベル表現 𝑠𝑦 (𝑖) の間の平均L2距離。 • • 予測器 ϕ とコンテキストエンコーダ 𝜃 のパラメータは、勾配ベースの最適化によって学習される。 ターゲットエンコーダ θത のパラメータは、コンテキストエンコーダパラメータの指数移動平均によって更新される。 Copyright (C) Present Square Co., Ltd. All Rights Reserved. 15

４. 実験：画像分類 ImageNet を用いた評価 • I-JEPAは、手動でのビューデータの拡張に頼らない他の 手法と比較して、ImageNet-1k（ViT H/16448は 448×448の解像度で事前学習済み）における線形 評価の性能を向上させた。 • また、この手法は優れたスケーラビリティを持つため、より 大規模なI-JEPAモデルはビューデータの拡張を必要とせ ずに、ビュー不変性のアプローチと性能が一致する。 Copyright (C) Present Square Co., Ltd. All Rights Reserved. 16

４. 実験：画像分類 ImageNet-1K 1％ を用いた評価 • ViT H/14アーキテクチャを用いたI-JEPAは、data2vecで事前 学習されたViT-L/16と同等の性能を示す。その計算量は大幅に 少ない。 • I-JEPAは解像度を上げることで、MSN、DINO、iBOTなど、 事前学習中に手動でデータ拡張を追加する結合埋め込み手法 を上回る結果を示した。 • 利用可能なラベルの1%だけを使用したImageNet-1Kの半教 師あり評価では、I-JEPAの事前学習は手動で作成したデータ拡 張に頼らないMAEを上回る性能を示し、さらにスケールの恩恵を 受ける。 • 具体的には、解像度448で訓練されたViT H/16は、手動による データ拡張を使用する従来の手法を凌駕した。 Copyright (C) Present Square Co., Ltd. All Rights Reserved. 17

４. 実験：画像分類 Linear-probe transfer • 線形プローブを使用した各種のダウンストリーム画像分類タ スクでの性能を示す。その結果、I-JEPAはデータ拡張を用 いない従来手法（MAE と data2vec）を大幅に上回 る性能を示した。 • さらに、I-JEPAはデータ拡張を活用する最良のビュー不変 性ベースの手法との差を縮めた。線形プローブを使用する CIFAR100とPlace205では、DINOを上回る結果を得た。 Copyright (C) Present Square Co., Ltd. All Rights Reserved. 18

４. 実験：Local Prediction Task Local Prediction task • I-JEPAの性能が、低レベルのタスクであるClevrデータセットの 物体カウントと奥行き予測において、ビュー不変性ベースの手 法（DINOやiBOT等）を上回っていることが示されている。 • これはI-JEPAが事前学習中に低レベルの画像特徴を効果的 に捉え、その結果として低レベルで高密度な予測タスクに優れ た性能を発揮することを確認するものである。 Copyright (C) Present Square Co., Ltd. All Rights Reserved. 19

４. 実験：Scalability モデルの効率性 • I-JEPAは以前の手法よりも少ない計算量で、データ拡張に頼 らずに高い性能を達成している。具体的には、再構成に基づく 手法（例：MAE）と比較して、I-JEPAは表現空間でターゲッ トを計算することにより余分なオーバーヘッドを導入するものの、 約5倍少ない反復で収束し、大幅に計算量を削減できる。 • さらに、データ拡張に依存し、各画像の複数のビューを作成・処 理するビュー不変性に基づく手法（例：iBOT）と比較すると、 I-JEPAは著しく高速に実行される。特に、大きなI-JEPAモデル （ViT-H/14）は、小さなiBOTモデル（ViT-S/16）よりも少 ない計算量で済む。 Copyright (C) Present Square Co., Ltd. All Rights Reserved. 20

４. 実験：Scalability データサイズによるスケーリング • 事前学習データセットのサイズを増やすと、意味的なタスクと低レベルのタスクの両方で転移学習の性能が向上することが 示されている。これは、より大きく多様なデータセットでの事前学習が、様々なタスクの学習において有益であることを示してい る。 モデルサイズによるスケーリング • IN22Kでのプレトレーニング時に、大きなモデルサイズ（具体的には、ViT-G/16）が有効であることを示している。 • ViT-G/16をプレトレーニングすると、ViT-H/14モデルに比べて画像分類タスク（例えば、Place205やINat18）の下 流性能が大幅に向上する。しかし、この大きなモデルサイズは低レベルのタスクの性能を向上させない。これは、ViTG/16がより大きな入力パッチを使用するため、局所的な予測タスクに対しては不利である可能性があるため。 Copyright (C) Present Square Co., Ltd. All Rights Reserved. 21

４. 実験：Predictor Visualizations 予測器の学習効果を評価 • 予測器がターゲットの位置の不確実性を正確に捉えるかどうかを調査するために、事前学習後に、予測器とコンテキスト・エ ンコーダの重みを固定し、予測器の出力の平均プーリングをピクセル空間にマップバックするために、RCDM （Representation Conditional Diffusion Model）フレームワークに基づいてデコーダを訓練する。 • 予測器の出力の視覚化は、予測器が位置の不確実性を正確に捕捉し、高レベルのオブジェクト部分（例えば、鳥の背 中や車の上部）を正確に生成できることを示している。しかし、正確な低レベルの詳細と背景情報を破棄する傾向がある。 Copyright (C) Present Square Co., Ltd. All Rights Reserved. 22

４. 実験：Ablations マスキング戦略の比較 • マルチブロックマスキングを、画像を4つの大きな象限に分割し、1つの象限をコンテキストとして他の3つの象限を予測すること を目的とするラスタライズドマスキングや、再構成ベースの手法で一般的に用いられる従来のブロックマスキングやランダムマスキ ングなどの他のマスキング戦略と比較。 • ViT-B/16を300エポック学習した後、利用可能なラベルの1%のみを使用したImageNet-1KでのLinear evaluationに おいて、このマルチブロックマスキング戦略の効果を評価。 ブロックマスキングでは、ターゲットは1つの画像ブロックであり、コンテキストは画像の補集合である。 ランダムマスキングでは、ターゲットはランダムなパッチの集合であり、コンテキストは画像の補集合である。 Copyright (C) Present Square Co., Ltd. All Rights Reserved. 23

４. 実験：Ablations 表現空間での予測 • ImageNet-1Kの1%のデータで行われた線形プローブを使用した評価により、I-JEPAのローショット性能は、損失がピク セル空間ではなく表現空間で計算されることにより向上することが明らかになった。 • これはターゲットエンコーダが抽象的な予測ターゲットを生成する能力を強化するためと推測される。ピクセル空間での予測 は線形プロービング性能を著しく低下させることが示されており、これは事前学習中のターゲットエンコーダの重要性を強調 している。 Copyright (C) Present Square Co., Ltd. All Rights Reserved. 24

５. まとめ 実証内容 ① I-JEPAは、手動によるデータ拡張を必要とせずに、高性能な表現を学習する。具体的には、ImageNet-1Kの線形プ ローブ、半教師ありの1% ImageNet-1K、及び意味転移タスクにおいて、ピクセル再構成手法（例：MAE）を上回る 結果を示す。 ② I-JEPAは、意味的なタスクにおいてビュー不変の事前学習アプローチと同等、さらには低レベルの視覚タスク（例：オ ブジェクトのカウントや深度の予測）においても優れたパフォーマンスを達成する。この結果は、より単純で、誘導バイア スの少ないモデルを用いることで、より幅広いタスクに対応可能であることを示している。 ③ I-JEPAはスケーラブルで効率的である。具体的には、ViT-H/14をImageNetで事前学習するのに必要な時間は、 iBOTで事前学習したViT-S/16より2.5倍以上速く、MAEで事前学習したViT-H/14よりも10倍以上効率的である。 予測を表現空間で行うことで、自己教師あり事前学習に必要な総計算量を大幅に削減することができる。 Copyright (C) Present Square Co., Ltd. All Rights Reserved. 25

Appendix 参考文献 [4] Mahmoud Assran, Mathilde Caron, Ishan Misra, Piotr Bojanowski, Florian Bordes, Pascal Vincent, Armand Joulin, Michael Rabbat, and Nicolas Ballas. Masked siamese networks for label-efficient learning. European Conference on Computer Vision, 2022. [7] Alexei Baevski, Arun Babu, Wei-Ning Hsu, and Michael Auli. Efficient self-supervised learning with contextualized target representations for vision, speech and language. arXiv preprint arXiv:2212.07525, 2022. [8] Alexei Baevski,Wei-Ning Hsu, Qiantong Xu, Arun Babu, Jiatao Gu, and Michael Auli. Data2vec: A general framework for self-supervised learning in speech, vision and language. arXiv preprint arXiv:2202.03555, 2022 [17] Mathilde Caron, Ishan Misra, Julien Mairal, Priya Goyal, Piotr Bojanowski, and Armand Joulin. Unsupervised learning of visual features by contrasting cluster assignments. arXiv preprint arXiv:2006.09882, 2020. [36] Kaiming He, Xinlei Chen, Saining Xie, Yanghao Li, Piotr Doll´ar, and Ross Girshick. Masked autoencoders are scalable vision learners. IEEE/CVF Conference on Computer Vision and Pattern Recognition, 2022. [49] Yann LeCun, Sumit Chopra, Raia Hadsell, M Ranzato, and Fujie Huang. A tutorial on energy-based learning. Predicting structured data, 1(0), 2006. [79] Jinghao Zhou, Chen Wei, Huiyu Wang, Wei Shen, Cihang Xie, Alan Yuille, and Tao Kong. Ibot: Image bert pretraining with online tokenizer. International Conference on Learning Representations, 2022. Copyright (C) Present Square Co., Ltd. All Rights Reserved. 26