注解版JEPA

# 注释版 JEPA（联合嵌入预测架构） 来源：https://elonlit.com/scrivings/the-annotated-jepa/ 译者：Elon Litman 花朵 本文是一份逐步、带注释、从头开始的 JEPA（联合嵌入预测架构）实战指南。目标是仿照《注释版Transformer》对Transformer所做的那样：构建完整对象，解释每个活动部件，最后以一个可工作的训练循环收尾。 JEPA 是 Yann LeCun 提出的，用于解决自监督学习中一个基本问题的答案：如何在没有标签、不坍缩到平凡解、不浪费容量在无关细节上的情况下，训练一个能理解世界的模型？ 答案，在原则上优雅而在实践中微妙，是在**表征空间或潜在空间中进行预测**。 为使讨论具体化，主要运行示例是 I-JEPA，即图像实例化。为什么选择图像和视频而非文本？LeCun 认为语言已经是高度压缩的离散知识表征；预测下一个词元需要建模人类交流模式，而非物理现实。相比之下，视觉预测要求理解持续性、遮挡和动态。JEPA 的设计针对那些像素级重建会浪费容量在无关细节上的领域，而离散词元不会以相同方式遇到此问题。我们将在末尾回到这一点。，它被引入为一种自监督方法，通过从可见区域预测掩码区域的表征来学习语义图像表征。参见 *Self-Supervised Learning from Images with a Joint-Embedding Predictive Architecture* (2023)。I-JEPA 是一种完全避免手工设计数据增强的非生成式方法。我们将从头构建 I-JEPA，然后讨论其到视频的扩展 V-JEPA 和 V-JEPA 2，参见 *V-JEPA: Latent Video Prediction for Visual Representation Learning* (2024) 以及 *V-JEPA 2: Self-Supervised Video Models Enable Understanding, Prediction and Planning* (2025)。最后审视 LeJEPA，后者是用分布正则化器取代工程启发式的最新尝试。参见 *LeJEPA: Provable and Scalable Self-Supervised Learning Without the Heuristics* (2025)。 以下内容旨在教学。实现省略了 FlashAttention、梯度检查点、混合精度以及使大规模训练可行的批处理策略。这些是会在生产代码库中占主导地位的工程选择，但可以轻易地与数学分离。 ## 问题 自监督表征学习问道：**如何在没有标签的情况下学习有用的特征？**你需要一个能捕捉有意义结构的目标函数，而在没有标签的情况下，找到一个真正有效的目标函数正是该领域的核心难点。 JEPA 的回答是：通过预测来训练，但在表征空间中进行预测。然而，这为何有效？ 假设你看到图像的一部分，即上下文 \(x\)，并想要学习表征。图像中某个地方有一个你看不到的目标区域 \(y\)。一个编码器将 \(y\) 映射为表征 \(s_y\)。一个预测器接收你对上下文的编码，并输出其对 \(s_y\) 的猜测 \(\hat{s}_y\)。训练最小化距离 \(D(\hat{s}_y, s_y)\)。 现在问：预测器何时能成功？仅当上下文编码 \(s_x\) 包含足够信息以确定 \(s_y\) 必须是什么。如果你看到了汽车引擎盖，预测车轮的表征就需要你的引擎盖编码捕捉到“这是一辆车”。如果你看到了一张脸，预测头发的表征就需要你的编码捕捉到身份、姿态和光照。预测器无法凭空捏造出上下文编码所缺乏的结构。 这就是强制函数。上下文编码器必须学会从 \(x\) 中提取能预测 \(y\) 表征的特征。这正是那些语义、结构特征：对象身份、空间关系、物理约束。\(x\) 中的像素级噪声无助于预测 \(s_y\)，因此编码器学会忽略它。剩下的是能泛化的东西。 目标编码器有互补的压力。其输出 \(s_y\) 必须能从上下文中预测。如果 \(s_y\) 编码了随机的高频纹理，再多上下文也无济于事。因此，目标编码器学会输出能捕捉 \(x\) 和 \(y\) 之间共享结构（即使得预测成为可能的结构）的表征，而非 \(y\) 本身的特殊细节。 通往自主机器智能之路 LeCun 的立场文件将其表述为基于能量的公式。参见 *A Path Towards Autonomous Machine Intelligence* (2022)，OpenReview。LeCun 将 JEPA 框架化为在潜在空间中进行规划的世界模型的基础。：将 \(x\) 和 \(y\) 编码为表征，从一个预测另一个，将能量定义为该抽象空间中的预测误差。架构分解为我们可以实现和分析的组件：两个编码器、一个预测器、一个距离函数。 ## JEPA 模板 一个 JEPA 从成对的、语义相关的世界视图开始。在最一般的形式中，想象三元组 \((x, y, z)\)：\(x\) 是你观察到的，\(y\) 是你想要预测的，\(z\) 是一个可选的潜在变量，捕捉使预测成为多模态的未知因素。成对的 \((x, y)\) 来自观察的某个联合分布。在图像预训练中，\(x\) 可能是可见块而 \(y\) 是同一图像的掩码块；在视频中，\(x\) 是片段前缀而 \(y\) 是后续内容；在跨模态设置中，\(x\) 可以是音频而 \(y\) 是对应的视频。唯一的结构性要求是，知道 \(x\) 应当约束 \(y\) 可能的值。 LeCun 的公式增加了一个可选的潜在变量 \(z\) 来处理预测关系中的多模态性。当多个 \(y\) 值与同一个 \(x\) 一致时，预测器以 \(z\) 为条件来选择其中之一。这对于时间预测很重要，因为未来是真正不确定的；而对于掩码图像建模则不那么重要，因为上下文通常决定了直到噪声的目标。 模板有三个部分。**编码**两个观察值到共享的表征空间： $$ s_x = f_\theta(x), \quad s_y = f_{\bar\theta}(y) $$ 编码器 \(f_\theta\) 和 \(f_{\bar\theta}\) 可以在架构或参数共享上有所不同。当 \(x\) 和 \(y\) 处于不同模态时，它们必须不同；当它们属于同一模态时，权重共享是一个设计选择，它在归纳偏置和灵活性之间权衡。 **预测**从上下文表征中预测目标表征： $$\widehat{s}_y = g_\phi(s_x, z)$$ **最小化**表征空间中的预测误差： $$ \mathcal{L}(\theta, \bar\theta, \phi) = D(\widehat{s}_y, s_y) $$ 本文的其余部分实例化了这些抽象：为图像具体定义 \(x\) 和 \(y\)、为编码器和预测器选择架构、指定距离函数 \(D\)，以及实现防止所有表征坍缩到常数的平凡解的机制。 ## 预测 vs. 重建 为什么在表征空间中预测，而不是在像素空间中？ 考虑两个图像块，它们在语义上相同，即同一对象，同一含义，但由于光照、纹理或 JPEG 伪影而在数百个像素上不同。一个重建目标必须解释所有这些差异。模型浪费容量在高熵噪声的建模上。 使用 JEPA，目标编码器可以学会输出丢弃无关细节的表征。当你回归到那个表征时，你推动模型去捕捉编码器所保留的内容，而不是像素恰好包含的内容。I-JEPA 明确了这一点：架构类似于生成模型，但损失存在于嵌入空间，而非输入空间。类似于感知损失 (LPIPS) 与像素损失 (MSE) 的类比。表征空间预测使模型能够专注于语义上有意义的特征。 ## 坍缩 最小化匹配损失最简单的方法是什么？为所有输入输出相同的向量。然后，损失变为零。但表征变得无用。这是所有联合嵌入方法的核心失败模式，而 I-JEPA 直接解决它：他们使用了编码器之间的不对称设计。 文献提供了许多反坍缩策略。SimCLR 使用对比负样本。VICReg 使用方差-不变性-协方差正则化。BYOL 和 SimSiam 使用停止梯度。DINO/MoCo 使用 EMA 教师。每种策略的权衡不同。：对比负样本（昂贵）、显式协方差约束（难以调优）、停止梯度（出奇有效），或带有指数移动平均更新的师生编码器（I-JEPA 的选择）。I-JEPA 通过梯度下降训练上下文编码器和预测器，但通过上下文编码器权重的指数移动平均 (EMA) 来更新目标编码器。这就是它防止坍缩的原因。 考虑如果两个编码器都通过梯度下降训练会发生什么。损失是 \(D(\hat{s}_y, s_y)\)。如果目标编码器可以自由改变，它将学会输出易于预测的表征，即常数。上下文编码器随之而去。两者收敛到为所有输入输出相同的向量：损失为零，表征无用。EMA 打破了这种共同适应。目标编码器更新为 $$\theta_{\text{target}} \leftarrow m \theta_{\text{target}} + (1-m) \theta_{\text{context}}$$ 其中 \(m\) 接近 1。这意味着目标编码器滞后于上下文编码器。目标表征在梯度步的时间尺度上是稳定的，因此上下文编码器无法利用快速移动的目标来寻找退化解。上下文编码器必须真正学会预测（缓慢演化的）目标编码器的输出，这迫使它提取有意义的特征。 稍后我们将讨论 LeJEPA，它试图使防止坍缩从启发式变得有原则。目前，我们专注于经典的 I-JEPA 设计。 ## 将 JEPA 特化到图像 一个 JEPA 需要共享语义的成对视图 \(x\) 和 \(y\)。对于 I-JEPA，两者来自同一图像：\(x\) 是块的一个大上下文块，而 \(y\) 是图像中其他位置的几个较小的目标块。目标是从上下文块预测目标块的表征。在编写任何代码之前，您应该理解 I-JEPA 中的三个关键设计选择。 第一个选择是像视觉Transformer一样，将图像表示为一系列补丁词元。参见 Dosovitskiy 等人，《An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale》(2020)。ViT 架构将图像分割成固定大小的补丁，并将其作为序列处理。一个 224×224 图像，补丁大小为 16×16，产生 196 个词元，排列成 14×14 网格。 第二个选择是掩码策略。论文指出，为了引导 I-JEPA 走向语义表征，预测尺度足够大的目标块至关重要，并且要使用一个信息量足够且空间分布广泛的上下文块。在方法部分，论文指定了典型的掩码采样范围：通常使用 \(M = 4\) 个目标块，宽高比范围 \((0.75, 1.5)\)，尺度范围 \((0.15, 0.2)\)。上下文块以尺度范围 \((0.85, 1.0)\) 和单位宽高比采样，并且移除目标和上下文之间的任何重叠，以使预测任务非平凡。 第三个选择是目标掩码发生的位置。目标块是通过掩码目标编码器的输出获得的，而不是输入。目标编码器看到完整图像并产生高层表征；然后我们选择要预测哪些表征。论文说这个区别至关重要，以确保目标表征具有高层语义。这在计算上比同时掩码两个视图的输入更重，但对于正确性是必要的。由于 ViT 使用全局注意力，一个补丁的表征依赖于完整的图像上下文。仅对目标补丁运行目标编码器会改变这些表征的含义。 这些不是随意的实现细节。它们是 I-JEPA 归纳偏置的全部要点，用结构化的预测任务取代手工设计的多视图增强。 ## 完整的图像-JEPA 架构 I-JEPA架构。我们将实现三个网络。 上下文编码器（也称为**学生**）是一个 Vision Transformer，它只处理可见的上下文补丁，并产生一组上下文嵌入。 目标编码器（也称为**教师**）是另一个 Vision Transformer，它处理完整图像并为所有补丁产生嵌入。其权重不通过梯度更新；而是通过上下文编码器权重的指数移动平均来更新。 预测器是一个较小的 Transformer，它接受上下文嵌入和一组对应于目标补丁位置的掩码词元，并为这些目标位置产生预测的嵌入。方法描述说，预测器的输入是上下文编码器输出加上每个要预测的补丁的掩码词元，掩码词元由一个共享的可学习向量和添加的位置嵌入参数化。 损失是目标块预测的补丁表征与目标编码器的补丁表征之间的均方误差。这是 JEPA 模板中的距离函数 \(D\)： $$D(\hat{s}_y, s_y) = \|\hat{s}_y - s_y\|^2_2$$ ## PyTorch 实现 该实现被分解为多个部分；您应该能够从头到尾阅读。通篇我们将使用这些约定：图像的形状为 \((B, C, H, W)\)，补丁词元的形状为 \((B, N, D)\)，其中 \(N\) 是补丁数量，\(D\) 是嵌入维度。每当我们选择补丁的子集时，我们将该子集表示为到 \(N\) 个补丁的索引向量，并且我们将为简单起见实现绝对位置嵌入。 ### 导入和小工具 ``python import math import random from dataclasses import dataclass from typing import List, Optional, Tuple import torch import torch.nn as nn import torch.nn.functional as F `` ### 补丁嵌入 一个标准的 Vision Transformer 使用步长卷积创建补丁词元。每个补丁变成 \(\mathbb{R}^D\) 中的一个向量。 ``python class PatchEmbed(nn.Module): """ 将图像转换为补丁嵌入序列。 输入： (B, C, H, W) 输出： (B, N, D) 其中 N = (H / P) * (W / P) """ def __init__( self, img_size: int = 224, patch_size: int = 16, in_chans: int = 3, embed_dim: int = 768, ): super().__init__() assert img_size % patch_size == 0, "img_size 必须能被 patch_size 整除" self.img_size = img_size self.patch_size = patch_size self.grid_size = img_size // patch_size self.num_patches = self.grid_size * self.grid_size self.proj = nn.Conv2d( in_channels=in_chans, out_channels=embed_dim, kernel_size=patch_size, stride=patch_size, bias=True, ) def forward(self, x: torch.Tensor) -> torch.Tensor: x = self.proj(x) # (B, D, G, G) x = x.flatten(2).transpose(1, 2) # (B, N, D) return x `` 唯一非显而易见的行是重塑。