PredRNN: A Recurrent Neural Network for Spatiotemporal Predictive Learning

Yunbo Wang; Haixu Wu; Jianjin Zhang; Zhifeng Gao; Jianmin Wang; Philip S. Yu; Mingsheng Long

Source-linked AI summary

PredRNN: A Recurrent Neural Network for Spatiotemporal Predictive Learning

Yunbo Wang, Haixu Wu, Jianjin Zhang, Zhifeng Gao, Jianmin Wang, Philip S. Yu, Mingsheng Long

arXiv:2103.09504v4Mar 17, 2021cs.LGcs.CV

TL;DR

Spatiotemporal predictive learning seeks to generate future images from historical context while capturing complex visual dynamics. PredRNN uses decoupled memory cells, zigzag cross-layer memory flow, and reverse scheduled sampling, achieving state-of-the-art performance across five datasets. The approach spans both action-free and action-conditioned predictive learning, but remains limited on sparse, high-intensity radar echoes.

Problem
Future-frame prediction must model both spatial deformations and long-term temporal dynamics, which stacked LSTM memory transitions may not represent optimally.
Method
PredRNN combines an LSTM temporal memory cell with a zigzag spatiotemporal memory flow, decoupling loss, and Reverse Scheduled Sampling.
Results
PredRNN achieves state-of-the-art performance on five synthetic and natural datasets covering action-free and action-conditioned prediction.
Takeaways & Limitations
The framework provides a unified recurrent approach for modeling short-term spatial deformations and long-term dynamics across predictive learning scenarios.
Takeaways & Limitations
PredRNN-V2 cannot perfectly model sparse high-intensity radar echoes arising from long-tail cases, despite better prediction of their future positions.

Abstract

from arXiv · show

The predictive learning of spatiotemporal sequences aims to generate future images by learning from the historical context, where the visual dynamics are believed to have modular structures that can be learned with compositional subsystems. This paper models these structures by presenting PredRNN, a new recurrent network, in which a pair of memory cells are explicitly decoupled, operate in nearly independent transition manners, and finally form unified representations of the complex environment. Concretely, besides the original memory cell of LSTM, this network is featured by a zigzag memory flow that propagates in both bottom-up and top-down directions across all layers, enabling the learned visual dynamics at different levels of RNNs to communicate. It also leverages a memory decoupling loss to keep the memory cells from learning redundant features. We further propose a new curriculum learning strategy to force PredRNN to learn long-term dynamics from context frames, which can be generalized to most sequence-to-sequence models. We provide detailed ablation studies to verify the effectiveness of each component. Our approach is shown to obtain highly competitive results on five datasets for both action-free and action-conditioned predictive learning scenarios.

1 INTRODUCTION

PredRNN addresses spatiotemporal prediction by coupling hierarchical spatial representations with temporal dynamics through zigzag memory flow and complementary memory cells. Its extensions add memory decoupling and reverse scheduled sampling, and the approach achieves state-of-the-art performance across five datasets.

Motivation: The framework targets predictive learning because future-frame generation requires preserving both fine-grained spatial details and high-level semantics.These requirements differ from tasks where high-level semantic features alone may suffice.
PredRNN framework: PredRNN extends LSTM memory transitions into a spatiotemporal memory flow that traverses recurrent layers in a bidirectional zigzag pattern.Low-level information moves bottom-up, while top-layer memory returns top-down at the next timestep.
ST-LSTM: ST-LSTM combines the proposed spatiotemporal memory cell with LSTM’s temporal memory cell to model short-term deformations and long-term dynamics.The spatiotemporal cell increases nonlinear depth between time-adjacent states, while the temporal cell supports long-term coherence.
Extensions: The extended PredRNN introduces a decoupling loss, action-conditioned prediction, and Reverse Scheduled Sampling for learning longer-term context dependencies.Reverse Scheduled Sampling randomly hides real observations with probabilities that change during training.
Results: PredRNN achieves state-of-the-art performance on five datasets spanning action-free and action-conditioned predictive learning.The datasets include Moving MNIST, KTH, radar precipitation forecasting, Traffic4Cast, and action-conditioned BAIR.

2 RELATED WORK

Prior spatiotemporal prediction methods encode different inductive biases through feed-forward, recurrent, and convolutional recurrent architectures. PredRNN builds on convolutional recurrent models with a zigzag memory transition, decoupled memory cells, and a sequence-to-sequence training procedure.

Architectural families: Existing spatiotemporal prediction architectures broadly include CNN-based feed-forward models, recurrent models, and Transformer- or flow-based alternatives.These groups encode different inductive biases for predictive learning.
Feed-forward models: CNN-based feed-forward models introduce spatial group-invariance biases and have been applied to next-frame, probabilistic, and traffic prediction.Examples include convolutional autoencoders, motion-kernel probabilistic models, and residual CNNs.
Recurrent models: Recurrent models are used to capture historical dynamics, extend prediction horizons, and represent temporal uncertainty.Reported designs include patch-cluster prediction, sequence-to-sequence video prediction, and hierarchical recurrent states.
Convolutional recurrent models: Convolutional recurrent models combine spatial appearance modeling with temporal dynamics by replacing standard recurrent matrix operations with convolutions.ConvLSTM is a foundational example, and later models add action conditioning, non-local connections, temporal snippets, or tensor factorization.
PredRNN: This paper improves convolutional recurrent prediction with zigzag spatiotemporal memory flow, decoupled memory cells, and a new sequence-to-sequence training procedure.The design combines convolutional and recurrent modeling advantages.

3 PRELIMINARIES

Spatiotemporal predictive learning predicts future spatially arranged observations from historical sequences, with video and radar echo maps as representative domains. The preliminaries identify three limitations of ConvLSTM involving hierarchical memory interaction, mixed temporal scales, and training dependencies.

Spatiotemporal predictive learning: Given tensor observations over T timesteps, spatiotemporal predictive learning predicts a most-probable future sequence of length K.Each observation is represented on an M × N spatial grid across J measurement channels.
Application domains: Video prediction uses RGB images, while precipitation nowcasting uses radar echo maps from a geographic region.Radar echoes are commonly visualized by mapping them to color images.
Convolutional LSTM: ConvLSTM represents inputs, memory states, and hidden states as spatial tensors and applies convolutions to recurrent transitions.Its local cell updates depend on the current frame and neighboring past states.
Convolutional LSTM: ConvLSTM’s gates update memory and hidden states using convolution and Hadamard-product operations with sigmoid activations.The equations define input, forget, cell, and output transformations.
Challenge I: In stacked ConvLSTM, memory states remain primarily within layers, limiting their dependence on hierarchical visual features from other layers.The bottom layer may therefore largely ignore information memorized by the top layer at the previous timestep.
Challenge II: ConvLSTM forces one memory cell to handle both long-term and short-term dynamics, potentially limiting modeling of complex spatiotemporal variations.The hidden state depends on this memory state and its output gate.
Challenge III: Standard sequence-to-sequence training feeds real observations during encoding but requires greater reliance on prior context during forecasting.Shared parameters and one-step encoding training may hinder learning jumpy frame dependencies for long-term prediction.

4 METHOD

PredRNN combines a zigzag spatiotemporal memory flow with paired, decoupled memory cells and reverse scheduled sampling to model spatial structure, short-term motion, and long-term dynamics. Its action-conditioned extension supports simulating future spatiotemporal variations in response to agent actions.

Spatiotemporal Memory Flow: PredRNN uses convolutional recurrent units and a zigzag memory flow to jointly model spatial correlations and temporal dynamics across RNN levels.The flow traverses recurrent nodes through bidirectional hierarchies, allowing memory states in different layers to communicate.
Spatiotemporal Memory Flow: The spatiotemporal memory flow updates states upward across layers before forwarding them to the next timestep, deepening transitions for short-term motion modeling.This path adds nonlinear neurons between horizontally adjacent states, but may make long-term dependencies harder to capture.
Spatiotemporal LSTM with Memory Decoupling: ST-LSTM combines the original temporal memory cell C with a spatiotemporal memory cell M to preserve long-term coherence while responding quickly to short-term dynamics.The two memory flows provide complementary transition paths within the recurrent unit.
Spatiotemporal LSTM with Memory Decoupling: A decoupling loss prevents C and M from learning redundant features by encouraging their memory increments to become orthogonal.The additional parameters are used only during training, so inference model size does not increase relative to the previous ST-LSTM.
Action-Conditioned PredRNN: Action-conditioned PredRNN learns to simulate the consequences of future action sequences over long periods and is validated on data collected with a real robot.The action-conditioned ST-LSTM extends the model beyond action-free prediction scenarios.
Training with Reverse Scheduled Sampling: Reverse scheduled sampling progressively hides true context frames during encoding, encouraging the model to learn long-term, non-Markovian dynamics.It can be combined with original scheduled sampling, which addresses the training-inference gap during forecasting; the second combined strategy performs slightly better empirically.

5 EXPERIMENTS

Experiments evaluate PredRNN and its variants across synthetic, real-world, action-free, and action-conditioned prediction settings. Ablations attribute gains to zigzag memory flow, ST-LSTM, memory decoupling, and reverse scheduled sampling, while results also identify a limitation on sparse high-intensity radar echoes.

Ablation studies: PredRNN-V2 combines memory decoupling and reverse scheduled sampling to improve the original PredRNN training process.These techniques add negligible model parameters and are shown to further improve models based on vanilla ST-LSTMs.
Moving MNIST: 14.8% MSE improvement reduces Moving MNIST error from 56.8 to 48.4 after adding memory decoupling and reverse scheduled sampling.The preceding stages reduce ConvLSTM MSE from 103.3 to 74.0 with spatiotemporal memory flow and then to 56.8 with ST-LSTM.
Ablation studies: PredRNN and PredRNN-V2 improve with longer training, while memory decoupling and reverse scheduled sampling further improve ST-LSTM-based models.The zigzag flow’s top-down transition contributes more to final Moving MNIST performance, and its long-term modeling effect is studied through saturation ratios.
Moving MNIST: PredRNN produces clearer long-term Moving MNIST predictions and more accurate digit trajectories and shapes after occlusions than compared models.The original PredRNN and PredRNN-V2 progressively improve prediction quality in the qualitative examples.
Traffic4Cast: On Traffic4Cast, the memory-decoupled ST-LSTM with reverse scheduled sampling predicts rapid traffic-flow changes most accurately.The evaluation predicts three future frames from nine context frames on Berlin traffic data.
Action-conditioned BAIR: Action-conditioned PredRNN variants outperform SVG, SV2P, and action-conditioned ConvLSTM, with PredRNN-V2 achieving state-of-the-art performance on BAIR.The models also predict commanded gripper positions more precisely and enrich object details.
Model limitations: PredRNN-V2 cannot perfectly model sparse, high-intensity radar-echo dynamics associated with long-tail cases.The paper identifies incorporating complex physical properties with deep learning as a future direction.

6 CONCLUSION

PredRNN combines a Spatiotemporal LSTM with zigzag memory flow to model short-term spatial deformations and long-term dynamics. Reverse scheduled sampling further encourages learning temporal dynamics from longer context periods, supporting strong results across action-free and action-conditioned settings.

PredRNN uses a Spatiotemporal LSTM to model short-term spatial deformations and long-term dynamics simultaneously.Its zigzag memory flow propagates across stacked recurrent layers and temporal states, allowing hierarchical memory representations to interact.
The zigzag memory flow connects hierarchical memory representations across recurrent layers and temporal states.The flow moves vertically across layers and horizontally through temporal states.
PredRNN decouples twisted memory states to learn distributed representations covering different aspects of spatiotemporal variation.The decoupling separates horizontal and zigzag pathway transitions.
Reverse scheduled sampling forces the encoding component to learn temporal dynamics from longer periods of context frames.The strategy is presented as a curriculum learning method for PredRNN.
PredRNN achieves state-of-the-art performance on synthetic and natural spatiotemporal datasets in action-free and action-conditioned prediction scenarios.