PredRNN: A Recurrent Neural Network for Spatiotemporal Predictive Learning

In plain words

Video prediction means showing a model a few past frames and asking it to draw the frames that come next. A common model for this is a stack of recurrent layers that each keep a memory of what they have seen. In the usual design, each layer's memory only moves forward in time, so the low-level detail seen by the bottom layer rarely reaches the top, and the top layer rarely shapes what the bottom layer remembers.

PredRNN changes the route the memory takes. It sends the memory up through the layers, then back to the bottom at the next time step, so the coarse and fine views of the scene can influence each other. It also keeps two memories in each unit and pushes them to store different things: one for slow changes and one for quick changes. The authors then train the model with a schedule that hides real frames from the encoder at random, so that the model learns longer-term patterns from the earlier frames rather than leaning on the most recent prediction.

The authors report that this beats earlier models on several datasets. The datasets are synthetic digits, human actions, radar echo maps, city traffic and robot pushing.

TL;DR

• A recurrent video predictor. Its memory flows in a zigzag, up across layers and forward in time, instead of only forward within each layer §4.1 fig. 1.
• Each unit has two memory cells. A cosine penalty pushes their increments apart, and the penalty is dropped at inference §4.2 eq. 5.
• Reverse scheduled sampling hides true input frames during encoding, with a probability that rises over training, so the encoder learns longer-term dynamics §4.4 eq. 9.
• Author-reported: PredRNN-V2 reaches MSE 48.4 on fixed-set Moving MNIST, against 103.3 for ConvLSTM tab. 2. It also gives the best KTH, radar and BAIR scores tab. 7 tab. 8 tab. 10.
• The official repo covers Moving MNIST, KTH and BAIR only. The radar and Traffic4Cast results have no code in it datasets_factory.py.

Problem

Future-frame prediction needs both fine detail (small, local motion) and high-level structure (where objects are going). Stacked ConvLSTM layers keep their memory cells within each layer and update them only along time, so the first layer at a step may ignore what the top layer remembered at the previous step §3.2, Challenge I. The output of a ConvLSTM cell also depends on its memory through the output gate, so one cell must handle both long-term and short-term change §3.2, Challenge II. Training follows the usual sequence-to-sequence scheme: the encoder always sees real frames, the forecaster sees its own earlier predictions, so the forecaster has to learn long-range dependence without seeing new frames §3.2, Challenge III §4.4.

Method

Notation

Frames are tensors of shape (J, M, N). The input window is X_1..X_T and the forecast is X_{T+1}..X_{T+K} §3.1. The paper uses * for convolution and ⊙ for the Hadamard product §3.2. The hidden state is H_t^l and the memory is C_t^l (temporal) or M_t^l (spatiotemporal), for layer l and step t §4.1 §4.2.

Spatiotemporal memory flow §4.1

1. The base unit is a ConvLSTM. Its memory is passed up the stack instead of along time fig. 1 eq. 3.
2. Each layer's gates read the hidden state of the layer below and the spatiotemporal memory that layer produced at the current step eq. 3.
3. The bottom layer reads the frame. Its memory comes from the top layer at the previous step §4.1.
4. Gates use convolutions instead of Hadamard products, which the authors say gives finer memory transitions and lets the model run on inputs of different resolutions §4.1.

# Spatiotemporal memory flow over L layers, one time step  [eq. 3] [§4.1]
# Input:  frame X_t; hidden[l] = H_{t-1}^l; memory = M_{t-1}^L
# Output: top-layer hidden state H_t^L and updated memory M_t^L

def memory_flow_step(frame, hidden, cells, memory):
    for layer in range(num_layers):
        # Bottom layer reads the frame; higher layers read the layer below
        # [eq. 3]
        layer_input = frame if layer == 0 else hidden[layer - 1]
        # Memory moves up through the layers within this step  [§4.1]
        hidden[layer], memory = cells[layer](layer_input, hidden[layer],
                                             memory)
    return hidden[num_layers - 1], memory

ST-LSTM unit with memory decoupling §4.2

The final unit is the Spatiotemporal LSTM (ST-LSTM). It keeps two memories: C runs along time within the layer, as in an LSTM, with no Hadamard terms from C in its input and forget gates; M runs up from the layer below eq. 4 fig. 2. The two get separate gates. The output gate sees both new memories, and the hidden state is a 1×1 convolution over [C, M] gated by the output gate eq. 4.

# ST-LSTM unit, layer l, step t  [eq. 4] [§4.2] [fig. 2]
# Input:  x = X_t; h_prev = H_{t-1}^l; c_prev = C_{t-1}^l; m_below = M_t^{l-1}
# Output: h = H_t^l; c = C_t^l; m = M_t^l; deltas for the decoupling loss

def st_lstm_cell(x, h_prev, c_prev, m_below):
    # Temporal memory C: gates read X_t and H_{t-1}^l  [eq. 4]
    input_gate = sigmoid(x_input(x) + h_input(h_prev))
    forget_gate = sigmoid(x_forget(x) + h_forget(h_prev))
    candidate = tanh(x_cand(x) + h_cand(h_prev))
    delta_c = input_gate * candidate
    c = forget_gate * c_prev + delta_c
    # Spatiotemporal memory M: its own gates, reading M_t^{l-1}  [eq. 4]
    input_gate_m = sigmoid(x_input_m(x) + m_input(m_below))
    forget_gate_m = sigmoid(x_forget_m(x) + m_forget(m_below))
    candidate_m = tanh(x_cand_m(x) + m_cand(m_below))
    delta_m = input_gate_m * candidate_m
    m = forget_gate_m * m_below + delta_m
    # Output gate reads both new memories  [eq. 4]
    out_gate = sigmoid(x_out(x) + h_out(h_prev) + c_out(c) + m_out(m))
    # 1x1 convolution over [C, M] keeps H the same size as the memories
    # [§4.2] [eq. 4]
    h = out_gate * tanh(conv_1x1(concat(c, m)))
    return h, c, m, delta_c, delta_m

The x_*, h_*, c_*, m_* functions are convolutions, one per gate and input. In the paper's ST-LSTM equations, the input gate and forget gate of C use X_t and H_{t-1} only, with no C term eq. 4 §4.2.

The authors' intuition is that M gives a deeper, zigzagging path for short-term change, while C runs on a slower timescale and gives a shorter gradient path between distant hidden states §4.2.

Memory decoupling loss §4.2 eq. 5

The increments are ΔC = W ∗ (i ⊙ g) and ΔM = W ∗ (i' ⊙ g'), where W is a single 1×1 convolution shared by all ST-LSTM units eq. 5 §4.2. The penalty is the absolute cosine similarity of ΔC and ΔM, computed per channel over flattened spatial maps, summed over channels, layers and steps eq. 5. The loss is training-only. The authors say it adds no parameters at inference §4.2.

# Memory decoupling penalty, training only  [eq. 5] [§4.2]
# Input:  delta_c, delta_m for every step t and layer l, each (B, C, H, W)
# Output: scalar penalty

def decouple_penalty(increments):
    total = 0.0
    for delta_c, delta_m in increments:
        # One 1x1 convolution shared by all units  [eq. 5] [§4.2]
        flat_c = flatten_spatial(shared_conv(delta_c))  # (B, C, H*W)
        flat_m = flatten_spatial(shared_conv(delta_m))  # (B, C, H*W)
        # Per-channel cosine similarity over space, then absolute value
        # [eq. 5]
        cos = cosine_over_last_axis(flat_c, flat_m)     # (B, C)
        # Summed over channels c (paper); see Paper vs code for the code
        total += abs(cos).sum()
    return total

Training objective eq. 6

The model is trained end to end on frame reconstruction. The loss is the sum of per-step squared errors over steps 2..T+K, plus the decoupling penalty with weight 1 eq. 6.

\mathcal{L}_\text{final} = \sum_{t=2}^{T+K} \lVert \widehat{X}_t - X_t \rVert_2^2 + \mathcal{L}_\text{decouple} eq. 6

Reverse scheduled sampling §4.4

RSS applies to the encoder steps (t ≤ T). At each step the model is fed the true frame with probability ε_k, otherwise its own previous prediction. ε_k rises from ε_s to ε_e over training eq. 9 §4.4. Scheduled sampling (SS) [bengio2015scheduled] applies to forecaster steps, with a probability η_k of the true frame that decays over training §4.4. Two combinations are described. Strategy 2 starts both schedules at similar values (the paper's preferred choice) fig. 6 tab. 6.

The three schedule shapes for ε_k are linear, exponential and sigmoid §4.4. The paper's best run uses the exponential form with ε_s = 0.5 and ε_e = 1.0 tab. 6. The exponential form is ε_k = ε_e − (ε_e − ε_s) exp(−k / α_e) §4.4.

# Input chosen for one step under RSS (encoder) or SS (forecaster)
# [eq. 9] [§4.4]
# Input:  step t, true frame, previous prediction, schedule values
# Output: the frame fed to the cell at step t

def choose_input(t, true_frame, prev_pred, eps, eta, input_length):
    if t <= input_length:
        # Encoder (RSS): true frame with probability eps  [eq. 9]
        keep_true = bernoulli(eps)
    else:
        # Forecaster (SS): true frame with probability eta  [§4.4]
        keep_true = bernoulli(eta)
    return true_frame if keep_true else prev_pred

Action-conditioned ST-LSTM §4.3

For action-conditioned prediction, the previous hidden state is fused with the action before it enters the cell. The paper uses a product of two convolutions, one on H_{t-1}^l and one on the action A_{t-1}. The fused result V takes the place of H_{t-1}^l in the cell eq. 7 fig. 4.

# Action fusion before the ST-LSTM unit  [eq. 7] [§4.3]
# Input:  hidden_prev = H_{t-1}^l; action = A_{t-1} (paper's index)
# Output: fused state V, passed to the cell as its previous hidden state

def action_fusion(hidden_prev, action):
    # Elementwise product of two convolutions  [eq. 7]
    return conv_hv(hidden_prev) * conv_av(action)

Hyperparameters

Values are the paper's unless marked "code". The code's shipped scripts are the three training runs for Moving MNIST, KTH and BAIR predrnn_v2_mnist_train.sh predrnn_v2_kth_train.sh predrnn_v2_bair_train.sh.

Results

Author-reported. Each number is from the paper's own runs §5. The baselines are the ones the paper names for each dataset §5.1 §5.2–§5.6.

Moving MNIST, fixed 10,000-sequence training set, averaged over 10 steps tab. 2:

Moving MNIST, dynamically generated data, 100,000 iterations tab. 3:

Starred models have three 64-channel layers. CrevNet numbers are taken from its paper tab. 3.

Ablation on the memory flow, MSE on fixed Moving MNIST tab. 4: PredRNN without RSS or decoupling 56.8; truncated bottom-up path 57.6; truncated top-down path 59.7.

Generality of the two techniques across backbones tab. 5, MSE:

RSS training schemes, Moving MNIST MSE tab. 6: PredRNN alone 57.3; with SS 56.8; strategy 1 at ε 0→1 gives 51.8 (linear), 53.9 (sigmoid), 51.8 (exponential); strategy 2 at ε 0.5→1 gives 51.9 (linear), 50.9 (sigmoid), 50.6 (exponential). Models in this table have no decoupling loss.

KTH, averaged over 20 future steps tab. 7: ConvLSTM PSNR 23.58, SSIM 0.712, LPIPS 0.231. Conv-TT-LSTM 27.62 / 0.815 / 0.196. PredRNN 27.55 / 0.839 / 0.204. PredRNN-V2 28.37 / 0.838 / 0.139. PSNR and SSIM are also given for MCnet+Residual (26.29 / 0.806), TrajGRU (26.97 / 0.790) and DFN (27.26 / 0.794) tab. 7.

Radar echo dataset tab. 8: MSE and CSI at 30, 40 and 50 dBZ. TrajGRU 68.3, 0.309, 0.266, 0.211. ConvLSTM 63.7, 0.381, 0.340, 0.286. MIM 39.3, 0.451, 0.418, 0.372. PredRNN 39.1, 0.455, 0.417, 0.358. PredRNN-V2 36.4, 0.462, 0.425, 0.378 (best on all CSI thresholds).

Traffic4Cast, MSE in units of 10⁻³ tab. 9: U-Net 6.992. U-Net + PredRNN with ST-LSTM 7.035, and with decoupling and RSS 5.135. CrevNet with ConvLSTM 6.789; with ST-LSTM 6.613; with decoupling and RSS 6.506.

BAIR robot pushing, averaged over 28 future steps tab. 10. Action-free: PredRNN SSIM 0.536, PSNR 16.21; PredRNN-V2 0.577, 16.80. Action-conditioned: ConvLSTM 0.807, 20.52; SVG 0.821, 19.97; SV2P 0.837, 21.32; PredRNN 0.843, 21.93; PredRNN-V2 0.861, 22.74.

Paper vs code

Official repository: https://github.com/thuml/predrnn-pytorch, at commit 49d3b0f3a2757ad9c299716ec455a96868455d67. The paper links it in §1 and §5 §1 §5. The README describes the repo as the PredRNN and PredRNN-V2 (TPAMI 2022) implementation and links this arXiv paper README.md. That is the basis for treating it as the authors' repository. The owner thuml is not identified as an author's lab in the repo itself; it is the Tsinghua group of corresponding author Mingsheng Long [author block]. No LICENSE file is in the repository, so code_license is recorded as none. No code is quoted in this card; all citations point to the code by path and line.

Differences between paper and code (checked against the code at the commit above):

• Loss scale and weight. Paper: sum of squared L2 errors over steps plus the decoupling term with weight 1 eq. 6. Code: nn.MSELoss (a mean over all elements) plus decouple_beta times the decoupling term, with β 0.1 (MNIST, BAIR) or 0.01 (KTH) predrnn_v2.py run.py predrnn_v2_kth_train.sh.
• Decoupling reduction. Paper: sum over steps, layers and channels of absolute per-channel cosine eq. 5. Code: mean over steps and layers of the batch-and-channel mean of absolute cosine predrnn_v2.py.
• Reverse input augmentation. Not in the paper. Default on: each iteration makes two optimiser steps, one on the batch and one on its time-reversed copy trainer.py model_factory.py.
• Forget-gate bias. Not in the paper's equations. The code adds +1.0 to both forget gates SpatioTemporalLSTMCell_v2.py SpatioTemporalLSTMCell_v2.py SpatioTemporalLSTMCell_v2.py.
• ST-LSTM equations. Match eq. 4: the C gates have no Hadamard terms, the output gate reads C_t and M_t, and the hidden state is a 1×1 convolution over [C, M] SpatioTemporalLSTMCell_v2.py.
• Memory flow. Match eq. 3: the bottom layer's memory input is the top layer's memory from the previous step predrnn_v2.py.
• RSS and SS combined. Match the paper's Strategy 2 in shape: one mask per sample and step, with ε-probability true frames in the encoder and η-probability in the forecaster run.py predrnn_v2.py.
• Action-conditioned fusion. The code multiplies two convolutions and applies the result to the hidden-state gates i, f, g and o SpatioTemporalLSTMCell_v2_action.py. The paper's eq. 7 has the same product, with V taking the place of H_{t-1}. The paper pairs H_{t-1} with A_{t-1}. The code pairs H_{t-1} with the action at the same loop index as the frame, actions[:, t] action_cond_predrnn_v2.py. Whether the two agree depends on how the BAIR loader aligns actions, which was not checked.
• Action path, BAIR script. conv_on_input downsamples the frame and the action by two stride-2 convolutions before the recurrent layers, which the paper does not describe action_cond_predrnn_v2.py predrnn_v2_bair_train.sh.
• Metric scale. Paper's MSE is reported per frame. The code's test loop sums squared error over all pixels of a frame and averages over sequences, so its MSE scale is per-frame sums, not per-pixel means trainer.py trainer.py. The Moving MNIST values in tab. 2 are at that scale.
• Reproducibility scope. The repository has data providers for Moving MNIST, KTH and BAIR only datasets_factory.py. No radar or Traffic4Cast loader, no U-Net or CrevNet backbone and no training script for those results is in the repository.
• Defaults. run.py defaults differ from the scripts: learning rate 1e-3, layer norm on, 64 channels, decoupling weight 0.1 run.py run.py. Use the scripts, not the defaults, to reproduce the paper's settings.

The repository does not contain the radar or Traffic4Cast pipelines, so those results cannot be reproduced from it.

Pitfalls

• Unstated values. Patch size, layer norm, forget-gate bias, the RSS schedule's step counts, the η decay range and the action dimension are not stated in the paper. The code sets all of them. An implementer working from the paper alone must choose values, and the results depend on them.
• Decoupling weight. Eq. 6 has weight 1. The code uses β in 0.01–0.1 and a mean-based MSE, so the decoupling term is much smaller relative to reconstruction than the equation suggests.
• Two gradient steps per batch. With the default reverse_input=1, each batch is used twice per iteration. This doubles the optimiser steps against the paper's stated iteration count trainer.py.
• First input is always real. Eq. 9 is written for t ≤ T, but the code never replaces the first input frame with a prediction, which has no prediction to fall back on predrnn_v2.py.
• Index of the action. The paper writes H_{t-1} with A_{t-1} eq. 7. The code pairs H_{t-1} with the action at index t action_cond_predrnn_v2.py. Check the BAIR loader before relying on either.
• Test horizon. The paper trains on 10 steps and tests on 20 (KTH) or 28 (BAIR) §5.3 §5.6. The shipped scripts set total_length for training only. The test-time configuration is not in the repository.
• Metric scale. MSE values in tab. 2 are per-frame sums of squared error over pixels, as the code computes them trainer.py. Do not compare them to per-pixel MSE from another paper.
• Dataset preprocessing. The paper says Traffic4Cast values are in [0, 1] §5.5. The KTH loader divides by 255 kth_action.py. Preprocessing for Moving MNIST and BAIR is not stated in the paper.
• Backbone-specific claims. The Traffic4Cast and radar results depend on backbones (U-Net, CrevNet) that are not in the repo tab. 9 §5.5.
• Radar limitations. The authors say the model still struggles with sparse, high-intensity echoes in the long tail of the radar data §5.7 fig. 18.
• Gradient and gate analyses. Gradient plots fig. 9 and the forget-gate saturation analysis fig. 10 are averaged over 100 sequences from one trained model each. Single-run variance is not reported.
• Single-run results. The paper reports no seeds or variance for any table.
• Best-of-100 for baselines. For SVG and SV2P the authors draw 100 samples and report the best-SSIM one §5.6. The paper does not say whether PredRNN-V2 was compared on the same footing, so the BAIR gap is not a like-for-like sampling comparison.
• Stated numbers. The paper says the memory-decoupling and RSS gain on Moving MNIST is 14.8% (56.8 to 48.4) §5.2. Table 5 lists the same 14.8% for PredRNN tab. 5.

Lineage

• Builds on: ConvLSTM (Shi et al., 2015) for the base unit §3.2. The conference version of PredRNN, which introduced the ST-LSTM and the spatiotemporal memory flow §1 §4.1. No arXiv id for it is in this card's sources.
• Uses: scheduled sampling (Bengio et al., 2015) for the forecaster §4.4. The memory-network view of the flow §4.1.
• Compared against: ConvLSTM (primary baseline), CrevNet, Conv-TT-LSTM, TrajGRU, SV2P, SVG and the other models in tab. 2 and tab. 3.
• Related extensions from the same group: PredRNN++ and MIM, which the paper combines with its two techniques tab. 5 §5.2.
• Superseded by: not stated in this card's sources.