用于扩散世界模型的记忆专家组合

# 记忆专家的组合用于扩散世界模型  
来源：https://arxiv.org/html/2605.18813  

Sebastian Stapf & Pablo Acuaviva Huertos & Aram Davtyan & Paolo Favaro  
计算机视觉组，计算机科学系，伯尔尼大学  
{sebastian.stapf, pablo.acuavivahuertos, aram.davytan, paolo.favaro}@unibe.ch  

###### 摘要  

世界模型旨在预测与过去观测一致的合理未来，这一能力是强化学习中规划和决策的核心。然而，现有架构面临一个根本性的记忆权衡：Transformer 保留了局部细节，但受二次注意力机制的瓶颈限制；循环神经网络和状态空间模型扩展性更好，但以保真度为代价压缩历史信息。为了克服这一权衡，我们建议将未来-过去一致性从单一架构中分离出来，转而利用一组专门的专家。我们引入了一个基于扩散的框架，通过对比专家乘积公式整合异构记忆模型。我们的方法实例化了三个互补角色：一个短期记忆专家，捕捉精细的局部动态；一个长期记忆专家，通过轻量级的测试时微调将情景历史存储在外部扩散权重中；以及一个空间长期记忆专家，强制几何和空间一致性。这种组合设计避免了模式崩溃，并且在不产生二次成本的情况下扩展到长上下文。在模拟和真实世界基准测试中，我们的方法提高了时间一致性、对过去观测的回忆以及导航性能，为构建和运行记忆增强扩散世界模型建立了一种新范式。  
项目页面：github.io/composition-of-memory-experts (https://wiqzard.github.io/composition-of-memory-experts/)  

图 1：我们的记忆增强扩散世界模型概述。专门的专家捕捉短期动态、长期情景记忆和空间一致性，通过专家乘积公式组合，生成与过去观测一致的未来。  

## 1 引言  

世界模型（WMs）是基于过去观测预测世界合理未来状态的有力工具 [Ha and Schmidhuber (2018a (https://arxiv.org/html/2605.18813#bib.bib82);b (https://arxiv.org/html/2605.18813#bib.bib34)); Hafner et al. (2019 (https://arxiv.org/html/2605.18813#bib.bib45)); Hu et al. (2023 (https://arxiv.org/html/2605.18813#bib.bib72)); Bruce et al. (2024 (https://arxiv.org/html/2605.18813#bib.bib73))]。通过学习对观测环境的分布进行建模，WMs 隐含地捕捉了其底层规则和动态。这种能力使其对下游任务中的决策特别有吸引力 [Alonso et al. (2024 (https://arxiv.org/html/2605.18813#bib.bib36)); Ha and Schmidhuber (2018b (https://arxiv.org/html/2605.18813#bib.bib34)); Hafner et al. (2019 (https://arxiv.org/html/2605.18813#bib.bib45))]。最近的进展，特别是基于扩散的世界模型，在生成高质量的未来轨迹方面取得了显著进步 [Hassan et al. (2024 (https://arxiv.org/html/2605.18813#bib.bib35))]。至关重要的是，这种预测能力使 WMs 非常适合导航和交互，在这些场景中，代理必须预测未来状态才能有效地规划和行动。在此类场景中，代理在不确定性下进行探索、感知物体并选择轨迹，方法是在行动之前模拟候选未来并评估其结果。这些想象的展开只有在与先前观测保持一致时才有用；例如，之前访问过的房间在后续访问时不应自发改变其内容。维护这种跨时间一致性是可靠预测和决策的关键。  

不幸的是，预测性 WMs 面临结构性的权衡：更丰富的时间上下文提高了保真度，但计算量爆炸增长。Transformer 架构可以产生高质量展开，但其二次注意力机制扩展性差且限制了上下文长度 [Vaswani et al. (2017 (https://arxiv.org/html/2605.18813#bib.bib53))]。循环网络和状态空间模型在上下文扩展方面更优，但会将历史压缩到隐藏状态，不可避免地随时间丢失细节。这里没有银弹：每种架构都在某个领域胜出，在另一个领域落败。简单地“增加”上下文不是可持续的解决方案，因为训练会变得不稳定且昂贵，推理成本很快会超出实际预算。  

我们主张一种不同的立场：记忆应该分布在系统中，而不是局限于单一架构。人类认知通过分离快速、容量有限的短期记忆（STM）和较慢但持久的长期记忆（LTM）来说明这一原理。这些记忆形式在机制和目的上都有所不同，正是这种划分使它们共同有效 [Liu et al. (2025 (https://arxiv.org/html/2605.18813#bib.bib93))]。我们对世界模型（WMs）采用相同的理念：不是让单个骨干网络承担所有时间需求，而是将记忆构建为专门专家的组合。这种分布式方法使 WMs 能够在效率与保真度之间取得平衡，支持对过去经验的可扩展和可靠利用。  

除了组合之外，我们还解决了在训练和推理中扩展上下文的硬现实问题：成本高昂且脆弱。即使采用线性时间设计，无限扩展时间范围也是不可行的。为了保持成本稳定，我们添加了一个长期记忆通道，直接将情节知识存储在外部扩散专家的权重中。少量有针对性的“记忆化”更新摊销了未来的回忆，使得在每一步无需将完整历史拖入核心模型的情况下，实现过去经验的常数时间重用。扩散模型特别适合这个想法，因为它们允许在无需重训练的情况下进行推理时异构专家的原则性组合 [Du et al. (2023 (https://arxiv.org/html/2605.18813#bib.bib17))]。我们利用这一点将大型预训练骨干网络与轻量级适配器和辅助专家整合到一个统一的框架中。  

然而，简单的组合是不够的：当专家仍然共享与过去不一致的模式时，正如我们所示，标准的专家乘积往往会过度锐化共同区域并破坏一致性。因此，我们引入了一种对比机制，在组合过程中剔除冗余模式，在关键之处保持一致，同时避免专家只是相互呼应的过度自信。  

最后，对于导航而言，世界不仅仅是时间性的，也是空间性的。代理的位置以及观测如何与地点关联，对于一致性至关重要。我们假设将记忆锚定到空间先验（例如，姿态/拓扑线索或地图对齐的键）可以改善检索和组合，特别是在变视角下重访位置的 RL 风格任务中。因此，我们将记忆模块置于空间结构之上，使得想象的未来既尊重看到的内容，也尊重看到的位置。  

总结起来，我们的主要贡献是：  

1. 我们旨在解决记忆问题，通过使用多个在建模不同尺度记忆方面表现出色的专家模型，并将其表述为专家乘积（PoE）问题，从而在视频世界模型中实现记忆整合的概率性原则视图；  
2. 我们引入了对比专家乘积（PoCE），这是一种针对此场景定制的新策略，确保异构记忆专家能够融合以进行一致预测。我们提供了理论和实证证据来支持其必要性；  
3. 我们提出了利用外部扩散模型作为长期记忆（LTM）的方法，并配合微调策略来适应预训练先验同时保留领域通用能力，并将此框架扩展到空间长期记忆模型（SLTM），以进一步提高生成序列的准确性和空间一致性。  

## 2 相关工作  

**视频生成与世界模型**  
最近的视频扩散模型在建模时间连贯的视频序列方面表现出强大能力，但确保长期一致性仍然是一个开放挑战 [He et al. (2022 (https://arxiv.org/html/2605.18813#bib.bib64)); Henschel et al. (2024 (https://arxiv.org/html/2605.18813#bib.bib62)); Bar-Tal et al. (2024 (https://arxiv.org/html/2605.18813#bib.bib48)); Ma et al. (2024 (https://arxiv.org/html/2605.18813#bib.bib50)); Lin et al. (2024 (https://arxiv.org/html/2605.18813#bib.bib51)); Davtyan et al. (2023 (https://arxiv.org/html/2605.18813#bib.bib28))]。这推动了几条工作线，包括架构修改，例如旨在时间稳定性的 Transformer 变体 [Tay et al. (2020 (https://arxiv.org/html/2605.18813#bib.bib46)); Lu et al. (2024 (https://arxiv.org/html/2605.18813#bib.bib63))] 或不同的建模框架 [Fuest et al. (2025 (https://arxiv.org/html/2605.18813#bib.bib33)); Ge et al. (2022 (https://arxiv.org/html/2605.18813#bib.bib66))]。另外还探索了替代采样策略，即指导生成朝着基于先前上下文的时间对齐预测 [Yin et al. (2023 (https://arxiv.org/html/2605.18813#bib.bib65)); Chen et al. (2024 (https://arxiv.org/html/2605.18813#bib.bib15)); Song et al. (2025 (https://arxiv.org/html/2605.18813#bib.bib81)); Davtyan et al. (2023 (https://arxiv.org/html/2605.18813#bib.bib28))]。在世界建模方面，早期工作从像素学习动态；较新的方法结合生成式记忆以在更长范围内进行推理 [Feng et al. (2023 (https://arxiv.org/html/2605.18813#bib.bib67)); Deng et al. (2023 (https://arxiv.org/html/2605.18813#bib.bib68)); Samsami et al. (2024 (https://arxiv.org/html/2605.18813#bib.bib70)); Alonso et al. (2024 (https://arxiv.org/html/2605.18813#bib.bib36))]。  

**专家乘积与组合式适应**  
将大型预训练扩散模型适应到新领域因其规模和领域特异性而具有挑战性。一种方法，概率适应 [Yang et al. (2024a (https://arxiv.org/html/2605.18813#bib.bib19))，训练一个较小的扩散模型与冻结的预训练模型一起，组合它们的分数以灵活地适应新领域，同时保留原始模型的鲁棒性。组合多个条件专家 [Du et al. (2023 (https://arxiv.org/html/2605.18813#bib.bib17)) 已应用于图像生成。在图像之外，组合性已被探索用于任务和视觉规划 [Liu et al. (2022 (https://arxiv.org/html/2605.18813#bib.bib59))，以及分解语言指令以指导视频生成模型 [Ajay et al. (2023 (https://arxiv.org/html/2605.18813#bib.bib87))]。  

**视频模型的记忆**  
SlowFast [Hong et al. (2025 (https://arxiv.org/html/2605.18813#bib.bib7)) 引入了一种两阶段记忆机制，对预训练扩散模型应用 LoRA 适配器 [Hu et al. (2022 (https://arxiv.org/html/2605.18813#bib.bib22)) 以快速适应最新序列（快速学习），同时维护一个独立的慢速循环来合并未来预测的变化（慢速学习）。与我们的方法不同，这种方法遵循元学习风格的记忆范式，需要额外训练可能很大的预训练网络。其他方法通过显式存储相关过去帧并学习如何检索并将其注入模型上下文窗口来实现记忆 [Xiao et al. (2025 (https://arxiv.org/html/2605.18813#bib.bib80))。然而，这引入了与存储需求增长和检索复杂性相关的挑战。其他工作线利用空间记忆模型，基于相机位置检索相关过去上下文 [Xiao et al. (2025 (https://arxiv.org/html/2605.18813#bib.bib80)); Yu et al. (2025 (https://arxiv.org/html/2605.18813#bib.bib91)); Wu et al. (2025 (https://arxiv.org/html/2605.18813#bib.bib94))。最后还有方法在其世界模型中实现循环网络如状态空间模型 [Savov et al. (2025 (https://arxiv.org/html/2605.18813#bib.bib89)); Po et al. (2025 (https://arxiv.org/html/2605.18813#bib.bib90))。  

## 3 记忆适应  

### 3.1 背景  

#### 去噪扩散概率模型  

扩散模型 [Ho et al. (2020 (https://arxiv.org/html/2605.18813#bib.bib20)) 旨在通过学习数据分布 \(q(\mathbf{x}_0)\)，引入一个序列 \(T\) 个逐步更嘈杂的版本 \(\{\mathbf{x}_t\}_{t=1}^T\) 的数据 \(\mathbf{x}_0\)。这是通过一个前向马尔可夫链完成的，其中每个转移概率 \(q(\mathbf{x}_t|\mathbf{x}_{t-1}) = \mathcal{N}(\mathbf{x}_t; \sqrt{1-\beta_t}\,\mathbf{x}_{t-1}, \beta_t\mathbf{I})\) 是高斯分布，具有预定义的噪声调度 \(\{\beta_t\}_{t=1}^T\)，其中 \(0 < \beta_t \leq 1\)。生成过程学习逆转这个扩散过程，建模 \(p_\theta(\mathbf{x}_{t-1}|\mathbf{x}_t) = \mathcal{N}(\mathbf{x}_{t-1}; \mu_\theta(\mathbf{x}_t, t), \tilde{\beta}_t\mathbf{I})\)，这近似了逆条件 \(q(\mathbf{x}_{t-1}|\mathbf{x}_t)\)。均值 \(\mu_\theta(\mathbf{x}_t, t)\) 由一个输出噪声估计 \(\epsilon_\theta(\mathbf{x}_t, t)\) 的神经网络预测。该模型通过最小化 Ho et al. (2020) 的简化目标进行训练：
\[
\mathcal{L}(\theta) = \mathbb{E}_{\mathbf{x}_0, t, \epsilon \sim \mathcal{N}(0, \mathbf{I})} \left[ \left\| \epsilon - \epsilon_\theta(\mathbf{x}_t, t) \right\|^2 \right],
\]
这隐含地学习了数据分布的得分函数 \(\nabla_{\mathbf{x}_t} \log p_\theta(\mathbf{x}_t) \approx -\frac{\beta_t}{\sqrt{1-\bar{\alpha}_t}} \epsilon_\theta(\mathbf{x}_t, t)\)。训练后，采样迭代地进行 Langevin 式更新：
\[
\mathbf{x}_{t-1} = \frac{1}{\sqrt{\alpha_t}} \left( \mathbf{x}_t - \frac{\beta_t}{\sqrt{1-\bar{\alpha}_t}} \epsilon_\theta(\mathbf{x}_t, t) \right) + \sigma_t \eta, \quad \eta \sim \mathcal{N}(0, \mathbf{I}),
\]
其中 \(\alpha_t, \bar{\alpha}_t, \sigma_t\) 由噪声调度定义。  

#### 通过专家乘积进行组合扩散  

在许多生成场景中，希望组合多个模型编码的知识或约束。一个原则性的方法是通过专家乘积（PoE）框架 [Hinton (2002 (https://arxiv.org/html/2605.18813#bib.bib60))，该框架将复合分布定义为单个（归一化）的乘积。每个模型可以看作一个专家 \(p_i(\mathbf{x})\)，查询本身可以定义一个额外的约束 \(p_{\text{query}}(\mathbf{x})\)。PoE 然后计算一个联合分布 \(p(\mathbf{x}) \propto p_{\text{query}}(\mathbf{x}) \prod_i p_i(\mathbf{x})\)，尖锐地聚焦于与所有证据源一致的潜在状态。这允许选择性、内容可寻址的生成，检索或合成满足条件的模式。