[Paper Review] 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 Analysis4 citations

TL;DR

This paper proposes SwinLSTM, a novel recurrent cell that replaces ConvLSTM's convolutional operations with Swin Transformer blocks to better capture global spatial dependencies in spatiotemporal data. By integrating self-attention mechanisms from Swin Transformers with a simplified LSTM, SwinLSTM achieves state-of-the-art performance on Moving MNIST, TaxiBJ, Human3.6m, and KTH datasets, significantly outperforming ConvLSTM in prediction accuracy without using specialized training tricks.

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.

Motivation & Objective

• To address the limitation of CNN-based models in capturing long-range spatial dependencies due to their local receptive fields.
• To improve spatiotemporal prediction accuracy by replacing convolutional operations in LSTM with self-attention mechanisms.
• To develop a new recurrent cell, SwinLSTM, that effectively models both spatial and temporal dynamics in sequential data.
• To establish a strong, generalizable baseline for future spatiotemporal prediction research.

Proposed method

• SwinLSTM replaces the convolutional operations in ConvLSTM with Swin Transformer blocks to model global spatial attention in 2D feature maps.
• The model processes input frames by dividing them into non-overlapping patches, which are embedded into learnable tokens for attention computation.
• A simplified LSTM unit is used to maintain temporal memory, with input, forget, and output gates adapted to work with the output of the Swin Transformer blocks.
• The Swin Transformer blocks employ shifted window self-attention to efficiently compute global dependencies while maintaining computational efficiency.
• The network architecture uses patch merging and expanding layers to downsample and upsample feature maps across encoder and decoder stages.
• A reconstruction head—using transposed convolution, bilinear interpolation, or linear projection—generates the predicted future frame from the final hidden states.

Experimental results

Research questions

• RQ1 Can replacing convolutional operations in LSTM with self-attention mechanisms improve spatiotemporal prediction accuracy?
• RQ2 How does modeling global spatial dependencies via Swin Transformers affect the performance of recurrent models on spatiotemporal tasks?
• RQ3 What is the optimal configuration of patch size and number of Swin Transformer blocks for generalization across diverse datasets?
• RQ4 Does the proposed SwinLSTM architecture generalize well across different spatiotemporal prediction benchmarks?
• RQ5 How do the hidden and cell states in SwinLSTM encode spatial motion and trajectory information?

Key findings

• SwinLSTM achieves the best performance on the Moving MNIST dataset, with a test MSE of 0.390 and SSIM of 0.980, outperforming ConvLSTM and other SOTA methods.
• On the Human3.6m dataset, SwinLSTM achieves a test MSE of 33.2 (divided by 10) and SSIM of 0.913, significantly improving upon ConvLSTM.
• The ablation study shows that transposed convolutional reconstruction yields the best results on TaxiBJ and Human3.6m, outperforming bilinear interpolation and linear projection.
• Optimal patch size varies by dataset: patch size 4 achieves the best performance on Human3.6m, while patch size 2 performs best on TaxiBJ.
• The number of Swin Transformer blocks (STB) has a non-monotonic effect on performance, with 18 blocks yielding the best results on both TaxiBJ and Human3.6m.
• Visualization of feature maps confirms that cell states memorize digit trajectories, while hidden states encode current position, and deeper Swin blocks progressively learn global spatial correlations.