Latent Dynamics · 2026-10-15
Enhancing Latent World Models with Associative Memory
Yuwei Sun, Siyu Zhu
AI-assisted summary based on the paper. Consult the linked source for methods and limitations. Benchmark numbers appear after editor verification.
TL;DR
The authors propose a memory-enhanced JEPA-style latent world model for long-horizon planning. It augments latent dynamics with a content-addressable associative memory bank, retrieves high-level slot-conditioned priors via cross-attention, and conditions prediction using Adaptive Layer Normalization (AdaLN). Two memory losses encourage memory slot diversity and balanced temporal utilization to prevent collapse.
Why it matters
The paper targets representational collapse, computational overhead, and limited long-horizon capability in latent dynamics world models. By adding content-addressable long-term memory and training losses for slot diversity and temporal utilization, the authors aim to improve long-horizon action-sequence prediction while keeping the architecture lightweight and efficient.
Method
- Adds a content-addressable associative memory bank (M learnable low-rank slots) to a JEPA-style latent dynamics model; retrieves memory via cross-attention between observation latents and learned keys, producing a high-level retrieved memory m_t.
- Conditions latent prediction using AdaLN: a sequential conditioning design uses the memory condition first and then action conditioning within a shared predictor to forecast next latent states.
- Trains with a prediction loss (MSE plus SIGReg) and two associative memory losses: slot diversity (SIGReg on memory slots) and temporal diversity (entropy-maximizing objective on attention usage across a horizon).
Limitation
The authors state that training becomes unstable and performance degrades in latent world models as the planning horizon increases. They also note their method requires action-labeled observations for end-to-end training, which incurs a cost in collecting labels.
Abstract (from arXiv)
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.
Related papers
- JEPA-TTT: Persistent Test-Time Training of Latent World Models for Planning under Dynamics Shifts
- Structured Latent Dynamics for Long-Horizon Embodied World Models
- LRC-JEPA: Disentangling Dynamics and Residual Context for Efficient World Models
- H-JEPA: End-to-End Learning of Hierarchical World Models for Visual Planning
- Beyond a single latent space: a dual-latent world model for long-horizon planning
Publication record
Bibliographic evidence describes where and when a record appeared. It does not validate scientific claims.
First public release: 2026-10-14 · Indexed by Pub-AI: 2026-10-15 UTC
Enhancing Latent World Models with Associative Memory
Preprint posted: 2026-10-14 · Original record ↗
- arxiv
- 2610.19683
Explicit source versions (1)
- v1 · 2026-10-14
- Preprint available
Source date: 2026-10-14 · First detected: 2026-10-25 UTC · Historical baseline
Source evidence (1)
- Official evidence ↗ · atom:published (official metadata)
Abstract · arxiv
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.