When Does LeJEPA Learn a World Model?

The phrase “world model” gets used a lot in AI, but it often hides a hard question:

If a model learns an internal representation from observations, how do we know that representation preserves the real degrees of freedom of the world?

That is the question behind When Does LeJEPA Learn a World Model? by David Klindt, Yann LeCun, and Randall Balestriero. The paper studies LeJEPA, a latent version of JEPA-style representation learning, and proves when the learned representation is not just useful, but mathematically tied to the underlying latent variables.

The short version:

LeJEPA can linearly recover the world’s latent variables, up to rotation, if and only if those latents are Gaussian.

That sounds narrow at first. It is actually a very strong statement. It says that under a specific but meaningful class of worlds, alignment plus Gaussian regularization is enough to recover the structure that matters for planning and compositional generalization.

Code availability

Yes, code is available.

The project page links to the repository:

github.com/klindtlab/lejepa-identifiability

The repository is not just a placeholder. It includes:

That matters because the paper makes both empirical and theoretical claims. The public repo gives readers a way to inspect the simulations, rerun core experiments, and look at the machine-checked proof structure.

What LeJEPA is trying to learn

The setup is simple but deep.

There is a hidden world state:

z

The world evolves over time:

z' = m(z) + noise

But the learner does not observe z directly. It observes a nonlinear transformation:

x = g(z)

So the learner sees scrambled observations and must learn a representation that recovers the real latent structure.

LeJEPA trains an encoder with two pressures:

1. Alignment: nearby or temporally related states should have predictable representations.
2. Gaussian regularization: the embedding distribution should look like a standard Gaussian.

In the paper’s simplified objective, the encoder minimizes representation change between positive pairs while satisfying a Gaussian embedding constraint:

minimize E[ || h(z') - h(z) ||^2 ]
subject to h(z) ~ N(0, I)

In practice, the Gaussianity constraint is enforced through SIGReg, the Sketched Isotropic Gaussian Regularizer.

The core result: linear identifiability

The main theoretical result is linear identifiability.

If the world’s latent variables are independent Gaussian variables and evolve through stationary additive-noise transitions, then an optimal LeJEPA encoder recovers the latent variables up to an orthogonal rotation:

h(z) = Qz

where Q is an orthogonal matrix.

That “up to rotation” part is important. It means the model may not recover the exact coordinate system used by the original world, but it recovers an equivalent linear structure. For many downstream tasks, especially planning with rotation-invariant costs, that is enough.

This is a stronger claim than “the representation performs well on a benchmark.” It says the representation is structurally aligned with the true latent world.

Why Gaussian latents are not just a convenience

The most interesting part of the paper is the converse result.

The authors do not only prove that Gaussian latent worlds work. They prove that within their class of stationary additive-noise worlds, the Gaussian is the unique latent distribution where the guarantee holds.

That changes how I read the result.

It is not:

Gaussian latents are mathematically convenient.

It is closer to:

The LeJEPA recipe has a precise identifiability story, and that story depends on the world being Gaussian.

The empirical ablation makes the point visually. When the latent distribution is swept through a generalized-normal family, recovery peaks sharply at the Gaussian case. Heavy-tailed, Laplace-like, and uniform-like worlds break the guarantee.

For practical AI systems, this is a useful warning. Representation learning objectives can look robust while secretly depending on distributional assumptions.

The Hermite polynomial argument in plain language

The proof uses a spectral decomposition under the Gaussian measure, based on Hermite polynomials.

You do not need to live inside the proof to understand the intuition.

Under a Gaussian distribution, functions can be decomposed into components by degree:

linear part
quadratic part
cubic part
...

The alignment objective penalizes nonlinear degrees more strongly than the linear degree. If the embedding must remain Gaussian, the optimal way to satisfy both alignment and Gaussianity is to keep the linear part and discard the nonlinear scrambling.

That is why the learned encoder becomes a rotation of the true latent variables.

This is the mathematical bridge between the objective and the world-model claim. The model is not simply encouraged to be smooth. It is pushed toward the one representation class that keeps the Gaussian latent structure intact.

Why this matters for planning

A world model is only useful if it supports action.

The paper proves that when the learned latent space is linearly identifiable up to an orthogonal transformation, planning can be optimal for a class of finite-horizon control problems whose costs are invariant under that transformation.

In simpler terms:

If the learned representation is a rotated version of the true latent world, and the control problem does not care which rotated coordinate system you use, then planning in the learned space can be as good as planning in the true latent space.

That is a big deal because it connects representation learning to control.

The Reacher pixel experiment makes the point concrete. A Gaussian-OU-trained encoder learns an identifiable latent space, so straight-line latent interpolation tracks the oracle joint-space path. A trajectory-trained encoder that is not identifiable deviates.

This is the kind of result world-model research needs more often: not only prettier representations, but a clear explanation of when those representations preserve the structure required for planning.

The Lean 4 verification is part of the story

One detail I really like: the theoretical results are formally verified in Lean 4.

The project page says the Lean build has zero sorry obligations, meaning the proof chain has no admitted gaps at the Lean level. Some background components are axiomatized because the required Hermite polynomial and related mathematical infrastructure is not all in Mathlib yet, but the authors are explicit about that boundary.

This matters for machine learning theory.

ML papers often combine dense measure-theoretic arguments, optimization claims, and informal proof sketches. Formal verification does not make a theorem automatically important, but it changes the trust surface. It forces the authors to specify assumptions, proof dependencies, and logical steps more sharply.

For a paper about identifiability, that is appropriate. The whole point is to know exactly when the guarantee holds.

What the experiments show

The empirical section checks the theory from several angles:

The scaling result is especially useful: SIGReg and VICReg maintain very high recovery across dimensions up to 1024 in the reported setup, while InfoNCE degrades at scale under the fixed kernel-width configuration.

I would not read that as “InfoNCE is bad.” I would read it as a reminder that contrastive objectives and Gaussianity-enforcing objectives have different failure modes.

What I would be careful about

This is a theory-forward paper, so the assumptions matter.

The guarantee applies to a broad but specific class of worlds:

• latent variables are Gaussian
• transitions are stationary
• transitions include additive noise
• the representation is constrained to be Gaussian
• identifiability is linear up to rotation

Those assumptions are not weaknesses. They are the contract.

The practical question is how often real-world AI problems can be transformed into a regime where this contract is approximately true. The paper helps by proving approximate identifiability and validating it empirically, but the gap between controlled latent worlds and messy real environments is still where engineering judgment lives.

For applied systems, I would treat this result as a design principle:

If you want a representation that supports planning, do not only ask whether it predicts well. Ask whether the objective can identify the latent degrees of freedom under assumptions you can defend.

Why this paper is worth reading

The value of this paper is not that it declares LeJEPA solved.

The value is that it gives a precise answer to a question that is usually vague:

When can a self-supervised representation be trusted as a world model?

The answer is conditional, but useful:

• Gaussian latent worlds give LeJEPA a linear identifiability guarantee.
• Non-Gaussian latent distributions break the exact guarantee.
• Approximate guarantees degrade gracefully.
• Linear identifiability is enough for a meaningful class of planning problems.
• The code and Lean proofs are available for inspection.

That combination of theory, experiments, code, and formal verification makes the work stand out.

For anyone building or evaluating world models, the paper is a reminder that representation quality is not just about downstream accuracy. It is about whether the learned space preserves the causal and geometric structure needed to plan.