When Does LeJEPA Learn a World Model? — part 1/9

Cleaned-text code points [0, 15643). Bibliography omitted here; source.md includes it.

# When Does LeJEPA Learn a World Model?

###### Abstract

A representation that scrambles the true degrees of freedom of the world cannot support reliable planning or compositional generalization. We prove that LeJEPA (alignment plus Gaussian regularization) linearly recovers the world’s latent variables from nonlinear observations, a property known as linear identifiability , in a broad class of worlds where latents evolve under stationary, additive-noise transitions. Our main result is that among all such worlds, the Gaussian is the unique latent distribution for which this guarantee holds. The forward direction rests on a spectral decomposition in which each degree of nonlinearity is strictly penalized by alignment, making the linear map the optimum; the converse rules out every non-Gaussian alternative. We further prove an approximate identifiability result where the guarantee degrades gracefully, and show that linear, orthogonal identifiability enables optimal latent-space planning . We validate the theory with experiments ranging from 2D examples to 1024-dimensional latents, including distributional ablations and pixel-based robotic control. Our theory turns an empirically successful recipe into a mathematical guarantee, providing the foundation for building World Models that provably recover the structure of the world.

## 1 Introduction

The promise of self-supervised learning (SSL) is that we can learn useful representations of the world without labeled data, just by observing and predicting. Joint-Embedding Predictive Architectures [ 1 , JEPAs,] pursue this vision by training a representation to produce similar embeddings for related views of the same input, while a regularizer prevents the representation from collapsing to a trivial constant [ 2 , 3 ] . The resulting representations have proven remarkably effective across image [ 4 ] , video [ 5 , 6 ] , and latent-space planning [ 7 , 8 , 9 ] . However, a deeper question remains: When is a learned representation a World Model, i.e., a faithful map of the world’s latent structure?

Our answer: When it linearly recovers the world’s latent variables (Fig. 1 ). A representation that entangles unrelated latent variables, scrambling an object’s position with its color, or mixing velocity with texture, may score well on narrow tasks but will fail when the world changes [ 10 ] . What we need is linear identifiability , i.e., a mathematical guarantee that the learned representation recovers the underlying latent variables, up to simple symmetries [ 11 ] . Practitioners already test for this routinely. Every time a representation is evaluated via linear probing [ 12 ] , the implicit question is whether the model has learned a linear representation of the latent variables [ 13 , 14 , 15 , 16 , 17 , 18 ] . Without this, linear probes cannot recover latent variables exactly. Linear identifiability is thus a necessary , albeit not sufficient [ 19 ] , condition for faithful linear probing.

Yet no identifiability results exist for JEPAs. This has remained out of reach because prior methods rely on implicit collapse prevention [ 20 , 21 , 22 ] with unspecified embedding distributions. A recent shift changes the picture: LeJEPA [ 2 ] prevents collapse by explicitly regularizing the embedding distribution toward an isotropic Gaussian via Sketched Isotropic Gaussian Regularization (SIGReg). Along with an alignment loss this enables stable end-to-end training from raw pixels [ 9 ] , but whether the resulting representation actually recovers the world’s latent variables still remains open.

##### Contributions.

We close this gap with the first identifiability result for JEPAs (Fig. 2 ). We consider a broad class of worlds with Gaussian latent variables and positive pairs from a stationary process with independent, additive-noise transitions, which are maximum-entropy choices [ 23 ] . We prove that LeJEPA learns a linearly identifiable representation if and only if the latent variables are Gaussian . The forward direction rests on a spectral decomposition penalizing every degree of nonlinearity; the converse rules out every non-Gaussian alternative. We further prove that identifiability degrades gracefully when the objectives are only approximately satisfied, and that linear identifiability suffices for optimal latent-space planning . We validate empirically across 2D mixings, 1024-dim latents, distributional ablations, pixel-based robotic control, and the approximate bound (Sec. 6 ).

## 2 Related Work

##### Representation Learning.

JEPAs [ 1 , 4 , 5 , 6 ] predict in representation space rather than pixel space, avoiding capacity wasted on irrelevant detail. More generally, SSL pulls positive views of the same content together while preventing collapse. Contrastive methods [ 24 , 25 ] use negatives explicitly [ 26 , InfoNCE,] ; non-contrastive methods substitute stop-gradient teachers [ 21 , 20 ] , covariance regularization [ 3 , 27 , VICReg,] , or self-distillation with feature clustering [ 22 , 28 ] . LeJEPA [ 2 ] adds an explicit Gaussianity regularizer (SIGReg); LeWorldModel [ 9 ] scales the recipe to action-conditioned control. InfoNCE, VICReg, and LeJEPA span a hierarchy of Gaussianity constraints, from implicit [ 29 , 30 , 31 , 32 ] through second-moment to full; we test all three (Tab. 1 , App. H.7 ).

##### World Models.

The internal-model concept spans cognitive science, control theory, and modern ML. Cognitive scientists framed mental simulation as the substrate of reasoning, perception, and motor control [ 33 , 34 , 35 , 36 , 37 ] , with free-energy [ 38 ] and probabilistic-program [ 39 ] formalizations of the brain as a generative model. Cybernetics extended this to engineered systems [ 40 , 41 ] , while classical control developed dynamics on a given state space [ 42 , 43 ] . Neural world models progressed from recurrent controller-model pairs [ 44 , 45 , 46 ] through latent dynamics from pixels [ 47 , 48 , 49 , 50 , 51 ] , value-equivalent prediction [ 52 ] , generative [ 53 ] and joint-embedding predictive [ 1 , 9 ] architectures, to large generative video simulators [ 54 , 55 ] . The JEPA argument is that pixel-perfect prediction wastes capacity [ 1 ] . We address the encoder side of this picture; classical control applies in the learned coordinates (Thm. 5.4 ).

##### Identifiability.

Nonlinear ICA is unidentifiable without additional structure [ 56 , 57 ] . Identifiability becomes possible when the world supplies it: non-stationarity [ 58 ] , temporal dependence [ 59 , 60 ] , auxiliary variables [ 61 , 62 ] , contrastive learning [ 63 ] , augmentation [ 64 ] , mechanism sparsity [ 65 ] , interventions [ 66 , 67 ] , or supervision [ 68 , 69 ] . These results typically constrain the representation to a smooth diffeomorphism; our Hermite approach works for arbitrary measurable maps. Most closely related is slow feature analysis (SFA) [ 70 ] , which JEPA objectives empirically recover [ 71 ] ; App. F contrasts in depth with SFA theory [ 72 ] . Similar spectral analysis in representation space has previously characterized other SSL objectives [ 73 ] , and our Thm. 5.3 connects to recent quantitative stability work [ 74 , 75 ] . The common lesson is that identifiability is always a joint statement about the World (‘data generating process’) and the Learner (‘learning objective’, LeJEPA).

## 3 The World and the Learner

What does it mean to Learn the World Model? Suppose the world has latent variables z ∈ ℝ n z\in\mathbb{R}^{n} , for instance, position, velocity, color or lighting. These degrees of freedom are called latent variables / sources in the ICA literature, or factors of variation [ 76 ] in representation learning. We never observe z z directly. Instead, an unknown process g g generates the data we see: x = g ⁡ ( z ) x=g(z) . 1 1 1 More philosophically, we only ever observe the shadows ( x ) (x) of reality ( z ) (z) , projected ( g ) (g) onto the walls of Plato ’s cave. Think of g g as rendering a 3D scene into an image, or mapping physical states to sensor readings. The process g g can be highly nonlinear, it scrambles the clean latent structure into complicated, entangled observations. We train a representation f f that maps observations back to representations: y = f ⁡ ( x ) y=f(x) . The ideal outcome is that f f undoes g g : the composed map h = f ∘ g h=f\circ g should recover the original latent variables z z . Of course, perfect recovery is too much to hope for, there are symmetries that cannot be resolved like the rotation invariance of a Gaussian. We will show that h h must be a linear function of the true latents: h ⁡ ( z ) = Q ​ z h(z)=Qz . This is a necessary condition for linear probes to work.

### 3.1 The World

The world specifies a joint distribution p ⁡ ( z , z ′ ) p(z,z^{\prime}) over positive pairs ( g ⁡ ( z ) , g ⁡ ( z ′ ) ) (g(z),g(z^{\prime})) . In SSL, these are two views of the same underlying content, e.g., two frames of a video, two augmentations of an image, two nearby time steps of a trajectory. A broad class of worlds is defined by three assumptions:

Independent latent variables are the standard assumption shared by ICA and disentangled representation learning [ 11 , 76 ] . Stationarity means the generative process does not change between views. Additive noise is the simplest perturbation model: random fluctuations are added on top of a deterministic signal, as in sensor noise, Brownian motion, or data augmentation by jitter. Together, these assumptions define a general class of World Models. Our forward result (Sec. 5.1 ) will specialize this class by choosing a specific latent variable distribution (Gaussian). Our converse result (Sec. 5.2 ) shows that this the unique choice yielding linear identifiability in this class of worlds.

#### 3.1.1 The Gaussian World

We now make a specific distributional choice within the framework above. We assume Gaussian latents , i.e., z ∼ 𝒩 ⁡ ( 0 , I n ) z\sim\mathcal{N}(0,I_{n}) . This is the maximum-entropy distribution for a given mean and covariance [ 23 ] ; it assumes as little structure as possible. Moreover, task-relevant latents are typically aggregates of many micro-variables, which, by the central limit theorem, tend toward Gaussianity. Gaussian latents and Assumptions ( 3.1 ) entail Gaussian transitions : Stationarity requires z ′ ∼ 𝒩 ⁡ ( 0 , I n ) z^{\prime}\sim\mathcal{N}(0,I_{n}) , and the only additive-noise perturbation of a Gaussian that preserves the distribution is the Ornstein–Uhlenbeck (OU) transition [ 77 , 78 ] : z ′ = ρ ​ z + 1 − ρ 2 ​ η , η ∼ 𝒩 ⁡ ( 0 , I n ) , η ⟂ z , z^{\prime}=\rho\,z+\sqrt{1-\rho^{2}}\;\eta,\qquad\eta\sim\mathcal{N}(0,I_{n}),\quad\eta\perp z, (1) where ρ ∈ ( 0 , 1 ) \rho\in(0,1) controls the correlation between views. One can verify: 𝔼 ⁡ [ z ′ ] = 0 \mathbb{E}[z^{\prime}]=0 , Var ⁡ ( z ′ ) = ρ 2 ​ I n + ( 1 − ρ 2 ) ​ I n = I n \mathrm{Var}(z^{\prime})=\rho^{2}I_{n}+(1-\rho^{2})I_{n}=I_{n} , and Cov ⁡ ( z , z ′ ) = ρ ​ I n \mathrm{Cov}(z,z^{\prime})=\rho I_{n} . The Gaussian is the unique distribution for which a channel of this form preserves the marginal, a consequence of the Gaussian being a fixed point of convolution up to rescaling. The independence assumption is satisfied because the components of z z are independent and the noise η \eta has diagonal covariance, so each ( z i , z i ′ ) (z_{i},z^{\prime}_{i}) evolves independently.

### 3.2 The Learner: LeJEPA

The representation is characterized by the composed map h = f ∘ g : ℝ n → ℝ n h=f\circ g:\mathbb{R}^{n}\to\mathbb{R}^{n} with g g the unknown generative function of the world and f f our learned representation. The two components of LeJEPA [ 2 ] training are an invariance loss pulling positive pairs together and a regularizer that shapes the embedding distribution to prevent collapse, i.e., h h (specifically, f f in h = f ∘ g h=f\circ g ) is trained to min h ⁡ ℒ ⁡ ( h ) = 𝔼 ⁡ [ ‖ h ⁡ ( z ′ ) − h ⁡ ( z ) ‖ 2 ] ⏟ Alignment s.t. h ⁡ ( z ) ∼ 𝒩 ⁡ ( 0 , I n ) ⏟ Gaussianity \min_{h}\mathcal{L}(h)=\underbrace{\mathbb{E}\!\left[\|h(z^{\prime})-h(z)\|^{2}\right]}_{\textbf{Alignment}}\qquad\text{s.t.}\qquad\underbrace{h(z)\sim\mathcal{N}(0,I_{n})}_{\textbf{Gaussianity}} (2) We model the situation where SIGReg has succeeded, i.e., h ⁡ ( z ) h(z) matches the target Gaussian; in practice this holds approximately (Sec. 5.3 ). We require nothing else about h h except measurability so expectations are defined. The result thus applies to any neural network, with output dimension taken to equal the latent dimension n n throughout (mismatched regimes m ≠ n m\neq n are discussed in Sec. 7 ). In addition, whitening (i.e., Cov ⁡ ( h ⁡ ( z ) ) = I n \mathrm{Cov}(h(z))=I_{n} ) fixes 𝔼 ⁡ [ ‖ h ⁡ ( z ) ‖ 2 ] = 𝔼 ⁡ [ ‖ h ⁡ ( z ′ ) ‖ 2 ] = n \mathbb{E}[\|h(z)\|^{2}]=\mathbb{E}[\|h(z^{\prime})\|^{2}]=n , so ℒ ⁡ ( h ) = 2 ​ n − 2 ​ ∑ i = 1 n 𝔼 ⁡ [ h i ​ ( z ′ ) ​ h i ​ ( z ) ] . \mathcal{L}(h)=2n-2\sum_{i=1}^{n}\mathbb{E}\!\left[h_{i}(z^{\prime})\,h_{i}(z)\right]. (3) Thus, minimizing distance is equivalent to maximizing correlation between the two views. The question becomes: among all measure-preserving maps h h with h ⁡ ( z ) ∼ 𝒩 ⁡ ( 0 , I n ) h(z)\sim\mathcal{N}(0,I_{n}) , which achieves the highest correlation between positive pairs h ⁡ ( z ) , h ⁡ ( z ′ ) h(z),h(z^{\prime}) ?

## 4 Spectral Analysis of the World

Before proving anything, let us build intuition for the mathematical tools behind our main results.

##### Transition Operator.

The world defines a transition from z z to z ′ z^{\prime} ( 3.1 ). This transition induces an operator on functions: given any scalar function h i ​ ( z ) h_{i}(z) , define the transition operator T T by ( T ​ h i ) ​ ( z ) = 𝔼 ⁡ [ h i ​ ( z ′ ) ∣ z ] (Th_{i})(z)=\mathbb{E}[h_{i}(z^{\prime})\mid z] , the expected value of h i h_{i} at the next view, given the current state. This is a linear operator with a spectral decomposition: a set of eigenfunctions φ k \varphi_{k} satisfying T ​ φ k = λ k ​ φ k T\varphi_{k}=\lambda_{k}\varphi_{k} , with eigenvalues 1 = λ 0 > λ 1 ≥ λ 2 ≥ ⋯ ≥ 0 1=\lambda_{0}>\lambda_{1}\geq\lambda_{2}\geq\cdots\geq 0 . The eigenfunctions with the largest eigenvalues are most correlated across positive pairs, i.e., the most predictable features of the latent variables. This spectral perspective is the common thread behind both our forward and converse results.

##### Gaussian World: Hermite Polynomials.

For Gaussian worlds (Sec. 3.1.1 ), the eigenfunctions are known in closed form: they are the Hermite polynomials { He k } k ≥ 0 \{\mathrm{He}_{k}\}_{k\geq 0} , the natural orthogonal basis for functions of Gaussian variables, analogous to Fourier modes for periodic functions. The eigenvalue of a degree- d d Hermite polynomial is exactly ρ d \rho^{d} , a consequence of Mehler’s formula [ 79 ] . This means any function h i ​ ( z ) h_{i}(z) with zero mean and unit variance can be decomposed into a linear part (degree 1), a quadratic part (degree 2), a cubic part (degree 3), and so on, with variance fractions w 1 , w 2 , w 3 , … w_{1},w_{2},w_{3},\ldots summing to 1. The correlation across positive pairs decomposes: 𝔼 ⁡ [ h i ​ ( z ′ ) ​ h i ​ ( z ) ] = w 1 ⋅ ρ + w 2 ⋅ ρ 2 + w 3 ⋅ ρ 3 + ⋯ ≤ ρ , \mathbb{E}\!\left[h_{i}(z^{\prime})\,h_{i}(z)\right]=w_{1}\cdot\rho+w_{2}\cdot\rho^{2}+w_{3}\cdot\rho^{3}+\cdots\;\leq\;\rho, (4) with equality if and only if w 1 = 1 w_{1}=1 , i.e., h i h_{i} is linear. In words: any nonlinear distortion of the representation strictly reduces the correlation between positive pairs . This is the key intuition.

##### General Case: Sturm-Liouville Theory.