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

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

For a general latent variable distribution (not necessarily Gaussian) evolving under constant diffusion, the eigenfunctions of the transition operator are characterized by a classical Sturm–Liouville (SL) equation [ 72 ] . The first non-constant eigenfunction φ 1 \varphi_{1} is always monotonic, providing identifiability up to a monotonic transformation. But linear identifiability requires φ 1 \varphi_{1} to be affine, which places a strong constraint on the latent variable distribution. This is the engine behind our converse result (Sec. 5.2 ): we show that only the Gaussian satisfies this constraint. Full details in App. A.2 (Gauss/Hermite) and App. F (SL connection).

## 5 Theory

We state four results that together characterize when LeJEPA learns the World Model. Thm. 5.1 shows that a Gaussian world is linearly identifiable with LeJEPA. Thm. 5.2 establishes that, within the class of worlds defined in Sec. 3.1 , the Gaussian is the unique distribution with this property. Thm. 5.3 bounds the recovery error when objectives are not fully satisfied. Thm. 5.4 shows that linear identifiability enables optimal planning in latent space. All proofs have been verified in Lean 4 theorem prover modulo standard background lemmas axiomatized from the literature (App. G ).

### 5.1 Forward Direction: LeJEPA Learns the World Model

##### Proof Sketch.

Minimizing ℒ \mathcal{L} is equivalent to maximizing ∑ i 𝔼 ⁡ [ h i ​ ( z ′ ) ​ h i ​ ( z ) ] \sum_{i}\mathbb{E}[h_{i}(z^{\prime})h_{i}(z)] ( 3 ). The spectral bound ( 4 ) gives 𝔼 ⁡ [ h i ​ ( z ′ ) ​ h i ​ ( z ) ] ≤ ρ \mathbb{E}[h_{i}(z^{\prime})h_{i}(z)]\leq\rho for each component, with equality if and only if h i h_{i} is linear. At equality, h ⁡ ( z ) = Q ​ z h(z)=Qz for a matrix Q Q with unit-norm rows; Gaussianity forces Q ​ Q ⊤ = I n QQ^{\top}=I_{n} , so Q Q is orthogonal. The transition then follows by direct substitution: h ⁡ ( z ′ ) = ρ ​ h ​ ( z ) + 1 − ρ 2 ​ Q ​ η h(z^{\prime})=\rho\,h(z)+\sqrt{1-\rho^{2}}\,Q\eta , and since Q Q is orthogonal, Q ​ η ∼ 𝒩 ⁡ ( 0 , I n ) Q\eta\sim\mathcal{N}(0,I_{n}) independently of h ⁡ ( z ) h(z) . (Proof in App. A .)

##### Interpretation.

The representation has no choice but to learn the full World Model. Any representation satisfying the two LeJEPA objectives must recover a rotation/reflection of the true latent variables and the true transition dynamics. The only remaining ambiguity is a global rotation, which is inherent to the isotropic Gaussian. We provide an alternative proof via Dirichlet energy and the Mazur–Ulam theorem in App. E , operating in a more theoretical ( ρ → 1 \rho\to 1 ) regime.

### 5.2 Converse Direction: The Gaussian is Unique

Thm. 5.1 uses only that h ⁡ ( z ) h(z) has whitened covariance ( Cov ⁡ ( h ⁡ ( z ) ) = I n \mathrm{Cov}(h(z))=I_{n} , entailed by SIGReg, Fig. 10 ). Does the specific choice of Gaussian matter, or would any white distribution work? The optimal representation extracts the slowest features of the latent process [ 71 ] . Under additive noise, Sturm–Liouville theory orders these as increasingly oscillatory eigenfunctions [ 72 ] ; the first is always monotonic, giving identifiability up to a monotonic transformation for any latent distribution. Linear identifiability demands this eigenfunction be affine, and only the Gaussian satisfies this.

##### Proof Sketch.

Demanding an affine eigenfunction forces the score function ( log ⁡ p ) ′ (\log p)^{\prime} of the latent distribution to be linear. If φ ⁡ ( z i ) = a ​ z i + b \varphi(z_{i})=az_{i}+b is an eigenfunction, then φ ′ \varphi^{\prime} is constant, and the eigenvalue equation collapses to a linear ODE for ( log ⁡ p ) ′ (\log p)^{\prime} in z i z_{i} . Solving it yields log ⁡ p ⁡ ( z i ) ∝ − ( z i − μ ) 2 \log p(z_{i})\propto-(z_{i}-\mu)^{2} (the sign fixed by normalizability), which is Gaussian. The argument is per-component; independence of the latent variables lifts the conclusion to the joint distribution. (Proof in App. B .)

### 5.3 Approximate Identifiability

The preceding results assume exact optimality. In practice, both objectives are only approximately satisfied: alignment reaches a value near but not equal to its minimum, and the regularizer enforces approximate whitening. Identifiability degrades gracefully in both quantities:

##### Interpretation.

The bound has two pieces: a nonlinearity term D D measuring how far h h is from linear, and a distortion term ( ε + D ) 2 (\varepsilon+D)^{2} measuring how far the linear part is from orthogonal. Both vanish in the exact case ( δ = ε = 0 \delta=\varepsilon=0 ), recovering Thm. 5.1 . In practice the first dominates, so recovery error scales as δ / 2 ​ ρ ​ ( 1 − ρ ) \delta/2\rho(1-\rho) : alignment is hard, whitening essentially free. (Proof in App. C .)

### 5.4 Optimal Latent Planning

Theorems 5.1 – 5.3 characterize what the encoder recovers. We now make explicit what orthogonal identifiability buys for a key motivation of World Models: planning actions in latent space.

##### Interpretation.

Linear identifiability turns the learned representation into a useful World Model for planning: trajectories planned in the learned latent are mathematically identical to trajectories planned in the true world, with the same actions and the same value. (Proof in App. D .)

## 6 Experiments

We validate each result: linear identifiability on Gaussian latent variables (Sec. 6.1 , Thm. 5.1 ); converse on a latent-distribution sweep and RL-policy [ 9 ] latents (Sec. 6.2 , Thm. 5.2 ); approximate bound across all runs (Sec. 6.3 , Thm. 5.3 ); and near-optimal latent planning [ 80 ] (Sec. 6.4 , Thm. 5.4 ).

### 6.1 Forward: Linear Identifiability from Gaussian Latent Variables

We first verify Thm. 5.1 in a controlled 2D setting. We sample latents z ∼ 𝒩 ⁡ ( 0 , I 2 ) z\sim\mathcal{N}(0,I_{2}) and apply four nonlinear mixing functions (Figures 1 , 3 ): a norm-dependent rotation g ⁡ ( z ) = R ⁡ ( π ​ ‖ z ‖ 2 ) ​ z g(z)=R(\pi\|z\|_{2})\,z producing a spiral [ 74 , c.f.] , a sinusoidal shear g ⁡ ( z 1 , z 2 ) = ( z 1 + sin ⁡ ( 1.5 ​ z 2 ) , z 2 ) g(z_{1},z_{2})=(z_{1}+\sin(1.5\,z_{2}),\;z_{2}) , a parabolic shear g ⁡ ( z 1 , z 2 ) = ( z 1 , z 2 + z 1 2 ) g(z_{1},z_{2})=(z_{1},\;z_{2}+z_{1}^{2}) , and a RealNVP-style coupling layer [ 81 ] . All four are diffeomorphisms. In addition, the spiral is measure-preserving, demonstrating that Gaussianity of the observations alone does not suffice for identifiability; alignment is essential. We train a 4-layer MLP with the LeJEPA loss ℒ = λ ​ ℒ SIG + ( 1 − λ ) ​ ℒ inv \mathcal{L}=\lambda\,\mathcal{L}_{\mathrm{SIG}}+(1-\lambda)\,\mathcal{L}_{\mathrm{inv}} , where positive pairs are generated via the OU transition in ( 1 ). Fig. 3 shows that the learned representation inverts each nonlinear mixing up to rotation, consistent with Thm. 5.1 . A grid search over λ \lambda and ρ \rho (App. H.6 ) shows: too much Gaussianity ( λ = 0.5 \lambda=0.5 ) collapses the representation, while the best recovery occurs at low λ \lambda and high ρ \rho .

##### Scaling to High Dimensions.

We next sweep the latent dimension N ∈ { 2 1 , … , 2 10 } N\in\{2^{1},\ldots,2^{10}\} (beyond, e.g., DINOv3’s 768 768 embedding dimensions), using a RealNVP mixing and matched encoder [ 81 ] so that any failure is attributable to optimization rather than encoder expressivity. We test all three Gaussianity-enforcing classes on the same setup: SIGReg [ 2 ] , VICReg [ 3 ] , and InfoNCE [ 26 ] . Table 1 : batch-statistic estimators (SIGReg, VICReg) maintain R 2 > 0.999 R^{2}>0.999 up to N = 1024 N=1024 ; InfoNCE degrades at scale under fixed kernel width. Thm. 5.1 is thus empirically supported across mixing functions, dimensions, and hyperparameters when the Gaussian-latent assumptions hold; the practical gap between methods appears only off the theoretical ideal.

### 6.2 Converse: Non-Gaussian Latent Variables Break Linear Identifiability

Next, we validate Thm. 5.2 ’s predictions that non-Gaussian latents break linear identifiability.

##### Latent-Distribution Sweep.

We sweep the latent variable through the generalized normal family with shape parameter α \alpha ( α → 0 \alpha\to 0 heavy-tailed, α = 1 \alpha=1 Laplace, α = 2 \alpha=2 Gaussian, α → ∞ \alpha\to\infty uniform). Linear recovery peaks sharply at α = 2 \alpha=2 across all three objectives (Fig. 4 b, App. H.7 ), illustrating Thm. 5.2 ; SIGReg and InfoNCE retain a wider plateau than VICReg for heavy-tailed latents.

##### Pixel-Based RL Trajectories.

The DMC Reacher [ 80 ] has two joints, giving a 2D latent state z = ( θ 0 , θ 1 ) z=(\theta_{0},\theta_{1}) (Fig. 11 ). We train a CNN encoder with LeJEPA (App. H.11 ) under two conditions sharing the same rendering pipeline but different distributions: (i) OU: Gaussian samples z ∼ 𝒩 ⁡ ( 0 , I 2 ) z\sim\mathcal{N}(0,I_{2}) as before ( 1 ); (ii) Trajectory: joint-angle pairs ( z t , z t + δ ) (z_{t},z_{t+\delta}) with δ \delta frames separation from 10 10 k RL episodes [ 9 ] . Table 2 (left) shows that OU pairs attain R 2 = 0.95 R^{2}=0.95 at ρ = 0.99 \rho=0.99 , with the two joint dimensions linearly recovered. In contrast, real trajectories break the Gaussian assumption (App. H.11 , Figs. 12 , 13 ). Table 2 (right) shows per-dimension R 2 R^{2} is anisotropic, and total R 2 R^{2} never exceeds 0.5 0.5 , consistent with Thm. 5.2 . The trajectory condition violates several theory assumptions at once (non-Gaussian marginals, anisotropic ρ 0 ≠ ρ 1 \rho_{0}\neq\rho_{1} , joint-limit wrapping).

### 6.3 Approximation Bound and Loss Predict Identifiability

Thm. 5.3 bounds the recovery error by the whitening error ε = ‖ Cov ⁡ ( h ⁡ ( z ) ) − I ‖ F \varepsilon=\|\mathrm{Cov}(h(z))-I\|_{F} and the alignment gap δ = ℒ ⁡ ( h ) − 2 ​ ( 1 − ρ ) ​ tr ​ ( Cov ⁡ ( h ⁡ ( z ) ) ) \delta=\mathcal{L}(h)-2(1-\rho)\,\mathrm{tr}(\mathrm{Cov}(h(z))) . For each run we compute ε \varepsilon , δ \delta , the bound D + ( ε + D ) 2 D+(\varepsilon+D)^{2} , and the actual recovery error min Q ∈ O ⁡ ( n ) ⁡ 𝔼 ⁡ [ ‖ h ⁡ ( z ) − Q ​ z ‖ 2 ] \min_{Q\in O(n)}\mathbb{E}[\|h(z)-Qz\|^{2}] (Fig. 4 a). The bound holds across grid search, 2D mixings, scaling, and the latent-distribution sweep, supporting Thm. 5.3 empirically. As a practical corollary, training loss is a reliable proxy for identifiability (App. H.9 ).

### 6.4 Linear Identifiability Enables Latent-Space Planning

Thm. 5.4 predicts that any planner using a rotation-invariant cost attains identical performance in the learned and true latent. We test a natural instance: goal-reaching, where the straight line z ^ 0 → z ^ ∗ \hat{z}_{0}\to\hat{z}^{*} is the cost-minimizing trajectory. For an encoder satisfying h ⁡ ( z ) ≈ Q ​ z h(z)\approx Qz , this latent straight line decodes to a near-straight path in the true latent; without, the same plan induces a curved path (Fig. 17 ). Figs. 4 c,d; 5 confirm both. The Gaussian encoder’s straight-latent plans decode to oracle-quality joint-space trajectories (left, middle), while the Trajectory encoder inflates control cost. Across all models, control cost tracks linear identifiability R 2 R^{2} monotonically (right). Thus, linear identifiability is the structural property that turns a faithful World Model into a useful planner.

## 7 Limitations

##### Are the latents Gaussian?

The Gaussian is the maximum-entropy distribution for a given mean and covariance [ 23 ] , making it the least-assumption prior. Whether real-world latents are Gaussian is unknowable from observations alone, but the same applies to the non-Gaussianity assumptions of classical ICA [ 82 ] (maybe a structuralist perspective [ 83 ] is more agnostic). There is a scale-of-description argument in favor: individual micro-variables may be non-Gaussian, but task-relevant latent variables are often aggregates that tend toward Gaussianity by the central limit theorem.

##### What if the dimension is wrong?

Our theorem assumes the encoder output dimension matches the true latent dimension ( m = n m=n ). When m < n m<n , the Gaussianity constraint does not determine which subspace is selected or whether the system resorts to superposition [ 15 , 84 , 18 ] ; when m > n m>n , extra dimensions must collapse or encode redundancy. Understanding this interaction is an important open problem with direct consequences for JEPA design [ 18 ] .

##### Finite samples and optimization.

Our result is a population-level statement about the global optimum. Thm. 5.3 shows the guarantee degrades continuously with respect to alignment gap and covariance deviation, but does not address how these scale with sample size [ 85 ] or training dynamics [ 86 ] . The few bound violations we observe empirically (Fig. 4 a) are consistent with finite-sample estimation noise in ε \varepsilon and δ \delta .

## 8 Discussion

##### Outlook.

Linear identifiability addresses the state side of a World Model; the action-conditioned transition p ^ ​ ( z ^ ′ ∣ z ^ , a ) \hat{p}(\hat{z}^{\prime}\mid\hat{z},a) must still be learned. Lifting identifiability to that setting connects naturally to interventional causal representation learning, where actions act as interventions on the latent variables [ 66 , 67 ] , and onward to causal graphs [ 87 ] . Our assumptions may not hold exactly in practice, but Thm. 5.3 shows the guarantee degrades gracefully.

##### Practical Implications.

Two implications, one for data and one for objective. Data: the Reacher result shows that the same physical system supports identifiability when sampled isotropically (OU) but not under a goal-directed policy, whose marginals collapse onto a low-entropy region of latent space (Tab. 2 ); for self-supervised pretraining, exploration approximating an isotropic random walk keeps the data in the regime our theory covers. Objective: SIGReg, VICReg, and InfoNCE all yield linear identifiability when their assumptions hold, and our experiments suggest they fail in different ways: pair-based estimators are sensitive to kernel choice at scale (Tab. 1 ), moment-based estimators to non-Gaussian latents (App. H.7 ). Which estimator is preferable in practice, and how the choice interacts with batch size, architecture, and data regime, is left to future work.

##### A Theoretical Foundation for World Models.

Our theory inverts the classical ICA narrative: in linear ICA the Gaussian is the one distribution where source separation fails [ 82 ] ; in our nonlinear setting it is exactly what enables it. The payoff is structural: linear identifiability turns a learned representation into a usable state for a control system, so any orthogonally-invariant cost transfers from the true world to the learned latent without modification, and simple planning (Thm. 5.4 ) and linear probing come for free. This is what it means to provably learn a World Model.

## Acknowledgments

We thank Aapo Hyvärinen, Patrik Reizinger, Attila Juhos, Heejeong (Hazel) Nam and Christian Internò for their feedback and inputs. DK acknowledges the CSHL GPU cluster with assistance from the US National Institutes of Health Grant S10OD028632-01.

## Appendix Overview

Contents Page Appendix A Proof of Thm. 5.1 (linear identifiability), incl. Hermite background A Appendix B Proof of Thm. 5.2 (Gaussian uniqueness) B Appendix C Proof of Thm. 5.3 (approximate identifiability) C Appendix D Proof and discussion of Thm. 5.4 (optimal latent planning) D Appendix E Alternative proof via Dirichlet energy E Appendix F Prior work: connection to Slow Feature Analysis F Appendix G Lean 4 formal verification G Appendix H Experimental details and additional results H