2023
Abstract
本文展示了一种学习高度语义图像表示的方法，而不依赖于手工制作的数据增强。本文介绍了基于图像的联合嵌入预测架构 (I-JEPA)，这是一种用于从图像进行自监督学习的非生成方法。 I-JEPA 背后的想法很简单：从单个上下文块，预测同一图像中各种目标块的表示。引导 I-JEPA 生成语义表示的核心设计选择是 掩码策略；具体来说，至关重要的是 (a) 预测图像中的几个目标块，(b) 对具有足够大尺度的目标块进行采样（占据图像的 15%–20%），以及 (c) 使用足够信息（空间分布）的上下文块。根据经验，当与 Vision Transformers 结合使用时，本文发现 I-JEPA 具有高度可扩展性。例如，本文在 ImageNet 上使用 32 个 A100 GPU 在 38 小时内训练了 ViT-Huge/16，以在需要不同抽象级别的各种任务（从线性分类到目标计数和深度预测）中实现强大的下游性能。
1. Introduction
在计算机视觉中，有两种常见的图像自监督学习方法：基于不变性的方法 [9,16,17,22,33,35,70] 和 生成方法 [7,26,34,55]。
基于不变性的方法：
基于不变性的预训练方法 优化编码器，为同一图像的两个或多个视图生成相似的嵌入 [14、19]，图像视图通常使用一组手工制作的数据增强来构建，例如随机缩放、裁剪和颜色抖动 [19]，等等 [33]。这些预训练方法可以产生高语义级别的表示 [3、17]，但它们也引入了强烈的偏差，这可能对某些下游任务甚至具有不同数据分布的预训练任务不利 [1]。通常，不清楚如何为需要不同抽象级别的任务概括这些偏差。例如，图像分类和实例分割不需要相同的不变性 [10]。此外，将这些特定于图像的增强推广到其他模式（例如音频）并不简单。
生成方法：
认知学习理论表明，生物系统中表示学习背后的驱动机制是 适应内部模型 以 预测感官输入反应 [29, 57]。这个想法是自监督生成方法的核心，它删除或破坏部分输入 并 学习预测损坏的内容 [8、34、55、64、65、68]。特别是，掩码去噪方法 通过 在像素或tokens级别 从输入重建随机屏蔽的patches 来学习表示。 掩码的预训练任务 比 视图不变性方法 需要更少的先验知识，并且很容易泛化到图像模态之外 [7]。然而，由此产生的表示通常具有较低的语义级别，并且在现成的评估（例如，线性探测）和 语义分类任务监督有限的迁移设置中 表现弱于 基于不变性的预训练 [3]。因此，需要更复杂的适应机制（例如，端到端微调）来充分利用这些方法的优势。
在这项工作中，探索了如何在 不使用通过图像变换编码的额外先验知识的情况下 提高自监督表示的语义水平。为此，本文引入了基于图像的联合嵌入预测架构(I-JEPA) [46]。图 3 提供了该方法的说明。I-JEPA 背后的想法是预测抽象表示空间中的缺失信息；例如，给定单个上下文块，预测同一图像中各种目标块的表示，其中目标表示由学习的目标编码器网络计算。
与 在 像素/token空间 中进行预测的生成方法相比，I-JEPA 利用抽象的预测目标，可能会消除不必要的像素级细节，从而使模型学习更多的语义特征。
另一个引导 I-JEPA 产生语义表示的核心设计选择是 提出的多块掩码策略。具体来说，本文证明了 使用 信息丰富的（空间分布的）上下文块 预测图像中的几个目标块（具有足够大的比例）的重要性。
通过广泛的实证评估，本文证明：
• I-JEPA 在 不使用手工的视图增强 的情况下 学习强大的现成的语义表示（参见图 1）。 I-JEPA 在 ImageNet-1K 线性探测、半监督 1% ImageNet-1K 和语义迁移任务上优于像素重建方法，例如 MAE [34]。
• I-JEPA 在语义任务上与视图不变的预训练方法相比具有竞争力，并且在目标计数和深度预测（第 5 节和第 6 节）等低级视觉任务上 取得了更好的性能。通过使用具有较少刚性归纳偏差的更简单模型，I-JEPA 适用于更广泛的任务集。
• I-JEPA 还具有可扩展性和高效性（第 7 节）。在 ImageNet 上预训练 ViT-H/14 需要大约 2400 个 GPU 小时，比使用 iBOT [74] 预训练的 ViTB/16 快 50%，比使用 MAE 预训练的 ViT-L/16 效率高 140%。在表示空间中进行预测 显著减少了自监督预训练所需的总计算量。
2. Background
自监督学习是一种表示学习方法，在这种方法中，系统 学习捕捉 其输入之间的关系。使用基于能量的模型 (EBM) [47] 的框架可以很容易地描述这个目标，其中自监督的目标是 将高能量分配给不兼容的输入，并将低能量分配给兼容的输入。许多现有的自监督学习的生成和非生成方法确实可以在这个框架中使用；参见图 2。
图2：自监督学习的通用架构，其中系统学习捕获其输入之间的关系。目标是为不兼容的输入分配高能量（大标度值），并为兼容的输入分配低能量（低标度值）。 (a) 联合嵌入架构学习为兼容的输入 x、y 输出相似的嵌入，为不兼容的输入输出不同的嵌入。 (b) 生成架构学习直接从兼容信号 x 重建信号 y，使用以附加（可能是潜在的）变量 z 为条件的解码器网络以促进重建。 (c) 联合嵌入预测架构 学习从兼容信号 x 预测信号 y 的嵌入，使用以附加（可能是潜在的）变量 z 为条件的预测网络来促进预测。
联合嵌入架构。（基于不变性的方法，对比方法）
可以使用联合嵌入架构 (JEA) 在 EBM 框架中进行基于不变性的预训练；参见图 2a。联合嵌入架构学习为兼容的输入 x、y 输出相似嵌入，为不兼容的输入输出不同嵌入。在基于图像的预训练环境中，兼容的 x、y 对通常是通过将手工制作的数据增强随机应用于同一输入图像 [19] 来构建的。
JEA 的主要挑战是 表征崩溃，其中能量形状是平坦的（即，无论输入如何，编码器都会产生恒定的输出）。在过去的几年中，已经研究了几种方法来防止表示崩溃，例如 显式推开负样本嵌入的对比损失 [14,22,35MoCo]，最小化嵌入之间的信息冗余的非对比损失 [9, 70 BarlowTwins]，以及 最大化平均嵌入熵的基于聚类的方法 [3, 4, 17]。还有一些启发式方法利用 x 编码器和 y 编码器之间的非对称架构设计来避免崩溃 [7、22 SimSiam、33 BYOL]。
Generative Architectures 生成架构
用于自监督学习的 基于重建的方法 也可以使用生成架构投射到 EBM 的框架中；参见图 2b。生成架构 学习直接从兼容信号 x 重建信号 y，使用以附加（可能是潜在的）变量 z 为条件的解码器网络以促进重建。在基于图像的预训练上下文中，计算机视觉中的一种常见方法是 使用掩码 [8, 36] 生成兼容的 x, y 对，其中 x 是图像 y 的副本，但其中一些patches被掩码。然后，条件变量 z 对应于一组（可能是可学习的）掩码和位置tokens，它向解码器指定要重建哪些图像块。只要 z 的信息容量比信号 y 低，这些架构就不会担心表示崩溃。
联合嵌入预测架构。
如图 2c 所示，联合嵌入预测架构 [46] 在概念上类似于生成架构；然而，一个关键的区别是损失函数应用于嵌入空间，而不是输入空间。 JEPA 学习从兼容信号 x 预测信号 y 的嵌入，使用以附加（可能是潜在的）变量 z 为条件的预测网络来促进预测。本文提出的 I-JEPA 在 使用掩码的图像上下文中 提供了该架构的实例；参见图 3。与联合嵌入架构相比，JEPA 不寻求对一组手工制作的数据增强不变的表示，而是寻求在以附加信息 z 为条件时 相互预测的表示。然而，与联合嵌入架构一样，表示崩溃也是 JEPA 的一个问题。本文 利用 x 和 y 编码器之间的非对称架构 来避免 I-JEPA 中的表示崩溃。
3. Method
现在描述提出的基于图像的联合嵌入预测架构 (I-JEPA)，如图 3 所示。总体目标如下：给定上下文块，预测同一图像中各种目标块的表示。本文使用 Vision Transformer [27] (ViT) 架构作为上下文编码器、目标编码器和预测器。 ViT 由一堆变换器层组成，每个变换器层由一个自注意力 [63] 操作和一个全连接的 MLP 组成。本文的编码器/预测器架构让人想起生成掩码自动编码器 (MAE) [34] 方法。然而，一个关键的区别是 I-JEPA 方法是非生成的，并且预测是在表示空间中进行的。
Targets.
首先描述我们如何在 I-JEPA 框架中生成目标：在 I-JEPA 中，目标对应于图像块的表示。给定输入图像 y，我们将其转换为 N 个非重叠块patch的序列，并将其通过目标编码器 
获得相应的块级表示 Sy = {Sy1, . . . ,SyN } 其中Syk 是与第 k 个patch关联的表示。为了获得损失的目标，我们从目标表示 sy 中随机抽取 M 个（可能重叠的）块block。我们用 Bi 表示第 i 个块对应的掩码，用 
表示 它的块级表示。通常，我们将 M 设置为 4，并以 (0.75, 1.5) 范围内的随机长宽比和 (0.15, 0.2) 范围内的随机比例 对块进行采样。请注意，目标块（target blocks）是通过掩码目标编码器的输出 而不是输入来获得的。这种区别对于确保高语义级别的目标表示至关重要；参见，例如，[7]。
Context.
回想一下，I-JEPA 背后的目标是从单个上下文块预测目标块表示。为了获得 I-JEPA 中的上下文，我们首先从图像中采样单个上下文块 x，其范围为 (0.85, 1.0) 和单位纵横比。我们用 Bx 表示与上下文块 x 关联的掩码。由于目标块是独立于上下文块采样的，因此可能存在明显的重叠。为了确保有效的预测任务，我们从上下文块中删除任何重叠区域。图 4 显示了实践中各种上下文和目标块的示例。接下来，掩码上下文块 x 通过上下文编码器 fθ 提供相应的patch级表示 
。
Prediction.
给定 上下文编码器的输出 Sx，我们希望 预测 M 个目标块表示 Sy(1), . . . , Sy(M) 。为此，对于 对应于目标掩码 Bi 的给定目标块 Sy(i)，预测器 gφ(·,·) 将上下文编码器的输出 Sx 和 我们希望预测的每个patch的掩码tokens 作为输入，
，并输出 patch级预测 
。掩码tokens由一个共享的可学习向量参数化，并添加了位置嵌入。由于我们希望对 M 个目标块进行预测，因此我们应用预测器 M 次，每次以对应于我们希望预测的目标块位置的掩码tokens为条件，并获得预测 
。
Loss.
损失简单地是 预测的patch级表示 与 目标patch级表示之间的平均 L2 距离；即，
预测器的参数 φ 和上下文编码器 θ 是通过基于梯度的优化学习的，而目标编码器的参数 ¯θ 是通过上下文编码器参数的指数移动平均值更新的。事实证明，使用指数移动平均目标编码器对于使用 Vision Transformers [17,23,74] 训练 JEA 至关重要，本文发现 I-JEPA 也是如此。
4. Related Work
大量工作通过预测缺失或损坏的感官输入的值来探索视觉表示学习。去噪自动编码器使用随机噪声作为输入损坏 [64]。上下文编码器根据其周围环境 [55] 回归整个图像区域。其他作品将图像着色作为去噪任务 [44、45、72]。
最近在掩码图像建模 [8、34、68] 的背景下重新审视了图像去噪的想法，其中使用 Vision Transformer [27] 来重建缺失的输入块。 Masked Autoencoders (MAE) [34] 的工作提出了一种高效的架构，它只需要编码器处理可见图像块。通过在像素空间中重建缺失的patches，MAE 在对大型标记数据集 进行 端到端微调时 实现了强大的性能，并表现出良好的缩放特性。其他类似的作品，例如 BEiT [8]，预测标记化空间中缺失补丁的价值；具体来说，使用冻结的离散 VAE 对图像块进行标记，该 VAE 在包含 2.5 亿张图像的数据集上进行训练 [56]。然而，像素级预训练已被证明在微调方面优于 BEiT [34]。另一项工作 SimMiM [68] 基于经典的梯度直方图 [25] 特征空间探索重建目标，并展示了优于像素空间重建的一些优势。与那些作品不同的是，本文的表示空间是在训练过程中通过联合嵌入预测架构学习的。本文的目标是学习不需要对下游任务进行大量微调的语义表示。
最接近本文工作的是 data2vec [7]，它学习预测 通过在线目标编码器计算的缺失patches的表示。 data2vec [7] 方法不依赖于手工增强，可以应用于多种模式，在视觉、文本和语音方面取得有希望的结果。然而，本文的实验表明，data2vec 需要在基于图像的预训练上下文中进行端到端微调，以便在下游任务中具有竞争力。
本文的方法和 data2vec 之间的主要区别是 1) 编码器/预测器架构设计 和 2) 提出的多块multi-block掩码策略。根据经验，本文证明 这些差异 对于 学习视觉中强大的现成表征至关重要。此外，I-JEPA 只需要通过上下文编码器 处理可见的patches，并且在计算效率上有显著提高。与本文的工作一致，data2vecv2 [6] 探索了不同模式下自监督表示学习的有效架构。
本文还将 I-JEPA 与基于联合嵌入架构的各种方法进行了比较；例如，DINO [17]、MSN [3] 和 iBOT [74]。这些方法依赖于预训练期间 手工制作的数据增强 来学习语义图像表示。 MSN [3] 的工作在预训练期间 使用掩码 作为额外的数据增强，而 iBOT 将 data2vec 样式的patch级重建损失 与 DINO 的视图不变性损失 相结合。这些方法的共同点是需要处理 每个输入图像的多个用户生成视图，从而阻碍了可扩展性。相比之下，I-JEPA 只需要处理每个图像的单个视图。我们发现使用 I-JEPA 训练的 ViT-Huge/14 比使用 iBOT 训练的 ViT-Small/16 需要更少的计算量。
5. Image Classification
为了证明 I-JEPA 在不依赖手工数据增强的情况下学习高级表示，本文使用线性探测和部分微调协议报告了各种图像分类任务的结果。在本节中，考虑在 ImageNet-1K 数据集 [58] 上预训练的自监督模型。预训练和评估实施细节在附录 A 中进行了描述。除非另有明确说明，否则所有 I-JEPA 模型均以 224×224 分辨率进行训练。
ImageNet-1K
表 1 显示了通用 ImageNet-1K 线性评估基准的性能。自监督预训练后，模型权重被冻结，并使用完整的 ImageNet-1K 训练集在顶部 训练线性分类器。与掩码自动编码器 (MAE) [34] 和 data2vec [7] 等流行方法相比，它们在预训练期间也不依赖于大量手工制作的数据增强，我们看到 I-JEPA 显著提高了线性探测性能，同时使用显著减少计算量（见第 7 节）。此外，I-JEPA 受益于规模。以 448 分辨率训练的 ViT-H/16 与 iBOT [74] 等视图不变方法的性能相匹配，无需额外的手工数据增强。
少样本 ImageNet-1K。
表 2 显示了 1% ImageNet 基准测试的性能。这里的想法是仅使用 1% 的可用 ImageNet 标签来调整用于 ImageNet 分类的预训练模型，对应于每个类别大约 12 或 13 个图像。通过微调或线性探测来调整模型，具体取决于哪种方法最适合每种方法。 当使用类似的编码器架构时，I-JEPA 优于 MAE，同时需要更少的预训练时期。 I-JEPA 使用 ViTH/14 架构，与使用 data2vec [7] 预训练的 ViT-L/16 的性能相匹配，同时使用的计算量要少得多（参见第 7 节）。通过 增加图像输入分辨率，I-JEPA 优于以前的方法，包括在预训练期间利用额外的手工数据增强的联合嵌入方法，例如 MSN [3]、DINO [16] 和 iBOT [74]。
迁移学习。
表 3 显示了使用线性探测在各种下游图像分类任务上的性能。 I-JEPA 明显优于以前不使用增强的方法（MAE 和 data2vec），并 缩小了与在预训练期间利用手工数据增强的最佳基于视图不变性的方法的差距，甚至在 CIFAR100 和 Place205 上用线性探测超越了流行的 DINO [17]。
6. 局部预测任务
如第 5 节所示，I-JEPA 学习语义图像表示，显著提高以前方法（如 MAE 和 data2vec）的下游图像分类性能。此外，I-JEPA 受益于规模，可以缩小差距，甚至超越基于视图不变性的方法，这些方法利用额外的手工数据增强。在本节中，发现 I-JEPA 还学习了局部图像特征，并在目标计数和深度预测等低级密集预测任务上超越了基于视图不变性的方法。
表 4 显示了使用线性探测在各种低级任务上的性能。特别是，在预训练之后，模型权重被冻结，并在顶部训练线性模型以在 Clevr 数据集 [41] 上执行目标计数和深度预测。与 DINO 和 iBOT 等视图不变性方法相比，I-JEPA 方法在预训练期间有效地捕获了低级图像特征，并在目标计数 (Clevr/Count) 和深度预测 (Clevr/Dist) 方面优于它们)。
7.可扩展性
模型效率。
与以前的方法相比，I-JEPA 具有高度可扩展性。图 5 显示了 1% ImageNet-1K 的半监督评估随 GPU 小时数的变化。 I-JEPA 比以前的方法需要更少的计算，并且在不依赖手工数据增强的情况下实现强大的性能。与直接使用像素作为目标的基于重建的方法（例如 MAE）相比，I-JEPA 通过在表示空间中计算目标引入了额外的开销（每次迭代的时间大约慢 7%）。然而，由于 I-JEPA 收敛的迭代次数减少了大约 5 倍，我们在实践中仍然看到显著的计算节省。与基于视图不变性的方法（例如 iBOT）相比，iBOT 依靠手工制作的数据增强来创建和处理每个图像的多个视图，IJEPA 的运行速度也快得多。特别是，最大的 I-JEPA 模型 (ViT-H/14) 比最小的 iBOT 模型 (ViT-S/16) 需要更少的计算。
缩放数据大小。
本文还发现 I-JEPA 可以从更大数据集的预训练中获益。表 5 显示了在增加预训练数据集的大小时（IN1K 与 IN22K）在语义和低级任务上的迁移学习性能。当在更大、更多样化的数据集上进行预训练时，迁移学习在这些概念上不同的任务上的性能会提高。
缩放模型大小。
表 5 还显示 I-JEPA 在 IN22K 上进行预训练时受益于更大的模型尺寸。与 ViT-H/14 模型相比，预训练 ViT-G/16 显著提高了 Place205 和 INat18 等图像分类任务的下游性能。 ViT-G/16 模型不会提高低级下游任务的性能。 ViT-G/16 使用较大的输入patch大小，这可能不利于局部预测任务。
8.预测可视化
I-JEPA 中 预测器 的作用是 获取上下文编码器的输出，并以位置掩码tokens为条件，预测掩码标记指定位置处 目标黑色的表示。一个自然的问题是，以位置掩码标记为条件的预测器 是否 正在学习 正确地捕获目标中的位置不确定性。为了定性地研究这个问题，本文将预测器的输出可视化。预训练后，我们冻结上下文编码器 和 预测器权重，并按照 RCDM 框架 [12] 训练解码器，将预测器输出的平均池化映射回像素空间。图 6 显示了各种随机种子的解码器输出。样本间共有的特性 表示着 包含在平均池化预测变量表征中的信息。 I-JEPA 预测器正确捕获位置不确定性 并 生成具有正确姿势的高级对象部分（例如，鸟的背部和汽车的顶部）。样本之间不同的特性表示着 表征中 未包含的信息。在这种情况下，I-JEPA 预测器会丢弃精确的低级细节以及背景信息。
9.消融
在表示空间中进行预测。
表 7 比较了在像素空间与表示空间中计算损失时 使用线性探测在 1% ImageNet-1K 上的少样本性能。本文推测 I-JEPA 的一个关键组成部分是 损失完全在表示空间中计算，从而使目标编码器能够生成抽象的预测目标，从而消除不相关的像素级细节。从表 7 可以清楚地看出，在像素空间中进行预测会导致线性探测性能显著下降。
掩码策略。
在表 8 中，本文在 I-JEPA 预训练期间 消融了目标块的数量 以及 所提出的多块屏蔽策略（如图 4 所示）中的上下文和目标块的规模。我们使用具有各种多块设置的 I-JEPA 训练 ViT-B/16 300 个时期，并使用线性探测比较 1% ImageNet-1K 基准测试的性能。简而言之，本文发现重要的是预测几个相对较大（语义）的目标块，并使用足够信息（空间分布）的上下文块。
表 6 在与其他掩码策略进行比较时进行了类似的消融。本文与光栅化掩码策略进行比较，其中图像被分成四个大象限，目标是使用一个象限作为上下文来预测其他三个象限。本文还与传统的块和随机掩码策略进行了比较，这些策略通常用于基于重建的方法。在块掩码中，目标是单个图像块，上下文是图像的补充。在随机掩码中，目标是一组随机（可能不连续）的图像块，上下文是图像的补集。
请注意，在所有考虑的掩码策略中，上下文和目标块之间没有重叠。事实证明，所提出的多块掩码策略 对于 I-JEPA 学习语义表示至关重要。即使切换到传统的块掩码，ImageNet 的性能也会降低超过 24% top-1。
10. 结论
本文提出了 I-JEPA，这是一种简单有效的学习语义图像表示的方法，无需依赖手工制作的数据增强。本文表明，通过在表示空间中进行预测，I-JEPA 比像素重建方法收敛得更快，并且可以学习高语义级别的表示。与基于视图不变性的方法相比，I-JEPA 强调了一种使用联合嵌入架构学习一般表示的途径，而不依赖于手工制作的视图增强。