worldmodel/JEPA/lejepa-identifiability/README.md

LeJEPA Identifiability

When Does LeJEPA Learn a World Model?

David Klindt, Yann LeCun and Randall Balestriero

Abstract: A representation that scrambles the true degrees of freedom of the world cannot support reliable planning or compositional generalization. We prove that LeJEPA (alignment plus Gaussian regularization) linearly recovers the world's latent variables from nonlinear observations, a property known as linear identifiability, in a broad class of worlds where latents evolve under stationary, additive-noise transitions. Our main result is that among all such worlds, the Gaussian is the unique latent distribution for which this guarantee holds. The forward direction rests on a spectral decomposition in which each degree of nonlinearity is strictly penalized by alignment, making the linear map the optimum; the converse rules out every non-Gaussian alternative. We further prove an approximate identifiability result where the guarantee degrades gracefully, and show that linear, orthogonal identifiability enables optimal latent-space planning. We validate the theory across 2D examples to 1024-dimensional latents, distributional ablations, and pixel-based robotic control. All theorems are formally verified in Lean 4.

[ Paper | Website | Colab | Video ]

If you find this work useful, please cite:

@article{klindt2026lejepa,
  title={When Does LeJEPA Learn a World Model?},
  author={Klindt, David and LeCun, Yann and Balestriero, Randall},
  journal={arXiv preprint arXiv:TODO},
  year={2026}
}

Quick Start

Try the 2D demo in your browser (~30s on a T4 GPU):

Repository Structure

lejepa-identifiability/
├── lean/                           # Lean 4 formal verification
│   ├── LeJEPA/
│   │   ├── Hermite.lean            # Forward direction (Hermite polynomial proof)
│   │   ├── Uniqueness.lean         # Converse (Gaussian uniqueness)
│   │   ├── Approx.lean             # Approximate identifiability bound
│   │   └── Dirichlet.lean          # Alternative proof (Dirichlet energy)
│   ├── LeJEPA.lean
│   ├── lakefile.lean
│   └── lean-toolchain              # Lean 4 v4.28.0
├── experiments/
│   ├── lejepa_id/                  # Shared library
│   │   ├── mixing.py               # Mixing functions (spiral, banana, sinusoid, coupling)
│   │   ├── models.py               # MLP and matched (inverse-NVP) encoders
│   │   ├── losses.py               # SIGReg, whitening, alignment, InfoNCE
│   │   ├── metrics.py              # R², orthogonality, bound quantities
│   │   ├── data.py                 # Gaussian / generalized-normal sampling, OU augmentation
│   │   ├── reacher.py              # Reacher pixel data utilities
│   │   └── engine.py               # Training loop (warmup + cosine LR, online data)
│   ├── run.py                      # Unified runner for 2D / scaling / gennorm / grid
│   ├── run_reacher.py              # Reacher pixel-observation runner
│   ├── prerender.py                # Render Reacher OU and trajectory frames
│   ├── analysis/                   # Post-hoc plotting and tables
│   ├── configs/                    # Experiment hyperparameters (YAML)
│   │   ├── 2d.yaml
│   │   ├── gennorm.yaml
│   │   ├── scaling.yaml
│   │   ├── grid.yaml
│   │   └── reacher.yaml
│   └── slurm/                      # SLURM launch scripts (CSHL cluster)
├── requirements.txt
└── README.md

Formal Verification (Lean 4)

All theoretical results are formalized in Lean 4 with Mathlib. The project compiles with zero sorry obligations — every logical chain from axiomatized premises to conclusions is machine-checked. Axiomatized components are standard results not yet available in Mathlib (Hermite polynomial infrastructure, Mazur–Ulam, AM–GM with uniform weights). See the paper appendix for the full verification inventory.

cd lean
lake build    # requires Lean 4 v4.28.0; fetches Mathlib automatically

Experiments

All experiments share the same training infrastructure (lejepa_id/engine.py) and read parameters from YAML configs. Training uses online data generation, a warmup + cosine LR schedule, and saves results as .json (scalars and curves); 2D and ablation runs additionally save .pt files with scatter arrays.

pip install -r requirements.txt
cd experiments

2D Illustrations

Four mixing functions (spiral, banana, sinusoidal shear, NVP) with MLP or matched encoders.

python run.py --config configs/2d.yaml --run spiral_lejepa --seed 1337
python analysis/plot_2d.py --results_dir results/2d/ --out figures/

Scaling (N = 2 to 1024)

Matched (inverse-NVP) encoder scaling with latent dimension, swept across SIGReg / VICReg / InfoNCE objectives. Each (N, seed) trains K=3 encoders in parallel for N ≤ 32 and picks the best by final loss.

python run.py --config configs/scaling.yaml --N 16 --seed 0
python run.py --config configs/scaling.yaml --N 16 --seed 0 --mode infonce
python analysis/plot_scaling.py --results_dir results/scaling/ --out figures/

Distributional Ablation (Generalized Normal)

Same mixings sweeping the latent shape parameter α (heavy-tailed → Laplace → Gaussian → uniform). Demonstrates that linear identifiability fails away from the Gaussian (α = 2).

python run.py --config configs/gennorm.yaml --run spiral_lejepa --alpha 2.0 --seed 1337
python analysis/plot_gennorm.py --results_dir results/gennorm/ --out figures/

Grid Search / Bound Verification

Sweep over regularization weight λ and OU correlation ρ on the 2D spiral mixing.

python run.py --config configs/grid.yaml --lamb 0.01 --rho 0.9 --seed 0
python analysis/plot_bound.py --results_dirs results/grid results/2d results/scaling --out figures/

Reacher (Pixel-Based RL)

CNN encoder on rendered DMC Reacher frames, comparing OU pairs against trajectory pairs from a learned policy.

python prerender.py ou --rho 0.95
python prerender.py traj --delta 16 --h5_path data/reacher.h5
python run_reacher.py --config configs/reacher.yaml --data_dir data/reacher/ou/rho=0.95

Cross-Experiment Analysis

python analysis/aggregate.py --results_dir results/ --recursive --out results/all.csv
python analysis/plot_scatter.py --results_dirs results/2d results/gennorm results/scaling results/grid --out figures/

Regenerate All Figures

bash analysis/run_all.sh

Requirements

• Lean: v4.28.0 + Mathlib v4.28.0 (managed by lake)
• Python: pip install -r requirements.txt

License

MIT