【论文解读】基于图的自监督学习联合嵌入预测架构

一、简要介绍

本文演示了一种学习高度语义的图像表示的方法，而不依赖于手工制作的数据增强。论文介绍了基于图像的联合嵌入预测架构（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学习了强大的现成的语义表示，而不使用手工制作的视图增强(图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）。基于不变性的预训练可以使用联合嵌入架构（JEA）在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，而不需要额外的手工数据扩充。

Low-Shot ImageNet-1K.表2显示了在1%的ImageNet基准测试上的性能。这里的方法是将预先训练好的模型用于ImageNet分类，只使用1%的ImageNet标签，对应于每个类大约12或13张图像。模型通过微调或线性探测来进行调整，这取决于每种方法的最佳效果。当使用类似的编码器架构时，I-JEPA优于MAE，同时需要更少的预训练时期。I-JEPA，使用ViTH/14架构，与使用数据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）

模型效率（Model Efficiency） 与以前的方法相比，I-JEPA具有高度的可缩放性。图5显示了对1% ImageNet-1K作为GPU hour的函数的半监督评估。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预训练过程中所提出的multi-block掩膜策略（如图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