160行代码拆解LeCun的JEPA：5个变体一次看懂世界模型核心

导语：160行代码，带你看懂JEPA的核心机制。

LeCun反复强调了多年的 JEPA，终于被开发者用极简代码“翻译”成了能读、能跑、能对照论文理解的版本。这个 GitHub 项目把 I-JEPA、V-JEPA、V-JEPA 2、C-JEPA 和 LeWorldModel 都做成了单文件 PyTorch 实现，代码量大多只在 160 到 278 行之间，依赖也只有 PyTorch 和 torchvision，普通笔记本就能启动实验。

为什么这份极简实现值得看

原版 JEPA 论文对应的官方实现，通常伴随分布式训练、大规模数据管线、EMA 调度和复杂配置，初学者很难快速抓住算法本身。这份极简版做的事情很直接：把工程包装全部剥离，只保留编码器、预测器、掩码采样、损失函数和目标网络更新，让论文里的 f_theta、g_phi、s_y 可以和代码逐行对应。

• 模型从 ViT-Huge 缩到 ViT-Tiny，参数量大幅下降
• 数据从 ImageNet/Kinetics 换成 CIFAR-10、Moving MNIST 等小数据集
• 关键机制没有删：multi-block masking、Smooth-L1 / MSE、EMA、warmup + cosine 学习率都保留下来

换句话说，它不是在追求最强结果，而是在追求“看得懂”。这也是它最大的教学价值：你不用先理解庞大的工程体系，就能直接看到 JEPA 的算法骨架。

I-JEPA：从图像补全理解核心机制

160 行的 ijepa.py 是全项目里最适合入门的一份。它把图像切成 patch，通过 multi-block masking 遮掉一部分区域，让模型只根据可见 patch 预测被遮区域的 embedding，而不是直接还原像素。训练中使用 EMA target encoder 生成目标表示，损失采用 Smooth-L1，并配合 warmup + cosine 学习率和权重衰减分离。

这个实现里几乎把 I-JEPA 的关键模块都装进去了：patch embedding、ViT encoder、EMA target encoder、predictor、mask 采样、loss 计算、目标网络更新，一个都不少。哪怕只看这一份文件，也能把论文中的主要流程串起来。

开发者用 CIFAR-10 跑 100 个 epoch，线性探测准确率达到 52.7%。虽然这和论文原版在 ImageNet 上的大规模结果还有明显差距，但机制是一致的，足够帮助读者理解 JEPA 为什么要在 embedding 空间里做预测，以及这种做法为什么比像素重建更适合学习表示。

训练过程还能“看得见”

项目还配套了 ijepa_extras.py，用来输出掩码动画、loss 曲线、t-SNE/PCA 降维图，以及按训练轮次保存的特征快照。开发者每训练 10 轮就保存一次测试集特征，再用 LDA 投影到二维平面，可以看到类别从一开始混在一起，到后期逐渐分开。对新手来说，这种可视化比单看 loss 数字更容易理解“表示学习”到底发生了什么。

四个扩展变体分别在解决什么

V-JEPA：把图像逻辑扩展到视频

V-JEPA 的实现大约 188 行，适配的是 Moving MNIST。它把二维 patch 扩展为三维 tubelet，同时设计短程和长程两组管状掩码，让模型从部分帧推断缺失帧的时空特征。核心思想仍然是“预测潜空间表示”，只是把空间维度扩展到了时间轴。

V-JEPA 2：引入动作条件

V-JEPA 2 增加了动作条件预测，代码约 278 行，采用两阶段训练：先像 V-JEPA 一样做预训练，再在冻结的 encoder latent 上训练 action-conditioned predictor。这样做的意义在于，把“看见了什么”和“下一步会发生什么”分开建模，更接近世界模型的思路。

C-JEPA：从 patch 走向物体轨迹

C-JEPA 约 174 行，任务换成了 3-digit bouncing 视频，重点不再是图像块，而是物体级轨迹预测。它保留 t=0 的身份锚点，后续时间步全部遮掉，再用双向 Transformer 在 object slot token 上做预测。因为目标已经是物体轨迹，这个版本更像是在学习“物体如何运动”，而不是单纯补全图像外观。

需要注意的是，这个教学版没有沿用原始物体发现预训练流程，而是用 frozen lookup 做占位，因此更适合看懂思路，不适合拿来和原论文精度直接比较。

LeWorldModel：最接近端到端世界模型

LeWorldModel 约 233 行，是这套代码里最像“世界模型”的版本。它不再依赖 EMA、stop-grad 或显式 masking，而是让编码器和 action-conditioned autoregressive predictor 端到端联合训练。换句话说，它把 JEPA 的表示预测思想进一步推向了“基于动作预测未来状态”的方向。

这份代码为什么能把复杂论文讲明白

这套项目最强的地方，不是它的实验规模，而是它把算法蒸馏到了只剩本质。每个实现文件都能独立运行，没有共享工具代码，没有复杂抽象，打开就是算法。注释还尽量对应论文符号，读者可以一边看论文，一边对照源码里的编码器、预测器、掩码采样和损失函数。

• 独立文件结构清晰：ijepa.py、vjepa.py、vjepa2.py、cjepa.py、leworldmodel.py 都能直接运行
• 兼容 CUDA、MPS 和 CPU，数据集支持自动下载
• 每个实现都配套可视化脚本，方便观察训练过程和表征变化
• 适合想弄懂 JEPA “到底怎么 work” 的学习者，而不是只想看结论的人

项目作者也坦率说明了教学版与论文原版之间的差异：例如 I-JEPA 的 52.7% 线性探测，和论文里在 ImageNet 上的结果不能直接横向对比；C-JEPA 跳过了 VideoSAUR 的物体发现预训练，使用 frozen lookup 作为占位；V-JEPA 2-AC 在玩具数据上 conditioning gap 较小，信号不够强。这些限制并不影响它作为入门材料的价值，反而更清楚地提醒读者：它的目标不是复现 SOTA，而是把原理讲透。

安装与运行方式

如果想快速上手，可以直接按项目提供的基础流程准备环境：

git clone [email protected]:keon/jepa.git && cd jepa
python -m venv .venv
source .venv/bin/activate
pip install -r requirements.txt

环境支持 CUDA、MPS 和 CPU，数据集也会自动下载。运行方式同样非常直接：只训练 I-JEPA 就执行 python ijepa.py；如果想看训练过程和可视化，就运行 python ijepa_extras.py。其他变体同理，换成 python vjepa.py、python vjepa2.py、python cjepa.py 或 python leworldmodel.py 即可。

结语：JEPA 终于从概念变成了可读代码

过去，JEPA 对很多人来说是“知道它很重要，但不知道该从哪学起”的典型代表。现在这份极简复刻把它重新压回算法本体，让读者能真正看清：模型不是在拼命还原像素，而是在潜空间里学习结构、关系和未来变化。尽管它和论文原版在模型规模、数据复杂度和实验精度上都存在差异，但作为教学项目，它已经非常接近“把复杂理论讲明白”的理想形态。

如果你想快速建立对 JEPA、视频表示学习和世界模型的直觉，这个仓库会是一个很好的起点。