When Does LeJEPA Learn a World Model?

Overview

Canonical JEPA schematic for Theory. The input is split into a visible context and hidden targets (token-level, masked blocks). The context encoder $f_\theta$ embeds what is visible; the target encoder $\bar f_\theta$ (an EMA copy, gradient stopped) embeds the targets; the predictor $g_\phi$ maps context to the target embeddings; training minimises the latent distance.

Problem. JEPAs empirically learn useful latents, but it has been unclear when such a model actually recovers the true generative world variables rather than an arbitrary tangle. This work gives identifiability conditions for LeJEPA.

Mechanism. Under stated assumptions — independent Gaussian latents, stationary additive-noise transitions, and successful Gaussian regularization of the embedding (LeJEPA's anti-collapse / SIGReg-style isotropy constraint) — the authors prove that a JEPA, with its context encoder and latent predictor, linearly recovers the latent world variables up to a rotation. Identifiability up to rotation means the learned latent space is an orthogonal transform of the true one, so latent transitions are faithful and the geometry supports optimal latent-space planning, $z_{t+1}=P(z_t,a_t)$.

Contribution. A rigorous link between JEPA's self-supervised objective and world-model identifiability, explaining why Gaussian-regularized latent prediction can yield plannable, structured representations rather than collapsed or rotationally arbitrary ones.

Significance. For world modeling, it grounds JEPA planning in identifiability theory and motivates isotropic Gaussian regularization as more than a collapse fix. Crucially, this is design guidance, not a guarantee. Real systems — especially biology — routinely violate independent-Gaussian latents and stationary additive-noise transitions: gene programs are correlated and nonlinear, interventions are multiplicative and state-dependent. For a biological or drug-discovery world model the theorem indicates which inductive biases to aim for when recovering plannable latent state under interventions as actions, while flagging that recovery is only approximate where assumptions break.