Enhancing Latent World Models with Associative Memory
cs.LG, cs.CV
Submitted: 2026-10-14
Updated: 2026-10-14
License: http://creativecommons.org/licenses/by/4.0/
The gist: World models trained via latent dynamics offer a scalable path for embodied planning, yet existing methods suffer from representational collapse, computational overhead, and limited long-horizon
Terminology
Abstract
World models trained via latent dynamics offer a scalable path for embodied planning, yet existing methods suffer from representational collapse, computational overhead, and limited long-horizon capability. We introduce a memory-enhanced latent world model that augments JEPA-style prediction with a content-addressable memory bank. Our approach learns high-level representations directly from observations, operating at multiple temporal scales within a shared latent space, and conditions predictions via adaptive layer normalization. To prevent collapse and ensure effective memory usage, we design two complementary losses that encourage slot diversity and balanced temporal utilization. Experiments on PushT, Tworoom, Reacher, and OGBCube demonstrate that our method achieves state-of-the-art performance with a lightweight architecture, outperforming conventional latent world models on long-horizon planning tasks.
Sources
- ThinkJEPA: Empowering Latent World Models with Large Vision-Language Reasoning Model
- VL-JEPA: Joint Embedding Predictive Architecture for Vision-language
- LeWorldModel: Stable End-to-End Joint-Embedding Predictive Architecture from Pixels
- DINO-WM: World Models on Pre-trained Visual Features enable Zero-shot Planning
- Hierarchical Planning with Latent World Models
- Mastering Diverse Domains through World Models
- stable-worldmodel-v1: Reproducible World Modeling Research and Evaluation
Related papers
- Polynomial-Augmented Neural Networks (PANNs) with Weak Orthogonality Constraints for Enhanced Function and PDE Approximation
- AIRL-S: Unifying Reinforcement Learning and Search-Based Test-Time Scaling via Adversarial Inverse Reinforcement Learning
- Transformers as Bayesian In-Context Experimenters: Smoothness-Adaptive Efficient ATE Estimation
- Convergence issues in Relational Concept Analysis based on AOC-posets
- Beliefs Beyond Posteriors: The Predictive Cost of ELBO Binding
- Understanding Diffusion Models via Ratio-Based Function Approximation with SignReLU Networks