Projects / Self-supervised world models
JEPA & Saccade JEPA
I-JEPA built from the paper before the official release, trained on ImageNet with one GPU, plus Saccade JEPA: an original variant that learns by predicting small eye movements.
An implementation of I-JEPA, written from the paper about a month before Meta released its official code, and trained for 21 epochs on ImageNet with a single GPU.
What’s here
- A barebones I-JEPA with vmapped, per-sample context and target masks, plus a variant built on an energy transformer.
- Evaluation beyond linear probes and kNN: an interactive attention-map dashboard, correlation-dimension estimates, and UMAP views of the embedding.
- Saccade JEPA (original). Small rotations and translations stand in for the eye’s constant jitter. A predictor, told the shift through NeRF-style embeddings, predicts the shifted representation, with cycle-consistency and VICReg terms keeping the representations well-behaved.
- Masked autoencoders, including a self-distillation variant.