Research map / Architectures and efficiency / RNNs and normalization

LeWorldModel: Stable End-to-End Joint-Embedding Predictive Architecture from Pixels

NeurIPS 2026 · Architectures and efficiency › RNNs and normalization · approach: Differentiable optimization layers · 0 citations (OpenAlex)

Read the paperCheck your idea against 108,967 papers

AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.

Proposes LeWorldModel, an end-to-end JEPA trained from pixels with a next-embedding prediction loss plus a SIGReg Gaussian regularizer, enabling stable latent world models with ~15M parameters that plan much faster than DINO-WM.

Problem

Existing JEPA world models collapse easily and either rely on heuristics (EMA, stop-gradient), multi-term losses, or frozen pretrained encoders, making end-to-end pixel-based latent world modeling fragile and costly.

Approach

Train encoder and predictor jointly from raw pixels with only two loss terms: MSE next-embedding prediction plus SIGReg, which projects embeddings onto random directions and applies an Epps-Pulley normality test to push latents toward an isotropic Gaussian, preventing collapse without stop-gradient or EMA.

Key results

• LeWM plans up to 48x faster than DINO-WM while remaining competitive on control tasks, with full planning under one second.
• LeWM achieves 18% higher success rate than PLDM on PushT and outperforms DINO-WM on PushT even without proprioception.
• With ~15M parameters, LeWM reduces tunable loss hyperparameters from six (PLDM) to one (lambda), trainable on a single GPU in a few hours.

Surprising findings

• LeWM underperforms PLDM and DINO-WM on the simplest environment, TwoRoom, which the authors attribute to low intrinsic dimensionality conflicting with the Gaussian prior.
• Planning performance is the weaker side on OGBench-Cube, where DINO-WM slightly outperforms LeWM.

Limitations

• Planning restricted to short horizons.
• SIGReg is weaker in low-diversity, low-intrinsic-dimensionality environments such as TwoRoom.
• Relies on offline datasets with action labels and sufficient coverage.

Experimental setup

Datasets: TwoRoom PushT OGBench-Cube DeepMind Control Reacher

Benchmarks: PushT OGBench-Cube TwoRoom Reacher

Metrics: success rate planning time probing MSE probing Pearson r surprise (paired t-test)

Baselines: DINO-WM PLDM GCBC GCIVL GCIQL

Builds on: JEPA DINO-WM PLDM VICReg Cramér-Wold theorem

Closest papers by meaning

• Drive vs. Decay: On the Training Dynamics of Joint-Embedding Predictive Architectures (NeurIPS 2026 · similarity 0.89)
• seq-JEPA: Autoregressive Predictive Learning of Invariant-Equivariant World Models (NeurIPS 2025 · similarity 0.87)
• Denoising with a Joint-Embedding Predictive Architecture (ICLR 2025 · similarity 0.86)
• How JEPA Avoids Noisy Features: The Implicit Bias of Deep Linear Self Distillation Networks (NeurIPS 2024 · similarity 0.85)
• Connecting Joint-Embedding Predictive Architecture with Contrastive Self-supervised Learning (NeurIPS 2024 · similarity 0.85)
• M3-JEPA: Multimodal Alignment via Multi-gate MoE based on the Joint-Embedding Predictive Architecture (ICML 2025 · similarity 0.83)