[Paper Review] CMS-LSTM: Context Embedding and Multi-Scale Spatiotemporal Expression LSTM for Predictive Learning

Zenghao Chai, Zhengzhuo Xu+3 more|arXiv (Cornell University)|Feb 6, 2021

Hydrological Forecasting Using AI4 citations

TL;DR

This paper proposes CMS-LSTM, a novel spatiotemporal recurrent network that enhances predictive learning through a Context Embedding (CE) block for iterative context interaction and a Multi-Scale Spatiotemporal Expression (SE) block for fine-grained, multi-scale feature representation. By integrating these modules into ConvLSTM, the model achieves state-of-the-art performance with fewer parameters on Moving MNIST and typhoon prediction benchmarks, significantly improving PSNR, SSIM, and reducing MSE and MAE.

ABSTRACT

Spatiotemporal predictive learning (ST-PL) is a hotspot with numerous applications, such as object movement and meteorological prediction. It aims at predicting the subsequent frames via observed sequences. However, inherent uncertainty among consecutive frames exacerbates the difficulty in long-term prediction. To tackle the increasing ambiguity during forecasting, we design CMS-LSTM to focus on context correlations and multi-scale spatiotemporal flow with details on fine-grained locals, containing two elaborate designed blocks: Context Embedding (CE) and Spatiotemporal Expression (SE) blocks. CE is designed for abundant context interactions, while SE focuses on multi-scale spatiotemporal expression in hidden states. The newly introduced blocks also facilitate other spatiotemporal models (e.g., PredRNN, SA-ConvLSTM) to produce representative implicit features for ST-PL and improve prediction quality. Qualitative and quantitative experiments demonstrate the effectiveness and flexibility of our proposed method. With fewer params, CMS-LSTM outperforms state-of-the-art methods in numbers of metrics on two representative benchmarks and scenarios. Code is available at https://github.com/czh-98/CMS-LSTM.

Motivation & Objective

• Address the challenge of increasing uncertainty and blurred details in long-term spatiotemporal sequence prediction.
• Improve correlation between input frames and hidden state representations to maintain spatiotemporal consistency across long sequences.
• Capture fine-grained, multi-scale spatiotemporal dynamics by expressing dominant changes while suppressing negligible regions.
• Design modular components (CE and SE blocks) that are transferable to existing spatiotemporal models to boost performance.
• Demonstrate the effectiveness and generalization capability of the proposed modules across diverse datasets and architectures.

Proposed method

• Introduce the Context Embedding (CE) block, which uses stacked lightweight CNN layers to iteratively re-weight input and hidden states for enhanced context interaction.
• Implement the Multi-Scale Spatiotemporal Expression (SE) block using self-attention to extract features at three spatial scales, emphasizing dynamic regions and suppressing static or unimportant areas.
• Integrate CE and SE blocks into a modified ConvLSTM architecture, where CE processes input and hidden states before LSTM gates, and SE refines the final output state.
• Use a three-stage processing flow: (1) CE block refines input and context states, (2) standard LSTM gates compute candidate cell and hidden states, (3) SE block enhances the final output with multi-scale attention.
• Apply residual connections and normalization in CE and SE blocks to stabilize training and improve gradient flow.
• Ensure modularity by enabling CE and SE blocks to be plugged into other models like PredRNN and SA-ConvLSTM without architectural overhaul.

Experimental results

Research questions

• RQ1 Can iterative context interaction between input and hidden states improve long-term spatiotemporal prediction accuracy?
• RQ2 Does multi-scale feature expression in latent states enhance the modeling of fine-grained dynamic changes in sequences?
• RQ3 To what extent do CE and SE blocks reduce uncertainty and blur in predicted frames compared to standard ConvLSTM?
• RQ4 Can the CE and SE modules be effectively transferred to other SOTA spatiotemporal models to improve their performance?
• RQ5 How do different configurations of the SE block (e.g., 1-scale vs. 3-scale) affect prediction quality and robustness?

Key findings

• CMS-LSTM achieves state-of-the-art performance on both Moving MNIST and typhoon prediction datasets, with a PSNR of 21.955 on Moving MNIST and 28.891 on the typhoon dataset, outperforming baseline ConvLSTM by +3.432 and +2.538 dB respectively.
• The ablation study confirms that removing the CE block reduces PSNR by 3.432 dB on Moving MNIST and 2.538 dB on the typhoon dataset, demonstrating the critical role of context interaction.
• The 3-scale SE block yields the best performance, reducing MSE by 36.8 on Moving MNIST and 4.16 on the typhoon dataset compared to baseline, showing the importance of multi-scale modeling.
• When transplanted into PredRNN and SA-ConvLSTM, the CE and SE blocks consistently improve metrics, with PSNR gains of up to +3.607 dB and +1.227 dB respectively, proving their generalization capability.
• The model achieves superior results with fewer parameters than SOTA methods, indicating improved parameter efficiency in capturing complex spatiotemporal dynamics.
• Qualitative results show sharper predicted frames with clearer motion details and reduced blur, especially in regions with high dynamic change, confirming improved feature expression.