StateSpaceDiffuser: Bringing Long Context to Diffusion World Models
作者: Nedko Savov, Naser Kazemi, Deheng Zhang, Danda Pani Paudel, Xi Wang, Luc Van Gool
分类: cs.CV
发布日期: 2025-05-28 (更新: 2025-10-30)
🔗 代码/项目: PROJECT_PAGE
💡 一句话要点
提出StateSpaceDiffuser，为扩散世界模型引入长时上下文建模能力
🎯 匹配领域: 支柱二：RL算法与架构 (RL & Architecture)
关键词: 扩散模型 世界模型 状态空间模型 长时上下文 时间一致性 视觉预测 序列建模
📋 核心要点
- 现有的基于扩散的世界模型依赖短期观测，缺乏长期记忆，导致生成视频的时间一致性差。
- StateSpaceDiffuser 将状态空间模型融入扩散模型，利用状态空间模型捕获长期交互历史，增强时间连贯性。
- 实验表明，StateSpaceDiffuser 在保持视觉上下文一致性方面显著优于仅扩散模型，尤其是在长 rollout 场景下。
📝 摘要（中文）
世界模型最近在复杂环境中基于动作的视觉预测方面获得了显著关注。然而，仅依赖于少量近期观测会导致模型丢失长期上下文信息。因此，在几个步骤之后，生成的场景会偏离先前观察到的内容，从而损害时间一致性。这种局限性在基于扩散的先进世界模型中很常见，其根源在于缺乏持久的环境状态。为了解决这个问题，我们引入了StateSpaceDiffuser，通过集成来自状态空间模型的特征，使扩散模型能够执行长上下文任务，该状态空间模型代表了整个交互历史。这种设计恢复了长期记忆，同时保留了扩散模型的高保真合成能力。为了严格衡量时间一致性，我们开发了一种评估协议，用于探测模型在扩展的 rollout 中重新实例化已见内容的能力。综合实验表明，StateSpaceDiffuser 显著优于强大的仅扩散基线，将连贯的视觉上下文保持的步数提高了一个数量级。它在 2D 迷宫导航和复杂的 3D 环境中都提供了连贯的视图。这些结果表明，将状态空间表示引入扩散模型在展示视觉细节和长期记忆方面非常有效。
🔬 方法详解
问题定义：现有基于扩散的世界模型在处理长序列预测时，由于缺乏长期记忆，容易出现时间不一致性问题。模型生成的场景会逐渐偏离历史观测，无法保持长期连贯的视觉上下文。这限制了它们在需要长期规划和记忆的任务中的应用。
核心思路：StateSpaceDiffuser 的核心思路是将状态空间模型（SSM）与扩散模型相结合。SSM 用于捕获和维护长期交互历史的状态表示，然后将这些状态信息融入到扩散模型的生成过程中，从而为扩散模型提供长期上下文信息，增强其时间一致性。
技术框架：StateSpaceDiffuser 的整体架构包含两个主要模块：状态空间模型（SSM）和扩散模型。SSM 负责编码历史观测和动作序列，生成环境的潜在状态表示。扩散模型则利用这些状态表示作为条件，生成未来的视觉场景。具体流程如下：1) 使用 SSM 编码历史信息；2) 将 SSM 的输出作为条件输入到扩散模型；3) 扩散模型生成下一帧图像。
关键创新：StateSpaceDiffuser 的关键创新在于将状态空间模型与扩散模型有效结合，从而为扩散模型引入了长期记忆能力。与传统的仅依赖短期观测的扩散模型相比，StateSpaceDiffuser 能够更好地保持时间一致性，生成更连贯的视觉序列。
关键设计：具体的状态空间模型和扩散模型的选择未知，论文可能使用了常见的 SSM 变体（如 LSTM 或 Transformer）和扩散模型架构（如 DDPM 或 DDIM）。关键在于如何将 SSM 的输出有效地融入到扩散模型的条件作用机制中。损失函数的设计也至关重要，可能包含重建损失、对抗损失以及鼓励时间一致性的正则化项。具体的参数设置和网络结构细节需要在论文原文中查找。
🖼️ 关键图片
📊 实验亮点
实验结果表明，StateSpaceDiffuser 在 2D 迷宫导航和 3D 环境中都显著优于仅扩散模型。在保持视觉上下文一致性方面，StateSpaceDiffuser 能够维持连贯的视觉上下文的步数比基线模型提高了一个数量级。论文还设计了一种新的评估协议来衡量模型在长 rollout 中重新实例化已见内容的能力，为时间一致性评估提供了新的思路。
🎯 应用场景
StateSpaceDiffuser 在机器人导航、游戏 AI、自动驾驶等领域具有广泛的应用前景。它可以帮助智能体更好地理解和记忆环境信息，从而做出更明智的决策和规划。例如，在机器人导航中，StateSpaceDiffuser 可以帮助机器人记住已经探索过的区域，避免重复路径，提高导航效率。在自动驾驶中，它可以帮助车辆预测周围环境的长期变化，提高驾驶安全性。
📄 摘要（原文）
World models have recently gained prominence for action-conditioned visual prediction in complex environments. However, relying on only a few recent observations causes them to lose long-term context. Consequently, within a few steps, the generated scenes drift from what was previously observed, undermining temporal coherence. This limitation, common in state-of-the-art world models, which are diffusion-based, stems from the lack of a lasting environment state. To address this problem, we introduce StateSpaceDiffuser, where a diffusion model is enabled to perform long-context tasks by integrating features from a state-space model, representing the entire interaction history. This design restores long-term memory while preserving the high-fidelity synthesis of diffusion models. To rigorously measure temporal consistency, we develop an evaluation protocol that probes a model's ability to reinstantiate seen content in extended rollouts. Comprehensive experiments show that StateSpaceDiffuser significantly outperforms a strong diffusion-only baseline, maintaining a coherent visual context for an order of magnitude more steps. It delivers consistent views in both a 2D maze navigation and a complex 3D environment. These results establish that bringing state-space representations into diffusion models is highly effective in demonstrating both visual details and long-term memory. Project page: https://insait-institute.github.io/StateSpaceDiffuser/.