simillusion/ijepa-vit-tiny-cifar10

I-JEPA ViT-Tiny on CIFAR-10

A demo checkpoint from Simulusion, an open-source PyTorch library that puts world models behind one API. It implements the paper Self-Supervised Learning from Images with a Joint-Embedding Predictive Architecture (Assran et al., CVPR 2023).

I-JEPA learns image features without labels and without reconstructing pixels: it predicts the representations of masked image blocks from a visible context block. In Simulusion it is a representation-learning component for world models.

Results

A ViT-Tiny (patch size 4) trained for 30 epochs on CIFAR-10, about an hour on a 4 GB RTX 3050 laptop GPU. Classifying test images by their 20 nearest training images in feature space scores 41.1%, against 29.6% for the same network at random initialisation (chance is 10%).

This is a short demo run sized for a laptop GPU, meant to show what the model does. It is not a reproduction of the paper's results. More clips and learning curves are in the gallery.

Usage

pip install simulusion huggingface_hub

from simulusion.models.jepa_agent import JEPAAgent

agent = JEPAAgent.from_pretrained("simillusion/ijepa-vit-tiny-cifar10")

Checkpoints are plain state dicts and load with torch.load(..., weights_only=True). They include optimizer state, so you can continue training from them.

Training

Reproduce this checkpoint from a clone of the repository:

python demos/jepa_demo.py train --epochs 30

Citation

If you use this model, please cite the original paper:

@inproceedings{assran2023ijepa,
  title     = {Self-Supervised Learning from Images with a Joint-Embedding Predictive Architecture},
  author    = {Assran, Mahmoud and Duval, Quentin and Misra, Ishan and Bojanowski, Piotr and Vincent, Pascal and Rabbat, Michael and LeCun, Yann and Ballas, Nicolas},
  booktitle = {Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition},
  year      = {2023},
  url       = {https://arxiv.org/abs/2301.08243},
}

and Simulusion:

@software{simulusion,
  title  = {Simulusion: a PyTorch library for world models},
  author = {Thakkar, Param},
  year   = {2026},
  url    = {https://github.com/ParamThakkar123/simulusion},
}