When Does LeJEPA Learn a World Model?

Source

• Raw Markdown: paper_lejepa-identifiability-2026.md
• PDF: paper_lejepa-identifiability-2026.pdf
• Preprint: arXiv:2605.26379
• Official project page: When Does LeJEPA Learn a World Model?
• Official code: klindtlab/lejepa-identifiability
• Official X thread: David Klindt announcement thread

Status And Credibility

This is a 2026-05-25 arXiv preprint by David Klindt, Yann LeCun, and Randall Balestriero, announced by David Klindt on X on 2026-05-27 with an official project page and code repository. It is credible enough to track as an important theory source because it extends the existing LeJEPA and LeWorldModel line, includes released code and Lean proofs, and is authored by the same research cluster. It is not yet peer reviewed, so benchmark and theory-boundary claims should stay tied to the paper’s stated assumptions.

Core Claim

LeJEPA learns a linearly identifiable state representation when the world has independent Gaussian latent variables, stationary isotropic additive-noise transitions, a successful Gaussian or whitening constraint, and alignment near the population optimum. Under those conditions, the learned representation recovers the true latent variables up to an orthogonal transform, and that linear identifiability is sufficient for rotation-invariant latent-space planning.

Author Narrative Context

The official X thread frames the paper as an answer to “what does JEPA actually learn”: LeJEPA recovers the latent variables of the world, planning in the learned world model gives the same shortest path, Gaussianity is why LeJEPA works, and identifiability is the right definition of learning a world model. The thread also emphasizes that the formal proof stack is checked in Lean 4 with zero sorry obligations, while standard background lemmas such as Mehler-style Hermite facts and AM-GM are axiomatized because the needed Mathlib infrastructure is missing.

The paper supports the narrower state-side version of that narrative. It proves linear identifiability for Gaussian latent worlds and shows that approximate recovery degrades gracefully when alignment and whitening are imperfect. It also shows that a linearly identifiable latent is enough for a planner with orthogonally invariant costs to recover the same optimal plan after rotating the state. It does not yet prove a complete action-conditioned world model: the discussion explicitly says the action-conditioned transition $\hat{p}({\hat{z}}^{\prime }\mid \hat{z},a)$ still has to be learned, and the X thread names adding action conditioning, as in LeWorldModel, as the natural next step.

The thread’s non-Gaussian discussion is also narrower in the paper: David Klindt notes that Laplace-like latents can still work empirically, but sufficiently non-Gaussian latents can break recovery. The paper’s converse theorem and Reacher trajectory experiments make that boundary concrete: outside the Gaussian/OU-style regime, the model may recover a monotonic or distorted transform rather than a linearly usable state.

Key Contributions

• Proves a forward identifiability theorem: with Gaussian latents, OU-style positive pairs, and a Gaussian/whitened embedding constraint, any optimum of the LeJEPA alignment objective recovers the true latent variables up to an orthogonal transform.
• Proves a converse theorem: within the stationary additive-noise class, Gaussian latents are the unique case that guarantees linear identifiability.
• Gives an approximate-identifiability bound where recovery error scales with alignment gap and covariance deviation rather than requiring exact optimization.
• Proves that orthogonal identifiability preserves optimal latent-space planning for rotation-invariant costs.
• Validates the theory with nonlinear 2D mixings, scaling to 1024 latent dimensions, generalized-normal distribution sweeps, and DMC Reacher pixel-observation plus latent-space planning experiments.
• Releases code and Lean 4 proof artifacts for the theory stack.

Method Notes

The setup separates the world, the observation map, and the learned encoder:

flowchart LR
  Z["latent state z"] --> G["unknown nonlinear observation map g"]
  G --> X["observation x = g(z)"]
  X --> F["encoder f"]
  F --> H["learned representation h(z) = f(g(z))"]
  Z -. "OU positive pair" .-> Z2["latent state z'"]
  Z2 --> G2["same observation map g"]
  G2 --> X2["observation x'"]
  X2 --> F2["same encoder f"]
  F2 --> H2["h(z')"]
  H <-- "alignment + Gaussian/whitening constraint" --> H2

The central objective is alignment under a Gaussianity or whitening constraint:

$$\underset{h}{\min }\mathbb{E}\left[\Vert h(z^{\prime })-h(z){\Vert }^2\right]\text{s.t.}h(z)\sim \mathcal{N}(0,I_n).$$

For the Gaussian world, $z^{\prime }=\rho z+\sqrt{1-{\rho }^2}\eta$ with $\eta \sim \mathcal{N}(0,I_n)$. Hermite polynomials diagonalize this transition: degree-$d$ nonlinear components are attenuated by ${\rho }^d$, so the linear part has the highest cross-view correlation when $0<\rho <1$. That is why any nonlinear component loses alignment relative to a linear orthogonal map.

The converse uses a Sturm-Liouville view: for general stationary additive-noise latents, the slowest non-constant eigenfunction is monotonic, but it is affine only when the latent density has a linear score, which characterizes a Gaussian. This is the paper’s cleanest bridge between JEPA, nonlinear ICA, and slow-feature analysis.

Evidence And Results

• In 2D nonlinear mixing experiments, LeJEPA recovers the latent coordinates up to rotation across spiral, sinusoidal, parabolic, and RealNVP-style observation maps.
• In the high-dimensional RealNVP sweep, SIGReg and VICReg maintain $R^2(h\to z)>0.999$ up to $N=1024$ latent dimensions, while fixed-kernel InfoNCE degrades at larger dimensions.
• In generalized-normal distribution sweeps, linear recovery peaks at the Gaussian case and degrades away from it, matching the Gaussian-uniqueness theorem.
• In DMC Reacher, OU-sampled Gaussian joint-angle latents reach much higher and more symmetric recovery than policy-trajectory latents; the paper reports that trajectory total $R^2$ never exceeds 0.5 because the policy-induced marginals, autocorrelations, and joint-limit wrapping violate the theory assumptions.
• In latent-space planning, the Gaussian/OU encoder’s straight latent interpolations decode close to oracle joint-space straight paths, while the trajectory encoder produces curved paths and higher control cost.
• The Lean appendix reports zero sorry obligations for the machine-checked reasoning chains, while standard mathematical results not yet available in Mathlib are included as axioms.

Limitations

• The source is a recent arXiv preprint, not a peer-reviewed paper.
• The theorem is a population-level global-optimum statement; it does not prove finite-sample or training-dynamics guarantees.
• The encoder output dimension is assumed to match the true latent dimension. Undercomplete or overcomplete representations may select a subspace, use superposition, collapse dimensions, or encode redundancy.
• Real-world latent Gaussianity is not directly observable. The paper argues Gaussianity is plausible as a maximum-entropy or aggregated-latent prior, but it remains an assumption.
• The planning theorem is about the learned state coordinates and rotation-invariant costs. It does not by itself learn or validate an action-conditioned transition model.
• The Reacher trajectory experiments are useful because they violate the assumptions, but they also mix several violations at once: non-Gaussian marginals, anisotropic autocorrelation, and joint-limit wrapping.

Foundation TSFM Relevance

Agenda slot	Verdict	Evidence	Missing pieces
Anti-collapse regularization	partially closes outside time series	Gives a sharper theory for why Gaussian/whitened LeJEPA representations can be identifiable rather than merely non-collapsed.	Needs rare-regime, long-tail, and numeric time-series stress tests.
Latent-state prediction	adjacent	Defines a condition under which representation-space prediction recovers true latent state up to rotation.	No numeric time-series benchmark and no streaming state-maintenance evidence.
Control and counterfactuals	adjacent	Shows that linear identifiability can preserve a rotation-invariant planner in learned coordinates.	The action-conditioned transition $\hat{p}({\hat{z}}^{\prime }\mid \hat{z},a)$ is explicitly future work.
Data diversity and long tail	warning	Non-Gaussian trajectory latents degrade recovery; sampling distribution matters for whether the theory applies.	Need exploration/data-collection policies that preserve intervention-relevant and rare states.

Links Into The Wiki

• LeJEPA
• JEPA
• World Models
• Representation Collapse
• Distribution Priors In Self-Supervised Learning
• Latent-Space Predictive Learning
• Latent-State Time-Series Modeling
• Self-Supervised Representation Learning
• Foundation Time-Series Model Research Agenda
• Dynamic Curriculum Learning For JEPA

Open Questions

• Which real-world state variables are well modeled as Gaussian aggregate latents, and which are intrinsically non-Gaussian, bounded, discrete, periodic, or graph-structured?
• Can the identifiability theorem be extended to action-conditioned world models where actions or interventions change the latent transition law?
• How should exploration or dataset curation keep SSL pretraining near the isotropic, informative-transition regime without erasing rare but important states?
• What happens when the learned representation is intentionally overcomplete or undercomplete relative to the true latent state?