用状态空间模型提升视频世界模型的长期记忆能力
Long-Context State-Space Video World Models
- 采用分块状态空间扫描+局部注意力,兼顾长时记忆与计算效率
- 在Memory Maze和Minecraft上实现更优的长期记忆表现
- 适合需要持续推理的交互式视频建模任务
视频扩散模型通过以动作条件的自回归帧预测展现世界建模潜力,但受限于注意力层处理长序列带来的高计算开销,难以维持长期记忆。为此,我们提出一种新型架构,利用状态空间模型(SSMs)扩展时间记忆,同时保持计算效率。不同于以往将SSMs用于非因果视觉任务的改造方式,本方法充分挖掘了SSMs在因果序列建模中的固有优势。核心是分块式SSM扫描策略,在牺牲部分空间一致性的同时换取更长的时间记忆能力,并结合密集局部注意力保证连续帧间的连贯性。我们在记忆迷宫(Memory Maze)和Minecraft数据集上评估了模型的长期记忆能力,通过空间检索和推理任务验证其在长时程上的表现优于基线方法,且推理速度足够支持交互式应用。
原文摘要 · Abstract (English)
Video diffusion models have recently shown promise for world modeling through autoregressive frame prediction conditioned on actions. However, they struggle to maintain long-term memory due to the high computational cost associated with processing extended sequences in attention layers. To overcome this limitation, we propose a novel architecture leveraging state-space models (SSMs) to extend temporal memory without compromising computational efficiency. Unlike previous approaches that retrofit SSMs for non-causal vision tasks, our method fully exploits the inherent advantages of SSMs in causal sequence modeling. Central to our design is a block-wise SSM scanning scheme, which strategically trades off spatial consistency for extended temporal memory, combined with dense local attention to ensure coherence between consecutive frames. We evaluate the long-term memory capabilities of our model through spatial retrieval and reasoning tasks over extended horizons. Experiments on Memory Maze and Minecraft datasets demonstrate that our approach surpasses baselines in preserving long-range memory, while maintaining practical inference speeds suitable for interactive applications.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。