用多个专家协作记忆,让扩散模型更准地预测未来。
Composition of Memory Experts for Diffusion World Models

- 用多个专用记忆模块组成记忆系统,分别处理短期、长期和空间信息。
- 在模拟与真实场景中,预测更连贯,回忆历史更准确,导航性能提升。
- 适合做长时序预测、强化学习中的世界模型构建者。
世界模型旨在根据过去观测预测合理未来的状态,这是强化学习中规划与决策的核心能力。然而现有架构面临根本性记忆权衡:变压器保留局部细节但受二次注意力瓶颈限制,循环与状态空间模型虽可高效扩展,却以牺牲历史保真度为代价压缩记忆。为此,我们提出将未来-过去一致性解耦于单一架构,转而利用一组专业化专家。引入基于扩散的框架,通过对比性的专家乘积形式整合异构记忆模型。该方法实现三种互补角色:短期记忆专家捕捉精细局部动态;长期记忆专家通过轻量级测试时微调,在外部扩散权重中存储情景化历史;空间长期记忆专家保障几何与空间一致性。这种组合式设计避免模式坍塌,可在不产生二次复杂度的前提下扩展至长上下文。在模拟与真实世界基准上,本方法显著提升时间连贯性、过往观测的召回率及导航表现,建立了一种新型记忆增强型扩散世界模型构建与运行范式。
原文摘要 · Abstract (English)
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.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。