keon/jepa 源地址: https://github.com/keon/jepa # jepa JEPA 系列模型的最小化、单文件 PyTorch 重新实现，附带配套教程。 | 文件 | 方法 | 数据集 | 代码行数 (LOC) | 教程 | |—|—|—|—:|—| | ijepa.py | I-JEPA | CIFAR-10 | 160 | ijepa_tutorial.md | | vjepa.py | V-JEPA | Moving MNIST | 188 | vjepa_tutorial.md | | vjepa2.py | V-JEPA 2 + V-JEPA 2-AC | 合成运动数字 | 278 | vjepa2_tutorial.md | | cjepa.py | C-JEPA | 3位数字弹跳视频 | 174 | cjepa_tutorial.md | 每个算法文件都是独立的——仅依赖 torch 和 torchvision，不共享任何工具类。对应的 _extras.py 文件添加了可视化功能（掩码网格、损失曲线、PCA/LDA/t-SNE 演化过程、线性探针）。 ## 快速开始 bash git clone [email protected]:keon/jepa.git cd jepa python -m venv .venv && source .venv/bin/activate pip install -r requirements.txt # 固定版本，见下文 python ijepa.py # 仅训练 I-JEPA（无绘图） python ijepa_extras.py # 训练 + 生成所有可视化图表 + 线性探针 支持 CUDA、MPS 或 CPU 运行。CIFAR-10 / MNIST 数据集会自动下载至 ./data/。 ### 可复现性 仓库在 requirements.txt 和 pyproject.toml 中固定了确切版本： python >= 3.10 (已在 3.13.5 上测试) torch == 2.11.0 torchvision == 0.26.0 matplotlib == 3.10.9 scikit-learn == 1.8.0 # ijepa_extras 中使用 scikit-learn 进行 t-SNE numpy == 2.4.4 pillow == 12.2.0 也可以将其作为包安装，而不是直接安装依赖项： bash pip install -e . ## 文件结构 . ├── ijepa.py / ijepa_extras.py # 基于 CIFAR-10 的 I-JEPA ├── vjepa.py / vjepa_extras.py # 基于 Moving MNIST 的 V-JEPA ├── vjepa2.py / vjepa2_extras.py # 基于合成数据的 V-JEPA 2 + V-JEPA 2-AC ├── cjepa.py / cjepa_extras.py # 基于 3位数字弹跳视频的 C-JEPA ├── ijepa_tutorial.md # 与代码对应的逐步教程 ├── vjepa_tutorial.md ├── vjepa2_tutorial.md ├── cjepa_tutorial.md ├── papers/ # 四篇源论文的 PDF 文件 ├── samples/ # 掩码网格、损失曲线、PCA/LDA/t-SNE 图表 └── figs/ # 教程中引用的论文插图 ## 方法概述（每段介绍一种方法） I-JEPA (Assran et al. 2023 (https://arxiv.org/abs/2301.08243)) —— 根据可见图像块的嵌入预测被保留图像块的嵌入。采用指数移动平均 (EMA) 目标编码器、多块掩码机制以及平滑 L1 损失。这是经典的自监督 JEPA。 V-JEPA (Bardes et al. 2024 (https://arxiv.org/abs/2404.08471)) —— 配方相同，但针对视频使用 3D 管状块 (tubelet patches)。包含两个掩码组（短程 + 长程管状结构），使用 L1 损失，EMA 系数从 0.998 逐渐增加到 1.0。 V-JEPA 2 (Assran et al. 2025 (https://arxiv.org/abs/2506.09985)) —— 分为两个阶段：先进行 V-JEPA 预训练，随后进行 V-JEPA 2-AC，这是一种基于冻结编码器隐空间、采用教师强制 (teacher forcing) 和展开 (rollout) 训练的动作条件预测器。第二阶段编码器保持冻结；不使用 EMA。 C-JEPA (Nam et al. 2026 (https://arxiv.org/abs/2602.11389)) —— 在 t=0 时刻使用身份锚点 (identity anchor) 的对象级轨迹掩码。不使用 EMA。在扁平化的槽位令牌 (slot tokens) 上使用双向 Transformer。构建于预训练的对象中心编码器之上（论文中为 VideoSAUR；我们使用固定的嵌入查找作为文档记录的替代方案）。 ## 注意事项 这些是教育性质的重新实现： - 使用 ViT-tiny 而非 ViT-Huge。数据集为 CIFAR-10 / Moving MNIST / 合成视频，而非 ImageNet / Kinetics。 - I-JEPA 在 CIFAR-10 上训练 100 个 epoch 后，线性探针准确率约为 52.7%。论文中的数据来自在 ImageNet 上训练 300 个 epoch 的 ViT-H/14 —— 计算规模完全不同。 - C-JEPA 跳过了槽位发现过程（使用已知位置）。真正的 C-JEPA 需要在冻结的 DINOv2 特征基础上，使用 VideoSAUR 进行预训练（约 10 万步）。 - 在我们的玩具数据集中，由于数据过于简单，V-JEPA 2-AC 的动作条件差距保持较小；机制是正确的，但需要更丰富的数据才能显现出信号。 每个教程都披露了与其源论文的具体偏差。 ## 许可 MIT 许可证。