keon/jepa

jepa

Minimal, single-file PyTorch reimplementations of the JEPA family, with paired tutorials.

Each algorithm file is standalone — only depends on torch and torchvision, no shared utilities. The matching <algo>_extras.py adds visualization (mask grids, loss curves, PCA/LDA/t-SNE evolution, linear probe).

See FAITHFULNESS.md for the load-bearing details each minimal implementation preserves and the educational substitutions it makes.

Quick start

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

python ijepa.py                     # train I-JEPA only (no plots)
python ijepa_extras.py              # train + write all visualizations + linear probe

Runs on CUDA, MPS, or CPU. CIFAR-10 / MNIST datasets auto-download to ./data/.

Reproducibility

The repo pins exact versions in requirements.txt and pyproject.toml:

python >= 3.10  (tested on 3.13.5)
torch == 2.11.0
torchvision == 0.26.0
matplotlib == 3.10.9
scikit-learn == 1.8.0   # used by ijepa_extras for t-SNE
numpy == 2.4.4
pillow == 12.2.0

Install as a package instead of installing requirements directly:

pip install -e .

What's where

.
├── ijepa.py / ijepa_extras.py                       # I-JEPA on CIFAR-10
├── vjepa.py / vjepa_extras.py                       # V-JEPA on Moving MNIST
├── vjepa2.py / vjepa2_extras.py                     # V-JEPA 2 + V-JEPA 2-AC (synthetic)
├── vjepa2_1.py / vjepa2_1_extras.py                 # V-JEPA 2.1 (toy image+video co-training)
├── cjepa.py / cjepa_extras.py                       # C-JEPA on 3-digit bouncing video
├── leworldmodel.py / leworldmodel_extras.py         # LeWorldModel (end-to-end JEPA, SIGReg)
├── ijepa_tutorial.md                                # walk-throughs that match the code
├── vjepa_tutorial.md
├── vjepa2_tutorial.md
├── vjepa2_1_tutorial.md
├── cjepa_tutorial.md
├── leworldmodel_tutorial.md
├── FAITHFULNESS.md                       # preserved details + deliberate simplifications
├── papers/                              # source PDFs bundled with the repo
├── samples/                             # mask grids, loss curves, PCA/LDA/t-SNE plots
└── figs/                                # paper figures referenced by tutorials

The methods, in one paragraph each

I-JEPA (Assran et al. 2023) — predict embeddings of held-out image patches from embeddings of visible patches. EMA target encoder, multi-block masking, smooth-L1 loss. The canonical self-supervised JEPA.

V-JEPA (Bardes et al. 2024) — same recipe, but 3D tubelet patches over video. Two mask groups (short-range + long-range tubes), L1 loss, EMA 0.998 → 1.0.

V-JEPA 2 (Assran et al. 2025) — two-phase: V-JEPA pretraining followed by V-JEPA 2-AC, an action-conditioned predictor trained on frozen-encoder latents with teacher forcing + rollout. The encoder is frozen in phase 2; no EMA.

V-JEPA 2.1 (official release notes) — updated pretraining recipe for dense, temporally consistent image/video features. The load-bearing changes are dense predictive loss (predict visible tokens too), deep self-supervision over multiple encoder depths, and image+video co-training through a shared backbone.

C-JEPA (Nam et al. 2026) — object-level trajectory masking with an identity anchor at $t=0$. No EMA. Bidirectional transformer over flattened slot tokens. Built on top of a pretrained object-centric encoder in the paper; here we use a frozen oracle position-slot embedding as a documented educational stand-in.

LeWorldModel (Maes et al. 2026) — end-to-end JEPA world model from pixels. No EMA, no stop-grad, no masking. The encoder and an action-conditioned AR predictor are jointly trained with two loss terms: next-embedding MSE plus a Sketch Isotropic Gaussian Regularizer (SIGReg) that prevents collapse by pushing the embedding marginals toward $\mathcal{N}(0, 1)$.

Caveats

These are educational reimplementations:

• ViT-tiny, not ViT-Huge. CIFAR-10 / Moving MNIST / synthetic videos, not ImageNet / Kinetics.
• I-JEPA hits ~52.7% linear probe on CIFAR-10 after 100 epochs. The paper's numbers come from ViT-H/14 on ImageNet for 300 epochs — different planet of compute.
• C-JEPA skips slot discovery (uses oracle positions). Real C-JEPA requires VideoSAUR/SAVi-style object-centric pretraining on top of visual features.
• V-JEPA 2-AC is a small block-causal, action/state-conditioned latent predictor, not Meta's 300M-parameter robot-action model; it preserves the teacher-forcing + rollout training shape.
• V-JEPA 2.1 uses toy MNIST + Moving MNIST co-training rather than the official large image/video mixture, but preserves dense token loss, multi-layer self-supervision, and multimodal pretraining shape.
• LeWorldModel includes the two-term objective and projection heads needed for SIGReg, but omits the paper's control/planning layer.

Each tutorial discloses the specific deviations from its source paper and keeps code snippets aligned with the minimized implementation.

License

MIT.