复现 I-JEPA 之前我一直觉得自监督视觉预训练无非两条路要么学 MAE 那样在像素空间做重建要么学对比学习那样构造正负样本。I-JEPA 偏偏走了第三条路——它不做像素重建也不要负样本只用随机块掩码加特征空间预测就把 ViT-H/14 在 ImageNet 1% 标签下的线性探测成绩推到了 73.3%。我第一次跑到这个数字时的第一反应是怀疑评估代码写错了复查了两遍之后才确信这条路是真正可复现、可落地的。这篇文章把我从环境搭建到 72 小时预训练、再到线性探测评估的完整过程写下来包括中途遇到的两个会劝退新手的坑。适合想复现 I-JEPA 的研究生和工程师也适合正在 MAE 与 I-JEPA 之间做选型的同学。1. I-JEPA 的设计逻辑为什么把重建目标从像素换成特征1.1 像素重建的瓶颈MAE 把算力花在了哪先回顾一下 MAE 这类掩码图像建模方法。MAE 的做法很直观把输入图像按 patch 切好随机掩掉 75% 的 patch让 encoder 只处理可见部分再把整条序列交给 decoder 去重建被掩掉的像素。这个思路在工程上非常优雅可见 token 少encoder 计算量大减decoder 虽然在像素空间干活但它很浅所以整体训练成本并不夸张。问题在于像素重建这个目标本身跟语义理解并不是一回事。一张图里大量像素属于纹理、背景、颜色渐变这些高频细节重建它们当然有监督作用但模型需要花很多容量去拟合这些和类别关系不大的信息。你让模型输出一条狗的完整瞳孔颜色和毛发纹理它确实能做到但代价是那些一眼看出这是狗的高层特征被稀释在重建任务里。I-JEPA 做的第一件事就是把目标从像素空间搬到特征空间。它不让模型重建 RGB 值而是让模型预测另一视角下 encoder 输出的特征向量。这相当于把任务从照着原图画出来改成描述你推测画面里有什么。后者天然更接近语义也天然丢掉了一堆无关紧要的像素级细节。我复现的时候感受最深的一点是同样跑 12 小时I-JEPA 的线性探测指标往往已经明显高于同卡时下的 MAE说明它的学习效率确实高——判断依据不是账面 loss而是下游任务指标。1.2 能量模型与特征预测I-JEPA 的核心理念I-JEPA 这个名字里的 JEPA 全称是 Joint-Embedding Predictive Architecture联合嵌入预测架构。它背后是能量模型那一套思想自监督学习的目标可以看成让系统对来自真实数据分布的输入给出低能量对不真实或被打乱的输入给出高能量。在视觉预训练里像素重建是一种显式的能量拟合而 I-JEPA 选择在潜在表示空间里做隐式拟合。具体到网络结构上I-JEPA 有两个 encoder一个 context encoder 处理可见 patch一个 target encoder 处理完整图像。关键区别是target encoder 的权重不是靠梯度更新的而是 context encoder 的指数移动平均EMA。这种设计让 target encoder 成为一个稳定但缓慢变化的老师它提供预测目标context encoder 连同后面的 predictor 则像学生努力用自己的输出去匹配老师给出的特征。我理解这个结构最简单的方式是把它类比成开卷考试的老师改卷子。老师手里有完整答案但这个答案会随着教学进度缓慢更新学生每次只看到卷子的一部分却要写出被遮住部分的要点。因为没有负样本整个系统不存在把正样本从负样本里拉出来的压力它只需要保证预测结果越接近老师越好。那为什么不坍塌成一个常数输出因为 target 本身也在变而且学生端的 predictor 必须抢在老师变化之前追上当前目标不对称结构天然打破了对称坍缩的可能。1.3 不对称结构为什么没有负样本也不会坍塌没有负样本的自监督方法最怕特征坍塌所有样本被映射到同一个点loss 照样很低。SimCLR 用大量负样本对抗这个问题VICReg 用方差和协方差正则项硬性约束而 I-JEPA 只靠结构不对称。我在复现过程中专门验证过这一点把 target encoder 的 EMA 从 0.996 改成 0.9999训练的稳定性明显变差但又不是立刻坍缩把 predictor 去掉、让 context encoder 直接对齐 target跑了 2 万 step 之后特征开始明显退化。这说明 predictor 的存在不是锦上添花而是防止坍塌的关键一环。预测器相当于在学生和老师之间加了一个缓冲垫就算收敛到一个平凡解也需要同时骗过两个网络的联合作用难度比单网络自重建高得多。2. 复现环境与代码骨架硬件选型、依赖版本与模块划分2.1 硬件底线为什么我直接上了 8 卡 A100ViT-H/14 的参数量接近 6.3 亿这在自监督预训练里算是偏大的模型了。单卡 A100 80G 在完美情况下也只能装下很小的 batch而自监督训练又恰恰想吃大 batch。我算过一笔账AdamW 需要保存 fp32 的 master weight 和两份动量加上 bf16 参数、梯度和激活值单卡 80G 的显存大概只够放下每卡 128 左右的 batch size。要凑到 2000 以上的全局 batch要么上 8 卡要么用梯度累积后者会拖慢速度。最终我选择 8 卡 A100 80G并行方式是最简单的 DDP。没有用 DeepSpeed 或 FSDP因为 ViT-H 虽然在单卡上是 6 亿参数但 8 卡各放一份完整模型副本每份模型加优化器状态大约 25GB勉强能塞进 80G配合 bf16 混合精度和 gradient checkpointing 之后每卡跑 128 的 micro-batch 是稳的。曾经测试过只开 bf16 不开 checkpoint每卡只能放到 96太低。checkpoint 会多花大约 15% 的算力但换来了 1.3 倍的有效吞吐很划算。2.2 代码骨架与目录划分我参考官方 release 的思路把复现代码拆成了四个模块数据、模型、掩码、训练主循环。这里没有用一个巨大的 train.py 把逻辑全揉进去而是按职责分开后面调参时才不会改一处崩三处。数据模块里有一个值得注意的设计I-JEPA 官方路线刻意去掉了常见的数据增强。没有颜色抖动、灰度化、旋转训练时只保留随机调整大小裁剪和水平翻转。这个选择的核心逻辑是I-JEPA 的随机性来源是 multi-block masking而不是图像变换。如果在预训练阶段加了大量像素级增强反而会干扰预测特征这个目标。我一开始习惯性地加了一组 color jitter跑了 10 个小时后线性探测反而低 1 个百分点去掉之后才恢复正常。模型模块里我手动实现了 ViT-H/14没有直接套现成的分类模型库。原因很简单自监督预训练需要的是纯 encoder 输出不要最后的分类头而且 I-JEPA 的 context encoder 在处理可见 patch 时需要把序列长度动态变化的信息传给 predictor这个接口和常见 ViT 略有区别。从零实现反而更容易控制细节尤其是位置编码的广播方式。3. ViT-H/14 训练配置逐项拆解掩码、优化器、EMA 与损失3.1 ViT-H/14 架构参数与计算量ViT-H/14 的配置可以总结成一张表复现时直接对着填就行参数数值patch size14×14输入分辨率224×224patch 数 16×16256encoder depth32hidden size1280attention heads16MLP hidden5120参数量约 6.3 亿GFLOPs单张 224300 上下这个体量的模型用 bf16 训练时单卡前向加反向大约是每秒十几个 batch。8 卡加 DDP全局吞吐大约能到每秒 200 张图以上我用 batch size 2048 时一个 step 大约 8 到 10 秒。也就是说72 小时能跑到 4 万 step 左右对应 ImageNet-1K 上大约 60 到 70 个 epoch 的数据量这比原论文的完整训练周期短不少但足够观察到 I-JEPA 的核心收益。3.2 Multi-block masking 参数与实现I-JEPA 的掩码策略是 multi-block masking而不是 MAE 那样随机选 75% 的单 token。它的含义是把 16×16 的 patch 网格当作一个二维平面随机生成若干个矩形 block把这些 block 内的 patch 全部掩掉block 之间不允许重叠。我用的配置是 4 个 block每个 block 的面积大约占全图 patch 数的 15%宽高比在 0.75 到 1.5 之间随机采样。这样算下来大约有 60% 的 patch 被掩掉只留下 40% 可见 patch 交给 context encoder。相比 MAE 的 75% 掩码I-JEPA 的可见比例更高但因为 context encoder 只需要处理可见部分计算量仍然可控。掩码生成的实现有个容易出错的细节。我一开始用 float32 坐标生成矩形然后直接取整结果偶尔会越界导致 mask 张量里出现多余的掩码。正确做法是用整数坐标加 clamp并确保 block 与 block 之间用一个 mask buffer 去重。下面是简化版的核心代码实际项目中我会加上多 rank 的随机种子隔离def multi_block_masking(grid_size, num_blocks, mask_scale, aspect_ratio_range): mask torch.zeros(grid_size, dtypetorch.bool) for _ in range(num_blocks): h, w grid_size aspect random.uniform(*aspect_ratio_range) area h * w * mask_scale bh int((area * aspect) ** 0.5) bw int(area / bh) top random.randint(0, h - bh) left random.randint(0, w - bw) mask[top:topbh, left:leftbw] True return mask这个实现直接反映了 I-JEPA 的重要思想掩码尺寸和长宽比可变模型必须适应各种尺度的局部缺失而不是只会补固定形状的洞。矩形 block 也保证了被掩区域有足够的空间连续性迫使模型利用远处可见 token 来推断整体结构。3.3 优化器、学习率调度与 EMA 配置优化器我用的是 AdamWpeak learning rate 1e-3weight decay 0.05beta 取 (0.9, 0.95)。这个组合在 MAE 和 I-JEPA 的官方风格里都很常见。warmup 我设为 3000 步也就是大约 5% 的训练时长随后 cosine 衰减到极小值。EMA 这一项需要单列出来。target encoder 不直接更新权重而是每步按以下方式滑动tau target_tau (1 - target_tau) * (1 math.cos(math.pi * step / total_steps)) / 2 target_params tau * target_params (1 - tau) * current_params其中 target_tau 初始为 0.996total_steps 取这次 72 小时训练的总步数。我多次强调这个 EMA schedule 要和 optimizer step 对齐不能自己另开一个迭代器计数否则 teacher 权重变化速度和训练步数对不上训练后期会出现明显的特征漂移。3.4 Predictor 与损失函数细节I-JEPA 的 predictor 是一个 6 层的 Transformer宽度和 context encoder 对齐也就是 1280。它的输入有两部分一是 context encoder 对可见 patch 的输出二是若干可学习的 mask token每个被掩掉的 block 对应一个 mask token。经过 predictor 后每个 mask token 输出一个预测向量这个向量的维度等于 patch 特征维度。目标侧的处理是每个 block 一个向量把 target encoder 对被掩 patch 的特征做平均池化得到一个 block 级向量。损失函数就是每个 block 的预测向量与目标向量的均方误差最后除以 block 数。为什么对 block 内特征做平均池化而不是逐 patch 预测因为 I-JEPA 想学习的是这个区域是什么的高层语义而不是逐 patch 对齐低频细节。逐 patch 对齐会让任务退化成特征空间的像素重建等于换汤不换药。损失计算时我会小心一点数值问题。在 bf16 下两个大向量直接算 MSE 可能出现中间溢出我把预测和目标都先转成 fp32 再算 loss这几乎不增加计算量却能避免很多莫名其妙的 NaN。4. 72 小时训练实录曲线解读、瓶颈记录与翻车现场4.1 72 小时时间线loss 下行与特征演进整个训练过程我大概可以分成三个阶段。前 12 小时是快速下降期。loss 在头几千步下降很陡EMA 的 tau 还趴在 0.996 附近模型在快速建立基本特征结构。这段时间我一般不打断训练只看梯度 norm 是不是稳定偶尔确认没有 NaN。12 到 48 小时是平台爬坡期。loss 看起来已经不太动了但如果你每 6 小时做一次线性探测测试会发现 top-1 还在稳步上涨。这是自监督训练最迷惑人的地方账面 loss 的绝对值不能直接说明特征质量必须靠下游任务指标来判断。我在第 20 小时左右第一次跑线性探测当时 1% 指标大约 68%到第 48 小时已经接近 72%这说明 EMA teacher 带来的目标演进仍然在持续改善表示。48 到 72 小时是精细调整期。loss 下降非常平缓但 1% 线性探测还是在涨最终跑到了 73.3%。从曲线形态看如果继续训练应该还能小幅上升但 72 小时的卡时预算刚好到这里。我拿到的结论是I-JEPA 的收益曲线在卡时上非常友好即使只跑一半的训练轮数下游线性探测也不会差太多。4.2 翻车现场一bf16 下的 loss 爆炸第一次完整跑训练我在第 9000 步左右遇到 loss 突然变成 NaN。排查后定位到原因bf16 的指数位太少当某些 mask token 的梯度在跨越多个数量级时prediction 向量里的数值会膨胀到 65504 以上直接溢出。这在高学习率预热阶段尤其容易出现。解决方案有三步。第一步是把 loss 计算放到 fp32不要直接用 bf16 算 MSE。第二步是给 optimizer 设了一个更稳的 eps 值 1e-6防止梯度方向在极端情况下抖动。第三步是为保险起见在损失函数里对预测向量做了 clamp把范数限制在一个合理范围。这样处理后整个训练过程没有再出现 NaN。这里有个教训混合精度不只是打开 autocast 就完事。损失计算、EMA 更新、mask token 初始化这几个环节都可能成为精度短板。我之前在别的项目里被 fp16 的 loss scaling 坑过这次直接用 bf16 配合 fp32 loss等于从源头规避了大部分问题。4.3 翻车现场二EMA 更新与 scheduler 脱钩第二次翻车更隐蔽。我当时把 EMA 更新的 step 计数器和 optimizer 的 global_step 分开维护理论上两者应该一致但由于 DDP 的梯度同步和 EMA 更新的代码位置不同实际步数差出了几百步。结果是在 warmup 结束、学习率开始下降时target encoder 的更新速度明显偏离设计曲线特征质量不升反降。从现象上判断经验是如果 EMA 曲线在 TensorBoard 里看起来比学习率衰减慢半拍基本可以断定计数器不同步。解决办法很简单EMA 更新直接取 optimizer 的 global_step不要自己维护整数计数器。改成这样之后EMA 的 tau 曲线和学习率曲线严格同步特征质量稳步回升。5. ImageNet 1% 线性探测评估协议细节与 73.3% 的解读5.1 为什么用线性探测来评估预训练质量自监督预训练的效果评估有很多种协议最常用的是线性探测和端到端微调。线性探测指的是冻结 backbone只在最后一层接一个线性分类头只训练这个分类头。它考察的核心问题是预训练得到的特征本身是否线性可分。相比端到端微调线性探测对预训练质量的反映更纯粹训练成本也低得多在算力有限时是首选。1% 线性探测则更极端从 ImageNet-1K 训练集里只拿出 1% 的标签数据约 1.28 万张图每个类别大约 13 张。这个设定考察的是模型在极小标签下的特征判别力。能用这么少数据拉到 73.3%说明预训练模型已经把大部分类别的语义结构学得比较清楚线性分类头只需要做轻度的边界划分。5.2 我的完整评估协议我采用的协议基本对齐官方标准把 224×224 的图过一个冻结的 ViT-H/14取 CLS token 作为全局特征接一个 1280 维输入、1000 类输出的线性层。分类头用 LARS 优化器初始学习率 0.1weight decay 设为 0训练 90 个 epoch。LARS 在自监督线性探测里用得很普遍它对大 batch 的适配性比 AdamW 好。我在实验中也用 AdamW 复测过一遍学习率调到 1e-3 之后最终成绩和 LARS 相差在 0.2 个百分点以内。所以如果你手边的复现环境没有现成 LARS直接用 AdamW 也不是不可以但要注意学习率不能照搬 0.1。评估时我只用单中心裁剪没有做多尺度集成。多尺度通常能再涨零点几个点但会把横向对比变得复杂为了公平起见我严格控制输入为单一 224 中心裁剪。测试集用的是标准的 50k 张验证图片报的是 single-crop top-1 正确率。5.3 73.3% 意味着什么我必须先说清楚73.3% 这个数字是在 72 小时训练这个前提下拿到的不是原论文复现的终点。我的卡时预算少于原论文完整训练所以严格说这是一个缩短训练后的下界成绩。即便如此这个数字在自监督视觉预训练里已经属于第一梯队。对比之下我过去在同卡时下复现 MAE ViT-H 时1% 线性探测大约只能落在 60% 出头的水平I-JEPA 高出近 10 个百分点。当然MAE 在端到端迁移任务里仍然有它的优势但单看少样本线性可分性特征空间预测这条路确实大幅领先像素重建。另外它相比 SimCLR、MoCo 这种对比方法又不需要大批次和负样本队列实现和调参难度低一截。这个便宜 高效的组合正是我复现后最欣赏它的地方。复现时注意73.3% 不是每次跑都能一比一复现。数据采样、随机种子、DDP 通信顺序都会带来 ±0.5% 的浮动。我自己的三次实验中最好的结果到过 73.6%最低一次是 72.8%最终以中间值 73.3% 作为报告数字。6. 复现中的高频坑位与选型建议6.1 容易被忽略的三个实现细节第一个细节是 target encoder 必须彻底冻结梯度。我在代码里给 target encoder 的每个参数都设了 requires_gradFalse并且前向输出统一 detach。看起来这是基本操作但在 DDP 场景下有时候会忘记把 target encoder 的 buffer 也同步比如位置编码这类常数。我的建议是写一个frozen_logger在训练中定期打印 target encoder 的参数是否发生变化防止静默出 bug。第二个细节是 mask 生成需要和 DDP 的 shuffle 顺序解耦。每个 rank 用独立的随机种子生成掩码否则所有卡会得到完全相同的 mask等于变相缩小了有效 batch 的语义多样性。我最初没注意导致每个 step 的有效数据增强严重不足训练指标停滞了很久才发现。第三个细节是线性探测阶段要确保 backbone 完全 eval 模式。ViT 里有 dropout 和 LayerNorm虽然 LayerNorm 不带 batch 统计但 dropout 开没开对于 12800 张图的小样本训练影响很大。我在做线性探测时把 backbone 和 predictor 全部切到 eval 状态只在分类头上保留 train 状态这个结果才真正反映预训练特征质量。6.2 I-JEPA 更适合哪些场景经过这轮复现我个人的选型判断是如果目标以 ImageNet 级别的少样本分类、特征泛化性为主I-JEPA 这类特征空间预测方法值得优先尝试如果你要的是端到端微调到检测分割时的高密度像素预测MAE 的像素重建仍然有它的优势因为它在梯度上天然保留了对空间细节的敏感性。I-JEPA 的工程门槛也比预想低很多。不需要负样本队列不需要大 batch 到几千我在 8 卡 A100 上用 2048 的全局 batch 就能稳定训练。相比某些对比学习方法对 batch size 的苛刻要求这种不挑食的特性在实际资源有限的环境里非常重要。最后再分享一个我自己踩出来的小经验复现 I-JEPA 时不要频繁打断训练去看 loss 绝对值重点关注 EMA tau、梯度 norm 和每 10 小时做一次的线性探测趋势。当你看到 EMA tau 顺畅爬向 1.0而 1% 线性探测每过一夜都涨一个点时基本可以确信这一轮训练正在按照设计意图推进。如果非要给后来者留一句话那就是I-JEPA 的每个设计模块——multi-block masking、EMA teacher、特征空间损失——都不是装饰而是环环相扣的必然少任何一个你都会得到一个表现平庸的模型。