Masked Modeling 与 JEPA

一台机械臂的腕部相机每秒产生几十帧图像。要让它学会抓取，最直接的办法是监督学习：人工标注每张图里物体在哪、该从什么角度抓，再训练一个网络（16.1 节）。问题是标注很贵，而机器人演示数据更贵：一条遥操作演示要人花几十秒到几分钟（22.3 节），一个任务通常只有几十到几百条。从零训练的视觉网络在这么少的数据上很容易过拟合（16.2 节）。

可没有标签的图像几乎是免费的：机器人自己乱动时拍的视频、互联网上的视频、人做家务的第一人称视频。自监督学习的思路是：从数据本身造出一个“题目”，让网络在海量无标签数据上做题，学到一个好的表示（representation），再用少量标注数据把它接到具体任务上（21.1 节讲微调）。18.3 节的对比学习是一种题目：“这两张图是不是同一个东西的两种增强？”本节讲另一类，也是今天大模型最常用的一类：遮住一部分，猜回来（masked modeling）。

它从语言开始。BERT 把句子里 15% 的词遮住，让 Transformer 根据两侧的上下文猜它们；GPT 一类的自回归模型（19.4 节）只看左侧，猜下一个词。图像上，MAE（masked autoencoder）把 75% 的图块遮住，只把剩下的图块送进编码器，再用一个小解码器重建全部像素。JEPA（joint-embedding predictive architecture）则认为重建像素是在浪费力气：被遮住的那块地毯上每根绒毛的样子本来就猜不出来。它改成在表示空间里预测：用一个目标编码器把被遮住的区域编码成向量，让预测器去猜这个向量。这样做省掉了像素细节，却引入了一个新危险：如果编码器把所有东西都编码成同一个常数，“猜”就永远是对的，这叫表示塌缩。I-JEPA、V-JEPA，以及 BYOL、DINO 等方法都靠停止梯度、指数滑动平均（EMA）目标编码器和预测器这一类设计来避免它。

对机器人来说，结果是一批可以直接拿来用的预训练视觉编码器：MVP 系列工作在大量第一人称视频上用 MAE 预训练 ViT，冻结后接一个小策略网络，比 CLIP 或 ImageNet 监督预训练的编码器更适合操作任务；V-JEPA 2 进一步把 JEPA 编码器和一个动作条件预测器组合成世界模型，用来规划抓放。本节最后（FIG 6）会在一个小例子里亲手体会这一点：预训练让一个原本难以从像素读出的量（工具的朝向）变得可以用几十条标注读出。

先从语言开始。下面是一个只有 180 句话的小语料库，全是给机器人的指令，形如“grab the red cube and put it in the bowl”。它有几条隐藏规则：颜色和物体可以随便组合；容器 bowl、box 用介词 in，tray、shelf 用 on；动词 drop 只能配 in。你点一个词把它遮住（变成 [MASK]），图里就给出在这个语料库上能达到的最好预测：数一数所有与未遮住的词一致的句子，被遮住的位置上各个词出现的比例。这正是一个无限大、训练完美的 BERT 会给出的分布。左栏用两侧的上下文，右栏只用左侧（像 GPT 那样）。每栏下面的 \(\hl{H}{H}\) 是这个分布的熵（4.11 节），单位是比特：0 表示确定，\(\log_2 3\approx1.58\) 表示三选一完全不确定。

反过来，在动词不是 drop 的句子里把介词和后面的容器一起遮住，再看介词：两侧上下文也帮不上忙了，\(H\) 回到 1 比特。遮得越多，题越难。再遮一下颜色词：无论看哪边都是 1.58 比特，因为语料里没有任何东西决定颜色。这种不可约的不确定性很重要：一个被训练去预测颜色的模型，最好也只能输出三分之一、三分之一、三分之一，损失永远降不到零，也不应该降到零。语言里这样的词不多，因为词本身已经是高度压缩的符号；可图像里到处都是这样的东西：一片纹理里每个像素的精确值、传感器噪声、反光的形状。下面看图像。

FIG 2 用的是本节贯穿始终的玩具图像：一台俯视桌面的机器人相机拍到的 12×12 像素画面，上方是背景，下方是桌面（边缘的高度和倾斜随机），桌上放着一根棒状工具（位置和朝向随机）。图像切成 36 个 2×2 的图块（patch），这和 ViT 把 224×224 的图像切成 16×16 的图块是同一个做法（17.3 节）。遮住一些图块，最笨的猜法是插值：用附近可见像素的加权平均填进去（权重与距离的四次方成反比）；另一个笨办法是不看图，直接填上整个数据集的平均图像。

现在训练一个真正学结构的模型。它是一个小小的“masked autoencoder”：输入是遮挡后的图像（被遮的像素置 0）加上 36 个标志位，告诉网络哪些图块被遮住了；网络是 180 → 64 → 64 → 144 的 MLP（GELU 激活，用本课程的 nn.js 实现），输出完整的 144 个像素。损失只算在被遮住的像素上：可见像素原样抄过去毫无意义。每一步从生成器里新抽 32 张图、每张随机遮挡，用 Adam（学习率 0.002）更新一次。数据无穷无尽，所以这里没有过拟合，测试集只是固定的 200 张图和固定的遮挡，用来公平地画曲线。

真正的 MAE 有两点不同。第一，编码器是 ViT，只处理可见图块（遮 75% 时只有四分之一的 token），解码器很轻，在编码结果中插入“mask token”后重建全部图块，所以训练又快又省（FIG 5）。第二，预测目标通常是每个图块内归一化后的像素。我们的 MLP 没有 token 的概念，只能用“置零加标志位”的方式告诉它哪里被遮住了，但训练目标与 MAE 完全一样。

FIG 3 暴露了像素重建的两个问题。一是浪费：模型的容量和梯度有很大一部分花在噪声、纹理这类本来就猜不出来的细节上（FIG 1 的颜色词、FIG 3 的噪声地板）。二是目标不对：机器人要的是“工具在哪、朝哪”，而不是每个像素的亮度，最小化像素误差并不直接奖励前者。Yann LeCun 在 2022 年的立场论文里提出的 JEPA 换了一个目标：把输入分成上下文 \(x\) 和目标 \(y\)，用编码器分别把它们变成表示，让预测器从 \(x\) 的表示预测 \(y\) 的表示：

\(f_\theta\) 是上下文编码器（被训练的那个），\(g_\phi\) 是预测器，\(f_\xi\) 是目标编码器：它不被梯度训练，而是上下文编码器参数的指数滑动平均（EMA，\(\tau\) 接近 1）。\(\operatorname{sg}\) 表示停止梯度（stop-gradient）：反向传播时把目标当成常数。LN 是层归一化（17.3 节），把每个目标向量减去均值、除以标准差。I-JEPA 中 \(x\) 是图像里一大块上下文区域的图块，\(y\) 是另外 4 个目标块，预测器还被告知目标块的位置；V-JEPA 把这套搬到视频上。我们的玩具做了简化：上下文是 FIG 3 那样遮掉 75% 的图像，目标是整幅图像的一个 8 维表示，编码器与 MAE 的前两层同构（180 → 64 → 64 → 8），预测器是 8 → 32 → 8。

危险在于：上面的损失有一个平凡的全局最小值。若编码器对任何输入都输出同一个向量，预测器只要输出这个常数，损失就是零。这叫完全塌缩，得到的表示没有任何信息。18.3 节的对比学习靠负样本防止它：InfoNCE 要求不同图像的表示彼此分开。JEPA 没有负样本，它靠的是结构上的不对称：只有一侧有预测器，另一侧停止梯度、缓慢地跟随（EMA）。下面的实验比较“有这套不对称”和“把梯度也传进目标一侧”两种情况。衡量塌缩的指标是展开度 \(\hl{spread}{s}\)：150 张测试图的目标表示在每一维上的标准差的平均。完全塌缩时 \(s=0\)。

这个小实验和文献的结论一致：把梯度传进目标一侧，网络立刻找到平凡解；停止梯度加预测器就能避开它（SimSiam 显示，甚至 \(\tau=0\) 也行，我们的玩具里它先差点塌缩、再恢复）；EMA 让目标慢慢变化，训练更稳。四档 \(\tau\) 在 4000 步内的最低展开度分别约为 0.04（\(\tau=0\)）、0.11（0.9）、0.45（0.99）和 0.55（0.999，从未低于初值 0.53）。真实系统用的 \(\tau\) 更接近 1：I-JEPA 从 0.996 线性增加到 1.0。要注意，为什么这套不对称能防止塌缩，目前没有完整的理论保证（第 5 节）；塌缩的解依然存在，只是训练动力学通常不会走到那里。

遮挡策略决定了题目的难度，也决定了计算量。下图在一个 14×14 的 token 网格上（224×224 的图像、16×16 的图块）画出三种策略。MAE 随机遮 75% 的图块，编码器只处理剩下的 49 个。I-JEPA 先随机取 4 个目标块（每块占图像的 15% 到 20%，宽高比 0.75 到 1.5），再取一个占 85% 到 100% 的方形上下文块，并从上下文里删去与目标重叠的部分：目标块要大到有“语义”，上下文要足够分散。V-JEPA 把同样的空间块在所有时间步上重复，形成贯穿整段视频的“管道”，否则模型只需从相邻帧抄过来；它有两种掩码：8 个占 15% 的块的并集（短程），和 2 个占 70% 的块的并集（长程）。

在我们的采样器里（块的位置均匀随机，大小取整到网格），I-JEPA 的上下文平均保留约 90 个 token（46%），4 个目标的并集平均约 94 个（48%）；V-JEPA 的长程掩码平均遮住约 84%，短程约 63%。论文报告的平均遮挡比例约为 90%，与我们的简化采样器不完全一样。不论细节，三者的共同点是：被预测的部分很大，而且不能靠局部复制得到。

设输入由 \(n\) 个 token（词或图块）组成，\(\hl{mask}{\mathcal M}\subset\{1,\dots,n\}\) 是随机抽出的遮挡集合，\(x_{\setminus\mathcal M}\) 是可见部分。遮挡建模最小化

BERT 的 token 是词，\(\ell\) 是交叉熵（16.1 节），\(|\mathcal M|\approx0.15n\)；被选中的位置 80% 换成 [MASK]、10% 换成随机词、10% 保持不变，因为下游使用时输入里不会出现 [MASK]。MAE 的 token 是图块，\(\ell\) 是平方误差，\(|\mathcal M|=0.75n\)。两者的最优预测器都可以写出来。在交叉熵下，最优输出就是条件分布 \(p(x_i\mid x_{\setminus\mathcal M})\)，最小损失是条件熵：

这就是 FIG 1 显示的数：两侧上下文时介词的条件熵为 0，颜色为 \(\log_23\)。在平方误差下，最优输出是条件均值，最小损失是条件方差：

证明只需一行：对任意 \(c\)，\(\mathbb E[(c-x)^2]=(c-\mathbb E x)^2+\operatorname{Var}x\)，在 \(c=\mathbb E x\) 时最小（4.3 节）。两个后果在 FIG 3 里都看到了：被遮内容有多种可能时，最优重建是它们的平均，所以是模糊的；若像素带有与一切无关的噪声 \(\varepsilon\sim\mathcal N(0,\sigma^2)\)，条件方差至少是 \(\sigma^2\)，损失永远不会低于这条地板。

一层宽度为 \(D\) 的 Transformer 处理 \(n\) 个 token 时，QKV 和输出投影约 \(4nD^2\) 次乘加，MLP（隐藏层 \(4D\)）约 \(8nD^2\)，注意力矩阵及其加权约 \(2n^2D\)；按每次乘加 2 次运算算，一层约 \(24nD^2+4n^2D\) 次浮点运算（17.3 节）。MAE 编码器只处理 \(\hl{ctx}{n_{\mathrm{vis}}}=(1-r)n\) 个 token：

其中 \(C(n)=24nD^2+4n^2D\)，近似在 token 数远小于 \(6D\) 时成立。图像只有 196 个 token、\(D=1024\) 时，注意力只占一层计算的约 3%，所以编码器的计算量几乎与可见比例成正比：遮 75% 时约为 24%（FIG 5）。解码器很浅很窄，MAE 论文报告每个 token 的计算量不到编码器的 10%。视频的 token 多得多（V-JEPA 是 \(8\times14\times14=1568\) 个），注意力的 \(n^2\) 项更重，高遮挡比例省下的就更多。

本节玩具 JEPA 的损失是 \(\frac1d\|g_\phi(f_\theta(x))-\operatorname{sg}(\bar{\mathbf z})\|^2\)，\(d=8\)。I-JEPA 对每个目标图块的表示做同样的 L2 回归并取平均，V-JEPA 换成了 L1。EMA 更新 \(\xi_{t+1}=\tau\xi_t+(1-\tau)\theta_{t+1}\) 展开后是

第二个式子是权重的平均“年龄”：目标编码器是过去约 \(1/(1-\tau)\) 步在线参数的加权平均（初值的影响按 \(\tau^t\) 消失）：\(\tau=0.99\) 约 100 步，\(\tau=0.996\) 约 250 步。为什么不塌缩？先看塌缩为什么会发生。若不停止梯度，目标 \(\bar{\mathbf z}\) 对 \(\theta\) 的导数也进入更新，梯度下降同时把预测拉向目标、把目标拉向预测；而让所有表示都相等是同时满足两者的最省事办法。FIG 4 里这只用了不到 50 步。停止梯度后，目标在每一步都是常数，编码器只能通过让上下文的表示更好地预测目标来降低损失。BYOL 的作者给出过一个启发式解释（V-JEPA 的论文把它改写成 L1 损失的版本，那时最优预测器是条件中位数）：对平方误差，若预测器始终接近最优，即 \(g^\star=\mathbb E[\bar{\mathbf z}\mid f_\theta(x)]\)，那么损失就等于给定上下文表示时目标的条件方差，编码器的梯度会推动上下文表示携带尽可能多关于目标的信息；EMA 让目标变得比预测器慢，使预测器能一直保持在最优附近。Tian 等（2021）在线性网络上分析了这套动力学，说明预测器和停止梯度共同使塌缩解不再吸引训练。这些都是解释，不是对一般非线性网络的证明；实践中通常还要监控展开度或表示的有效秩。

与对比学习的关系：18.3 节的 InfoNCE 可以拆成两项，一项把正样本对拉近（“对齐”），一项把所有样本推开（“均匀”）。非对比方法只保留了对齐项，用结构上的不对称（BYOL、SimSiam、JEPA）、输出分布的中心化与锐化（DINO），或显式的方差与协方差正则（VICReg，Bardes 等，2022）去代替均匀项。下面是 DINO 的损失，\(\mathbf z_t,\mathbf z_s\) 是教师和学生的输出，\(\mathbf c\) 是中心，\(m\) 是它的动量，教师温度 \(T_t\) 低于学生温度 \(T_s\)：

评价一个表示最常用的办法是线性探针：冻结编码器，取特征 \(\phi(x)\in\R^{p}\)，只用 \(N\) 个标注样本拟合一个线性映射。FIG 6 用岭回归（1.13 节最小二乘加 L2 正则），\(\lambda=1\)，偏置不惩罚：

式中 \(\Phi\) 的每一行是一个样本的特征，先减去 \(N\) 个样本的均值（于是偏置单独求出、不受惩罚）。工具是一根直棒，转 180° 后样子不变，所以它的朝向 \(\alpha\) 只在模 180° 意义下有定义；用 \((\cos2\alpha,\sin2\alpha)\) 表示朝向，可以避开 0° 与 180° 处的跳变，预测后再用 \(\tfrac12\operatorname{atan2}\) 取回角度。线性探针测的是“信息是否以线性可读的形式存在”。全量微调（21.1 节）会改变编码器本身，往往得分更高；MAE 的特征就以“线性探针一般、微调很强”著称，而 I-JEPA 的论文报告它的线性探针明显好于 MAE。

MVP 的配方很简单：先在大量无标签图像上用 MAE 预训练一个视觉编码器，冻结它，再在它的特征上用很少的机器人数据训练一个小策略。下图在玩具世界里重演这个配方。机器人要抓桌上的工具，需要知道两件事：抓取点（工具中心，像素坐标）和夹爪的偏航角（工具的朝向）。先不用任何标签，各训练 2000 步的 MAE 和 JEPA；然后冻结，取每个编码器两层隐藏层的 128 个激活作为特征，用 \(N\) 张有标注的图拟合线性探针，在另外 300 张图上测误差。对照组是原始像素（144 维）和一个结构相同、但没训练过的随机编码器。

这个小结果包含了真实系统里的两条经验。第一，预训练的价值在于把原始输入里非线性地纠缠在一起的量（朝向需要比较像素之间的关系）变成可以用少量标签读出来的形式；300 个标注时，像素探针的偏航角误差约 19°，MAE 和 JEPA 约 11°。第二，冻结的通用特征不一定保留任务需要的每个细节：精确的位置信息在压缩时可能丢失一部分。真实的机器人系统因此常常把冻结的视觉特征和本体感受（关节角、夹爪状态）一起输入策略，或者在机器人数据上再做微调。