Sub-JEPA: a simple fix to LeCun group's LeWorldModel that consistently improves performance [P]

· r/MachineLearning ·

Sub-JEPA applies Gaussian regularization inside frozen random orthogonal subspaces instead of the full latent space, fixing a geometry mismatch in LeCun group's LeWorldModel and consistently improving performance across tasks.

Categories: Research

Excerpt

**World models** learn compact latent representations for planning without pixel reconstruction. LeWorldModel (LeWM), from LeCun's group at NYU, achieves stable end-to-end JEPA training by enforcing an isotropic Gaussian prior over the full latent space. **The flaw:** real environment dynamics live on low-dimensional manifolds, so a global high-dimensional Gaussian is an overly rigid prior — mismatched to the task geometry. LeWM itself struggles most on low-intrinsic-dimension tasks like Two-Room. **Our fix (Sub-JEPA):** apply the Gaussian regularization inside multiple frozen random orthogonal subspaces instead. This relaxes the global constraint while keeping the anti-collapse benefit. No new hyperparameters, same two-term objective. Sub-JEPA consistently outperforms LeWM across all four benchmarks, with up to +10.7 pp on Two-Room. We also observe straighter latent trajectories and better physical state decodability as emergent benefits. ![](https://kaizhao.net/images/projects/sub-jepa/overview.png) ![](https://kaizhao.net/images/projects/sub-jepa/cube.gif) 🌐 Project: [https://kaizhao.net/sub-jepa](https://kaizhao.net/sub-jepa) 💻 Code: [https://github.com/intcomp/sub-jepa](https://github.com/intcomp/sub-jepa) 📄 Paper: [https://arxiv.org/pdf/2605.09241](https://arxiv.org/pdf/2605.09241)

Discussions