Composition of Memory Experts for Diffusion World Models

• URL: http://arxiv.org/abs/2605.18813v1
• Date: Tue, 12 May 2026 09:43:10 GMT
• Title: Composition of Memory Experts for Diffusion World Models
• Abstract summary: World models aim to predict plausible futures consistent with past observations, a capability central to planning and decision-making in reinforcement learning.<n>Existing architectures face a fundamental memory trade-off: transformers preserve local detail but are bottlenecked by quadratic attention.<n>We suggest decoupling future-past consistency from any single architecture and instead leveraging a set of specialized experts.
• Score: 23.602187351200183
• License: http://creativecommons.org/licenses/by/4.0/
• Abstract: World models aim to predict plausible futures consistent with past observations, a capability central to planning and decision-making in reinforcement learning. Yet, existing architectures face a fundamental memory trade-off: transformers preserve local detail but are bottlenecked by quadratic attention, while recurrent and state-space models scale more efficiently but compress history at the cost of fidelity. To overcome this trade-off, we suggest decoupling future-past consistency from any single architecture and instead leveraging a set of specialized experts. We introduce a diffusion-based framework that integrates heterogeneous memory models through a contrastive product-of-experts formulation. Our approach instantiates three complementary roles: a short-term memory expert that captures fine local dynamics, a long-term memory expert that stores episodic history in external diffusion weights via lightweight test-time finetuning, and a spatial long-term memory expert that enforces geometric and spatial coherence. This compositional design avoids mode collapse and scales to long contexts without incurring a quadratic cost. Across simulated and real-world benchmarks, our method improves temporal consistency, recall of past observations, and navigation performance, establishing a novel paradigm for building and operating memory-augmented diffusion world models.

Related papers

• Memory for Large Language Models [18.001409529004807]
  We present a systematic, architecture-centric taxonomy of memory in large language models.<n>We formalize the granular mechanisms dictating memory writing, routing, state transitions, and consolidation.<n>This unified perspective elucidates the conceptual boundaries between computation-coupled and independently addressable memory.
  arXiv  Detail & Related papers  (2026-07-28T07:34:53Z)
• Learning Spatio-Temporal Foundation Models from Pure Synthetic Data [85.23980441738811]
  NeoST is a scalable synthetic-training corpus to real-world bias, a latent-space reasoning architecture that generates and iteratively refines multiple future trajectories, and latent-space objectives that emphasize structural dynamics and enable inference-time correction under distribution shifts.<n>Extensive experiments across diverse real-world benchmarks show that NeoST consistently outperforms existing STFMs in diverse real-worldtemporal systems.
  arXiv  Detail & Related papers  (2026-06-27T02:56:46Z)
• MemoryWAM: Efficient World Action Modeling with Persistent Memory [80.90899269062128]
  We introduce MemoryWAM, a world action model with efficient persistent memory.<n>A tailored attention mechanism enables retrieval of both detailed short-term context and compressed long-term context.<n>MemoryWAM outperforms strong vision-language-action (VLA) and WAM baselines in both simulation and the real world.
  arXiv  Detail & Related papers  (2026-06-18T17:59:51Z)
• Memorize When Needed: Decoupled Memory Control for Spatially Consistent Long-Horizon Video Generation [22.26820693283945]
  We propose a decoupled framework that separates memory conditioning from generation.<n>We employ a lightweight, independent memory branch to learn precise spatial consistency from historical observation.<n>Our approach achieves state-of-the-art performance in terms of both visual quality and spatial consistency.
  arXiv  Detail & Related papers  (2026-04-20T13:00:17Z)
• RainDiff: End-to-end Precipitation Nowcasting Via Token-wise Attention Diffusion [64.49056527678606]
  We propose a Token-wise Attention integrated into not only the U-Net diffusion model but also the radar-temporal encoder.<n>Unlike prior approaches, our method integrates attention into the architecture without incurring the high resource cost typical of pixel-space diffusion.<n>Our experiments and evaluations demonstrate that the proposed method significantly outperforms state-of-the-art approaches, robustness local fidelity, generalization, and superior in complex precipitation forecasting scenarios.
  arXiv  Detail & Related papers  (2025-10-16T17:59:13Z)
• Memory Forcing: Spatio-Temporal Memory for Consistent Scene Generation on Minecraft [45.363427511806385]
  Memory Forcing is a learning framework that pairs training protocols with a geometry-indexed spatial memory.<n>We show that Memory Forcing achieves superior long-term spatial consistency and generative quality across diverse environments.
  arXiv  Detail & Related papers  (2025-10-03T17:35:16Z)
• Video World Models with Long-term Spatial Memory [110.530715838396]
  We introduce a novel framework to enhance long-term consistency of video world models.<n>Our framework includes mechanisms to store and retrieve information from the long-term spatial memory.<n>Our evaluations show improved quality, consistency, and context length compared to relevant baselines.
  arXiv  Detail & Related papers  (2025-06-05T17:42:34Z)
• StateSpaceDiffuser: Bringing Long Context to Diffusion World Models [52.92249035412797]
  We introduce StateSpaceDiffuser, where a diffusion model is enabled to perform long-context tasks by integrating features from a state-space model.<n>This design restores long-term memory while preserving the high-fidelity synthesis of diffusion models.
  arXiv  Detail & Related papers  (2025-05-28T11:27:54Z)
• Long-Context State-Space Video World Models [66.28743632951218]
  We propose a novel architecture leveraging state-space models (SSMs) to extend temporal memory without compromising computational efficiency.<n>Central to our design is a block-wise SSM scanning scheme, which strategically trades off spatial consistency for extended temporal memory.<n>Experiments on Memory Maze and Minecraft datasets demonstrate that our approach surpasses baselines in preserving long-range memory.
  arXiv  Detail & Related papers  (2025-05-26T16:12:41Z)
• FACTS: A Factored State-Space Framework For World Modelling [24.08175276756845]
  We propose a novel recurrent framework, the textbfFACTored textbfState-space (textbfFACTS) model, for spatial-temporal world modelling.<n>The FACTS framework constructs a graph-memory with a routing mechanism that learns permutable memory representations.<n>It consistently outperforms or matches specialised state-of-the-art models, despite its general-purpose world modelling design.
  arXiv  Detail & Related papers  (2024-10-28T11:04:42Z)

This list is automatically generated from the titles and abstracts of the papers in this site.

This site does not guarantee the quality of this site (including all information) and is not responsible for any consequences.