🤖 machine learning

Composition of Memory Experts for Diffusion World Models

本文提出了一种基于扩散的世界模型，该模型通过对比性专家乘积公式组合专用的短期、情景和空间记忆专家，克服了现有架构在记忆保真度上的权衡，从而在避免二次计算成本的同时实现了可扩展的高保真未来预测。

原作者： Sebastian Stapf, Pablo Acuaviva Huertos, Aram Davtyan, Paolo Favaro

发布于 2026-05-20

📖 1 分钟阅读☕ 轻松阅读

原作者： Sebastian Stapf, Pablo Acuaviva Huertos, Aram Davtyan, Paolo Favaro

原始论文采用 CC BY 4.0 许可（http://creativecommons.org/licenses/by/4.0/）。 ✨ 这是对下方论文的AI生成解释。它不是由作者撰写或认可的。如需技术准确性，请参阅原始论文。阅读完整免责声明

想象一下，你正在讲述一次穿越迷宫的漫长旅程。你希望这个故事完美无缺：转弯合乎逻辑，数小时前造访的房间在返回时看起来应分毫不差，且景色的细节清晰锐利。

当前人工智能“讲故事者”（被称为世界模型）的问题在于它们存在记忆缺陷。

“短期”人工智能就像一个对刚听到的几句话记忆极佳的人，但会忘记一小时前发生的事。它们能完美描述紧接着的下一步，却丢失了整个旅程的情节。
“长期”人工智能则像一位记得整个旅程的历史学家，但对紧接着下一步的细节却模糊不清。它们知道迷宫的大致形状，却可能会搞混你此刻所站房间的墙壁颜色。

试图构建一个能完美记住“所有”事物的巨型人工智能，就像试图背着图书馆跑马拉松；它太重、太慢，最终你会被书本绊倒。

解决方案：专家团队

本文作者提出了一种构建这些人工智能讲故事者的新方法。他们不使用一个巨型大脑，而是利用一个协同工作的专家团队来讲述故事。他们称之为记忆专家组合（CoME）。

这就像电影制作团队：

“短期”专家（编剧）： 这位专家专注于当下的场景。他们查看最后几帧视频并说：“好的，角色刚向左转，所以下一帧应该显示右侧有一堵墙。”他们擅长处理精细细节和即时动作，但无法记住 100 帧之前发生的事。
“长期”专家（档案管理员）： 这位专家拥有整个旅程的庞大图书馆。他们不会实时查看每一帧（那样太慢了）。相反，他们“研读”旅程的历史，并更新自己的内部笔记（即它们的“权重”）以记住大局。如果你问：“五分钟前那个房间是什么样？”他们可以完美地回忆起来，因为他们已经记住了这一集。
“空间”专家（制图师）： 这位专家负责追踪事物“在哪里”。如果人工智能在一个循环中行走（比如一条处处看起来都一样的走廊），其他专家可能会感到困惑。制图师会说：“等等，我们在北入口，而不是南入口”，从而确保故事保持在正确的位置。

它们如何协同工作：“专家乘积”

通常，如果你向三个人征求意见，你可能会取平均值。但在人工智能中，取平均值有时会产生模糊、混乱的结果。

作者使用了一种巧妙的技巧，称为专家乘积。想象这三位专家都举着写有他们预测的牌子。

如果编剧、档案管理员和制图师都同意下一帧应该是什么样子，人工智能就会非常有信心地画出那一帧。
如果他们意见不一，人工智能不会只是猜测，而是寻找“共同点”。

然而，这里有一个陷阱。有时专家们之所以在错误的事情上达成一致，仅仅是因为他们在互相附和。为了解决这个问题，作者发明了一种**“对比”**方法。

类比： 想象专家们正在唱一首歌。有时他们都会哼出同一个错误的音符，因为他们太在意彼此了。“对比”方法就像一位指挥，说道：“停止哼唱背景噪音！只专注于那些不同于你们习惯的独特音符。”这确保了人工智能不会陷入重复错误的循环，并保持故事的鲜活与准确。

结果

该论文在以下方面测试了这种团队方法：

虚拟迷宫： 智能体必须在其中导航三维迷宫并记住它去过的地方。
现实世界视频： 比如在房子里行走或驾驶机器人。

他们发现，这种团队方法比试图使用一个巨型人工智能要好得多。

更好的记忆： 人工智能可以记住数百帧之前的细节而不会混淆。
更快： 它不需要在每一步都携带庞大的数据“图书馆”；它只需咨询其专家。
更一致： 当人工智能走回它之前造访过的房间时，房间看起来分毫不差，而不是一团模糊。

总结

该论文认为，与其试图构建一个包罗万象的超级大脑，不如建立一个由专业化大脑组成的委员会。通过让短期专家处理细节，长期专家处理历史，空间专家处理位置，然后让它们利用特殊的“对比”规则对未来进行投票，我们可以创造出既能完美记住过去又能准确预测未来的智能体——而且这一切都不需要超级计算机来完成。

技术摘要：用于扩散世界模型的记忆专家组合

问题陈述

世界模型（WMs）对于强化学习（RL）和规划至关重要，因为它们能够预测与过去观测一致且合理的未来状态。然而，现有架构在记忆保真度与计算效率之间面临根本性的权衡。基于 Transformer 的模型保留了局部细节，但受限于二次方注意力成本，导致上下文长度受限。相反，循环网络和状态空间模型（SSMs）具有高效的扩展性，但将历史压缩为隐藏状态，不可避免地随时间推移丢失细粒度细节。单纯增加单一架构的上下文长度会导致训练不稳定和推理成本过高。此外，多个模型的朴素组合往往导致“模式崩溃”，即组合分布过度锐化常见区域并丧失多样性。

方法论

作者提出了记忆专家组合（Composition of Memory Experts, CoME），这是一种基于扩散的框架，将未来与过去的一致性从单一的整体架构中解耦。相反，CoME 利用一组专门的专家，每个专家处理特定的记忆角色，并在推理时以概率方式融合它们。

1. 对比专家乘积（Product of Contrastive Experts, PoCE）

为了解决标准专家乘积（Product of Experts, PoE）可能抑制多样性并放大虚假模式的问题，作者引入了对比专家乘积公式。

机制： 对于一组异构专家 $p_i(x)$ ，其组合定义为 $p(x) \propto \prod \tilde{p}_i(x)$ ，其中 $\tilde{p}_i(x) \propto p_i(x)^{\alpha_i} \bar{p}_i(x)^{1-\alpha_i}$ 。
对比基线： $\bar{p}_i(x)$ 代表专家的无条件基线分布。系数 $\alpha_i > 1$ 充当对比因子。
理论依据： 与朴素退火（会缩小核方差并扭曲局部几何结构）不同，这种对比混合在保持核形状的同时重新加权混合分量。它有效地抑制了虚假模式（即专家在没有强证据的情况下分配高概率的区域），同时保留有效的次要模式，前提是条件专家对有效模式分配的概率高于基线。

2. 专用记忆专家

该框架实例化了三个互补的专家：

短期记忆（STM）： 一个以近期上下文窗口（10–100 帧）为条件的扩散模型。它利用具有扩展上下文长度的架构（例如滑动窗口注意力）来捕捉精细的局部动态和即时时间依赖关系。
长期记忆（LTM）： 一个独立的扩散模型，将情节历史（100–1000 帧）直接存储在其权重中。这是通过在长期上下文上对轻量级 LoRA 适配器进行测试时微调来实现的。这种方法摊销了回忆的成本，允许以恒定时间重用过去的经验，而无需在每一步都将整个历史输入模型。通过冻结骨干网络并仅更新适配器，它避免了灾难性遗忘。
空间长期记忆（SLTM）： LTM 的扩展，融合了从历史中提取的空间先验（例如相机姿态、点云图）。该专家强制几何和空间一致性，在重复或视觉相似的环境中消除位置歧义。

3. 统一生成

最终分布是预训练先验与三个对比专家（STM、LTM、SLTM）的乘积。在采样过程中，这些专家的得分函数以对数空间相加的方式结合，以引导去噪过程，确保生成的未来既尊重近期动态，又符合长期历史和空间一致性。

主要贡献

解耦的记忆架构： 本文提出通过将记忆分布在专用专家而非单一骨干网络上来解决记忆权衡问题，并将其表述为专家乘积问题。
对比专家乘积（PoCE）： 一种融合异构专家的新策略。作者提供了理论和实证证据，证明该方法通过剔除冗余模式，防止了模式崩溃和过度自信，这与标准 PoE 不同。
外部扩散记忆： 引入了一种 LTM 模块，利用 LoRA 适配器的测试时微调将情节知识存储在权重中。该模块扩展为 SLTM 以纳入空间结构，提高了导航任务的准确性。

实验结果

该方法在合成（Memory Maze、DMLab、Minecraft）和现实世界（RealEstate10K、RECON）基准上进行了评估。

重建质量： 在 Memory Maze 数据集上，CoME 在 LPIPS（0.097）、SSIM（0.892）和 PSNR（23.07）方面取得了最先进（SOTA）的性能，显著优于包括全注意力 Transformer、SSM 和滑动窗口模型在内的基线。
可扩展性： 与需要比 STM 基线多 60 倍计算量的全注意力模型不同，CoME 以一小部分训练和推理成本实现了更优越的一致性。
上下文长度： 该模型在长达 480 帧的上下文中未出现性能一致性的饱和，证明了 LTM 组件的有效性。
规划与导航： 在基于目标的导航任务（RECON 基准）中，CoME 将绝对轨迹误差（ATE）降低至 0.96，优于专用导航世界模型（NWM）和其他基线。
回忆能力： 在 RealEstate10K 的反向轨迹实验中，CoME 成功地在向前滚动后回忆起了初始帧，而没有记忆的基线则生成了合理但不一致的场景。
消融研究： 移除对比公式导致性能崩溃，证实了 PoCE 机制的必要性。增加 LTM 的 LoRA 秩和上下文长度始终改善了结果，且未发生过拟合。

意义与主张

本文声称建立了一种构建记忆增强型扩散世界模型的新范式。通过将记忆分布在专用专家上并使用对比组合策略，CoME 在效率与高保真时间一致性之间取得了平衡。作者认为，这种方法使世界模型能够扩展到长上下文，而无需承担二次方成本或遭受循环压缩固有的细节丢失。该方法被呈现为一个通用框架，适用于各种基于扩散的世界模型，增强了其在强化学习设置中用于规划和决策的可靠性。这项工作强调，记忆应是一种分布式系统属性，而非单一架构的约束。