Causal-JEPA: Causal World Modeling

• Causal-JEPA is a self-supervised framework that integrates causal inductive bias into JEPA via targeted masking in latent space.
• It utilizes motion-centric (IA-JEPA) and object-level masking strategies to drive improved prediction of physical interactions and counterfactual inference.
• Empirical results demonstrate enhanced causal reasoning accuracy, data efficiency, and planning speed compared to standard patch-masked JEPA methods.

Causal-JEPA encompasses a family of world modeling architectures that inject causal inductive bias into Joint Embedding Predictive Architectures (JEPA) through targeted masking strategies in latent space. These methods supersede standard patch-masked JEPA by enforcing interaction-centric prediction requirements, thereby internalizing physical dynamics necessary for causal reasoning, counterfactual inference, and data-efficient planning. The principal mechanisms include motion-centric masking at the pixel level (“Interaction-Aware JEPA,” IA-JEPA) and object-level masking that functions as a latent intervention (Causal-JEPA, C-JEPA), forming a new paradigm for self-supervised, causally grounded video world models (Paidi, 14 May 2026, Nam et al., 11 Feb 2026).

1. Foundations: JEPA and Its Limitations

JEPA is a representation learning framework for video prediction that operates by masking spatiotemporal regions—typically non-overlapping “tubelets”—of the input and requiring the latent reconstruction of these regions using a context encoder and transformer-based predictor. Standard JEPA employs random or uniform patch masking, resulting in models that predominantly capture semantic classification features and fail to internalize the dynamics of physical causality. This failure arises from a “static bias”: uniform masking leads the model to explain static backgrounds at the expense of learning the rare, informative events (e.g., collisions, momentum transfers) that encode the causal dynamics essential for downstream reasoning (Paidi, 14 May 2026).

JEPA variant	Masking strategy	Targeted causal bias	Main limitation
Patch-masked	Uniform/random patches	None	Static bias, low entropy
IA-JEPA	High-motion patch masking	Kinematics/events	No entity structure
C-JEPA	Object-level/intervention	Object interactions	Dependent on object encoder

2. Interaction-Aware Masking: IA-JEPA

IA-JEPA introduces a data-driven, motion-aware masking strategy designed to counteract static bias. The core algorithm computes a per-pixel interaction score via the second-order temporal difference:

$G(p, t) = \|x(p, t+1) - 2x(p, t) + x(p, t-1)\|$

This score is average-pooled to the patch grid to yield a saliency map $S_{i,j}$, from which a mask probability $p_{i,j}$ is derived. IA-JEPA samples the top proportion (typically 40%) of patches with the highest $p_{i,j}$, targeting regions containing collisions or momentum imparting events. The self-supervised pre-training objective remains the minimization of $L_2$ distance in latent space between predicted and target tubelet representations. By masking and predicting only in high-motion regions, IA-JEPA incentivizes the context encoder to encode rich kinematic abstractions—position, velocity, and proxies for mass—and discourages copying static context (Paidi, 14 May 2026).

Empirical analysis demonstrates IA-JEPA’s efficacy:

• Entropy of latent representations increases by 10% (latent std deviation rises from 0.856 to 0.938).
• Latent states exhibit a linear relationship with physical motion energy ($R^2=0.43$), indicating energy is approximately linearized in latent space.
• On the CLEVRER benchmark, IA-JEPA achieves 14.26% accuracy on causal reasoning, up from 3.22% with patch masking.
• Transfer to real-world video (Something-Something V2) yields 40.60% accuracy (baseline: 34.40%), and zero-shot performance on the PHYRE-Lite physics puzzle benchmark rises to 79% (baseline: 60%).

3. Object-Level Latent Interventions: The Causal-JEPA Method

Causal-JEPA (C-JEPA) generalizes the causal inductive bias by performing masking at the object level. It operates as follows:

1. Object Extraction: Each video frame $X_t$ is parsed into $N$ slot vectors $S_t = \{s_t^1, \dots, s_t^N\}$ via an encoder (e.g., Slot Attention over DINOv2 features).
2. Masked Latent Prediction: Object tokens, potentially augmented with auxiliary variables (actions, proprioception), are masked per object—i.e., selected slots have their information withheld across temporal windows.
3. Bidirectional Masked Transformer: The sequence of masked and unmasked object slots is input to a masked transformer $f$ for prediction. The loss is the mean-squared error between predicted and masked true object slots over both the history and the (fully masked) future:

$S_{i,j}$0

Object-level masking constitutes a “latent intervention”: by hiding all access to $S_{i,j}$1 at timestep $S_{i,j}$2, the architecture is compelled to infer this state from the “influence neighborhood”—the minimal set of other object states necessary for accurate imputation, formalized by conditional independence in the latent space. Theoretical analysis shows that MSE-optimal masked prediction yields $S_{i,j}$3, and the object-slot attention patterns induced by this process are stable under interventions, providing a causal inductive bias (Nam et al., 11 Feb 2026).

4. Empirical Evidence: Reasoning, Planning, and Ablations

Causal-JEPA has been validated on diverse tasks:

• Visual Question Answering (CLEVRER):

With four objects masked per batch, average QA accuracy improves by $S_{i,j}$4 pp and counterfactual QA by up to $S_{i,j}$5 pp relative to non-masked object-centric JEPA.

• Comparison to Baselines:

C-JEPA, without explicit slot trajectory reconstruction, outperforms SlotFormer and other sequence prediction baselines.

• Model-Based Planning (Push-T):

With only $S_{i,j}$6 of latent tokens (object-centric slots versus patch tokens), C-JEPA achieves planning success rates comparable to DINO-WM (88.67% vs. 91.33%) and $S_{i,j}$7 speedup (673 s vs. 5763 s for 50 trajectories).

• Masking Strategy Ablation:

Object-level masking confers more stable gains across masking budgets compared to token or tube masking. In control tasks, tube masking collapses at moderate budgets, while object masking remains robust.

Task	Patch-based baseline	IA-JEPA/C-JEPA	Absolute gain
CLEVRER causal QA	3.22% (patch JEPA)	14.26% (IA-JEPA)	+11.04 pp
CLEVRER counterfactual	47.68% (OC-JEPA)	68.81% (C-JEPA, 4/7)	+21.13 pp
Something-Something V2	34.40%	40.60% (IA-JEPA)	+6.2 pp
PHYRE-Lite	60.0%	79.0% (IA-JEPA)	+19.0 pp
Push-T planning	91.33% (DINO-WM)	88.67% (C-JEPA)	–2.66 pp (~1% tokens)

5. Inductive Biases and Theoretical Implications

Causal-JEPA’s principal innovation is enforcing an inductive bias that compels the model to simulate interactions in latent space. In IA-JEPA, high-motion masking compels the predictor to internalize kinematics; in C-JEPA, object-level masking enforces reasoning about inter-object influence and latent counterfactuals. The “where” of masking is shown to be as critical as the “what”: masking interaction interfaces selectively breaks shortcut solutions, increases latent entropy (+10%), and aligns latent structure with physical energy.

Formally, repeated exposure to masking as a latent intervention induces predictor attention patterns matched to the influence neighborhoods of each object, providing a mechanism analogous to invariant predictors in causal discovery. Because training uses bidirectional attention over history and futures, the model’s reasoning patterns become symmetric in time, enabling robust backward completion and forward prediction (Paidi, 14 May 2026, Nam et al., 11 Feb 2026).

6. Limitations and Prospects for Extension

Causal-JEPA is constrained by the quality of its object-centric encoder; the encoder–predictor interface is not jointly trained, which may limit optimality. Influence neighborhoods are formalized theoretically but not systematically validated against ground-truth causal graphs, suggesting future work could include evaluation on synthetic causal dynamical systems and exploration of fully learnable, dynamic entity discovery. Performance in more complex environments—deformable objects, fluids, and real-world robotics—remains untested.

Possible extensions include combining interaction-aware masking with object-centric slots for simultaneous entity discovery and dynamic interface prioritization, augmenting masking heuristics beyond motion (e.g., incorporating learned flow or depth cues), and probing counterfactuals by systematic latent interventions in discovered entities’ trajectories. These directions point toward a scalable, fully self-supervised “Causal-JEPA” framework for video world modeling rooted in causal mechanism learning rather than mere visual recognition (Paidi, 14 May 2026, Nam et al., 11 Feb 2026).

7. Significance and Broader Impact

By drastically reducing the number of tokens required for accurate prediction, Causal-JEPA achieves state-of-the-art data efficiency and planning speed in model-based control while improving causal and counterfactual reasoning performance in visual domains. Its object-centric reasoning and intrinsic efficiency may accelerate applications in robotics, simulation, and autonomous systems. The method is foundational and agnostic to downstream tasks, and at present no additional societal risks are identified beyond those common to learned models (Nam et al., 11 Feb 2026).

The Causal-JEPA paradigm establishes that strategic masking—targeted either at high-velocity regions (IA-JEPA) or at the level of latent entities (C-JEPA)—is a general tool for bridging representation learning with causal abstraction and simulation. This suggests a path toward world models that move beyond static visual summaries to embody mechanistic understanding of physical dynamics, supporting both prediction and robust reasoning under intervention.