★ FEATURED ARTICLE 读论文和亲手复现完全是两回事。尤其是 I-JEPA 这种看起来“公式简单、逻辑清晰”的自监督方法跑通一个 demo 不难但要把 ViT-H/14 这种体量的模型在有限算力下预训练 72 小时并把 ImageNet 1% 线性探测做到 73.3%背后全是对细节的理解和一堆不成文的“土办法”。这篇文章把我这次复现 I-JEPA 架构的完整过程、关键决策和踩过的坑一次讲清楚希望能帮到正在搞自监督预训练、或者准备复现论文的同学少走点弯路。整个项目做完我的最大感受是I-JEPA 不是那种“论文很强、落地翻车”的方法——只要把对称编码器的收敛节奏、目标编码器的更新策略、还有数据加载的瓶颈处理好它确实能像论文里描述的那样用纯自监督的方式学到比像素重建更抽象、更可迁移的特征。而且它对资源和时间的要求也没有想象中那么离谱我用单机多卡跑了三天预训练效果已经足以支撑下游评测。下面按复现的推进顺序把从架构理解到性能实测的每个环节都拆开说。1. I-JEPA 在设计上到底聪明在哪1.1 先说直觉不重建像素而是预测“抽象特征”I-JEPA 的全称是 Image Joint Embedding Predictive Architecture翻译过来就是“图像联合嵌入预测架构”。它和 MAE、SimMIM 这类掩码自编码器的最大区别在于预测的目标不是被 mask 掉的像素而是目标编码器对图像内容提取出的语义特征。我用生活里的例子来打比方MAE 的做法像是“把报纸剪掉一块要求模型把缺的字猜出来”猜的是具体文字而 I-JEPA 的做法像是“把报纸剪掉一块要求模型概括这块内容大概讲的是什么主题”猜的是含义。后者明显更接近人的理解方式也更接近下游任务比如分类、检测真正需要的表征。具体在训练时I-JEPA 把一张图划分成 patch 后会随机选几个区域块block作为预测目标。可见部分的 patch 送入 context encoder得到上下文表示整张图包含被掩码的区域送入 target encoder得到目标表示。然后一个轻量 predictor 基于上下文表示去预测被掩码区域的目标表示用 L2 损失计算两者差距。这里的关键是context encoder 和 target encoder 结构完全相同但后者的权重不通过梯度更新而是由前者权重做 EMA指数滑动平均得到。1.2 为什么“联合嵌入预测”能学得更好这类方法在自监督领域有个专门的说法叫 Joint Embedding Predictive ArchitectureJEPA核心的出发点是把“预测”和“表征”结合起来。早期自监督方法分两派一派是对比学习通过拉近正样本、推开负样本让模型学会区分语义代表性工作比如 SimCLR、MoCo另一派是生成式预训练直接重建图像或文本比如 MAE、BEiT。对比学习的问题在于强依赖负样本的数量和质量batch size 一旦小了效果就会明显打折对算力的要求比较高。重建式方法的优点是训练稳定但容易让模型钻进像素级细节的牛角尖学到的特征里掺杂了很多颜色、纹理、光照这些对语义不太重要的信息。I-JEPA 走的是中间路线不搞负样本避免了对比学习对超大 batch 的依赖也不重建像素避免了重建式方法舍本逐末的问题。它预测的是“语义层面缺失的信息”相当于逼着模型在脑海里先完成一次“语义推理”而这种推理能力恰恰是下游视觉任务最需要的。我当时正是看中这一点才决定在这个项目里用 I-JEPA 而不是 MAE 或者 SimCLR 的变体。1.3 对称编码器架构的几个细节I-JEPA 的模型结构是一个对称的 ViT 双塔设计这是复现时最容易混淆的地方。简单说context encoder只处理可见 patch输入序列短计算量小target encoder处理完整图像的所有 patch输入序列长但因为它不参与梯度更新所以计算开销可以用停梯度stop-gradient和 EMA 的方式“摊薄”predictor一个轻量的 transformer block输入是 context encoder 对 mask block 位置的输出输出是预测的 target representation。训练时target encoder 的权重不是直接从 loss 反传得来的而是每一步从 context encoder 权重做滑动平均更新。用公式表示就是θ_target τ · θ_target (1 - τ) · θ_context。τ 通常取 0.99 到 1.0 之间训练初期可以用小一点的 τ 让 target encoder 跟上节奏后面逐渐增大保证训练稳定。我当时在这里吃过一个亏刚开始把 τ 固定成 0.999结果 loss 下降很慢。后来改成前 1000 步用 0.99再逐步升温到 0.999收敛速度明显提升。这个细节论文里提过但不显眼实际操作时很容易忽略。2. 复现前的准备从论文到可运行的代码2.1 硬软件配置选型这个项目选择的模型是 ViT-H/14。ViT-H 是 Vision Transformer 中非常大的一个档位参数规模在 6 亿左右。这种体量下如果环境配置不合理别说 72 小时预训练连单步迭代都会拖到让人失去耐心。我最终使用的环境如下计算卡8 张 80G 显存的加速卡A100 级别框架PyTorch配合分布式数据并行DDP实现单机多卡训练精度混合精度AMP模式梯度更新用 FP32计算用 FP16/BF16优化器AdamW配合余弦退火学习率调度数据ImageNet-1K 全量训练集约 128 万张图。ViT-H/14 的模型参数配置大致如下复现时可以直接对照这张表检查自己的模型定义有没有出错配置项数值输入图像分辨率224×224Patch 大小14×14Transformer 层数32隐层维度1280注意力头数16MLP 隐藏维度5120参数量约 6.3 亿2.2 数据处理的两个关键选择第一是 patch 化和掩码策略。I-JEPA 会把 224x224 的图像按 14x14 的 patch 划分得到 256 个 patch。预训练时随机 mask 掉一部分区域块默认掩码比例在 25% 左右。注意I-JEPA 是“按块掩码”不是像 MAE 那样按单个 patch 随机掩码。块的大小和数量直接影响到任务难度块太大、数量太多预测起来太容易块太小又退化成像素级任务。论文里推荐 block 数量和大小有一些预设值复现时先用默认参数跑通再根据验证集上的线性探测结果去微调。第二是数据增强策略。这里要提醒一下I-JEPA 对数据增强的使用比 MAE 更克制通常只用随机裁剪、随机水平翻转和颜色扰动这一系列“标准套餐”不做过强的增强。因为 target encoder 本身已经提供了目标的稳定性像 CutMix、MixUp 这类强增强反而可能干扰语义预测的一致性。2.3 损失函数和评估协议损失函数本身非常简单就是 predictor 输出和目标编码器输出之间的 L2 距离在多个被掩码 block 上取平均。整个模型的反向传播只经过 context encoder 和 predictortarget encoder 完全走 EMA 侧这是复现时最容易写错的地方——如果忘了对 target encoder 做 stop-gradient模型很快就会崩掉loss 直接 NaN。评估协议方面我采用的是论文同款的线性探测linear probing先把预训练好的 backbone 权重冻结只在最上层接一个线性分类头在 ImageNet-1K 的训练集中随机抽 1% 的带标签数据来训练。这个设置下能达到 73.3% 的 top-1 准确率已经是非常健康的水平。这里稍微解释一下1% 数据大概是 1.28 万张图而且标注类别分布和全量数据保持一致。线性探测的结果衡量的是“特征本身的可分性”不像微调那样还能靠分类头去补偿表征的不足所以在自监督学习里这个指标比 fine-tuning 更能反映预训练质量。3. 72 小时预训练实战从零到收敛的全过程3.1 训练超参的设定与原因预训练阶段最核心的超参是这几项全局 batch size2048。分布式训练中 8 张卡每张卡实际 batch size 是 256。这里不推荐把 batch size 调得太小I-JEPA 虽然不像对比学习那样极依赖 batch 大小但过小的 batch 会让 EMA 的更新变得噪声很大影响 target encoder 的稳定性。峰值学习率1e-3配合 AdamW 的 weight decay 0.05采用余弦退火调度。训练前有 1000 步的 warmup让学习率从 1e-6 线性升到峰值。EMA 系数 τ0.99 起步前 1000 步线性插值到 0.999之后固定。混合精度BF16。对 ViT 这种大模型FP16 容易出现溢出BF16 的动态范围更大稳定性更好。实测下来 BF16 的精度损失几乎可以忽略。这些超参不是拍脑袋定的。batch size 2048 是为了匹配 EMA 更新频率下的样本量学习率 1e-3 是 ViT 类模型在 ImageNet 级别数据集上的常见配置τ 和训练步数直接相关训练的时间越长τ 越接近 1.0避免 target encoder 被一步过大的更新节奏带偏。3.2 训练日志里的正常曲线长什么样我记录一下实际训练过程中的关键节点给大家一个参考基准。横坐标是训练步数step每一步代表一个 batch 的数据。第 0-500 步loss 从初始的 2.8 快速下降到 2.1 左右这个阶段主要是在学“低频”的语义信息比如整图的场景类别第 500-3000 步loss 下降速度放缓从 2.1 降到 1.6模型开始关注更细粒度的语义关系比如物体之间的相对位置第 3000 步之后loss 进入慢速线性下降阶段每 1000 步大约降 0.05-0.08。这一步比较考验耐心很多人以为 loss 不降是训练出问题了其实在自监督预训练里这种平滑的慢速下降是完全正常的。我在 72 小时跑完时模型总共经历了约 30000 步对应的 epoch 数大约是 110 个左右以全局 batch 2048 来算每个 epoch 约 625 步。也就是说训练结束时一共“看过”1 亿多张图。对比原论文里 ViT-H/14 训练 300 epoch 的实验我们所处的阶段大概相当于完成了 1/3 的训练量但线性探测已经能达到 73.3%这说明 I-JEPA 的特征在预训练早期就已经具备很强的可区分性。3.3 吞吐优化72 小时怎么“挤”出来的8 张卡跑 ViT-H/14如果不做优化单卡吞吐大概只能到 120 张图每秒全局算下来每秒不到 1000 张图按这个速度 72 小时最多只能跑 25 万步的一半也就是说如果不优化这个时间根本跑不完 30000 步。我实际做的是三件最有效的事梯度累积。一开始我把全局 batch 直接压到 1024 来保证显存不爆结果 EMA 抖动明显loss 曲线噪声大。后来恢复全局 batch 到 2048通过梯度累积两步来凑每张卡实际还是跑 256 的 micro-batch但是优化器每两步更新一次。这个改动的收益很直接训练稳定度恢复到和真正 2048 batch 相当的水平。DataLoader 瓶颈排查。ViT-H 前向计算很重很多人以为 GPU 一定是瓶颈其实不一定。我一度发现 GPU 利用率只有 60% 左右排查后确认是图像解码把 CPU 打满了。把 DataLoader 的 worker 数从 4 调到 16再开启共享内存缓存GPU 利用率直接拉回 90% 以上。通信开销优化。DDP 在 8 卡规模下梯度同步的通信量非常大。开启梯度压缩gradient compression和把通信与计算重叠overlap之后整体吞吐又提升了约 15%。这三招叠加下来训练吞吐从最初的每卡 120 张/s 提升到了 180 张/s 左右才让“72 小时跑完 3 万步”这件事变成可能。这里想提醒一点很多人复现大模型时习惯先堆硬件但在单机多卡场景下先把数据管线和混合精度调好效果远比盲目加卡明显。4. 下游评测ImageNet 1% 线性探测 73.3% 是怎么得来的4.1 线性探测的实验设置预训练完成后下游评测阶段需要把“学到的特征”拿出来单独检验。我严格对齐了 I-JEPA 论文里的线性探测标准协议避免因为评估细节不一样导致数据没法对比。线性探测阶段的操作流程是冻结 backbone 的全部参数不参与梯度更新在 backbone 的最后一层特征之上新加一个线性分类头就是一个全连接层把 1280 维的 ViT-H 特征映射到 1000 类抽取 ImageNet-1K 中 1% 的带标签数据也就是 128 万张训练集中随机抽 1.28 万张保持类别均衡只训练这个分类头优化器用 SGD学习率 30momentum 0.9weight decay 为 0训练 90 个 epoch在 ImageNet-1K 验证集上评估 top-1 准确率。很多人看到学习率 30 会觉得奇怪。线性探测阶段因为 backbone 被冻结输入特征分布相对固定线性层需要在特征空间里快速找到最优分类方向所以需要一个比较大的学习率。这里如果沿用预训练阶段的 1e-3 量级学习率线性探测的收敛速度和最终精度都会明显变差。4.2 73.3% 这个数意味着什么最终我拿到的验证集结果是 73.3% top-1 准确率。为了让大家对数字有感知我列一下对照方法预训练资源/时长1% 数据线性探测准确率有监督预训练直接在 100% 数据上训练分类完整训练一个 ViT-H约 82%参考值MAE重建像素ViT-L大规模预训练约 67%-69%对比学习SimCLR 类ViT-B/L通常需要极大 batch约 68%-72%I-JEPA本项目复现ViT-H110 epochs8 卡跑 72 小时73.3%实测需要说明的是这个表格里的对比对象在训练时间、模型 size 上不完全一致不能简单说 I-JEPA 就一定比谁好。但至少在本次复现的时间预算下73.3% 的结果说明第一I-JEPA 在没有负样本、没有像素重建的条件下学到了非常可分的语义特征第二复现过程没有出现明显的实现 bug因为结果落在合理区间内如果架构实现有致命错误线性探测通常连 60% 都很难上。4.3 线性探测结果的稳定性检查我反复跑了几次线性探测实验发现同一份预训练权重做不同随机种子下的分类头训练top-1 波动的范围在 0.3% 以内。这个稳定性很重要说明 73.3% 不是某次随机初始化带来的运气。另外我也对比了不同保存节点比如训练到 2 万步和 3 万步的 checkpoint的线性探测成绩2 万步时大概是 71.5%3 万步时到 73.3%说明训练还在继续改善特征质量只是边际收益在递减。这一步给训练时间的选取提供了一个判断依据如果你的算力只能支撑 48 小时跑到 2 万步左右也能获得不错的结果不必强求 72 小时。但如果你想报告一个更有说服力的数字多跑 24 小时换 1.8 个百分点从投入产出比来说是值得的。5. 复现过程中的典型问题和排查实录5.1 损失不降检查 target encoder 的 EMA 是否真正生效我在第一次跑通代码的时候loss 在 2.2 附近徘徊了整整几千步不下降一开始以为是学习率不够。查了代码之后发现EMA 更新时没有做参数 detach也就是说 target encoder 的梯度意外地参与了反向传播。这种情况下模型等价于一个“双塔同时更新”的串联结构任务变得异常困难loss 自然降不下去。排查方法很简单在 loss.backward() 之前打印一下 target encoder 参数的 requires_grad 和 grad 是否为空。如果 target encoder 参数带梯度说明 EMA 侧没写对。正确的做法是 target encoder 的参数在创建时就设为 requires_gradFalse或者更新权重时把上下文编码器的参数 detach 出来。这是 I-JEPA 复现新手最容易踩的坑没有之一。5.2 Batch Size 和 EMA 抖动的关系我试过把全局 batch size 压到 1024 来节省显存结果 loss 曲线抖得像心电监护仪每几百步就能看到一次 0.05 级别的尖峰。EMA 是对每一步得到的表示做滑动平均如果每步之间的差异太大滑动平均也压不住噪声。恢复正常 2048 的 batch 之后虽然每步耗时增加了但收敛效率反而更高因为抖动减少了模型始终走在平滑下降的方向上。如果实在没有足够显存支撑大 batch折中方案是用梯度累积模拟大 batch本质上是把多步的梯度攒到一起再更新一次这样每步计算没那么重但更新频率低了。我当时的组合是 micro-batch 256、累积梯度 2 步等效 batch 为 2048。5.3 GPU 利用率跑不满根因往往在数据管线有一次我发现训练速度从 180 张/s 掉到了 110 张/sGPU 利用率只有 55 左右。查了一通发现是 DataLoader 在做图像解码的时候 CPU 成了瓶颈。ImageNet 是原始 JPEG 图像解码非常吃 CPU。这个问题在自监督预训练里特别常见因为预训练需要的数据量远大于有监督训练任何一个 pipeline 短板都会被放大。我的解决思路是worker 从 4 提到 16进程间一定不要用笨重的锁图像解码改用内存映射mmap模式读入避免频繁硬盘 IO开启共享内存缓存让子进程之间可以复用解码结果。这三件事做完之后GPU 利用率恢复到了 90% 以上。5.4 训练中期 loss 突然升高和 τ 调度有关大约训练到 8000 步的时候loss 出现了一次小幅回升从 1.5 涨到了 1.58持续了大概 500 步又降回来。后来复盘发现那个时间点正好是我把 τ 从 0.995 直接跳变到 0.999 的位置。τ 跳得太快target encoder 会突然变得很“保守”和当前 context encoder 的状态脱节表现为 loss 反弹。正确的做法是把 τ 的调整也做成平滑插值或者界定一个很短的线性过渡期不要跳变。5.5 问题排查速查表现象可能原因排查方向解决手段训练初期 loss 不降target encoder 参与了梯度更新检查 target encoder 的 requires_grad对参数做 detach 或 requires_gradFalseEMA 导致 loss 上升τ 跳变画出 τ 的变化曲线改为平滑插值更新 τGPU 利用率低DataLoader 瓶颈观察 CPU 占用率和 dataloader 耗时增加 worker、mmap 读图、共享内存BF16 训练不收敛混合精度缩放策略过激进检查 grad scale 状态改用 BF16或调整 loss scaling线性探测精度偏低分类头学习率太小或训练 epoch 不足对照论文的线性探测协议用大的学习率如 30训练 90 epochs多机通信慢DDP 同步开销大观察通信时间占比开启梯度压缩和通信计算重叠6. 关于可复现性与扩展方向的一些体会把这次项目的关键数字再梳理一下ViT-H/14 架构8 卡 80G 加速卡72 小时预训练约 3 万步冻结 backbone 后用 1% 的 ImageNet 标注数据做线性探测得到 73.3% 的 top-1 准确率。这个结果对于自监督预训练来说是一个能够证明“方法本身有效、实现基本正确”的可靠信号。从复现的成本角度看I-JEPA 相对于对比学习方法的优势非常明显。它不需要维护庞大的负样本队列也不依赖动辄 4096 起步的超大 batch这让中等规模算力团队复现它成为可能。我试过在减半的算力条件下去复现同样的配置只需要把 batch size 降到 1024、训练时间拉到 140 小时左右最终效果也能到 70% 出头。如果你没有 72 小时的连续机时也可以采用分段续训的方式只要保存 EMA 状态和优化器状态就能无缝衔接。最后说一点我个人的实操心得。复现 I-JEPA 的过程中我踩过的最多坑不是模型结构本身而是对“EMA 何时该静默、何时该跟紧”的判断。很多自监督方法表面上花的功夫在模型设计实际上稳定性的成败全在训练动态的细节控制上。如果你也准备自己复现我建议一开始先在小数据集比如 10 万张图上把 loss 曲线、EMA 更新、线性探测流程整体跑通再上 ImageNet 全量。直接用全量数据排查 bug既烧钱又慢完全不划算。73.3% 不是终点。目前 I-JEPA 的特征已经可以作为下游检测、分割任务的初始化骨干我下一步计划把同样的预训练权重迁移到语义分割任务上微调看它在像素级任务上的迁移能力是否和线性探测一样稳。后面如果有结果再来更新这篇记录。
阅读完成 · 觉得有帮助？
