[论文解读] SwinLSTM:Improving Spatiotemporal Prediction Accuracy using Swin Transformer and LSTM

Song Tang, Chuang Li+2 more|arXiv (Cornell University)|Aug 19, 2023

Human Mobility and Location-Based Analysis被引用 4

一句话总结

该论文提出SwinLSTM，一种新型循环单元，通过用Swin Transformer模块替代ConvLSTM中的卷积操作，更好地捕捉时空数据中的全局空间依赖性。通过将Swin Transformer的自注意力机制与简化的LSTM相结合，SwinLSTM在Moving MNIST、TaxiBJ、Human3.6m和KTH数据集上实现了最先进性能，预测精度显著优于ConvLSTM，且无需使用特殊训练技巧。

ABSTRACT

Integrating CNNs and RNNs to capture spatiotemporal dependencies is a prevalent strategy for spatiotemporal prediction tasks. However, the property of CNNs to learn local spatial information decreases their efficiency in capturing spatiotemporal dependencies, thereby limiting their prediction accuracy. In this paper, we propose a new recurrent cell, SwinLSTM, which integrates Swin Transformer blocks and the simplified LSTM, an extension that replaces the convolutional structure in ConvLSTM with the self-attention mechanism. Furthermore, we construct a network with SwinLSTM cell as the core for spatiotemporal prediction. Without using unique tricks, SwinLSTM outperforms state-of-the-art methods on Moving MNIST, Human3.6m, TaxiBJ, and KTH datasets. In particular, it exhibits a significant improvement in prediction accuracy compared to ConvLSTM. Our competitive experimental results demonstrate that learning global spatial dependencies is more advantageous for models to capture spatiotemporal dependencies. We hope that SwinLSTM can serve as a solid baseline to promote the advancement of spatiotemporal prediction accuracy. The codes are publicly available at https://github.com/SongTang-x/SwinLSTM.

研究动机与目标

• 解决基于CNN的模型因感受野局部化而难以捕捉长程空间依赖性的局限性。
• 通过用自注意力机制替代LSTM中的卷积操作，提升时空预测精度。
• 开发一种新型循环单元SwinLSTM，有效建模序列数据中的时空动态。
• 为未来时空预测研究建立一个强大且可泛化的基线。

提出的方法

• SwinLSTM用Swin Transformer模块替代ConvLSTM中的卷积操作，以在2D特征图中建模全局空间注意力。
• 模型通过将输入帧划分为非重叠的图像块，将这些块嵌入可学习的token以进行注意力计算。
• 采用简化的LSTM单元以保持时间记忆，其输入门、遗忘门和输出门经调整以适配Swin Transformer模块的输出。
• Swin Transformer模块采用移位窗口自注意力机制，高效计算全局依赖性，同时保持计算效率。
• 网络架构使用图像块下采样与上采样层，在编码器和解码器阶段实现特征图的下采样与上采样。
• 重建头通过转置卷积、双线性插值或线性投影，从最终隐藏状态生成预测的未来帧。

实验结果

研究问题

• RQ1 用自注意力机制替代LSTM中的卷积操作，能否提升时空预测精度？
• RQ2 通过Swin Transformer建模全局空间依赖性，对循环模型在时空任务上的性能有何影响？
• RQ3 在不同数据集上实现良好泛化性能时，图像块大小与Swin Transformer模块数量的最佳配置是什么？
• RQ4 所提出的SwinLSTM架构是否在不同时空预测基准上具有良好泛化能力？
• RQ5 SwinLSTM中的隐藏状态与细胞状态如何编码空间运动与轨迹信息？

主要发现

• 在Moving MNIST数据集上，SwinLSTM表现最佳，测试MSE为0.390，SSIM为0.980，显著优于ConvLSTM及其他SOTA方法。
• 在Human3.6m数据集上，SwinLSTM的测试MSE为33.2（除以10），SSIM为0.913，显著优于ConvLSTM。
• 消融实验表明，转置卷积重建在TaxiBJ和Human3.6m数据集上表现最佳，优于双线性插值与线性投影。
• 最优图像块大小因数据集而异：在Human3.6m上，块大小为4时性能最佳；在TaxiBJ上，块大小为2时表现最佳。
• Swin Transformer模块数量（STB）对性能的影响呈非单调性，TaxiBJ与Human3.6m上均以18个模块时达到最佳性能。
• 特征图可视化结果表明，细胞状态记忆了数字轨迹，隐藏状态编码了当前位置，而更深的Swin模块逐步学习到全局空间相关性。