Source-linked AI summary
SwinLSTM:Improving Spatiotemporal Prediction Accuracy using Swin Transformer and LSTM
Song Tang, Chuang Li, Pu Zhang, RongNian Tang
TL;DR
CNN-based recurrent models can struggle to capture global spatial dependencies because convolution emphasizes local information. The paper proposes SwinLSTM, combining Swin Transformer blocks with simplified LSTM processing, and reports state-of-the-art performance across four datasets, including substantial gains over ConvLSTM.
Problem
CNN-based approaches emphasize local spatial information, limiting their efficiency in capturing spatiotemporal dependencies and prediction accuracy.
Method
SwinLSTM integrates Swin Transformer blocks with a simplified LSTM and forms the core of a predictive network for spatiotemporal prediction.
Results
SwinLSTM achieves state-of-the-art performance on Moving MNIST, Human3.6m, TaxiBJ, and KTH, outperforming ConvLSTM in prediction accuracy.
Takeaways & Limitations
The results suggest that learning global spatial dependencies helps models capture spatiotemporal dependencies more effectively than local convolutional modeling.
Abstract
from arXiv · showhide
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.
1. Introduction
Spatiotemporal prediction is difficult because CNN-based recurrent models emphasize local spatial relations, limiting their ability to capture global dependencies. The paper introduces SwinLSTM and a predictive architecture that combines Swin Transformer blocks with simplified LSTM processing.
- 1. Introduction: CNN-based recurrent models can be inefficient at modeling global spatial information because convolution operators focus on local features and relations.The paper links this locality to restricted prediction accuracy when capturing spatiotemporal dependencies.
- 1. Introduction: SwinLSTM combines Swin Transformer blocks with a simplified LSTM to extract spatiotemporal representations.The proposed predictive network uses SwinLSTM as its core to model spatial and temporal dependencies.
- 1. Introduction: The SwinLSTM architecture splits input images into patches, embeds them, recurrently processes transformed patches and states, and reconstructs the next frame.This pipeline uses patch transformations and recurrent cell and hidden states from the previous time step.
- 1. Introduction: The model is evaluated on Moving MNIST, TaxiBJ, Human3.6m, and KTH, where experiments report excellent performance across all four datasets.The supplied introduction identifies these datasets as the evaluation scope and reports strong overall performance.
2. Related Work
Prior spatiotemporal prediction methods combine CNNs and RNNs through recurrent-state designs and convolutional variants. Vision Transformer research introduced transformer-based visual modeling, while later variants addressed data and task limitations.
- CNN-Based Models: ConvLSTM extends FC-LSTM by replacing fully connected operations with convolutions, while later models modify recurrent units or use 3D convolutions.PredRNN, PredRNN++, E3D-LSTM, and MIM target different spatiotemporal modeling or optimization challenges.
- CNN-Based Models: PredRNN transfers hidden states horizontally and vertically to jointly model spatiotemporal information.PredRNN++ adds a Gradient Highway unit to address vanishing gradients, while other variants alter convolutional or recurrent components.
- Vision Transformers: ViT introduced direct Transformer-based image classification, but its strong performance depended on large datasets, motivating variants such as DeiT for smaller-data settings.The related work positions these methods as developments in applying Transformers to vision.
3. Method
SwinLSTM replaces ConvLSTM’s convolutional state transitions with Swin Transformer blocks and a simplified LSTM to model global spatial and temporal dependencies. The resulting architecture supports base and deeper predictive networks with patch-based inputs, recurrent states, reconstruction, and warm-up/prediction training phases.
- 3.1. Overall Architecture: The model splits each input frame into non-overlapping patches, embeds them linearly, and feeds them with recurrent hidden and cell states into SwinLSTM cells.The hidden state is sent to reconstruction for prediction and carried with the cell state to the next recurrent step.
- 3.1. Overall Architecture: SwinLSTM-B uses one cell, whereas SwinLSTM-D stacks multiple cells and adds Patch Merging for downsampling and Patch Expanding for upsampling.Figure 2 distinguishes the recurrent-cell structure, the single-cell base model, and the multi-cell deep model.
- 3.1. Overall Architecture: Training uses warm-up inputs followed by autoregressive prediction, concatenating both phases before applying L2 or L1+L2 prediction loss.The loss compares predicted warm-up and prediction-phase frames with the corresponding input and future ground-truth frames.
- 3.2. SwinLSTM Module: SwinLSTM combines Swin Transformer blocks with a simplified LSTM, using self-attention instead of convolution to model global spatial dependencies.The simplified cell fuses the input, forget, and output gates into one filter gate, while Swin Transformer blocks process spatial information.
- 3.2. SwinLSTM Module: The recurrent cell updates hidden and cell states horizontally for temporal dependencies while Swin Transformer blocks learn spatial dependencies vertically.The cell equations use Linear Projection and Swin Transformer blocks to transform the current input together with the previous hidden state.
- 3.3. Swin Transformer Block: Each Swin Transformer block applies window-based or shifted-window self-attention with LayerNorm, residual connections, and a two-layer GELU MLP.The second block replaces W-MSA with SW-MSA to form the shifted-window variant.
4. Experiments
SwinLSTM is evaluated on four datasets using standard prediction metrics, quantitative comparisons, qualitative visualizations, ablations, and feature-map analysis. It achieves strong benchmark results, with particularly large gains over ConvLSTM on Moving MNIST and strong performance on KTH.
- 4.4. Main results: SwinLSTM achieves state-of-the-art results on Moving MNIST, Human3.6m, and KTH, while reducing TaxiBJ MSE for every predicted frame.The comparison covers four datasets and previous state-of-the-art models.
- 4.4. Main results: 103.3 to 17.7 MSE and 0.707 to 0.962 SSIM are reported for SwinLSTM versus ConvLSTM on Moving MNIST.Each model observes 10 frames and predicts the subsequent 10 frames.
- 4.4. Main results: 4.49 and 5.59 PSNR gains are reported on KTH for predicting 20 and 40 frames from 10 observations.Metrics are averaged over predicted frames, with higher PSNR indicating better prediction quality.
- 4.4. Main results: On Moving MNIST, SwinLSTM predictions maintain higher visual similarity than ConvLSTM as intertwined digits evolve over time.The qualitative comparison shows increasing blur in ConvLSTM predictions while SwinLSTM remains visually similar.
- 4.5. Ablation Study: The ablations show that transposed convolution outperforms bilinear interpolation and linear projection, while the best patch size varies by dataset.These studies evaluate reconstruction layers and patch sizes on TaxiBJ and Human3.6m.
- 4.6. Feature map Visualization: Feature maps indicate that hidden states track current digit positions, cell states memorize trajectories, and later Swin blocks learn broader spatial correlations.The visualization compares Swin Transformer blocks at different interaction depths.
5. Conclusion
SwinLSTM combines Swin Transformer blocks with simplified LSTM in a recurrent cell and predictive network for spatiotemporal prediction. It performs well across four datasets and outperforms ConvLSTM, supporting global spatial dependency learning.
- SwinLSTM achieves excellent performance on Moving MNIST, TaxiBJ, Human3.6m, and KTH datasets.
- SwinLSTM outperforms ConvLSTM in prediction accuracy.
- The results suggest that learning global spatial dependencies helps models capture spatiotemporal dependencies more effectively.