160行代码读懂LeCun的JEPA世界模型
闻乐 发自 凹非寺量子位 | 公众号 QbitAILeCun念叨了好几年的JEPA被160行代码给复刻了。GitHub上有个开发者用极简单文件形式用PyTorch把JEPA核心系列全部实现了一遍从I-JEPA到LeWorldModel五个变体一个没落就为了——教学。没错就为了让你看懂原理把核心机制全拆了出来代码行数从160到278不等依赖只有PyTorch和torchvision普通笔记本就能跑。五个模型速览一下iJEPA图像160行代码跑CIFAR-10核心是掩码块嵌入预测遮掉图像的一部分patch让模型从可见区域去预测被遮区域的embedding。EMA目标编码器、multi-block masking、smooth-L1 loss经典自监督JEPA的核心模块一个不少。V-JEPA视频188行适配Moving MNIST。把二维patch扩展成3D管块tubelet同时做短程和长程两组管状掩码让模型从部分帧去推测缺失帧的时空特征。机制跟I-JEPA同源但多了一维时间轴。V-JEPA 2278行支持动作条件预测。两阶段玩法先像V-JEPA一样预训练再在冻结的encoder latent上训练action-conditioned predictor。C-JEPA物体轨迹174行聚焦3-digit弹跳视频。不搞patch masking了改做物体级轨迹掩码t0保留身份锚点后面全遮掉用双向Transformer在物体slot token上做预测。LeWorldModel233行端到端JEPA世界模型。没有EMA、没有stop-grad、没有masking编码器和action-conditioned自回归预测器联合端到端训练。咋做的先说轻量化设计的思路。原版JEPA论文跑的是ViT-Huge在ImageNet或Kinetics上的规模动辄几百张GPU普通人连下载模型的带宽都不够。但这个教学版三条腿走路第一模型规模从ViT-Huge砍成ViT-Tiny参数量差了起码两个数量级第二数据集从ImageNet/Kinetics换成CIFAR-10和合成视频几十兆就能跑效果可验证但门槛极低第三核心机制一个没砍掩码策略、损失函数Smooth-L1、MSE、预测逻辑、EMA更新、warmupcosine学习率全部精准保留。拿最核心的ijepa.py来说160行里塞了啥从patch embedding、ViT encoder、EMA target encoder、multi-block masking到predictor、smooth-L1 loss、warmupcosine学习率调度还有权重衰减分离一个不少。用这套机制在CIFAR-10上跑100个epoch线性探测准确率能到52.7%。为了能直观看到模型的学习变化过程坐着每训练10轮就把测试集的特征保存一份快照。再用LDA把这些特征压缩投影到二维平面展示如下。当然论文原版是ViT-Huge在ImageNet上跑300个epoch那是另一个星球的计算量了但两者机制是完全一致的。再说拓展功能。每个xxx.py旁边都有个 xxx_extras.py跑完之后自动出全套可视化。掩码动画让你直观看到哪些patch被遮了、哪些参与预测loss曲线看收敛过程t-SNE和PCA降维图看embedding空间里各类别是怎么逐步分离的从第1个epoch混成一团到第100个epoch泾渭分明比看loss数字更直观。实操也不复杂git clone gitgithub.com:keon/jepa.git cd jepa python -m venv .venv source .venv/bin/activatepip install -r requirements.txtrequirements文档里的环境依赖如下项目目前涵盖的5个JEPA变体全部是独立文件没有共享的工具代码每个都能直接python xxx.py跑python ijepa.py # 纯训练不画图 python ijepa_extras.py # 训练 全量可视化 线性探测CUDA、MPS、CPU都支持数据集自动下载。想跑V-JEPA就换python vjepa.py想跑LeWorldModel就换python leworldmodel.py每个文件独立跑互相不依赖。JEPA终于从论文概念变成了可读代码JEPA这套东西其实一直有点“只闻其名”的那味儿。LeCun在各种场合推了好几年论文出了一篇又一篇大家都知道这是个重要的方向——在embedding空间做预测不碰像素效率高、泛化好。但真正打开Meta官方的V-JEPA仓库想学一下的时候第一反应往往是「我到底该先看哪」就比如Meta官方仓库里塞了分布式训练、数据流水线、EMA调度、日志系统、大规模视频预训练的完整pipeline光配置文件就一大堆。而这个极简仓库干的事本质上就是把JEPA重新压缩回算法本体。每个实现文件打开就是算法没有工程包装没有分布式调度没有多余抽象。Encoder、Predictor、mask采样、loss计算、EMA更新一个模块几十行注释里标明了对应论文的符号f_theta、g_phi、s_y论文跟代码可以左右对照着读。当然开源极简版作为教学版模型规模、数据集复杂度、实验精度跟原版论文必然存在差距。开发者自己也坦率地列了每个实现的偏差。比如I-JEPA的52.7%线性探测对比论文的ImageNet结果是降维打击C-JEPA跳过了VideoSAUR的物体发现预训练用frozen lookup做占位V-JEPA 2-AC在玩具数据上conditioning gap太小信号不够明显。但这个项目最大的目的也不是复现SOTA用开发者自己的话说就是——把算法蒸馏到只剩数学本质。对于初学者或者是想搞懂JEPA到底是怎么work的人来说这个仓库可能就是最好的起点。项目地址https://github.com/keon/jepa一键三连「点赞」「转发」「小心心」欢迎在评论区留下你的想法—完—中国AIGC产业峰会全阵容发布查看议程从底层架构到爆款应用从软件到硬件从创作到投资... 这一次我们希望聚齐AI赛道的实战派聊透AI到底怎么落地、怎么赚钱、怎么改变工作。5月20日北京·金茂万丽酒店所有人马上AI起来 报名参会一键关注 点亮星标科技前沿进展每日见
)
