状态空间模型为视频世界模型解锁长期记忆

视频世界模型通过根据动作预测未来帧，为人工智能在动态环境中进行规划和推理开辟了广阔前景。近期，特别是视频扩散模型的进步，在生成逼真的未来序列方面展现出了惊人的能力。然而，一个重大瓶颈依然存在：长期记忆的维持。由于使用传统的注意力层处理长序列会产生高昂的计算成本，当前模型难以记住过去较远时间的事件和状态，这限制了它们执行需要持续理解复杂场景的任务的能力。

一篇由斯坦福大学、普林斯顿大学和某机构研究人员合著的新论文《Long-Context State-Space Video World Models》提出了一种创新的解决方案。他们引入了一种新颖的架构，利用状态空间模型在不牺牲计算效率的前提下扩展了时间记忆能力。

核心问题在于注意力机制相对于序列长度的二次方计算复杂度。随着视频上下文的增长，注意力层所需的资源会爆炸式增长，使得长期记忆在实际应用中变得不切实际。这意味着在达到一定帧数后，模型会有效地“忘记”早期事件，从而损害其在需要长期连贯性或长时间推理的任务中的性能。

作者的核心见解在于充分利用状态空间模型在因果序列建模方面的固有优势。与以往将状态空间模型改造用于非因果视觉任务的尝试不同，这项工作充分利用了其在高效处理序列方面的优势。

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

• 分块状态空间模型扫描方案：这是其设计的核心。他们没有用一次状态空间模型扫描来处理整个视频序列，而是采用了一种分块方案。这在策略上以牺牲部分空间一致性（在一个块内）为代价，来换取显著扩展的时间记忆。通过将长序列分解为可管理的块，他们可以维持一个压缩的“状态”，该状态可跨块传递信息，从而有效扩展模型的记忆范围。
• 密集局部注意力：为了弥补分块状态空间模型扫描可能带来的空间连贯性损失，该模型加入了密集局部注意力。这确保了块内及跨块的连续帧之间保持强关联，从而保留了生成逼真视频所必需的细粒度细节和一致性。这种全局（状态空间模型）和局部（注意力）双重处理方法，使他们能够同时实现长期记忆和局部保真度。

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

• 扩散强制：这种技术鼓励模型基于输入的前缀部分生成帧，有效地迫使其学习在更长的时间跨度内维持一致性。通过有时不对前缀进行采样而让所有token保持噪声状态，训练就等同于扩散强制，这被认为是长上下文训练中前缀长度为零的一种特殊情况。这促使模型即使在初始上下文极少的情况下也能生成连贯的序列。
• 帧局部注意力：为了实现更快的训练和采样，作者实现了一种“帧局部注意力”机制。该机制利用FlexAttention，相比完全因果掩码实现了显著的加速。通过将帧分组（例如，每组5帧，帧窗口大小为10），组内的帧保持双向性，同时也能关注前一组中的帧。这在优化计算负载的同时，实现了有效的感受野。

研究人员在具有挑战性的数据集上评估了他们的长上下文状态空间视频世界模型，包括Memory Maze和Minecraft，这些数据集专门设计用于通过空间检索和推理任务来测试长期记忆能力。

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

该论文《Long-Context State-Space Video World Models》可在arXiv上查阅。FINISHED