Research map / Probabilistic and scientific modeling / Causal discovery
When does LeJEPA learn a World Model?
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Proves that LeJEPA (alignment plus SIGReg Gaussian regularization) recovers latent world variables linearly under additive-noise stationary worlds, and that the Gaussian is the unique latent distribution giving this guarantee.
Problem
When does a JEPA-style representation learn a faithful World Model, i.e., linearly recover the true latent variables from nonlinear observations? Prior JEPA methods lacked identifiability guarantees because embedding distributions were unspecified.
Approach
Under Gaussian latents with stationary additive-noise (Ornstein-Uhlenbeck) transitions, alignment plus whitening is maximized by linear orthogonal maps via a Hermite spectral decomposition in which each nonlinear degree d attenuates correlation by rho^d; a Sturm-Liouville argument shows Gaussianity is the unique latent distribution yielding linear identifiability.
Key results
- Linear identifiability with SIGReg, VICReg reaching R^2 > 0.999 up to latent dimension N=1024 under RealNVP mixing, while InfoNCE degrades at scale (e.g., 0.72 at N=1024).
- On Reacher, OU Gaussian pairs reach R^2 = 0.95 at rho=0.99, while RL trajectory pairs never exceed R^2 of about 0.5.
- Theory: Gaussian is the unique latent distribution for which linear identifiability holds in the additive-noise stationary class (converse theorem).
Surprising findings
- Real RL-policy trajectories (non-Gaussian, anisotropic) break linear identifiability, so goal-directed data can fail where isotropic random-walk data succeeds.
- InfoNCE, a pair-based estimator, degrades at high latent dimension under fixed kernel width while moment-based SIGReg and VICReg do not.
Limitations
- Assumes encoder output dimension equals latent dimension; mismatched m != n is not resolved.
- Population-level result about global optimum; finite-sample and optimization dynamics not addressed.
- Assumes Gaussian latents, which are unverifiable from observations; real latents may be non-Gaussian.
Experimental setup
Datasets: synthetic 2D Gaussian latents with nonlinear mixings DMC Reacher trajectories (10k RL episodes)
Benchmarks: DMC Reacher pixel-based control
Metrics: R^2 linear identifiability control cost alignment loss
Baselines: VICReg InfoNCE
Builds on: LeJEPA VICReg InfoNCE Mehler's formula Sturm-Liouville theory
Closest papers by meaning
- Learning Linear Causal Representations from General Environments: Identifiability and Intrinsic Ambiguity (NeurIPS 2024 · similarity 0.85)
- Beyond ICA: Identifiability by Symmetry Breaking (NeurIPS 2026 · similarity 0.85)
- Identifying latent state transitions in non-linear dynamical systems (ICLR 2025 · similarity 0.85)
- LoRA Training Provably Converges to a Low-Rank Global Minimum Or It Fails Loudly (But it Probably Won't Fail) (ICML 2025 · similarity 0.84)
- Nonparametric Identification of Latent Concepts (ICML 2025 · similarity 0.84)
- When do neural networks learn world models? (ICML 2025 · similarity 0.84)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.