Adobe 研究团队利用状态空间模型解锁视频世界模型中的长期记忆

视频世界模型能够基于动作预测未来帧，在人工智能领域具有巨大潜力，使智能体能够在动态环境中进行规划和推理。近期，视频扩散模型在生成逼真未来序列方面表现出色，但一个关键瓶颈仍然存在：长期记忆。由于传统注意力层在处理长序列时计算成本过高，现有模型难以记住远距离过去的事件和状态，限制了其在需要持续理解场景的复杂任务中的表现。

一篇由斯坦福大学、普林斯顿大学和 Adobe 研究团队合著的新论文《长上下文状态空间视频世界模型》提出了一种创新解决方案。他们引入了一种新型架构，利用状态空间模型（SSM）在不牺牲计算效率的前提下扩展时间记忆。

核心问题在于注意力机制相对于序列长度的二次计算复杂度。随着视频上下文的增长，注意力层的资源消耗急剧增加，使得长期记忆在实际应用中变得不可行。这意味着模型在经过一定帧数后实际上会“忘记”早期事件，从而影响其在需要长距离连贯性或长时间推理的任务中的性能。

作者的关键洞察是利用状态空间模型在因果序列建模方面的固有优势。与之前尝试将 SSM 改造用于非因果视觉任务不同，这项工作充分利用了其在高效处理序列方面的优势。

提出的长上下文状态空间视频世界模型（LSSVWM）包含了几个关键设计选择：

块状 SSM 扫描方案：这是设计的核心。他们并非用单个 SSM 扫描处理整个视频序列，而是采用块状方案。这策略性地牺牲了一定空间一致性（在块内），以显著扩展时间记忆。通过将长序列分解为可管理的块，模型可以在块之间保持压缩的“状态”，从而有效延长记忆范围。

密集局部注意力：为了补偿块状 SSM 扫描可能带来的空间连贯性损失，模型引入了密集局部注意力。这确保块内和块间的连续帧保持强关联，保留了逼真视频生成所需的细粒度细节和一致性。这种全局（SSM）和局部（注意力）的双重处理方法使得模型既能实现长期记忆，又能保持局部保真度。

论文还介绍了两种关键训练策略以进一步提升长上下文性能：

扩散强制：该技术鼓励模型基于输入前缀生成帧，从而迫使其学习在更长持续时间内保持一致性。通过有时不采样前缀并使所有标记保持噪声，训练相当于扩散强制，这被强调为前缀长度为零的长上下文训练的特例。这推动模型即使在最小初始上下文下也能生成连贯序列。

帧局部注意力：为了加快训练和采样速度，作者实现了“帧局部注意力”机制。它利用 FlexAttention 实现了相比完全因果掩码的显著加速。通过将帧分组为块（例如，块大小为5，帧窗口大小为10），块内的帧保持双向性，同时也能关注前一个块中的帧。这使得在优化计算负载的同时保持了有效的感受野。

研究人员在具有挑战性的数据集上评估了 LSSVWM，包括 Memory Maze 和 Minecraft，这些数据集专门设计用于通过空间检索和推理任务测试长期记忆能力。

实验结果表明，该方法在保持长期记忆方面显著超越了基线。定性结果（如补充图 S1、S2、S3 所示）显示，与仅依赖因果注意力或没有帧局部注意力的 Mamba2 相比，LSSVWM 能够在更长时间内生成更连贯且准确的序列。例如，在迷宫数据集的推理任务中，该模型在长时间范围内保持了更好的一致性和准确性。同样，在检索任务中，LSSVWM 在回忆和利用遥远过去帧的信息方面表现出改进的能力。至关重要的一点是，这些改进是在保持实际推理速度的同时实现的，使得模型适用于交互式应用。

论文《长上下文状态空间视频世界模型》可在 arXiv 上获取。

该文章最初发表于 Synced。