状态空间模型破解视频世界模型长期记忆难题
定性结果,如补充图表所示,说明了与仅依赖因果注意力甚至是不带帧局部注意力的Mamba2模型相比,长上下文状态空间视频世界模型能在更长时间跨度内生成更连贯、更准确的序列。然而,一个显著的瓶颈依然存在:维持长期记忆。这意味着在特定数量的帧之后,模型会有效地“忘记”较早的事件,从而阻碍其在需要长程一致性或长时间跨度推理的任务上的表现。研究人员在具有挑战性的数据集上评估了他们的长上下文状态空间视频世界模型
视频世界模型能够根据动作预测未来帧,对人工智能领域意义重大,它使智能体能够在动态环境中进行规划和推理。近期的进展,特别是视频扩散模型,在生成逼真的未来序列方面展现了强大能力。然而,一个显著的瓶颈依然存在:维持长期记忆。由于使用传统注意力层处理长序列的计算成本高昂,现有模型难以记住过去较远的事件和状态。这限制了它们执行需要持续理解场景的复杂任务的能力。
一篇由斯坦福大学、普林斯顿大学和某研究机构的研究人员合著的新论文《长上下文状态空间视频世界模型》提出了一个针对此挑战的创新解决方案。他们介绍了一种新颖的架构,该架构利用状态空间模型在不牺牲计算效率的前提下扩展时间记忆。
核心问题在于注意力机制相对于序列长度的二次方计算复杂度。随着视频上下文的增长,注意力层所需的资源会爆炸式增长,使得在实际应用中实现长期记忆变得不切实际。这意味着在特定数量的帧之后,模型会有效地“忘记”较早的事件,从而阻碍其在需要长程一致性或长时间跨度推理的任务上的表现。
作者的關鍵见解在于利用状态空间模型在因果序列建模方面的固有能力。与以往尝试将状态空间模型改造用于非因果视觉任务不同,这项工作充分发挥了其高效处理序列的优势。
提出的长上下文状态空间视频世界模型包含几个关键的设计选择:
- 分块式状态空间模型扫描方案:这是他们设计的核心。他们并未用单个状态空间模型扫描来处理整个视频序列,而是采用了一种分块方案。这策略性地在(块内的)空间一致性上做出了一些权衡,以换取显著扩展的时间记忆。通过将长序列分解为可管理的块,他们可以维持一个压缩的“状态”,该状态能跨块携带信息,从而有效地扩展模型的记忆范围。
- 密集局部注意力:为了弥补分块式状态空间模型扫描可能带来的空间连贯性损失,该模型纳入了密集局部注意力。这确保了块内和块间的连续帧能保持强关联性,从而保留了生成逼真视频所需的细粒度和一致性。这种全局(状态空间模型)与局部(注意力)处理的结合方式,使其能够同时实现长期记忆和局部保真度。
论文还介绍了两种关键的训练策略,以进一步提升长上下文性能:
- 扩散强制:这种技术鼓励模型在生成帧时以输入的前缀为条件,实际上强制其学习在更长的时间跨度内保持一致性。通过有时不对前缀进行采样并使所有token保持噪声状态,训练变得等同于扩散强制,这被认为是前缀长度为零时的一种特殊长上下文训练情况。这促使模型即使在初始上下文极少的情况下也能生成连贯的序列。
- 帧局部注意力:为了加快训练和采样速度,作者实现了一种“帧局部注意力”机制。该机制利用FlexAttention,与完全因果掩码相比,实现了显著的加速。通过将帧分组为块(例如,块大小为5,帧窗口大小为10),块内的帧保持双向性,同时也能关注前一个块中的帧。这使得在优化计算负载的同时,能够拥有有效的感受野。
研究人员在具有挑战性的数据集上评估了他们的长上下文状态空间视频世界模型,包括专门设计用于通过空间检索和推理任务来测试长期记忆能力的Memory Maze和Minecraft数据集。
实验表明,他们的方法在保持长程记忆方面显著超越了基线模型。定性结果,如补充图表所示,说明了与仅依赖因果注意力甚至是不带帧局部注意力的Mamba2模型相比,长上下文状态空间视频世界模型能在更长时间跨度内生成更连贯、更准确的序列。例如,在迷宫数据集的推理任务上,他们的模型在长距离上保持了更好的一致性和准确性。同样,在检索任务中,长上下文状态空间视频世界模型也显示出其回忆和利用遥远过去帧信息的能力有所提高。重要的是,这些改进是在保持实际推理速度的同时实现的,这使得该模型适用于交互式应用。
这篇论文《长上下文状态空间视频世界模型》可在arXiv上查阅。FINISHED
更多精彩内容 请关注我的个人公众号 公众号(办公AI智能小助手)或者 我的个人博客 https://blog.qife122.com/
对网络安全、黑客技术感兴趣的朋友可以关注我的安全公众号(网络安全技术点滴分享)
更多推荐

所有评论(0)