Causal-JEPA: Learning World Models through Object-Level Latent Masking | Heejeong Nam et al. | ResearchPod

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

Alex: Welcome to another episode of ResearchPod. Today we're looking at a paper called Causal-JEPA: Learning World Models through Object-Level Latent Masking.

Sam: So what's the central problem this is trying to solve?

Alex: World models trained on visual scenes tend to exploit spurious correlations rather than learning how objects actually interact. The model sees a ball moving and a wall nearby, and instead of learning that the wall causes the bounce, it learns a statistical shortcut — the ball's own past trajectory is usually enough to predict its future. The proposed fix is object-level latent masking, and the key design choice is subtle: during training, the model selectively masks entire object slots across a history window, but keeps a minimal identity anchor for each masked object. That anchor is enough to say "this object exists," but not enough to predict its trajectory from its own past states alone.

Sam: So the predictor is forced to infer the missing trajectory from the object's relationships with everything else in the scene.

Alex: Exactly. That's the causal inductive bias the title is gesturing at. It's not causal discovery in the structural sense — the authors are careful about this — but it does enforce predictive sufficiency through relational context rather than self-dynamics. The Bayes-optimal predictor under this objective converges to the minimal sufficient set of contextual variables needed to reconstruct the masked state. Which is a weaker claim than recovering the true causal graph, but arguably a more honest one given what's actually identifiable from observational video.

Sam: Where does that actually show up in the results?

Alex: They evaluate on two settings: CLEVRER for visual counterfactual reasoning, and Push-T for agent control. On CLEVRER, object-level masking yields roughly a twenty-point absolute improvement in counterfactual accuracy over patch-based baselines. That's the load-bearing finding — it's what the causal reasoning claim actually rests on. Push-T then tests whether this transfers to control, not just reasoning, and that's where an efficiency story emerges. C-JEPA achieves comparable task performance to DINO-WM while operating on just over one percent of the total latent input features. That compression translates to more than an eightfold speedup in model predictive control planning on a single GPU.

Sam: That's a striking efficiency gap. Where does it come from?

Alex: Object-centric representations are dramatically sparser than patch-based ones. When your world model only needs to track a handful of object slots rather than thousands of image patches, the predictor has far less to attend over at inference time. The object-centric representation isn't just cleaner conceptually — it's cheaper at every forward pass during planning.

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

Heejeong Nam, Quentin Le Lidec, Lucas Maes, Yann LeCun, Randall Balestriero

Feb 11, 2026·5:31·arxiv·Original Paper

World ModelsObject-Centric RepresentationsObject-Level Latent MaskingCausal-JEPA (C-JEPA)Joint Embedding Predictive Architecture (JEPA)Model Predictive Control (MPC)

Listen

6 min

Briefing

Research Question

This paper investigates how to improve object-centric world models so they can robustly capture interaction-dependent dynamics rather than relying on trivial object self-dynamics or incidental pixel-level correlations. The authors ask whether applying structured masking at the object latent level can induce a causal inductive bias that makes interaction reasoning functionally necessary for the model.

Approach

The authors propose Causal-JEPA (C-JEPA), an object-centric world model that extends masked joint embedding prediction from image patches to object-centric latent representations. Using Slot Attention to extract object slots from video frames, C-JEPA selectively masks object-level latents across the history window—retaining only a minimal identity anchor—and requires the predictor to recover these masked states from the unmasked context slots and auxiliary variables like actions. This setup creates counterfactual-like prediction queries during training that discourage shortcut solutions like temporal interpolation.

Main Findings

Empirically, C-JEPA achieves consistent performance gains across visual question answering and agent control tasks. On the CLEVRER benchmark, C-JEPA improves overall visual question answering accuracy and yields an absolute gain of about 20% in counterfactual reasoning compared to the same architecture without object-level masking. In model predictive control tasks on the Push-T environment, C-JEPA delivers task performance comparable to patch-based world models while utilizing only about 1% of their latent input feature size, enabling substantially faster planning.

Why It Matters

World models are vital for scalable planning, reasoning, and decision-making in intelligent agents. By demonstrating that object-level latent masking effectively forces a model to learn true interactions without requiring handcrafted relational graphs or complex architectural factorization, C-JEPA provides a simple, flexible, and computationally efficient foundation for robust video understanding and robotic control.

Key Terms and Definitions

Object-Centric Representation — A structured latent abstraction that groups high-dimensional visual observations into a set of discrete vectors, where each vector corresponds to a distinct entity or object in the scene.
Joint Embedding Predictive Architecture (JEPA) — A self-supervised learning paradigm that trains a predictor to map latent representations to other latent representations without relying on pixel-level reconstruction.
Object-Level Masking — A training strategy where specific object slots are artificially hidden across time steps, forcing the network to infer their states from interacting entities.
Identity Anchor — The preserved earliest time step representation of an object that provides necessary entity identity information to counteract the permutation equivariance of slot-based representations.
Counterfactual Reasoning — The capability to answer questions about what would have happened under hypothetical interventions or alternative scenarios in a dynamic environment.

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.

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