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

一、简要介绍

本文演示了一种学习高度语义的图像表示的方法，而不依赖于手工制作的数据增强。论文介绍了基于图像的联合嵌入预测架构（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