maes-2026-leworldmodel-stable-end-to-end-joint-embedding-predictive

LeWorldModel (LeWM)

Paper: LeWorldModel: Stable End-to-End Joint-Embedding Predictive Architecture from Pixels
Authors: Maes, Le Lidec, Scieur, LeCun, Balestriero (Mila / NYU / Samsung / Brown)
Year: 2026 | arXiv: 2603.19312
Source: maes-2026-leworldmodel-stable-end-to-end-joint-embedding-predictive

---

What Is It?

LeWM is the first JEPA (Joint-Embedding Predictive Architecture) that trains stably end-to-end from raw pixels without any training heuristics — no stop-gradient, no exponential moving averages, no pretrained encoders.

It learns a compact latent world model using only two loss terms:

1. MSE prediction loss — predict the next latent embedding
2. SIGReg — enforce isotropic Gaussian distribution on embeddings to prevent collapse

This gives a single tunable hyperparameter (λ), compared to 6 for the prior SOTA (PLDM).

---

Core Problem Solved

Prior JEPA methods suffer from representation collapse — the encoder maps all inputs to nearly identical representations to trivially satisfy the prediction objective. Existing fixes rely on:

• Heuristic tricks (EMA, stop-gradient) — no theoretical guarantees
• Pretrained frozen encoders (DINO-WM) — limits expressivity, not end-to-end
• Complex multi-term losses (PLDM, 7 terms) — unstable, hard to tune

LeWM addresses this with SIGReg: a principled, provably anti-collapse regularizer.

---

Architecture

Observation o_t
      ↓
  Encoder (ViT-Tiny, ~5M)
      ↓
  z_t  (192-dim CLS token → BatchNorm MLP projection)
      ↓
  Predictor (Transformer, ~10M, AdaLN action conditioning)
      ↓
  ẑ_{t+1}  (predicted next embedding)

Total: 15M parameters, trainable on a single GPU in a few hours.

---

SIGReg (Anti-Collapse Regularizer)

SIGReg projects all embeddings Z ∈ ℝ^{N×B×d} onto M=1024 random unit-norm directions, applies the Epps–Pulley normality test to each 1D projection, and minimizes the aggregate statistic.

By the Cramér–Wold theorem, matching all 1D marginals ≡ matching the full joint Gaussian distribution.

• Provable anti-collapse guarantee
• Scales to high dimensions (random projections)
• M (number of projections) has negligible effect on performance — only λ matters

---

Latent Planning

At inference, LeWM performs Model Predictive Control (MPC) in latent space:

1. Encode current observation and goal: z₁, z_g
2. Predictor rolls out: ẑ_{t+1} = pred(ẑ_t, a_t) for horizon H
3. Minimize terminal cost: C(ẑ_H) = ||ẑ_H − z_g||²
4. Solve with Cross-Entropy Method (CEM)
5. Execute first K actions, then replan (MPC)

Planning time: 0.98s vs. 47s for DINO-WM (48× faster).

---

Key Results

Task	LeWM	PLDM	DINO-WM
Two-Room	87%	100%	97%
Reacher	86%	78%	79%
Push-T	96%	78%	92% (w/ proprioception)
OGBench-Cube	74%	65%	86%

• LeWM best on 2 of 4 tasks; competitive on 3 of 4
• 48× faster planning under same FLOP budget: 96% vs 13% on Push-T

---

Physical Understanding

• Linear probing: latent encodes agent/block position and orientation — competitive with DINOv2-based DINO-WM
• Temporal straightening: latent trajectories become straight over training (emergent, no explicit loss)
• Violation-of-Expectation (VoE): higher surprise for object teleportation (physical violation) vs. unperturbed; p < 0.01 — model reliably detects physically implausible events

---

Limitations

• Short horizons only — error accumulates in autoregressive rollout
• SIGReg mismatch in low-complexity envs — forcing high-dim Gaussian on low intrinsic-dimensionality data hurts (Two-Room)
• Requires action labels
• Fine details lost — 192-dim latent misses fine-grained info (e.g., end-effector angle)

---

Positioning

Method	End-to-End	Pixels	Task-Agnostic	Anti-Collapse	# Hyperparams
PLDM	✓	✓	✓	Weak (VICReg)	6
DINO-WM	✗	✓	✓	Frozen encoder	Low
Dreamer	✓	✓	✗ (reward)	Reconstruction	Several
TD-MPC	✓	✗ (state)	✗	Reward signal	Several
LeWM	✓	✓	✓	Provable (SIGReg)	1

---

Related Pages

• jepa-world-models (to create)