CJEPA / README.md
HazelNam's picture
Improve model card for C-JEPA (#1)
d56bfc1
metadata
license: apache-2.0
paper_ids:
  - '2602.11389'
pipeline_tag: image-feature-extraction
datasets:
  - clevrer
  - pusht
tags:
  - object-centric
  - world-models
  - causal-inference
  - jepa
  - representation-learning
  - vision

C-JEPA: Causal-JEPA

This repository contains the weights and code for Causal-JEPA (C-JEPA), a simple and flexible object-centric world model architecture presented in the paper Causal-JEPA: Learning World Models through Object-Level Latent Interventions.

Summary

World models require robust relational understanding to support prediction, reasoning, and control. While object-centric representations provide a useful abstraction, they are not sufficient to capture interaction-dependent dynamics. C-JEPA is a simple and flexible object-centric world model that extends masked joint embedding prediction from image patches to object-centric representations. By applying object-level masking that requires an object's state to be inferred from other objects, C-JEPA induces latent interventions with counterfactual-like effects and prevents shortcut solutions, making interaction reasoning essential.

Empirically, C-JEPA demonstrates:

  • Improved Visual Reasoning: Consistent gains in visual question answering, with an absolute improvement of about 20% in counterfactual reasoning compared to the same architecture without object-level masking on benchmarks like CLEVRER.
  • Efficient Planning: Substantially more efficient planning in agent control tasks (e.g., Push-T), using only 1% of the total latent input features required by patch-based world models while achieving comparable performance.
  • Causal Inductive Bias: A formal analysis demonstrates that object-level masking induces a causal inductive bias via latent interventions.

Architecture

architecture

Setup and Usage

C-JEPA relies on object-centric encoders (like VideoSAUR or SAVi) to extract representations. For detailed environment setup, dataset preparation, and training/evaluation scripts, please refer to the official GitHub repository. The repository also provides model checkpoints and pre-extracted slot representations for various configurations.

Citation

If you find this work useful, please consider citing:

@article{nam2026causal,
  title={Causal-JEPA: Learning World Models through Object-Level Latent Interventions},
  author={Nam, Heejeong and Le Lidec, Quentin and Maes, Lucas and LeCun, Yann and Balestriero, Randall},
  journal={arXiv preprint arXiv:2602.11389},
  year={2026}
}