[论文解读] MS-RNN: A Flexible Multi-Scale Framework for Spatiotemporal Predictive Learning

Zhifeng Ma, Hao Zhang+1 more|arXiv (Cornell University)|Jun 7, 2022

Geographic Information Systems Studies被引用 10

一句话总结

本文提出 MS-RNN，一种通用的多尺度循环神经网络框架，通过在极低内存开销下整合多尺度感受野与多分辨率特征学习，显著提升 RNN 在时空预测中的性能。MS-RNN 在八种 RNN 模型和四个数据集上，将内存使用量最高降低 56.25%，同时提升预测准确率，展现出相较于更深或更宽模型的卓越效率与性能。

ABSTRACT

Spatiotemporal predictive learning, which predicts future frames through historical prior knowledge with the aid of deep learning, is widely used in many fields. Previous work essentially improves the model performance by widening or deepening the network, but it also brings surging memory overhead, which seriously hinders the development and application of this technology. In order to improve the performance without increasing memory consumption, we focus on scale, which is another dimension to improve model performance but with low memory requirement. The effectiveness has been widely demonstrated in many CNN-based tasks such as image classification and semantic segmentation, but it has not been fully explored in recent RNN models. In this paper, learning from the benefit of multi-scale, we propose a general framework named Multi-Scale RNN (MS-RNN) to boost recent RNN models for spatiotemporal predictive learning. We verify the MS-RNN framework by thorough theoretical analyses and exhaustive experiments, where the theory focuses on memory reduction and performance improvement while the experiments employ eight RNN models (ConvLSTM, TrajGRU, PredRNN, PredRNN++, MIM, MotionRNN, PredRNN-V2, and PrecipLSTM) and four datasets (Moving MNIST, TaxiBJ, KTH, and Germany). The results show the efficiency that RNN models incorporating our framework have much lower memory cost but better performance than before. Our code is released at \url{https://github.com/mazhf/MS-RNN}.

研究动机与目标

• 解决当前最先进 RNN 在时空预测中内存消耗过高的问题，该问题限制了模型的可扩展性与实际部署能力。
• 通过探索尺度这一尚未被充分挖掘的维度，克服模型性能与内存使用之间的权衡。
• 开发一种通用、即插即用的框架，无需对现有 RNN 进行架构重构或显著增加参数量即可提升性能。
• 通过降低内存占用，实现高分辨率视频预测，使模型在固定 GPU 内存限制下处理更大输入尺寸。
• 证明多尺度设计可增强感受野与时间记忆能力，从而更好地建模复杂时空动态。

提出的方法

• 提出一种多尺度架构，通过空间池化生成多分辨率特征图，使模型能够捕捉不同尺度下的物体与运动。
• 通过在共享或并行 RNN 单元前连接或融合来自不同尺度的特征，将多尺度特征整合进现有 RNN（如 ConvLSTM、PredRNN++）。
• 采用自粗到精的预测策略，即模型先预测低分辨率帧，再在更高分辨率上逐步优化，从而提升效率与准确性。
• 形式化分析内存与 FLOPs 的减少比例，理论上证明在理想条件下两者均可减少高达 56.25%。
• 设计一种灵活的框架，通过用多尺度组件替换或增强其输入与隐藏状态处理模块，支持多种 RNN 变体。
• 在四个数据集（Moving MNIST、TaxiBJ、KTH、Germany）上，将该框架应用于八种 RNN 模型（ConvLSTM、TrajGRU、PredRNN、PredRNN++、MIM、MotionRNN、PredRNN-V2、PrecipLSTM）。

实验结果

研究问题

• RQ1 多尺度设计是否能在不增加内存消耗的前提下提升 RNN 在时空预测中的性能？
• RQ2 与更深或更宽的 RNN 变体相比，MS-RNN 在内存使用和 FLOPs 方面能减少多少？
• RQ3 多尺度架构是否增强了感受野与时间记忆能力，从而实现对复杂运动模式的更好建模？
• RQ4 在准确率与内存效率方面，MS-RNN 相较于竞争性多尺度模型（如 CMS-LSTM、MoDeRNN、SimVP 和 Earthformer）表现如何？
• RQ5 MS-RNN 是否能使现有 RNN 在固定 GPU 内存限制下处理更高分辨率的输入？

主要发现

• MS-RNN 相较基线 RNN，内存使用量最高降低 56.25%，FLOPs 最高降低 56.25%，理论分析证实了这些上限。
• 在 Moving MNIST 数据集上，MS-PredRNN++ 将 MSE 降低 43.2%，内存使用从 15.20G 降至 8.18G，减少 48.7%。
• MS-MotionRNN 将内存从 20.65G 降至 10.48G（减少 90.5%），MSE 从 55.35 降至 51.55，性能提升达 37.1%。
• 该框架使 ConvLSTM 可处理最高达 220×220 的图像，而 MS-ConvLSTM 仅消耗 3.86G 内存，同时保持高性能，支持更大输入尺寸。
• 在所有数据集上，MS-RNN 在内存效率与预测准确率方面均优于竞争模型（如 CMS-LSTM、MoDeRNN、SimVP 和 Earthformer）。
• 多尺度设计扩大了有效感受野，并支持自粗到精的预测，解释了实验中观察到的性能提升。