【論文解讀】基於圖的自監督學習聯合嵌入預測架構

一、簡單介紹

本文示範了一種學習高度語意的圖像表示的方法，而不依賴於手工製作的資料增強。論文介紹了基於影像的聯合嵌入預測架構（I-JEPA），這是一種用於從影像中進行自監督學習的非生成性方法。 I-JEPA背後的idea很簡單：從單一上下文區塊中，預測同一影像中不同目標區塊的表示。指導I-JEPA產生語義表示的核心設計選擇是掩膜策略；具體來說，(a)預測圖像中的幾個目標塊，(b)採樣足夠大規模的樣本目標塊（佔圖像的15%- 20%），(c)使用足夠豐富的（空間分佈）上下文區塊，是至關重要的。根據經驗，當與視覺transformer結合時，論文發現I-JEPA具有高度的可縮放性。例如，論文在ImageNet上使用32個A100 GPU在38小時內訓練一個ViT-Huge/16，以在需要不同抽象級別的廣泛任務中實現強大的下游性能，從線性分類到對象計數和深度預測。

二、研究背景

在電腦視覺中，有兩種​​常見的圖像自監督學習方法。

基於不變性的方法和產生方法。透過基於不變性的預訓練方法優化編碼器，可以為相同影像的兩個或多個視圖產生相似的嵌入。通常，圖像視圖會使用一組手工製作的資料增強方法來構建，例如隨機縮放、裁剪、顏色抖動等。這些預訓練方法可以產生高語義等級的表示，但同時也會引入強烈的偏差，可能對某些下游任務甚至不同資料分佈的預訓練任務產生負面影響

認知學習理論認為，生物系統中表徵學習背後的一個驅動機制是如何適應一個內在模型來預測感官輸入反應。這個想法是自我監督生成方法的核心，它刪除或損壞部分輸入，並學習預測損壞的內容。特別是，掩膜去噪方法透過從輸入的像素或令牌層級重建隨機掩膜補丁來學習表示。與視圖不變性方法相比，掩膜的預訓練任務需要更少的先驗知識，並且很容易推廣到圖像模態之外。然而，所得到的表示通常具有較低的語義級別，並且在現成的評估（如線性探測）和對語義分類任務監督有限的轉移設定中缺乏基於不變性的預訓練。因此，需要一個更複雜的適應機制（例如，端到端微調）來獲得這些方法的充分優勢。

在這項工作中，論文探索如何提高自我監督表示的語義水平，而不使用額外的先驗知識編碼的圖像轉換。為此，論文引入了一種圖像聯合嵌入預測架構（I-JEPA）。圖3提供了該方法的一個說明。 I-JEPA背後的構思是預測抽象表示空間中的缺失資訊；例如，給定一個上下文區塊，預測同一影像中不同目標區塊的表示，其中目標表示是由學習到的目標編碼器網路計算的。

與在像素/標記空間中進行預測的生成方法相比，I-JEPA利用抽象的預測目標，可能消除不必要的像素級細節，從而導致模型學習更多的語義特徵。另一個指導I-JEPA產生語義表示的核心設計選擇是所提出的多塊掩膜策略。具體來說，論文演示了使用一個資訊豐富的（空間分佈的）上下文塊來預測圖像中的幾個目標塊（具有足夠大的規模）的重要性。 重寫後的內容：相較於在像素/標記空間中進行預測的生成方法，I-JEPA利用抽象的預測目標，可能消除不必要的像素級細節，從而使模型能夠學習更多的語義特徵。 I-JEPA的另一個核心設計選擇是採用多塊掩膜策略來產生語義表示。具體而言，該論文展示了使用資訊豐富的（空間分佈的）上下文塊來預測圖像中幾個目標塊（具有足夠大的規模）的重要性

根據大量的實證評估，研究顯示：

I-JEPA學習了強大的現成的語意表示，而不使用手工製作的視圖增強(圖1)。 I-JEPA優於像素重建方法，例如在ImageNet-1K線性探測上的MAE、半監督的1% ImageNet-1K和語義傳輸任務。

I-JEPA在語意任務上與視圖不變預訓練方法具有競爭力，並且在物件計數和深度預測等低階視覺任務上取得了更好的性能。透過使用一個更簡單的模型和更少的剛性歸納偏差，I-JEPA是適用於更廣泛的任務集。

I-JEPA也是可縮放的和高效的。在ImageNet上對ViT-H/14進行預訓練大約需要2400 GPU小時，比使用iBOT 預訓練的ViTB/16快50%，比使用MAE預訓練的ViT-L/16的效率快140%。在表示空間中的預測顯著減少了自我監督預訓練所需的總計算量。

自我監督學習是一種表徵學習的方法，其中一個系統學習捕捉其輸入之間的關係。這個目標可以很容易地使用基於能量的模型（EBMs）的框架來描述，其中自我監督的目標是為不相容的輸入分配高能，並為相容的輸入分配低能量。許多現有的生成式和非生成式的自我監督學習方法確實可以在這個框架中進行轉換；見圖2

#重寫後的內容：聯合嵌入式架構（Joint-Embedding Architectures）是一種基於不變性的預訓練方法，在EBM框架中使用可以進行強制轉換，見圖2a。聯合嵌入架構的學習目標是使相容的輸入x和y輸出相似的嵌入，而不相容的輸入則輸出不同的嵌入。在基於影像的預訓練中，通常會透過對相同的輸入影像隨機應用手工製作的資料增強來建立相容的x和y對

JEA的主要挑戰是表示坍縮，其中能量景觀是平坦的（即，編碼器產生一個恆定的輸出，而不管輸入如何）。在過去的幾年裡，已經研究了幾種方法來防止表示坍縮，例如明確推動負例子嵌入的對比損失，最小化嵌入的信息冗餘的非對比損失，以及基於聚類的方法來最大化平均嵌入的熵。也有一些啟發式的方法，利用x編碼器和y編碼器之間的非對稱架構設計來避免坍縮的。生成式架構（Generative Architectures）。基於重構的自監督學習方法也可以在使用生成式架構的EBM框架中進行強制轉換；見圖2b

生成式架構學習從一個相容的訊號x直接重建訊號y，使用一個附加的（可能是潛在的）變數z的解碼器網絡，以促進重建。在基於影像的預訓練中，電腦視覺中常見的方法是使用遮罩產生相容的x，y對，其中x是影像y的副本，但有一些補丁被遮罩。然後，條件反射變數z對應於一組（可能是可學習的）遮罩和位置標記，它指定了要重建的影像補丁的解碼器。只要z的資訊容量比訊號y低，這些體系結構就不會關注表示坍縮。

聯合嵌入式預測性架構（Joint-Embedding Predictive Architectures）。如圖2c所示，聯合嵌入預測架構在概念上類似於生成式架構；然而，一個關鍵的差異是，損失函數應用於嵌入空間，而不是輸入空間。 JEPA學習從相容的訊號x預測訊號y的嵌入，使用一個附加（可能是潛在）變數z的預測網路以促進預測。論文提出的I-JEPA在使用掩膜的圖像上下文中提供了該架構的實例化；參見圖3。與聯合嵌入架構相比，JEPA並不尋求對一組手工製作的資料增強不變的表示，而是尋求在附加資訊z條件時相互預測的表示。然而，與聯合嵌入架構一樣，表示坍縮也是JEPA關注的議題。論文利用x和y編碼器之間的非對稱架構，以避免I-JEPA中的表示坍縮。

 三、方法介紹

論文現在描述了所提出的基於圖像的聯合嵌入預測體系結構（I-JEPA），如圖3所示。總體目標如下：給定一個上下文區塊，預測同一影像中不同目標區塊的表示。論文使用視覺Transformer（ViT）架構作為上下文編碼器、目標編碼器和預測器。一個ViT由一堆Transformer層組成，每個Transformer層由一個自註意的操作和一個全連接的MLP組成。論文的編碼器/預測器架構讓人想起生成掩膜自動編碼器（MAE）方法。然而，一個關鍵的區別是，I-JEPA方法是非生成性的，並且預測是在表示空間中做出的。

#圖片分類

为了证明I-JEPA在不依赖于手工制作的数据增强的情况下学习高级表示，论文报告了使用线性探测和部分微调协议的各种图像分类任务的结果。在本节中，论文将考虑在ImageNet-1K数据集上预先训练过的自监督模型。预训练和评估实施细节见附录a。所有的I-JEPA模型都在分辨率224×224中进行训练，除非另有明确说明。

ImageNet-1K.表1显示了在通用的ImageNet-1K线性评估基准上的性能。经过自监督预训练后，冻结模型权值，并使用完整的ImageNet-1K训练集在顶部训练一个线性分类器。与流行的掩膜自动编码器（MAE）和data2vec 方法相比，它们在训练前也不依赖大量手工制作的数据增强，论文看到I-JEPA显著提高了线性探测性能，同时使用了更少的计算量。此外，I-JEPA还受益于规模。在分辨率为448时训练的ViT-H/16与视图不变方法的性能相匹配，如iBOT，而不需要额外的手工数据扩充。

低样本量的ImageNet-1K。表2展示了在1%的ImageNet基准测试上的性能。这些方法是利用预训练好的模型进行ImageNet分类，只使用1%的ImageNet标签，每个类别大约有12或13张图像。模型通过微调或线性探测进行调整，具体取决于每种方法的最佳效果。当使用相似的编码器架构时，I-JEPA表现优于MAE，并且需要更少的预训练时期。使用ViTH/14架构的I-JEPA与使用数据2vec预训练的ViT-L/16的性能相当，但计算量明显更少。通过提高图像输入分辨率，I-JEPA在比以前的方法更好，包括联合嵌入方法以及在训练前利用额外的手工数据增强方法，如MSN、DINO和iBOT

Transfer learning.表3显示了使用线性探头的各种下游图像分类任务的性能。I-JEPA显著优于以前不使用增强的方法（MAE和数据2vec），并减少了在训练前利用手工制作的基于观点不变的最佳方法的差距，甚至超过了CIFAR100和Place205上流行的DINO。

 五、本地预测任务（Local Prediction Tasks）

I-JEPA学习语义图像表示，显著提高了以往方法的下游图像分类性能，如MAE和data2vec。此外，I-JEPA受益于规模，并可以缩小差距，甚至超越，利用额外的手工制作的数据增强的基于视图不变性的方法。在本节中，论文发现I-JEPA也能学习局部图像特征，并在低水平和密集的预测任务中，如对象计数和深度预测中，超越了基于视图不变性的方法。

表4显示了使用线性探测的各种低级任务的性能。特别是，经过预训练后，模型的权值被冻结，并在顶部训练一个线性模型，对Clevr数据集进行对象计数和深度预测。与DINO和iBOT等视图不变性方法相比，I-JEPA方法在训练前有效地捕获了低级图像特征，并在对象计数（Clevr/Count）和（大幅度）深度预测（Clevr/Dist）方面优于它们。六、可缩放性（Scalability）

重写后的内容如下：根据以前的方法进行比较，I-JEPA在模型效率方面具有高度的可扩展性。图5展示了对1%的ImageNet-1K进行GPU小时评估的半监督结果。相较于以前的方法，I-JEPA所需的计算量更少，并且在没有依赖手工数据增强的情况下获得了强大的性能。与基于重建的方法（如直接使用像素作为目标的MAE等）相比，I-JEPA通过在表示空间中计算目标而引入了额外的开销（每次迭代的时间大约慢7%）

缩放数据大小（Scaling data size）。论文还发现I-JEPA受益于更大数据集的预训练。表5显示了在增加训练前数据集的大小（IN1KvsIN22K）时，在语义任务和低水平任务上的迁移学习性能。当对更大更多样化的数据集进行预训练时，这些概念不同的任务上的迁移学习性能会提高。缩放模型大小（Scaling model size）。表5还显示，在IN22K上进行预训练时，I-JEPA受益于更大的模型规模。与ViT-H/14模型相比，对ViT-G/16的预训练显著提高了对Place205和INat18等图像分类任务的下游性能。ViTG/16模型并不能提高对低级下游任务的性能。ViT-G/16使用更大的输入补丁大小，这可能对本地预测任务有害。

七、預測視覺化（Predictor Visualizations）可以進行重寫

I-JEPA中預測器的作用是取上下文編碼器的輸出，並以位置掩膜令牌為條件，預測掩膜令牌指定的位置上目標塊的表示。一個問題是，以位置掩模令牌條件的預測器是否正在學習正確捕獲目標中的位置不確定性。為了定性地研究這個問題，論文將預測器的輸出視覺化。在預訓練後，論文凍結了上下文編碼器和預測器的權值，並按照RCDM框架訓練一個解碼器，將預測器輸出的平均池映射回像素空間。圖6顯示了各種隨機種子的解碼器輸出。在樣本之間常見的特性表示包含在平均合併預測器表示中的資訊。 I-JEPA預測器能正確地捕捉到位置的不確定性，並產生具有正確姿態的高級物體部件（例如，鳥的背部和汽車的頂部）。不同樣本中不同的品質表示表示法中不包含的資訊。在這種情況下，I-JEPA預測器丟棄了精確的低階細節和背景資訊。

八、消融（Ablations）的重要性

在表示空間中的預測（ Predicting in representation space.）。表7比較了在像素空間和表示空間中計算1% ImageNet-1K時的low-shot效能。論文推測，I-JEPA的一個關鍵組成部分是，損失完全在表示空間中計算，從而使目標編碼器能夠產生抽象的預測目標，從而消除了無關的像素級細節。從表7中可以清楚看出，在像素空間中的預測會導致線性探測性能的顯著下降。

重寫後的內容：遮罩策略（Masking strategy）在表8中進行了修改。本研究減少了在I-JEPA預訓練過程中所提出的多重遮罩策略中的目標區塊數量，並調整了情境和目標區塊的規模，如圖4所示。我們使用了各種不同多塊設定的I-JEPA進行了300個週期的訓練，並使用線性探針在1%的ImageNet-1K基準測試上進行了效能比較。總結來說，我們發現預測幾個相對較大的（語意）目標區塊，並結合資訊豐富的（空間分佈的）上下文區塊是非常重要的

表6與其他掩膜策略進行比較時，也進行了類似的消融。論文與rasterized掩膜策略進行了比較，其中圖像被分割成四個大象限，其目標是使用一個象限作為上下文來預測其他三個像限。論文還比較了通常用於基於重建的方法的傳統的block和random掩膜策略。在block掩膜中，目標是單一影像區塊，上下文是影像補體。在random掩膜中，目標是一組隨機的（可能是不連續的）影像補丁，而上下文是影像的補體。請注意，在所有考慮的遮罩策略中，上下文區塊和目標區塊之間沒有重疊。所提出的multi-block掩膜策略是I-JEPA學習語意表示的關鍵。即使切換到傳統的block遮罩，ImageNet的效能也會降低超過24%。

九、結論論文

提出了一種名為I-JEPA的方法，用於學習語義圖像表示。此方法不依賴手工製作的數據增強。研究表明，透過在表示空間中進行預測，I-JEPA比像素重建方法收斂速度更快，並且能夠學習到高語義水平的表示。與基於視圖不變性的方法相比，I-JEPA強調了使用聯合嵌入架構學習一般表示的路徑，而不依賴於手工製作的視圖增強

附錄見原文，原文鏈接：https:/ /arxiv.org/abs/2301.08243     

以上是【論文解讀】基於圖的自監督學習聯合嵌入預測架構的詳細內容。更多資訊請關注PHP中文網其他相關文章！