When Does LeJEPA Learn a World Model?

arxiv:2605.26379

Original source · Paper JSON

Self-supervised learning and vision · 2026-05-25

Cites (5) · Cited by (0)

Complete cleaned Markdown · Source Markdown including bibliography

Read in parts

1. Part 1
2. Part 2
3. Part 3
4. Part 4
5. Part 5
6. Part 6
7. Part 7
8. Part 8
9. Part 9

Contents

• When Does LeJEPA Learn a World Model?
• Abstract
• 1 Introduction
• Contributions.
• 2 Related Work
• Representation Learning.
• World Models.
• Identifiability.
• 3 The World and the Learner
• 3.1 The World
• 3.1.1 The Gaussian World
• 3.2 The Learner: LeJEPA
• 4 Spectral Analysis of the World
• Transition Operator.
• Gaussian World: Hermite Polynomials.
• General Case: Sturm-Liouville Theory.
• 5 Theory
• 5.1 Forward Direction: LeJEPA Learns the World Model
• Proof Sketch.
• Interpretation.
• 5.2 Converse Direction: The Gaussian is Unique
• Proof Sketch.
• 5.3 Approximate Identifiability
• Interpretation.
• 5.4 Optimal Latent Planning
• Interpretation.
• 6 Experiments
• 6.1 Forward: Linear Identifiability from Gaussian Latent Variables
• Scaling to High Dimensions.
• 6.2 Converse: Non-Gaussian Latent Variables Break Linear Identifiability
• Latent-Distribution Sweep.
• Pixel-Based RL Trajectories.
• 6.3 Approximation Bound and Loss Predict Identifiability
• 6.4 Linear Identifiability Enables Latent-Space Planning
• 7 Limitations
• Are the latents Gaussian?
• What if the dimension is wrong?
• Finite samples and optimization.
• 8 Discussion
• Outlook.
• Practical Implications.
• A Theoretical Foundation for World Models.
• Acknowledgments
• Appendix Overview
• Appendix A Proof of Theorem 5.1 (Forward Direction)
• A.1 Intuition: Why Should This Force Linearity?
• A.2 Hermite Polynomial Background
• Definitions and orthogonality.
• Multivariate extension.
• Hermite expansion of measure-preserving maps.
• A.3 Proof
• Setup and notation.
• Step 1: The transition attenuates each Hermite degree by ρ d \rho^{d} .
• Step 2: Different Hermite components do not interact across views.
• Step 3: Nonlinearity strictly reduces correlation.
• Step 4: Assembly.
• Step 5: Recovery of the transition.
• Appendix B Proof of Theorem 5.2 (Gaussian Uniqueness)
• B.1 Background
• B.2 Proof
• Appendix C Proof of Theorem 5.3 (Approximate Identifiability)
• Setup and notation.
• C.1 Step 1: Alignment gap controls nonlinear energy
• C.2 Step 2: Whitening constraint forces M M near orthogonal
• C.3 Step 3: Pythagorean decomposition
• Role of the Gaussian on the latent side vs. the output side.
• Appendix D Proof and Discussion of Theorem 5.4
• Proof.
• D.1 Worked Applications
• Goal-reaching.
• Linear-quadratic regulation.
• General rotation-invariant objectives.
• D.2 Scope: encoder versus action-conditioned dynamics
• D.3 Connection to classical optimal control
• Appendix E Alternative Proof via Dirichlet Energy
• E.1 Setup
• E.2 Proof
• Step 1: Reduction to Dirichlet energy.
• Step 2: The expected log-determinant is zero.
• Step 3: Lower bound via AM-GM and Jensen.
• Step 4: Equality forces orthogonal Jacobian.
• Step 5: Orthogonal Jacobian implies linearity (Mazur–Ulam).
• Generalization.
• E.3 Comparison of the Two Proofs
• Appendix F Prior Work: Connection to Slow Feature Analysis
• SFA and its optimization problem.
• Nonlinear identifiability via SFA.
• Eigenfunction structure and Sturm–Liouville theory.
• The Gaussian specialization and Hermite polynomials.
• Sequential versus simultaneous extraction.
• The role of isotropic transitions.
• Beyond Gaussian latents.
• Relation to Mehler’s formula and maximal correlation.
• Practical differences: learned features versus fixed kernels.
• Connections to diffusion maps and spectral geometry.
• Summary of the relationship.
• Appendix G Lean Verification
• Scope and methodology.
• G.1 What Is Axiomatized and Why
• 1. Hermite polynomial infrastructure.
• 2. Sturm–Liouville spectral theory.
• 3. Standard analysis results requiring API software engineering.
• 4. Measure-theoretic and matrix-analytic facts.
• 5. Trajectory pushforward for expected costs.
• What is verified.
• G.2 Build Information
• Appendix H Experimental Details and Additional Results
• H.1 Nonlinear Mixing Functions
• H.2 Network Architecture and Training
• H.3 Training Procedure
• H.4 Hyperparameters
• H.5 Evaluation Metrics
• H.6 Grid Search
• H.7 Latent Variable Distribution Sweep Across Mixings
• H.8 Approximate Bound Decomposition
• H.9 Loss Predictivity
• H.10 Scaling Experiment
• H.11 Reacher Experiment
• Environment and rendering.
• Encoder architecture.
• Training.
• Evaluation.
• Latent distributions.
• Gaussianity and autocorrelation jointly determine identifiability.
• Additional results on OU and Trajectory conditions.
• H.12 Planning Experiment
• Setup.
• Metric.
• Interpretation.
• Instructions for reporting errors