When Does LeJEPA Learn a World Model?

TL;DR. LeJEPA — the self-supervised recipe of "pull positive pairs together, and regularize embeddings toward an isotropic Gaussian" — gets an identifiability theorem. If the world's latent variables are Gaussian and evolve under stationary additive-noise dynamics, any encoder that exactly satisfies both LeJEPA objectives must output a rotation of the true latents: h(z) = Qz with Q orthogonal. Better still, the Gaussian is the unique latent distribution in this class for which the guarantee holds — a delicious inversion of classical ICA, where Gaussianity is the one case where source separation fails. The paper adds a graceful-degradation bound, a proof that orthogonal identifiability suffices for optimal latent-space planning, Lean 4 verification of all proofs, and experiments up to 1024 dimensions and pixel-based robot control. The catch, which the paper honestly demonstrates: real robot trajectories from a trained policy violate the assumptions, and identifiability drops from R^2 = 0.95 to below 0.5.

The question nobody had answered for JEPAs

Self-supervised methods like JEPA train an encoder so that two related views of the world — two video frames, two augmentations — map to nearby embeddings, while some regularizer prevents everything from collapsing to a constant. This works empirically across images, video, and latent-space planning. But there's a foundational question hiding underneath: does the learned representation actually recover the world's degrees of freedom, or does it merely produce features that happen to score well on benchmarks while scrambling position with color and velocity with texture?

The formal version of this question is linear identifiability. Suppose the world has latent variables z \in \mathbb{R}^n (joint angles, object positions, lighting), and we only observe x = g(z) through some unknown, highly nonlinear rendering g — pixels, sensor readings, Plato's cave shadows. We learn an encoder f, and the object of interest is the composed map h = f \circ g. Perfect recovery h = \mathrm{id} is impossible (symmetries can't be resolved), but we can ask for h(z) = Qz for some fixed invertible — ideally orthogonal — matrix Q. This is precisely the condition under which linear probing, which the entire field uses as its evaluation ritual, can actually read out the true latents.

For contrastive and nonlinear-ICA-style methods, identifiability theory exists. For JEPAs, it didn't — because collapse prevention was implicit (stop-gradients, teacher networks), leaving the embedding distribution unspecified and the math with nothing to grip. LeJEPA changed the setup: its SIGReg regularizer explicitly pushes the embedding distribution toward \mathcal{N}(0, I_n). That explicit distributional constraint is what makes a theorem possible.

The setup

The world: latents z with independent components, and positive pairs (z, z') generated by a stationary transition with additive noise, z'_i = m_i(z_i) + \eta_i. Specializing to Gaussian latents z \sim \mathcal{N}(0, I_n), stationarity forces the transition to be the Ornstein–Uhlenbeck channel:

where \rho \in (0,1) controls how correlated the two views are. Think of this as the maximum-entropy model of "two nearby frames of a video": same marginal, correlated content, independent jitter.

The learner: minimize the alignment loss \mathcal{L}(h) = \mathbb{E}\|h(z') - h(z)\|^2 subject to h(z) \sim \mathcal{N}(0, I_n) (SIGReg having succeeded). Crucially, h can be any measurable function — no smoothness, no invertibility, no architectural restriction. Since the Gaussian constraint fixes \mathbb{E}\|h(z)\|^2 = n, minimizing alignment is exactly maximizing the cross-view correlation \sum_i \mathbb{E}[h_i(z')\,h_i(z)].

The aha: alignment is a low-pass filter on nonlinearity

The heart of the forward proof (Theorem 1) is a spectral decomposition that is genuinely pretty. Any square-integrable function of a Gaussian variable can be expanded in Hermite polynomials — the Gaussian world's analogue of Fourier modes. Degree 1 is the linear part, degree 2 the quadratic part, and so on. The magic property, Mehler's formula, is that the OU transition diagonalizes in this basis: a degree-d Hermite component of the embedding correlates across the two views at exactly \rho^d, and different degrees don't cross-correlate at all.

So write each embedding coordinate's variance budget as fractions w_1, w_2, w_3, \ldots across Hermite degrees (summing to 1, forced by the unit-variance constraint). The cross-view correlation is then

with equality iff w_1 = 1. Since \rho < 1, every unit of variance you spend on nonlinear structure earns strictly less correlation than the same unit spent linearly. Alignment is a low-pass filter over degrees of nonlinearity, and the optimum is forced to sit entirely at degree 1. Linearity then gives h(z) = Qz with unit-norm rows, and the Gaussianity constraint (\mathrm{Cov} = I_n) upgrades Q to orthogonal. As a bonus, the learned representation inherits the true transition dynamics exactly: h(z') \mid h(z) \sim \mathcal{N}(\rho\,h(z), (1-\rho^2)I_n).

Notably, this argument needs no smoothness of h and works at any noise level \rho \in (0,1) — much weaker regularity than the diffeomorphism assumptions typical in nonlinear ICA. An alternative proof via Dirichlet energy and Mazur–Ulam (Appendix E) offers a geometric view — the Jacobian must be a pointwise isometry — but requires C^1 and infinitesimal noise.

The converse: only the Gaussian works

Theorem 2 is the more surprising result. Within the whole class of stationary, independent, additive-noise worlds, suppose every minimizer of the LeJEPA objective is linear. Then the latents must be Gaussian.

The engine is Sturm–Liouville theory. For a general latent distribution under constant-diffusion dynamics, the transition operator's eigenfunctions are increasingly oscillatory; the slowest non-constant one is always monotonic — so you always get identifiability up to a monotonic per-coordinate warp. Linear identifiability demands that first eigenfunction be affine. Plugging an affine function into the eigenvalue equation collapses it to an ODE for the score (\log p)', whose only normalizable solution is the Gaussian: linear score \Leftrightarrow Gaussian density.

The framing here is worth savoring: in linear ICA, Gaussian sources are the pathological case where separation is impossible (rotational symmetry destroys identifiability). In this nonlinear temporal setting, that same rotational symmetry is exactly what pins the solution down to a rotation — the Gaussian flips from the failure mode to the only success mode.

Two more results round out the theory. Theorem 3 gives quantitative robustness: with alignment gap \delta and whitening error \varepsilon, recovery error is bounded by D + (\varepsilon + D)^2 where D = \delta / (2\rho(1-\rho)) — the denominator is the spectral gap between Hermite degrees 1 and 2. Empirically the \delta term dominates: whitening is essentially free; alignment quality is the binding constraint, which makes training loss a usable proxy for identifiability. Theorem 4 shows the payoff: for any control problem whose costs are rotation-invariant (goal-reaching, LQR, anything built from norms and inner products), planning in the learned latent yields identical value and identical optimal actions as planning in the true latent. The orthogonal ambiguity is harmless as long as encoder, dynamics model, and cost are consistent in the learned coordinates. All five results are machine-checked in Lean 4 (with classical facts like Mehler's formula axiomatized, since Mathlib lacks Hermite polynomials).

The evidence

On-assumption, the theory is airtight empirically. With RealNVP mixings and a matched encoder (so expressivity is not a confounder), SIGReg and VICReg hold R^2 > 0.999 from 2 to 1024 latent dimensions. InfoNCE matches at low dimension but collapses at scale — the paper traces this precisely to the fixed Gaussian kernel underflowing when \|z\|^2 \approx N, a kernel-tuning artifact rather than an objective-level failure.

The converse gets clean support too: sweeping the latent distribution through the generalized-normal family (Laplace → Gaussian → uniform), linear recovery peaks sharply at exactly the Gaussian, for all three regularizers.

Off-assumption, the theory breaks — informatively. The pixel experiment uses DeepMind Control Reacher: the true latent is 2D (two joint angles), rendered to 64×64 images by MuJoCo, encoder is a CNN trained with LeJEPA from pixels. When the joint angles are sampled as an OU process — exactly matching the assumptions — the encoder recovers the joints with R^2 = 0.95. When positive pairs instead come from real trajectories of a trained SAC policy on the same physical system with the same renderer, identifiability craters: total R^2 never exceeds 0.5, with strongly asymmetric per-joint recovery. The policy's marginals are non-Gaussian (the wrist is nearly bimodal), the two joints decorrelate at different rates (\rho_0 \neq \rho_1), and angles wrap — violating Gaussianity, isotropy, and topology at once.

The planning experiment closes the loop: straight-line interpolation in the OU-trained encoder's latent space decodes to joint-space paths statistically indistinguishable from the oracle straight line, while the trajectory-trained encoder's plans curve, and across all models control cost tracks identifiability R^2 monotonically.

What this changes, and what to doubt

What it changes. This is, to my knowledge, the first identifiability theorem for any JEPA-family objective, and it's a strong one: measurable maps, finite noise, a converse, a quantitative stability bound, and a formal link from representation to control. It gives a principled answer to why "alignment + explicit Gaussianization" is the right pairing (implicit collapse prevention leaves the distribution unpinned, and any non-Gaussian target provably breaks linearity), and it grounds the field's linear-probing ritual in something more than convention. The practical corollaries are actionable: training loss predicts identifiability; whitening is cheap while alignment is the bottleneck; and — most interesting for robotics pretraining — data collection matters as much as the objective. The Reacher result says an isotropic-random-walk exploration policy puts you inside the theory's guarantee, while goal-directed policy data pushes you out of it. That's a concrete, testable design principle for self-supervised world-model pretraining.

What to doubt. The assumptions are load-bearing everywhere. Independent latents, isotropic transitions (Appendix F shows the fastest latent must be less than twice as fast as the slowest, or eigenvalue interleaving makes the encoder grab the second Hermite harmonic of a slow variable instead of a fast one), and encoder output dimension exactly matching the true latent dimension n — the m \neq n regime, which is every real deployment, is explicitly open. The result is population-level, about the global optimum, silent on finite samples and SGD dynamics. The "LeJEPA" analyzed here is the objective in its purest form; real JEPA systems have predictors, masking, and asymmetric architectures the theory doesn't touch. And the theorem covers only the encoder: the action-conditioned dynamics model that makes a world model a world model still has to be learned, with no identifiability guarantee yet (the authors flag persistent-excitation-style conditions as ongoing work). Finally, whether real-world latents are Gaussian is unfalsifiable; the CLT hand-wave ("task-relevant latents are aggregates of micro-variables") is plausible but the paper's own trajectory experiment shows realistic data is decisively non-Gaussian. Read the RL-trajectory result as the theory's most honest moment: the guarantee is real, and so is the gap between its preconditions and data in the wild.

Where to spend your time. Appendix A — the Hermite proof is short, self-contained, and the generating-function computation of the \rho^d contraction is the whole paper in five equations. Then Appendix F, which situates the result within Slow Feature Analysis theory (Sprekeler et al.'s Sturm–Liouville analysis) and honestly delineates what is new here versus what is a Gaussian specialization of known spectral structure: the simultaneous-extraction formulation, the converse, the quantitative bound, and the planning link.