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

> Paper: LeWorldModel: Stable End-to-End Joint-Embedding Predictive Architecture from Pixels (arXiv:2603.19312, March 2026) > Authors: Lucas Maes*, Quentin Le Lidec*, Damien Scieur, Yann LeCun, Randall Balestriero > Institutions: Mila & Universite de Montreal, NYU, Samsung SAIL, Brown University > Code: github.com/lucas-maes/le-wm | Project page: https://le-wm.github.io

While generative world models like Sora demand massive GPU clusters and DreamerV3 requires careful tuning, LeWM demonstrates that a 15M-parameter model with just two loss terms, trained end-to-end from pixels on a single GPU, can match foundation-model-based world models in planning performance while being 48x faster (full planning completes in under 1 second).

Key points

The problem: representation collapse in JEPA

• JEPA (Joint-Embedding Predictive Architecture) encodes images into low-dim vectors and predicts the next embedding given an action — no pixel reconstruction, no diffusion, no VAE.
• Its fatal flaw: an encoder can collapse all inputs to the same vector, making the prediction loss trivially zero (the model learns nothing).
• Prior fixes are inelegant: PLDM uses 7 loss terms with 6 hyperparameters; DINO-WM freezes a DINOv2 encoder (giving up end-to-end learning); I-JEPA/V-JEPA rely on EMA + stop-gradient with unclear theoretical grounding.

The solution: SIGReg regularization

• Idea: if embeddings collapse to a point, their distribution is degenerate. Force the embedding distribution toward a standard isotropic Gaussian, and collapse becomes impossible.
• Method: by the Cramer-Wold theorem, a multivariate distribution is determined by its 1D projections. SIGReg projects embeddings onto K=1024 random unit directions, then applies the Epps-Pulley goodness-of-fit test (based on the empirical characteristic function, i.e., essentially Fourier analysis) to measure deviation from a standard normal on each projection, averaging over directions.
• Chosen over KL divergence because density estimation in high dimensions is intractable, while characteristic functions are cheap in any dimension.
• The implementation is ~36 lines of PyTorch; random projection directions are resampled each forward pass (a form of stochastic regularization).

Architecture (extremely minimal)

• Loss: L_LeWM = L_pred + lambda * SIGReg(Z). That's it — two terms, one hyperparameter (lambda = 0.09, found via binary search across 6 environments in just 20 trials).
• Encoder (~5M params): ViT-Tiny, 12 layers, 3 heads, hidden dim 192, 14x14 patches, CLS token + 1-layer MLP projection with BatchNorm (not LayerNorm — LayerNorm would normalize away the sample-wise variance SIGReg needs to shape).
• Predictor (~10M params): 6-layer Transformer, 16 heads, AdaLN-zero action injection, causal masking, 10% dropout.
• Each frame becomes a single 192-dim token vs ~200 tokens for DINO-WM — the main reason planning is 48x faster.

Results

• Planning: best on PushT (beating even DINO-WM with proprioception) and Reacher; second-best on OGBench-Cube (behind DINO-WM, whose DINOv2 pretraining provides richer visual priors) and Two-Room (where the low intrinsic dimensionality makes Gaussian matching in high dimensions slightly harmful).
• Speed: <1 second vs ~47 seconds for DINO-WM full planning.
• Physics understanding: linear/nonlinear probes recover physical quantities (position, velocity, angle) better than PLDM; Violation-of-Expectation tests show strong "surprise" signals for physically impossible events (teleportation) but not mere color changes (paired t-test, p < 0.001).
• Emergent temporal path straightening: latent trajectories become increasingly straight during training with no explicit loss encouraging it — the authors attribute this to SIGReg being applied only within time steps, leaving the temporal dimension free to "collapse" into linear paths, which benefits downstream performance.
• Ablations: nearly insensitive to SIGReg internals (K=256 vs 1024, 17 vs 33 integration nodes); encoder-agnostic (ResNet-18 performs similarly); smooth, monotonic loss curves versus noisy PLDM training.

Reproducibility

Repository at github.com/lucas-maes/le-wm includes jepa.py (model + planning), module.py (SIGReg, Transformer, AdaLN-zero), train.py/eval.py (Lightning + Hydra), and configs. Training runs in hours on one GPU (bf16, batch size 128, AdamW).

Limitations and outlook

• Underperforms on simple low-dimensional environments; still requires action labels from offline data; planning limited to short horizons.
• Promising future directions: hierarchical world models, large-scale video pretraining, and inverse dynamics models to remove action-label dependence.

Takeaway

LeWM argues that when the mathematics is right, engineering becomes simple: two loss terms, one hyperparameter, 15M parameters, and a single GPU produce a world model competitive with systems built on 1.24B-image pretrained encoders — evidence that simplicity, backed by solid statistical theory (Cramer-Wold + Epps-Pulley), can outcompete scale and patchwork complexity.