Combining CNNs or ViTs, with RNNs for spatiotemporal forecasting, has yielded unparalleled results in predicting temporal and spatial dynamics. However, modeling extensive global information remains a formidable challenge; CNNs are limited by their narrow receptive fields, and ViTs struggle with the intensive computational demands of their attention mechanisms. The emergence of recent Mamba-based architectures has been met with enthusiasm for their exceptional long-sequence modeling capabilities, surpassing established vision models in efficiency and accuracy, which motivates us to develop an innovative architecture tailored for spatiotemporal forecasting. In this paper, we propose the VMRNN cell, a new recurrent unit that integrates the strengths of Vision Mamba blocks with LSTM. We construct a network centered on VMRNN cells to tackle spatiotemporal prediction tasks effectively. Our extensive evaluations show that our proposed approach secures competitive results on a variety of tasks while maintaining a smaller model size. Our code is available at https://github.com/yyyujintang/VMRNN-PyTorch.
翻译:将CNN或ViT与RNN结合用于时空预测,在预测时间与空间动态方面取得了无与伦比的成果。然而,建模大规模全局信息仍是一项严峻挑战:CNN受限于其狭窄的感受野,ViT则因注意力机制带来密集计算需求而捉襟见肘。近期基于Mamba的架构因其卓越的长序列建模能力而备受关注,其在效率和准确性上超越成熟视觉模型,这激励我们开发一种适用于时空预测的创新架构。本文提出VMRNN单元——一种融合Vision Mamba模块与LSTM优势的新型循环单元,并构建以VMRNN单元为核心的网络以有效解决时空预测任务。广泛评估表明,所提方法在保持较小模型规模的同时,能在多种任务上取得具有竞争力的结果。我们的代码开源在https://github.com/yyyujintang/VMRNN-PyTorch。