Technical Summary: When Does LeJEPA Learn a World Model?

Problem Statement

The paper addresses a fundamental gap in self-supervised learning (SSL): while Joint-Embedding Predictive Architectures (JEPAs) have demonstrated empirical success in learning representations for vision, video, and planning, there is no theoretical guarantee that these representations constitute a "World Model." Specifically, it remains unclear under what conditions a learned representation faithfully recovers the underlying latent structure of the world (e.g., position, velocity, color) rather than entangling these factors.

A representation that scrambles the true degrees of freedom (e.g., mixing an object's position with its color) may perform well on narrow tasks but fails when the world changes or when used for compositional generalization. The core question is: When does a learned representation linearly recover the world's latent variables? The authors define this property as linear identifiability, a necessary condition for reliable linear probing and optimal latent-space planning.

Methodology and Theoretical Framework

The World Model

The authors formalize the "World" as a system with latent variables $z \in \mathbb{R}^n$ that are unobserved. The observed data $x$ is generated by an unknown, potentially highly nonlinear mixing function $g$, such that $x = g(z)$. The learning objective involves positive pairs $(z, z')$ generated by a stationary process with independent, additive-noise transitions.

The paper focuses on a specific class of worlds where:

1. Independence: Latent variables and their transitions are independent across dimensions.
2. Stationarity: The marginal distribution of the latent variables remains constant between views.
3. Additive Noise: Transitions follow $z' = m(z) + \eta$, where $\eta$ is independent noise.

Crucially, the authors specialize this class to the Gaussian World, where $z \sim \mathcal{N}(0, I_n)$. Under these assumptions, the transition dynamics correspond to an Ornstein–Uhlenbeck (OU) process: $z' = \rho z + \sqrt{1-\rho^2}\eta$.

The Learner: LeJEPA

The learning architecture is LeJEPA (LeCun's JEPA), which trains an encoder $f$ to produce embeddings $h(z) = f(g(z))$ by minimizing a loss function composed of two terms:

1. Alignment Loss: Minimizes the distance between embeddings of positive pairs, $E[\|h(z') - h(z)\|^2]$. This encourages the representation to capture slow, predictable features.
2. SIGReg (Sketched Isotropic Gaussian Regularization): Explicitly regularizes the embedding distribution to be an isotropic Gaussian, $h(z) \sim \mathcal{N}(0, I_n)$, preventing representation collapse.

The optimization problem is:
$$\min_h E[\|h(z') - h(z)\|^2] \quad \text{s.t.} \quad h(z) \sim \mathcal{N}(0, I_n)$$

Theoretical Approach

The authors employ spectral analysis of the transition operator $T$, defined as $(Th)(z) = E[h(z')|z]$.

• Hermite Expansion: For Gaussian latents, the eigenfunctions of the transition operator are the Hermite polynomials. The eigenvalues decay geometrically as $\rho^d$, where $d$ is the degree of the polynomial.
• Spectral Penalty: The alignment loss penalizes higher-degree (nonlinear) components more heavily than linear components because they decorrelate faster under the transition.
• Sturm-Liouville Theory: For general (non-Gaussian) distributions, the eigenfunctions are characterized by Sturm-Liouville theory. The first non-constant eigenfunction is always monotonic, but linearity is only guaranteed if the distribution is Gaussian.

Key Contributions

1. Forward Direction: Linear Identifiability for Gaussian Worlds

Theorem 1 proves that if the latent variables are Gaussian and the transition is an OU process, any minimizer of the LeJEPA objective (subject to Gaussianity constraints) must be a linear orthogonal transformation of the true latents: $h(z) = Qz$ for some orthogonal matrix $Q$.

• Mechanism: The proof relies on a spectral decomposition where any nonlinearity strictly reduces the correlation between positive pairs. Since the alignment loss seeks to maximize this correlation, the global optimum forces the representation to be purely linear (degree-1 Hermite polynomials).
• Result: The learned representation recovers the latent variables up to a global rotation, and the transition dynamics are recovered exactly in the latent space.

2. Converse Direction: Uniqueness of the Gaussian

Theorem 2 establishes that the Gaussian distribution is the unique latent distribution for which linear identifiability holds in this class of worlds.

• Mechanism: Using Sturm-Liouville theory, the authors show that demanding the first eigenfunction of the transition operator to be affine (linear) forces the score function of the latent distribution to be linear. This condition uniquely characterizes the Gaussian distribution.
• Implication: If the latent variables are non-Gaussian, LeJEPA cannot guarantee linear identifiability; the learned representation will necessarily be a nonlinear transformation of the latents.

3. Approximate Identifiability

Theorem 3 provides a quantitative bound on the recovery error when the objectives are only approximately satisfied (i.e., imperfect alignment and imperfect whitening).

• The error scales gracefully with the alignment gap $\delta$ and the whitening error $\epsilon$.
• The bound highlights that the alignment gap is the dominant factor in recovery error, while approximate whitening is relatively "free" to achieve.

4. Optimal Latent Planning

Theorem 4 connects linear identifiability to control theory. It proves that if the encoder is linearly identifiable ($h(z) = Qz$), then any optimal control problem with costs invariant under orthogonal transformations (e.g., Euclidean distance to a goal) can be solved optimally in the learned latent space. The optimal actions and value functions in the latent space map directly to the true world, requiring no modification of the cost function.

Experimental Results

The authors validate their theory across several experimental settings:

1. 2D Simulations: Using various nonlinear mixing functions (spiral, shear, RealNVP), LeJEPA successfully recovers the latent variables up to rotation when the latents are Gaussian, confirming Theorem 1.
2. High-Dimensional Scaling: Experiments up to $N=1024$ dimensions show that batch-statistic regularizers (SIGReg, VICReg) maintain near-perfect linear identifiability ($R^2 > 0.999$), whereas contrastive methods (InfoNCE) degrade at scale due to kernel width mismatches.
3. Distributional Ablations: Sweeping the latent distribution shape (from heavy-tailed to uniform) demonstrates that linear recovery peaks sharply at the Gaussian ($\alpha=2$) and degrades for non-Gaussian distributions, validating Theorem 2.
4. Pixel-Based Robotic Control: Using the DeepMind Control Suite Reacher, the authors compare training on synthetic Gaussian transitions (OU) versus real RL trajectories.
   • OU Condition: Achieves high linear identifiability ($R^2 \approx 0.95$).
   • Trajectory Condition: Real-world policies induce non-Gaussian, anisotropic marginals, resulting in significantly lower identifiability ($R^2 < 0.5$).
5. Planning Validation: In the Reacher task, straight-line interpolation in the learned latent space of the Gaussian-trained encoder decodes to near-straight paths in joint space (optimal planning). In contrast, the trajectory-trained encoder produces curved, suboptimal paths, confirming Theorem 4.

Significance and Claims

The paper claims to provide the first identifiability result for JEPAs, turning an empirically successful recipe into a mathematical guarantee.

• Foundational Guarantee: The work establishes that LeJEPA provably recovers the structure of the world (latent variables and dynamics) under specific, well-defined conditions (Gaussian latents, stationary additive noise).
• Role of Gaussianity: The paper inverts the classical ICA narrative. In linear ICA, Gaussianity is the distribution where source separation fails. In this nonlinear, self-supervised setting, Gaussianity is the unique distribution that enables linear identifiability.
• Practical Implications: The theory explains why exploration strategies that approximate isotropic random walks (keeping data in the Gaussian regime) are crucial for learning useful world models, whereas goal-directed policies (which collapse data into low-entropy regions) may break identifiability.
• Planning Foundation: By proving that linear identifiability enables optimal planning with rotation-invariant costs, the paper provides a theoretical justification for using learned latent spaces for control without needing to map back to the true physical coordinates.

The authors note limitations: the theory assumes the encoder dimension matches the latent dimension and relies on population-level optimality. However, they demonstrate that the guarantees degrade gracefully with finite samples and optimization errors, and they formally verify all main theorems using the Lean 4 theorem prover.