Meta 自监督学习 I-JEPA 论文笔记

Meta 宣布推出一个全新的 AI 模型 Image Joint Embedding Predictive Architecture (I-JEPA)，可通过对图像的自我监督学习来学习世界的抽象表征，实现比现有模型更准确地分析和完成未完成的图像。

简介

I-JEPA ( Image Joint Embedding Predictive Architecture ) 是一种自监督学习框架，结合了 联合嵌入架构 (Joint-Embedding Architecture) 和 生成架构 (Generative Architecture) 的特性提出。

结合了 Meta 首席 AI 科学家 Yann LeCun 所提倡的类人推理方式，帮助避免 AI 生成图像常见的一些错误，比如多出的手指。I-JEPA 在多项计算机视觉任务上表现出色，且计算效率比其他广泛使用的计算机视觉模型高得多。

论文链接：https://arxiv.org/pdf/2301.08243.pdf

摘要

本文展示了一种学习高度语义图像表示的方法，而不依赖于手工制作的数据增强。本文介绍了基于图像的联合嵌入预测架构 (I-JEPA)，这是一种用于从图像进行自监督学习的非生成方法。

I-JEPA 背后的想法很简单：从单个内容块，预测同一图像中各种目标块的表示。引导 I-JEPA 生成语义表示的核心设计选择是掩码策略；具体来说，至关重要的是：

(a) 预测图像中的几个目标块

(b) 对具有足够大尺度的目标块进行采样（占据图像的 15%–20%）

© 使用足够信息（空间分布）的内容块。

根据经验，当与 Vision Transformers 结合使用时，本文发现 I-JEPA 具有高度可扩展性。例如，本文在 ImageNet 上使用 32 个 A100 GPU 在 38 小时内训练了 ViT-Huge/16，以在需要不同抽象级别的各种任务（从线性分类到目标计数和深度预测）中实现强大的下游性能。

引言

在计算机视觉中，有两种常见的图像自监督学习方法：基于不变性的方法 和 生成方法 。

不变性的方法

模型接受两个输入，相似的两个输入模型输出相近的 嵌入信息，否则输出不同的嵌入信息。常用的方法是使用一组手工制作的数据增强来构建，例如随机缩放、裁剪和颜色抖动。这些预训练方法可以产生高语义级别的表示，但它们也引入了强烈的偏差，这可能对某些下游任务甚至具有不同数据分布的预训练任务不利。通常，我们是不清楚如何为需要不同抽象级别的任务概括这些偏差的。例如，图像分类和实例分割不需要相同的不变性 。此外，将这些特定于图像的增强推广到其他模式（例如音频）并不简单。

生成方法

认知学习理论表明，生物系统中表示学习背后的驱动机制是适应内部模型 以 预测感官输入反应 。这个想法是自监督生成方法的核心，它删除或破坏部分输入 并 学习预测损坏的内容。特别是，掩码去噪方法通过在像素或tokens级别从输入重建随机屏蔽的图像块来学习表示。 掩码的预训练任务 比 视图不变性方法 需要更少的先验知识，并且很容易泛化到图像模态之外。

然而，由此产生的表示通常具有较低的语义级别，并且在现成的评估（例如，线性探测）和 语义分类任务监督有限的迁移设置中 表现弱于 基于不变性的预训练。因此，需要更复杂的适应机制（例如，端到端微调）来充分利用这些方法的优势。

本文工作

在这项工作中，探索了如何在 不使用通过图像变换编码的额外先验知识的情况下 提高自监督表示的语义水平。为此，本文引入了基于图像的联合嵌入预测架构(I-JEPA) 。下图提供了该方法的说明。I-JEPA 背后的想法是预测抽象表示空间中的缺失信息；例如，给定单个内容块，预测同一图像中各种目标块的表示，其中目标表示由学习的目标编码器网络计算。

与在像素/token空间中进行预测的生成方法相比，I-JEPA 利用抽象的预测目标，可能会消除不必要的像素级细节，从而使模型学习更多的语义特征。

另一个引导 I-JEPA 产生语义表示的核心设计选择是 提出的多块掩码策略。具体来说，本文证明了使用信息丰富的（空间分布的）内容块 预测图像中的几个目标块（具有足够大的比例）的重要性。

通过广泛的实证评估，本文证明：

I-JEPA 在 不使用手工的视图增强 的情况下 学习强大的现成的语义表示（参见下图）。 I-JEPA 在 ImageNet-1K 线性探测、半监督 1% ImageNet-1K 和语义迁移任务上优于像素重建方法，例如 MAE。

I-JEPA 在语义任务上与视图不变的预训练方法相比具有竞争力，并且在目标计数和深度预测等低级视觉任务上 取得了更好的性能。通过使用具有较少刚性归纳偏差的更简单模型，I-JEPA 适用于更广泛的任务集。

I-JEPA 还具有可扩展性和高效性。在 ImageNet 上预训练 ViT-H/14 需要大约 2400 个 GPU 小时，比使用 iBOT 预训练的 ViTB/16 快 50%，比使用 MAE 预训练的 ViT-L/16 效率高 140%。在表示空间中进行预测 显著减少了自监督预训练所需的总计算量。

背景知识

自监督学习是一种表示学习方法，在这种方法中，系统 学习捕捉 其输入之间的关系。使用基于能量的模型 (EBM) [47] 的框架可以很容易地描述这个目标，其中自监督的目标是 将高能量分配给不兼容的输入，并将低能量分配给兼容的输入。许多现有的自监督学习的生成和非生成方法确实可以在这个框架中使用。

联合嵌入架构（Joint-Embedding Architecture）

基于不变性的方法，对比方法

(a) 联合嵌入架构学习为兼容的输入 x、y 输出相似的嵌入，为不兼容的输入输出不同的嵌入。

可以使用联合嵌入架构 (JEA) 在 EBM 框架中进行基于不变性的预训练；参见上图。联合嵌入架构学习为兼容的输入 x、y 输出相似嵌入，为不兼容的输入输出不同嵌入。在基于图像的预训练环境中，兼容的 x、y 对通常是通过将手工制作的数据增强随机应用于同一输入图像来构建的。

JEA 的主要挑战是表征崩溃，其中能量形状是平坦的（即，无论输入如何，编码器都会产生恒定的输出）。在过去的几年中，已经研究了几种方法来防止表示崩溃，例如 显式推开负样本嵌入的对比损失 ，最小化嵌入之间的信息冗余的非对比损失，以及 最大化平均嵌入熵的基于聚类的方法。还有一些启发式方法利用 x 编码器和 y 编码器之间的非对称架构设计来避免崩溃。

生成架构 （Generative Architectures）

(b) 生成架构学习直接从兼容信号 x 重建信号 y，使用以附加（可能是潜在的）变量 z 为条件的解码器网络以促进重建。

用于自监督学习的 基于重建的方法 也可以使用生成架构投射到 EBM 的框架中；参见上图。生成架构 学习直接从兼容信号 x 重建信号 y，使用以附加（可能是潜在的）变量 z 为条件的解码器网络以促进重建。在基于图像的预训练内容中，计算机视觉中的一种常见方法是 使用掩码生成兼容的 x, y 对，其中 x 是图像 y 的副本，但其中一些patches被掩码。然后，条件变量 z 对应于一组（可能是可学习的）掩码和位置tokens，它向解码器指定要重建哪些图像块。只要 z 的信息容量比信号 y 低，这些架构就不会担心表示崩溃。

联合嵌入预测架构（Joint-Embedding Predictive Architecture）

联合嵌入预测架构 学习从兼容信号 x 预测信号 y 的嵌入，使用以附加（可能是潜在的）变量 z 为条件的预测网络来促进预测。

如上图所示，联合嵌入预测架构在概念上类似于生成架构；然而，一个关键的区别是损失函数应用于嵌入空间，而不是输入空间。 JEPA 学习从兼容信号 x 预测信号 y 的嵌入，使用以附加（可能是潜在的）变量 z 为条件的预测网络来促进预测。本文提出的 I-JEPA 在 使用掩码的图像内容中 提供了该架构的实例。

与联合嵌入架构相比，JEPA 不寻求对一组手工制作的数据增强不变的表示，而是寻求在以附加信息 z 为条件时 相互预测的表示。然而，与联合嵌入架构一样，表示崩溃也是 JEPA 的一个问题。本文 利用 x 和 y 编码器之间的非对称架构 来避免 I-JEPA 中的表示崩溃。

方法（Method）

现在描述提出的基于图像的联合嵌入预测架构 (I-JEPA)。总体目标如下：给定内容块，预测同一图像中各种目标块的表示。本文使用 Vision Transformer(ViT) 架构作为内容编码器、目标编码器和预测器。 ViT 由一堆变换器层组成，每个变换器层由一个自注意力操作和一个全连接的 MLP 组成。本文的编码器/预测器架构让人想起生成掩码自动编码器 (MAE)方法。然而，一个关键的区别是 I-JEPA 方法是非生成的，并且预测是在表示空间中进行的。

目标（Targets）

首先描述我们如何在 I-JEPA 框架中生成目标：在 I-JEPA 中，目标对应于图像块的表示。给定输入图像 y，我们将其转换为 N 个非重叠块patch的序列，并将其通过目标编码器 $ f_{\bar{\theta}} $ 获得相应的块级表示 $S_y = {S_{y1}, . . . ,S_{yN} }$ 其中 $S_{yk}$ 是与第 k 个patch关联的表示。为了获得损失的目标，我们从目标表示 $s_y$ 中随机抽取 M 个（可能重叠的）块block。我们用 $B_i$ 表示第 i 个块对应的掩码，用 $ \boldsymbol{s}_{y}(i)=\left\{s_{y^{J}}\right\}_{j \omega B_{i}} $ 表示 它的块级表示。通常，我们将 M 设置为 4，并以 (0.75, 1.5) 范围内的随机长宽比和 (0.15, 0.2) 范围内的随机比例 对块进行采样。请注意，目标块（target blocks）是通过掩码目标编码器的输出 而不是输入来获得的。这种区别对于确保高语义级别的目标表示至关重要。

内容块（Context）

I-JEPA 背后的目标是从单个内容块预测目标块表示。为了获得 I-JEPA 中的内容，我们首先从图像中采样单个内容块 x，其范围为 (0.85, 1.0) 和单位纵横比。我们用 $B_x$ 表示与内容块 x 关联的掩码。由于目标块是独立于内容块采样的，因此可能存在明显的重叠。为了确保有效的预测任务，我们从内容块中删除任何重叠区域。下图显示了实践中各种内容和目标块的示例。接下来，掩码内容块 x 通过内容编码器 $f_θ$ 提供相应的 patch 级表示 $ \boldsymbol{s}_{x}=\left\{\boldsymbol{s}_{x_{j}}\right\}_{j \in B_{x}} $.

预测 (Prediction)

给定 内容编码器的输出 $S_x$，我们希望 预测 M 个目标块表示 $S_y(1), . . . , S_y(M) $。为此，对于 对应于目标掩码 $B_i$ 的给定目标块 $S_y(i)$，预测器 $g_φ(·,·)$ 将内容编码器的输出 $S_x$ 和 我们希望预测的每个 patch 的掩码 tokens 作为输入，$ \left\{\boldsymbol{m}_{j}\right\}_{j \in B_{i}} $ 并输出 patch 级预测 $ \hat{\boldsymbol{s}}_{y}(i)= \left\{\hat{s}_{y_{j}}\right\}_{j \in B_{i}}=g_{\phi}\left(\boldsymbol{s}_{x},\left\{\boldsymbol{m}_{j}\right\}_{j \in B_{i}}\right) $ 。 掩码tokens由一个共享的可学习向量参数化，并添加了位置嵌入。由于我们希望对 M 个目标块进行预测，因此我们应用预测器 M 次，每次以对应于我们希望预测的目标块位置的掩码 tokens 为条件，并获得预测 $ \hat{\boldsymbol{s}}_{y}(1), \ldots, \hat{\boldsymbol{s}}_{y}(M) $。

损失计算（Loss）

计算基于内容块预测的 patch 级的特征表示 $ \hat{s}_{y}(i) $ 和目标块 patch 级的特征表示 $ {s}_{y}(i) $ 的 $L_2$ 距离作为损失函数：

$$ \frac{1}{M} \sum_{i=1}^{M} D\left(\hat{s}_{y}(i), \boldsymbol{s}_{y}(i)\right)=\frac{1}{M} \sum_{i=1}^{M} \sum_{j \in B_{i}}\left\|\hat{s}_{y_{j}}-\boldsymbol{s}_{y_{j}}\right\|_{2}^{2} $$

预测器的参数 φ 和内容编码器 θ 是通过基于梯度的优化学习的，而目标编码器的参数 $ \bar{\theta} $ 是通过内容编码器参数的指数移动平均值更新的。事实证明，使用指数移动平均目标编码器对于使用 Vision Transformers 训练 JEA 至关重要，本文发现 I-JEPA 也是如此。

图像分类应用

文章展示了在 ImageNet 分类能力上的自监督模型性能，证明了该自监督学习框架的能力和结构合理。

可扩展性

模型效率

​ 与以前的方法相比，I-JEPA 具有高度可扩展性。下图显示了 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 可以从更大数据集的预训练中获益。下表显示了在增加预训练数据集的大小时（IN1K 与 IN22K）在语义和低级任务上的迁移学习性能。当在更大、更多样化的数据集上进行预训练时，迁移学习在这些概念上不同的任务上的性能会提高。

同时下表还显示 I-JEPA 在 IN22K 上进行预训练时受益于更大的模型尺寸。与 ViT-H/14 模型相比，预训练 ViT-G/16 显著提高了 Place205 和 INat18 等图像分类任务的下游性能。 ViT-G/16 模型不会提高低级下游任务的性能。 ViT-G/16 使用较大的输入patch大小，这可能不利于局部预测任务。

预测可视化

I-JEPA 中 预测器 的作用是 获取内容编码器的输出，并以位置掩码tokens为条件，预测掩码标记指定位置处 目标黑色的表示。一个自然的问题是，以位置掩码标记为条件的预测器 是否 正在学习 正确地捕获目标中的位置不确定性。为了定性地研究这个问题，本文将预测器的输出可视化。预训练后，我们冻结内容编码器 和 预测器权重，并按照 RCDM 框架训练解码器，将预测器输出的平均池化映射回像素空间。下图显示了各种随机种子的解码器输出。样本间共有的特性 表示着 包含在平均池化预测变量表征中的信息。 I-JEPA 预测器正确捕获位置不确定性 并 生成具有正确姿势的高级对象部分（例如，鸟的背部和汽车的顶部）。样本之间不同的特性表示着 表征中 未包含的信息。在这种情况下，I-JEPA 预测器会丢弃精确的低级细节以及背景信息。

结论

本文提出了 I-JEPA，这是一种简单有效的学习语义图像表示的方法，无需依赖手工制作的数据增强。本文表明，通过在表示空间中进行预测，I-JEPA 比像素重建方法收敛得更快，并且可以学习高语义级别的表示。与基于视图不变性的方法相比，I-JEPA 强调了一种使用联合嵌入架构学习一般表示的途径，而不依赖于手工制作的视图增强。

原始论文

### 参考资料

• https://arxiv.org/pdf/2301.08243.pdf
• https://blog.csdn.net/yoooool_/article/details/129011503
• https://baijiahao.baidu.com/s?id=1768843236426429078&wfr=spider&for=pc

文章链接：
https://www.zywvvd.com/notes/study/deep-learning/self-supervised/i-jepa/i-jepa/