PaperlayerView exact revision on arXiv

Source-linked AI summary

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

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

arXiv:2602.11389v2cs.AI

TL;DR

World models need representations that capture entities and their interactions, but object-centric representations alone do not ensure interaction-dependent dynamics. C-JEPA addresses this by masking object-level latents during JEPA prediction, yielding gains in visual reasoning and efficient control while formally characterizing the resulting inductive bias.

  • Problem

    World models require robust relational understanding, while existing object-centric and patch-based approaches do not by themselves enforce object-level interaction reasoning.

  • Method

    C-JEPA combines JEPA with object-level masking so masked object states are inferred from surrounding context under structured partial observability.

  • Results

    C-JEPA improves counterfactual visual question answering by about 20% absolute and achieves comparable predictive-control performance using 1.02% of the total input feature size.

  • Takeaways & Limitations

    Object-level masking provides an efficient mechanism for learning interaction-aware latent dynamics without reconstruction losses or task-specific supervision.

  • Takeaways & Limitations

    The method depends on object-centric encoder quality, and its influence neighborhoods are not directly validated on datasets with explicit temporal causal graphs.

Abstract

from arXiv · show

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.

1. Introduction

C-JEPA introduces object-level masking for object-centric world modeling, making interaction-dependent prediction necessary while retaining an efficient latent representation. It improves visual reasoning and enables efficient predictive control compared with relevant baselines.

  • C-JEPA: C-JEPA masks object-level latents so masked object trajectories must be inferred from other evolving object states.The approach creates counterfactual-like interventions during learning and supports optional auxiliary variables such as actions and proprioception.
  • C-JEPA: C-JEPA uses compact object-centric representations with JEPA prediction to reduce computational and memory overhead relative to patch-based predictors.The same quadratic attention mechanism operates over substantially fewer entity tokens.
  • Evaluation: About 20% absolute gain in counterfactual visual question answering is achieved over the same architecture without object-level masking.The evaluation uses CLEVRER, a video question-answering benchmark involving multi-object interactions.
  • Evaluation: Comparable predictive-control performance is achieved with only 1.02% of the total input feature size and more than 8× faster model predictive control.The comparison is against patch-based world models in the Push-T manipulation task.
  • Analysis: Object-level masked prediction is analyzed as an inductive bias that makes interaction reasoning functionally necessary for minimizing the training objective.The analysis links the bias to interventions on predictor observability.

2. Related Works

Related work develops object-centric and masked-prediction approaches for visual dynamics, but these lines of work apply masking and object representations in different ways. C-JEPA is positioned against methods that use patch masking, object conditioning, or specialized causal structure.

  • Object-centric world models: World-modeling methods use weak entity supervision, reinforcement-learning-specific architectures, or self-supervised object-centric representations.These approaches target structured representations for entities and their interactions in dynamical scenes.
  • Masked prediction: Masked image modeling began as scalable self-supervised learning for improving representation quality and later incorporated motion- or dynamics-informed masking.The cited progression connects general masked image modeling with guided masking strategies.
  • Masked prediction: Nakano et al. mask image patch tokens while using object representations only as conditioning signals, whereas causal-discovery methods use masking to identify latent structure.The passage distinguishes token-level masking from object-level representation use.

3. Preliminaries

JEPA predicts latent relationships without reconstructing pixels, while Slot Attention produces object-level representations by grouping features into slots. C-JEPA builds on these two foundations for object-centric dynamics modeling.

  • JEPA: JEPA learns predictive relationships between latent embeddings rather than reconstructing pixels.Its extensions include image-level masked prediction in I-JEPA and spatiotemporal tube-based prediction in V-JEPA.
  • Slot Attention: Slot Attention iteratively groups feature maps into a fixed set of slots through competitive attention without requiring supervision.The mechanism alternates between attention-based assignment and slot updates.
  • Object-centric representations: C-JEPA uses Slot Attention to encode each observation into a fixed-size set of object-centric latent states, with each slot corresponding to a distinct scene entity.The framework can use pretrained feature encoders such as DINO or DINOv2.

4. Method

C-JEPA is an object-centric latent world model that applies object-level masking during training to induce interaction-aware predictive representations. It predicts masked and future object states from partially observable histories, auxiliary variables, and contextual interactions.

  • Core design: C-JEPA applies object-level masking during training to induce interaction-aware predictive representations.The method treats masking as a latent intervention on predictor observability.
  • Problem setting: An object-centric encoder maps each video frame to a fixed set of slot representations used to model future object states.The slots are permutation-equivariant, with N slots and dimensionality d for each slot.
  • Masking scheme: During training, selected object states are replaced by mask tokens across the history window, while future entity tokens are masked for prediction.A minimal earliest-time identity anchor preserves entity identity despite slot permutation equivariance.
  • Conditioning variables: Auxiliary actions and proprioceptive signals condition predictions alongside object states, capturing object–object and action-mediated dependencies.These variables remain external to the object-level state representation.
  • Prediction objective: The predictor is a bidirectional masked transformer that jointly infers masked object tokens across history and future intervals.Training minimizes masked latent prediction over the selected tokens, while inference uses a fully observed history for forward prediction.
  • Inductive bias: The history reconstruction term suppresses reliance on trivial self-dynamics, while the future prediction term preserves forward world-modeling behavior.Together, the terms make interaction reasoning functionally necessary for minimizing the predictive objective.

5. Experiments

C-JEPA is evaluated for visual reasoning, predictive control, and interaction-aware prediction. Across these settings, object-level masking improves counterfactual reasoning and preserves strong control performance with substantially lower planning cost.

  • Visual reasoning: C-JEPA improves visual question answering, with larger gains on counterfactual questions than on overall accuracy or other question categories.The comparison with OC-JEPA isolates the learning objective from object-centric representations alone.
  • Visual reasoning: C-JEPA uses object-level masking to infer missing object trajectories from surrounding context, exposing the predictor to counterfactual-like queries.
  • Predictive control: C-JEPA achieves comparable performance to patch-based world models while using only 1.02% of latent tokens and more than 8× faster model predictive control.On identical single-L40s settings, planning 50 trajectories averaged 673 seconds for C-JEPA versus 5,763 seconds for DINO-WM.
  • Predictive control: Treating actions and proprioceptive signals as separate auxiliary entities consistently outperforms concatenating them into object latents.
  • Interaction analysis: Qualitative PHYRE analyses provide evidence that object-level masking encourages interaction-aware predictive dependencies in temporally extended physical environments.
  • Masking ablations: Object-level masking provides a more structured and controllable inductive bias than token- or tube-level masking, with more stable training and stronger interaction-heavy performance.

6. Theoretical Perspective: Causal Inductive Bias of Latent Masking

The theoretical analysis characterizes object-level masking as an observability intervention that makes contextual interaction information necessary for predicting masked object states. It formalizes influence neighborhoods as predictively sufficient sets while distinguishing them from true causal parents.

  • Assumptions: The analysis assumes temporally directed dependencies, a shared transition mechanism, object-aligned representations, and finite-history sufficiency.
  • Interaction necessity: Masked history prediction forces the predictor to use contextual variables beyond the target object when those variables reduce uncertainty.
  • Influence neighborhoods: An influence neighborhood is the minimal sufficient subset of contextual variables needed to recover a masked object state under partial observability.
  • Interaction necessity: A predictor that ignores variables in the influence neighborhood cannot attain the minimum achievable expected reconstruction error.
  • Inductive bias: Repeated object-level masking encourages state-dependent attention patterns aligned with intervention-stable influence neighborhoods.
  • Scope of interpretation: Influence neighborhoods are interpreted as predictively sufficient sets under masking, not estimates of true causal parents or causal mechanisms.
  • Implications: C-JEPA encourages these predictive neighborhoods for reasoning, planning, and model-based control without requiring an explicit causal graph.
  • Temporal direction: Bidirectional masked prediction yields a direction-agnostic interaction structure that abstracts away edge direction while capturing jointly informative variables.

7. Conclusion

C-JEPA combines joint embedding prediction with object-level masking to introduce a causal inductive bias directly through the learning objective. It improves visual reasoning and enables efficient predictive control, while key limitations concern encoder quality, causal-graph validation, and evaluation in richer environments.

  • C-JEPA combines joint embedding prediction with object-level masking to introduce a causal inductive bias directly through the learning objective.
  • C-JEPA yields strong gains in visual reasoning, especially on counterfactual questions.
  • C-JEPA enables highly efficient planning with orders of magnitude fewer tokens than patch-based world models while achieving comparable performance.
  • Performance depends on the quality of the object-centric encoder, which can limit the performance ceiling.
  • The formal influence-neighborhood analysis is not directly validated on datasets with explicit temporal causal graphs.
  • Future work includes refining object-centric encoders with pretrained backbones and evaluating C-JEPA in more complex environments with richer interactions.

Impact Statement

The paper presents a foundational method for efficient, interaction-aware, task-agnostic world modeling, with formal analysis centered on observability interventions. Its empirical settings include visual reasoning and predictive control benchmarks, while the causal interpretation remains predictive rather than graph-identifying.

  • The work targets efficiency and interaction-aware learning in task-agnostic world models and their applications.
  • Object-level masking changes predictor observability while leaving the underlying data-generating and transition mechanisms unchanged.
  • C-JEPA treats masking as a latent intervention that creates counterfactual-like prediction problems under controlled information removal.
  • Unlike methods requiring multiple environments or externally specified interventions, C-JEPA induces observability variation through masking within a single dataset.
  • Influence neighborhoods identify predictively sufficient observable variables under masking without assuming causal sufficiency or a true causal graph.
  • CLEVRER evaluates descriptive, counterfactual, explanatory, and predictive questions, while Push-T evaluates planar manipulation toward a target configuration.
  • CLEVRER results use the validation set because the evaluation server was unavailable, with no validation data used for training or model selection.

I. Full Results

Object-level masking consistently improves visual question answering over the unmasked baseline, especially for counterfactual questions. With a shared SAVi encoder, C-JEPA achieves the best overall performance across all question categories without reconstruction losses.

  • Object-level masking consistently improves VQA performance over the unmasked baseline, with especially large gains on counterfactual questions.
  • Masking object histories provides a meaningful training signal and strengthens interaction-dependent reasoning.
  • C-JEPA achieves the best overall performance across all question categories among object-centric baselines using a shared SAVi encoder.
  • C-JEPA shows particularly strong gains on counterfactual and predictive questions.
  • C-JEPA attains these improvements without relying on reconstruction losses.

J. Qualitative Analysis on PHYRE

Qualitative PHYRE analyses compare C-JEPA with its unmasked counterpart through imagined rollouts and cross-slot attention. C-JEPA produces more physically plausible futures and sharper interaction-relevant attention, but these observations do not establish formal causal graph recovery.

  • Qualitative Analysis on PHYRE: PHYRE tests temporally extended physical dependencies such as gravity, momentum transfer, and multi-body collisions without evaluating causal graph recovery.
  • Rollout plausibility: C-JEPA produces more physically plausible imagined rollouts than OC-JEPA, especially when early collisions should affect later object motion.
  • Attention-based dependency proxy: Cross-slot attention patterns are qualitative proxies for predictive dependencies, not recovered causal graphs.
  • Attention-based dependency proxy: C-JEPA concentrates cross-slot attention on the slot involved in the relevant interaction, whereas OC-JEPA attends more diffusely to less relevant slots.
  • Discussion: Overall, C-JEPA shows more physically plausible rollouts and sharper, more interpretable cross-slot attention than its unmasked counterpart in PHYRE.

K. Exploring Different Masking Strategies

The section compares object-, token-, and tube-level masking strategies, emphasizing how masking granularity affects interaction-aware learning and control performance.

  • Masking strategies: Object-level masking hides entire object-centric slots, requiring the model to infer each masked object's latent trajectory from the remaining objects.Its masking budget is the number of masked objects out of seven.
  • Masking strategies: Token-level masking randomly hides individual latent tokens, while tube-level masking hides contiguous spatio-temporal regions.Both strategies report their budgets as percentages of masked tokens.
  • Masking strategies: Equivalent token budgets and matched tube counts align token- and tube-level comparisons with the object-level masking index set.This controls for masking extent while preserving differences in granularity and shape.
  • Results: Masking performance generally improves as the budget increases across strategies, supporting masking as an effective regularization signal.The strategies can be viewed as variants differing primarily in granularity and mask shape.
  • Results: Object-level masking preserves robust Push-T planning at higher ratios, whereas tube-level masking degrades severely and token-level masking offers limited or inconsistent benefits.The comparison identifies whole-object masking as the critical granularity for meaningful interaction-aware learning in this control setting.

L. Assumptions

The analysis assumes temporally directed, shared object dynamics represented by coherent object-aligned slots and predictable from a finite history window. Its guarantees are bounded by representation quality, stationarity, and the chosen temporal context.

  • Temporal assumptions: Object-level transitions are assumed to follow time-directed predictive dependencies without instantaneous causal effects within the same timestep.The future state depends on past object observations and auxiliary variables.
  • Temporal assumptions: Same-timestep observations may complete masked states, but this is treated as partial-observability completion rather than instantaneous causal generation.Causal influence is defined operationally through time-lagged predictive dependencies.
  • Transition assumptions: Latent transitions are assumed to share an invariant conditional mechanism across trajectories, enabling one predictor to pool experience across episodes.Severe regime shifts would require context or environment conditioning outside the stated scope.
  • Representation assumptions: Each slot is assumed to encode a coherent object-level state sufficient for reasoning about object dynamics, so masking corresponds to removing an object variable.Perfect disentanglement and fixed semantics across scenes are not required, but temporal stability is.
  • Practical limitations: Object-centric model performance is bounded by encoder fidelity, and violations of object alignment can weaken the intended masking intervention.This limits how reliably object-level masking can represent the intended intervention.
  • History-window assumptions: A finite history window is assumed sufficient for predicting future object states, trading information completeness against computational cost.If the true system has longer memory than Th, accuracy may saturate and the predictor characterizes only the chosen-window regime.

M.2. Interpretation of Theorem 1.

Theorem 1 is interpreted as a result about predictive sufficiency under masked observability, not direct recovery of true causal interactions. Object-level masking can nevertheless favor stable attention toward variables that remain informative across masking interventions.

  • Interpretation: The theorem guarantees predictive sufficiency under masked observability rather than recovery of true causal interactions.Alignment with meaningful interactions requires well-aligned object latents, observed interaction variables within the history window, and insufficient spurious predictors.
  • Interpretation: Complete object-level trajectory masking removes the target object's history while preserving a minimal identity anchor, suppressing same-object temporal shortcuts.This makes contextual information from other objects more important than under token- or tube-level masking.
  • Interpretation: Under object-level masking, attention can concentrate on the influence neighborhood because variables outside it are conditionally uninformative for predicting the masked state.Repeated masking patterns act as latent interventions that favor stable, soft relational dependencies without explicitly estimating a causal graph.
Loading 2602.11389v2…