LeWorldModel: A 15M-Parameter World Model That Trains on One GPU

Forum topic · 小凯 · 2026-04-24

Summary

LeWorldModel (LeWM), introduced by LeCun's team from Mila, Universite de Montreal, NYU, Samsung SAIL, and Brown University, demonstrates that a minimal world model can match or surpass far larger baselines. The system uses only 15M parameters, trains end-to-end from pixels on a single GPU, and employs a two-term loss function. Its core innovation, SIGReg (Sliced Isotropic Gaussian Regularizer), replaces the complex multi-loss anti-collapse heuristics used in prior JEPA models. Grounded in the Cramer-Wold theorem and Epps-Pulley test, SIGReg forces embeddings to match an isotropic Gaussian via 1024 random projections, implemented in just 36 lines of code. LeWM matches or exceeds DINO-WM on PushT, Reacher, and OGBench-Cube while planning 48x faster (under one second). Probing experiments reveal that the latent space encodes physical quantities and exhibits natural temporal path straightening. The release provides open-source code with a single tunable hyperparameter (lambda=0.09), highlighting that principled statistical foundations can replace engineering complexity in self-supervised world modeling.

Overview

LeWorldModel (LeWM) is a JEPA-based world model introduced by researchers from Mila, Universite de Montreal, NYU, Samsung SAIL, and Brown University (arXiv:2603.19312). With just 15M parameters, training on a single GPU from pixels end-to-end, it matches or exceeds DINO-WM on planning tasks while running 48x faster. The result is achieved using a two-term loss function whose only hyperparameter is lambda.

Key points

The problem: representation collapse in JEPA

JEPA predicts embeddings rather than pixels, avoiding diffusion-style generation, but its MSE prediction loss is vulnerable to collapse: a constant embedding yields zero prediction error while encoding nothing.
Prior fixes (PLDM's 7-loss VICReg variant, frozen DINOv2 backbones in DINO-WM, EMA stop-gradient in I-JEPA/V-JEPA) trade simplicity for either instability, dependence on 124M-image pretraining, or weak theoretical grounding.

The solution: SIGReg

SIGReg (Sliced Isotropic Gaussian Regularizer) forces the encoder's output distribution to approximate a standard isotropic Gaussian. Collapse becomes impossible if the embedding distribution must span a Gaussian.
Leverages the Cramer-Wold theorem: a multivariate distribution is fully characterized by its 1D projections along every direction. Empirically, K=1024 random unit projections are sufficient.
For each projection, distance to N(0,1) is measured via the Epps-Pulley statistic, computed from the empirical characteristic function (Fourier-domain comparison, approximated by 17-node trapezoidal integration).
Implementation is 36 lines of Python: sample directions, project, apply cos/sin (Euler's formula), integrate against a Gaussian-windowed weight vector, average.
Internally validated: results are insensitive to choices of projection count (256 vs 1024) or integration knots (17 vs 33).

Architecture

Loss: L_LeWM = L_pred (MSE) + lambda * SIGReg(Z) with lambda = 0.09, recoverable via binary search in roughly 20 trials across 6 environments.
Encoder (~5M): ViT-Tiny, 12 layers, 3 heads, hidden dim 192, patch 14x14, CLS token + 1-layer MLP projection.
Predictor (~10M): 6-layer Transformer, 16 heads, AdaLN-zero action conditioning, causal mask, 10% dropout.
Critical detail: encoder uses BatchNorm, not LayerNorm, because LayerNorm would rescale per-sample embeddings and prevent SIGReg from matching a target distribution.

Planning

Encode start and goal frames to z_start, z_goal (single 192-dim token each).
Cross-Entropy Method (CEM) searches action sequences in latent space, scoring by MSE to z_goal.
A single embedding token per frame (vs ~200 tokens for DINOv2) accounts for much of the 48x speedup; full planning completes in under one second.

Experimental results

Planning: LeWM is best on PushT and Reacher, slightly trails DINO-WM on OGBench-Cube (3D), and is weakest on the simplest Two-Room (low intrinsic dimensionality).
LeWM surpasses DINO-WM on PushT even when DINO-WM uses proprioception.
Physical understanding: probing recovers object position, velocity, and angle; a violation-of-expectation test shows significantly higher model "surprise" for physically impossible jumps vs pure color changes (paired t-test, p < 0.001).
Temporal path straightening emerges without any explicit loss: average cosine similarity between consecutive velocity vectors z_{t+1} - z_t rises during training, suggesting time-axis collapse that benefits downstream tasks.
Replacing ViT-Tiny with ResNet-18 yields nearly identical results.
Both LeWM losses decrease smoothly and monotonically; PLDM's seven losses exhibit noisy, non-monotonic behavior.

Resources

Paper: arXiv:2603.19312 (March 2026).
Code: https://github.com/lucas-maes/le-wm
Project page: https://le-wm.github.io
Core files: jepa.py (model), module.py (SIGReg, Transformer, AdaLN-zero), train.py, eval.py, config/train/lewm.yaml.
Training uses bf16 precision, batch size 128, AdamW; completes in a few hours on a single GPU.

Why it matters

SIGReg generalizes: any self-supervised method threatened by collapse (contrastive learning, masked modeling) could benefit from Cramer-Wold plus random projection matching.
End-to-end training from pixels can outperform frozen-pretrained encoders when the training objective is well designed.
Limitations remain: poor performance on low-dimensional environments, dependence on offline action labels, and short-horizon planning. The authors point to hierarchical models, large-scale video pretraining, and inverse dynamics as future directions.

One-line summary

> When the math is right, engineering becomes simple: 15M parameters, two loss terms, one GPU, no pretraining required.

---

Paper | Code | Project