English

Causal-JEPA: Learning World Models through Object-Level Latent Masking

Artificial Intelligence 2026-05-29 v2

Abstract

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. We therefore propose C-JEPA, a simple and flexible object-centric world model that extends masked joint embedding prediction from image patches to object-centric representations. By masking object-level latents and requiring each masked object state to be inferred from the surrounding context, C-JEPA imposes structured partial observability during training, creating counterfactual-like prediction queries that discourage shortcut solutions and make interaction-dependent prediction necessary under the learning objective. Empirically, C-JEPA leads to consistent gains in visual question answering, with an absolute improvement of about 20% in counterfactual reasoning over the same architecture without object-level masking. On agent control tasks, C-JEPA enables substantially more efficient planning by using only 1% of the total latent input features required by patch-based world models, while achieving comparable performance. Finally, we provide a formal analysis demonstrating that object-level masking induces useful inductive bias by controlling observability. Our code is available at https://github.com/galilai-group/cjepa.

Keywords

Cite

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

Comments

Project Page: https://hazel-heejeong-nam.github.io/cjepa/ ICML 2026 Accepted

R2 v1 2026-07-01T10:32:44.377Z