1
2
3
4
5
6
7
英文名: Self-Supervised Learning from Images with a Joint-Embedding Predictive Architecture
中文名: 从图像中进行自监督学习的联合嵌入预测架构
地址: http://arxiv.org/pdf/2301.08243v3
作者: Mahmoud Assran, Quentin Duval, Ishan Misra, Piotr Bojanowski, Pascal Vincent, Michael Rabbat, Yann LeCun, Nicolas Ballas
机构: Meta AI (FAIR)，McGill University，Mila, Quebec AI Institute，New York University
日期: 2023-01-19
引用次数: 448

摘要

  • 目标: 通过无手工设计的数据增强方法学习具有高度语义性的图像表示。
  • 方法: 提出了图像基础的联合嵌入预测架构（I-JEPA），一种自监督学习的非生成性方法，核心思想是从单个上下文块预测同一图像中不同目标块的表示。
  • 结果: 与视觉变换器结合时，I-JEPA 表现出高度可扩展性，能够在不到 72 小时内在 ImageNet 上使用 16 个 A100 GPU 训练 ViT-Huge/14，并在多个任务上实现强劲的下游性能。

读后感

之前的方法要么学习本质，要么学习表象，I-JEPA 是一种自监督学习方法，旨在同时捕捉图像的本质（语义）和表象（细节）特征。与传统的生成式方法（如 MAE）不同，I-JEPA 的主要特点包括：

  • 非生成式：I-JEPA 不重建原始图像，而是预测目标区域在嵌入空间中的表示。
  • 嵌入空间的预测：损失函数在嵌入空间中计算，避免了对像素级细节的过度关注，从而更专注于语义特征的学习。

这种方法使模型能够学习更具语义性的表示，同时减少对低级细节的依赖。

1 介绍

在计算机视觉领域，主流的自监督学习方法大致分为两类：

  • 基于不变性的学习方法：这类方法（如 SimCLR、BYOL 等）通过对同一图像的不同视图进行编码，并使它们的表示尽可能接近，从而学习出具有语义一致性的嵌入。但这些方法往往依赖大量手工设计的数据增强（如随机裁剪、颜色扰动等），引入了较强的先验偏置，在面对不同数据分布的任务时可能会表现不佳。
  • 生成式方法：比如 MAE 等，它们直接重建被遮挡的图像区域。这种方式能较好地捕捉细节，但训练目标落在像素级，容易让模型过度关注低层特征而忽视语义信息。

I-JEPA（Image-based Joint-Embedding Predictive Architecture）提出了一个折中方案：它保留了生成式方法中“预测缺失信息”的思路，但关键区别在于——I-JEPA 并不重建图像本身，而是预测目标块在嵌入空间中的表示，从而跳过了像素级还原，直接引导模型学习更抽象的、语义层次更高的表示。

这个思想也可以放在能量模型（EBM）的框架下理解：

  • 传统能量模型旨在为“兼容的输入 - 输出对”分配低能量（即高相似度），为不兼容的输入分配高能量。
  • I-JEPA 将这一原则应用于表示空间：上下文和目标嵌入相似 ⇒ 表示一致 ⇒ 能量低。

2 背景

2.1.1 联合嵌入架构（Joint Embedding Architecture, JEA）

在联合嵌入架构中，模型的目标是学习一个嵌入空间，使得兼容的输入对（例如，同一图像的不同视图）在该空间中的表示相似，而不兼容的输入对（例如，不同图像）在该空间中的表示差异较大。

  • 输入：一对输入 \((x, y)\)，如同一图像的不同视图。
  • 编码器：分别将 \(x\) 和 \(y\) 映射到嵌入空间，得到 \(s_x\) 和 \(s_y\)。
  • 判别器：评估 \(s_x\) 和 \(s_y\) 的相似度，优化目标是使兼容对的相似度高，不兼容对的相似度低。

这种方法常用于对比学习中，如 SimCLR 和 MoCo。

2.1.2 生成式架构（Generative Architecture）

生成式架构的目标是从输入 \(x\) 重建输出 \(y\)，通常通过引入潜在变量 \(z\) 来捕捉数据的生成过程。

  • 输入：输入 \(x\)。
  • 编码器：将 \(x\) 编码为潜在表示 \(z\)。
  • 解码器：根据 \(z\) 生成输出 \(\hat{y}\)，并与真实的 \(y\) 进行比较，优化重建误差。

这种方法在自动编码器和变分自动编码器（VAE）中广泛应用。victor-explore.github.io

2.1.3 联合嵌入预测架构（Joint Embedding Predictive Architecture, JEPA）

JEPA 结合了上述两种架构的优点，其目标是从输入 \(x\) 的嵌入表示 \(s_x\) 预测输出 \(y\) 的嵌入表示 \(s_y\)，而不是直接重建 \(y\)。

  • 输入：输入 \(x\) 和目标 \(y\)。
  • 编码器：将 \(x\) 编码为 \(s_x\)，将 \(y\) 编码为 \(s_y\)。
  • 预测器：从 \(s_x\) 预测 \(s_y\)，优化目标是最小化预测的 \(s_y\) 与真实 \(s_y\) 之间的差异。

这种方法的优势在于，它关注于学习语义层面的表示，而不是像素级的重建，从而提高了模型对语义信息的捕捉能力。

图 2 直观地展示了三种架构的区别：

  • 联合嵌入架构：两个编码器分别处理输入对，输出的嵌入通过判别器进行相似度评估。
  • 生成式架构：编码器将输入映射到潜在空间，解码器从潜在表示生成输出，优化重建误差。
  • 联合嵌入预测架构：编码器将输入和目标分别映射到嵌入空间，预测器从输入的嵌入预测目标的嵌入，优化预测误差。

2.2 方法

I-JEPA 的架构包括三个主要组件：

  • Context Encoder：对输入图像的上下文块进行编码，生成上下文表示。
  • Target Encoder：对目标块进行编码，生成目标表示。
  • Predictor：基于上下文表示预测目标块的表示。

训练过程中，模型通过最小化预测表示与目标表示之间的 L2 距离来优化参数。此外，Target Encoder 的参数通过 Context Encoder 参数的指数移动平均（EMA）进行更新，以防止表示崩塌。

图 3 展示了 I-JEPA 的核心流程：

  • 模型从图像中抽取一个上下文块（例如中间区域），由上下文编码器（ViT）提取表示；
  • 然后，模型尝试预测该图像中多个目标块（被遮挡区域）的表示；
  • 为此，它使用一个预测器网络，结合上下文表示和位置信息，来输出对目标块的表示的预测；
  • 真正的目标表示来自另一个称为目标编码器的网络。为了训练稳定，目标编码器的参数是上下文编码器参数的指数移动平均（EMA）更新而来。

这个机制让模型在不看到目标块像素的情况下，通过上下文信息“想象”出这些位置的语义表示。

图 4 说明了如何选择上下文块和目标块：

  • 一张完整图像中，系统会随机选择 4~4 个目标块，这些块的面积占比为图像的 15%~20%，纵横比控制在 0.75~1.5 之间，确保形状接近自然物体；
  • 接着，系统从图像中采样一个较大的上下文块（面积约 85%~100%），作为输入的主要内容；
  • 所有与目标块有重叠的部分会从上下文中剔除，确保上下文和目标不泄漏信息；
  • 这样的策略保证了上下文块既信息丰富，又不会直接暴露目标区域，使预测变得有挑战性但又可行。

3 相关重要论文

  • EBMs: A tutorial on energy-based learning. Predicting structured data
  • JEPA: A path towards autonomous machine intelligence version