66 - The ConvLSTM Cell: A Memory That Keeps Geometry

In the previous lesson we put a CNN in front of an LSTM: the CNN reduces each frame to a vector, the LSTM reads the vectors in order. It works, but it has a hidden cost: the moment the frame becomes a vector, the geometry (what is where) is crushed into a list of numbers with no shape. The ConvLSTM, proposed by Shi and colleagues in 2015 to predict radar images of rain in the following hours, makes a different choice: it keeps the memory as a map, with the same shape as the image, and replaces every matrix product inside the cell with a convolution. The gates, the memory and the output are the same as in the LSTM; the only change is that every number has become a small image. In this lesson you write the equations, implement the cell from scratch, check it pixel by pixel, measure how many parameters it saves and discover two properties that make it suited to video: translation equivariance and a speed limit that is easy to violate without noticing.

The words of this lesson

This lesson reuses the definitions of memory cell, gate, sigmoid and element-wise product from the lesson on the LSTM, and those of convolution, kernel and feature map from the previous lesson. Here we add five terms.

The "same" padding adds zeros around the image before the convolution, so that the output has the same size as the input: with a 3×3 kernel one pixel of zeros is added on each side. The receptive field is the set of input pixels that can influence a given output pixel; with a single 3×3 kernel it is a 3×3 window, and it grows by stacking layers or by repeating over time. Translation equivariance says that, if you shift the input by k pixels, the output shifts by k pixels with the same content: a network that has it recognises the same pattern wherever it is.

Peephole connections are a historical variant of the LSTM in which the gates also look at the cell state; the original ConvLSTM contains them, most modern implementations omit them, and so do we. Finally spatio-temporal forecasting is the task the cell was born for: given a film of maps (radar, traffic, crowd density), predict the following maps.

The analogy: from a single meteorologist to a team of local observers

With the CNN + LSTM, a single meteorologist looks at the whole radar map, writes a three-line summary and passes it to the colleague of the following day, who writes another. The map has been lost: in the summary "a disturbance to the north-east" the exact position of the edges of the rain can no longer be recovered. With the ConvLSTM the organisation is different: a grid of observers, one for each point of the map. Each keeps a personal notebook (the local memory) and at every instant updates the notebook looking at only three things: what they see at their own point now, what was written in their notebook, and what is written in the notebooks of their immediate neighbours. All the observers use the same rules of update (the kernel weights, shared over the whole map). A disturbance that moves is "handed over" from one observer to the neighbour: the memory moves together with the phenomenon.

The problem: the classic LSTM on images

The fully connected LSTM (the "FC-LSTM") would require flattening each frame into a vector of H × W numbers and having a state with as many components, all connected to all. Let us count with a 16×16 map and a memory state equivalent to 8 channels (that is 8 × 16 × 16 = 2048 numbers).

# Parameters: a fully connected LSTM that has the SAME "memory size" as the ConvLSTM
H = W = 16
hidden_ch = 8
n_state = hidden_ch * H * W                  # 2048 numbers of state, as in the ConvLSTM
n_input = 1 * H * W                          # 256 pixels per frame
fc_lstm = 4 * n_state * (n_state + n_input + 2)
conv_lstm = 4 * hidden_ch * (1 + hidden_ch) * 3 * 3 + 4 * hidden_ch
print("FC-LSTM  :", f"{fc_lstm:,}")          # 18,890,752
print("ConvLSTM :", f"{conv_lstm:,}")        # 2,624
print("ratio    :", round(fc_lstm / conv_lstm))

FC-LSTM  : 18,890,752
ConvLSTM : 2,624
ratio    : 7199

Explanation. For the FC-LSTM, each of the four transformations has a matrix of state × (state + input) weights: 4 · 2048 · (2048 + 256 + 2) = 18,890,752. For the ConvLSTM, each of the four transformations is a 3×3 kernel connecting 1 + 8 input channels to 8 output channels: 4 · 8 · (1 + 8) · 9 weights plus 4 · 8 biases = 2,624. The ratio is 7,199 times. The saving does not come from a compression trick but from two assumptions about the structure of the data: (1) a pixel is influenced mostly by its neighbours (locality), (2) the rule linking a pixel to its neighbours is the same all over the image (weight sharing). If the assumptions hold for your problem, the architecture exploits them; if they do not (for example when distant objects influence each other in an essential way), it does not.

The mathematics: from the LSTM to the ConvLSTM

Recall the equations of the LSTM, with W · [h, x] denoting the product of a matrix with the concatenated vector:

f = σ( W_f · [h_{t-1}, x_t] + b_f )        i = σ( W_i · [h_{t-1}, x_t] + b_i )
g = tanh( W_g · [h_{t-1}, x_t] + b_g )     o = σ( W_o · [h_{t-1}, x_t] + b_o )
c_t = f ⊙ c_{t-1} + i ⊙ g                   h_t = o ⊙ tanh(c_t)

The ConvLSTM is exactly the same structure with three changes. First: x_t, h_t and c_t are no longer vectors but tensors (channels, height, width). Second: every matrix-vector product W · v becomes a convolution W ∗ V with a kernel. Third: ⊙ remains the pixel-by-pixel product. In compact form, without peephole connections:

f_t = σ( W_xf ∗ X_t + W_hf ∗ H_{t-1} + b_f )
i_t = σ( W_xi ∗ X_t + W_hi ∗ H_{t-1} + b_i )
g_t = tanh( W_xg ∗ X_t + W_hg ∗ H_{t-1} + b_g )
o_t = σ( W_xo ∗ X_t + W_ho ∗ H_{t-1} + b_o )

C_t = f_t ⊙ C_{t-1} + i_t ⊙ g_t
H_t = o_t ⊙ tanh( C_t )

The symbol ∗ is the convolution. In the original paper there are also terms W_ci ⊙ C_{t-1} (the peephole connections); we omit them because the most widespread implementations do not have them and the results are similar. Two convolutions added together (one on the input, one on the previous state) are equivalent to one convolution on the concatenation of the two along the channel axis: it is the trick used in the code to compute the four gates with a single call.

How to read the drawing, comparing it with that of the LSTM. The green line at the top is still the memory, but now each "box" is an entire map per channel. The product with the forget gate and the sum with the new content happen pixel by pixel: at every point of the map, the gate decides how much to keep and how much to write. The row at the bottom now represents the concatenation of X(t) and H(t−1) on the channels, on which a 3×3 convolution produces, for each pixel, the values of the four gates by looking at the window of nine pixels around it.

Code: the cell from scratch

We write the cell in PyTorch. The core is a single nn.Conv2d that takes the input channels plus those of the state and produces four times the state channels (one per gate).

import torch
import torch.nn as nn

class ConvLSTMCell(nn.Module):
    """LSTM cell whose gates are computed with convolutions. State = maps (B, hidden, H, W)."""
    def __init__(self, in_ch, hidden_ch, kernel=3):
        super().__init__()
        self.hidden_ch = hidden_ch
        # ONE convolution for the four gates, applied to [x, h] stacked on the channel axis
        self.conv = nn.Conv2d(in_ch + hidden_ch, 4 * hidden_ch, kernel, padding=kernel // 2)
        with torch.no_grad():
            self.conv.bias[hidden_ch:2 * hidden_ch] = 1.0            # forget-gate bias = 1, as in the LSTM lesson

    def forward(self, x, state):
        h, c = state                                                 # each (B, hidden, H, W)
        z = self.conv(torch.cat([x, h], dim=1))                      # (B, 4*hidden, H, W)
        i, f, g, o = torch.chunk(z, 4, dim=1)
        i, f, o, g = torch.sigmoid(i), torch.sigmoid(f), torch.sigmoid(o), torch.tanh(g)
        c = f * c + i * g                                            # the additive memory update, per pixel
        h = o * torch.tanh(c)
        return h, c

class ConvLSTM(nn.Module):
    def __init__(self, in_ch, hidden_ch, kernel=3):
        super().__init__()
        self.cell = ConvLSTMCell(in_ch, hidden_ch, kernel)
        self.hidden_ch = hidden_ch

    def forward(self, seq):                                          # seq: (B, T, C, H, W)
        B, T, _, H, W = seq.shape
        h = seq.new_zeros(B, self.hidden_ch, H, W)
        c = seq.new_zeros(B, self.hidden_ch, H, W)
        outs = []
        for t in range(T):
            h, c = self.cell(seq[:, t], (h, c))
            outs.append(h)
        return torch.stack(outs, dim=1), (h, c)                      # (B, T, hidden, H, W)

torch.manual_seed(0)                                         # fixed seed: same weights, same numbers on every run
layer = ConvLSTM(in_ch=1, hidden_ch=8)
x = torch.randn(2, 5, 1, 16, 16)
out, (h, c) = layer(x)
print(out.shape, h.shape, c.shape)          # (2, 5, 8, 16, 16)  (2, 8, 16, 16)  (2, 8, 16, 16)
print(sum(p.numel() for p in layer.parameters()))     # 4 * 8 * (1 + 8) * 3 * 3 + 4 * 8 = 2624

Explanation line by line. In the constructor, nn.Conv2d(in_ch + hidden_ch, 4 * hidden_ch, kernel, padding=kernel // 2) is the single convolution: input = the channels of x plus the channels of h concatenated, output = four blocks of hidden_ch channels. The padding of kernel // 2 (1 for a 3×3 kernel) keeps the spatial size. The forget-gate bias is set to 1, as in the lesson on the LSTM, to start with a cell that tends to remember. In the forward, torch.cat([x, h], dim=1) concatenates on the channel axis; torch.chunk(z, 4, dim=1) splits the result into the four blocks; then come the sigmoid and tanh and the two update lines, identical to those of the LSTM but on tensors. The ConvLSTM class repeats the cell for T frames, starts from zero states with the right shape and stacks the outputs into (B, T, hidden, H, W). The last line counts the parameters: 4·8·(1+8)·3·3 + 4·8 = 2624, exactly the value of the manual count.

torch.Size([2, 5, 8, 16, 16]) torch.Size([2, 8, 16, 16]) torch.Size([2, 8, 16, 16])
2624

Pixel-by-pixel check: the convolution really is a sum of 81 products

To be sure we have understood what that single self.conv(...) line does, we recompute by hand a single value: the pre-activation value of the input gate (channel 0) at the pixel in row 5, column 7, with the formula written as an explicit sum. For each pixel, the gate looks at 9 channels (1 from x and 8 from h) and for each the 3×3 window around the pixel: 81 products, plus the bias.

import numpy as np

# Check ONE gate value at ONE pixel against the formula written out with explicit sums.
torch.manual_seed(1)
cell = layer.cell
x1 = torch.randn(1, 1, 16, 16)
h0 = torch.randn(1, 8, 16, 16)
z = cell.conv(torch.cat([x1, h0], dim=1))                  # (1, 32, 16, 16): [i | f | g | o] pre-activations
gate_i_lib = z[0, 0, 5, 7].item()                          # input-gate unit 0 at pixel (5, 7), before the sigmoid

W = cell.conv.weight.detach().numpy()                      # (32, 1+8, 3, 3): output ch, input ch, ky, kx
b = cell.conv.bias.detach().numpy()
inp = torch.cat([x1, h0], dim=1)[0].numpy()                # (9, 16, 16)
total = b[0]
for ch in range(9):                                        # sum over input channels (1 from x, 8 from h)
    for ky in range(3):
        for kx in range(3):
            total += W[0, ch, ky, kx] * inp[ch, 5 + ky - 1, 7 + kx - 1]     # the 3x3 neighbourhood of (5, 7)
print(gate_i_lib, total, abs(gate_i_lib - total))          # identical up to float precision

-0.7144114375114441 -0.7144115 5.9604645e-08

The first number is what the library computes, the second our explicit sum over 9 channels × 3 × 3 weights, the third the difference: 6·10⁻⁸, that is 32-bit numerical precision. Now the convolution is no longer a black box: every gate of every pixel is a weighted sum of the 81 neighbouring numbers. The index 5 + ky − 1 with ky from 0 to 2 runs through rows 4, 5, 6 around row 5; the same for the columns.

Two properties that can be measured

The receptive field grows by one radius per step

An output pixel at time t depends on its own 3×3 neighbourhood at time t and on the previous state, which in turn depended on the neighbourhood of the step before. So, going back in time, the neighbourhood widens by one pixel per side at every step. Let us measure it: we compute the gradient of one pixel of the final output (the centre) with respect to all the pixels of the input frames; where the gradient is not zero, that input pixel can influence it.

# How far does one input pixel spread, and what happens if the picture moves?
torch.manual_seed(0)
layer = ConvLSTM(1, 4)

# (1) Receptive field: which pixels of the FIRST frame influence ONE output pixel (the centre) at the end?
T = 4
x = torch.zeros(1, T, 1, 21, 21, requires_grad=True)
out, _ = layer(x)
out[0, -1, :, 10, 10].sum().backward()                     # the centre pixel of the last output map
reach = (x.grad[0, 0, 0] != 0)                              # frame 0: which input pixels have non-zero gradient?
print("frame-0 pixels that influence the centre (row 10):", int(reach[10].sum()), "wide")
print("same for the LAST frame (t=T-1):", int((x.grad[0, T - 1, 0] != 0)[10].sum()), "wide")

# (2) Translation equivariance: shift the input by 2 pixels, the output shifts by 2 pixels (away from borders)
img = torch.zeros(1, 3, 1, 21, 21)
img[:, :, :, 8:11, 8:11] = torch.rand(1, 3, 1, 3, 3)       # a small blob far from the borders
shifted = torch.roll(img, shifts=2, dims=4)                 # move right by 2
o1, _ = layer(img); o2, _ = layer(shifted)
a = torch.roll(o1, shifts=2, dims=4)[..., 6:-6, 6:-6]       # compare only the interior: borders break the symmetry
b = o2[..., 6:-6, 6:-6]
print("interior: max |shift(output) - output(shift)| =", f"{(a - b).abs().max().item():.1e}")

frame-0 pixels that influence the centre (row 10): 9 wide
same for the LAST frame (t=T-1): 3 wide
interior: max |shift(output) - output(shift)| = 0.0e+00

With T = 4 frames, the pixels of the first frame that influence the centre of the last output form a window 9 wide (radius 4), whereas for the last frame the window is 3 wide (radius 1): exactly radius 1 for each of the steps back. In the code, out[0, -1, :, 10, 10].sum().backward() picks a single output pixel; x.grad[0, t, 0] != 0 is the mask of the pixels of frame t with non-zero gradient; we count how many there are in the central row.

The practical consequence is a speed limit. If in your video an object moves by more pixels per frame than the kernel covers (radius 1 for a 3×3), the memory cannot "keep pace" with the object: information propagates at most one radius per step. Remedies: larger kernels (5×5, 7×7), more stacked layers, or reducing the spatial resolution (pooling, strided convolutions) so that the same movement corresponds to fewer pixels. In real data this becomes a design rule: estimate the maximum speed in pixels per frame and size the kernel accordingly.

Translation equivariance

In the second part of the experiment we take a small blob, shift it by 2 pixels to the right and compare: shifting the output must give the same result as computing the output of the shifted input. Only the interior is compared ([6:-6]), because at the borders the zero padding breaks the symmetry. Maximum difference: 0.0. The cell recognises the same phenomenon wherever it is, because it uses the same weights at every position: an FC-LSTM would have to learn the phenomenon separately for each position.

From block to model: how it is used in practice

• Stack of layers. Several ConvLSTMs are stacked: the output H of the first layer is the input sequence of the second. More layers increase the receptive field and the abstraction.
• Encoder-decoder for forecasting. A classic scheme: an encoder reads the past frames and produces the final state; a decoder starts from that state and produces the future frames. The original paper uses exactly this structure for the radar forecasts.
• Final head. After the ConvLSTM, a 1×1 convolution reduces the channels of H to those of the image to be predicted (for example 1), with a sigmoid if the pixels are in [0, 1].
• Variants. The ConvGRU applies the same idea to the GRU; the kernels can be different for input and state; layer normalisation can be added.

The four points compose a complete model in the order in which you use them: the stack builds the abstraction, the encoder-decoder scheme decides whether you only read or also generate, the head brings the channels back to those of the image, the variants are refinements to try only after a base version that works. Here we have not trained any complete model: we do it in the next lesson, with a single ConvLSTM.

Trade-offs

• Compute memory. The states are maps. An estimate (not measured): with batch 16, 20 frames of 64×64 and 64 channels, a single gate tensor has 16·20·64·64·64 ≈ 84 million numbers; during training several of them have to be kept per layer (the four gates, c, h and intermediates), about 7, that is on the order of 2.3 gigabytes for a single layer in 32-bit floating point. The real constraint of ConvLSTMs is often GPU memory, not the number of parameters.
• Sequentiality. Like every RNN, reading over time cannot be parallelised.
• Limited receptive field. It can link distant pixels only through many steps or many layers; for long-range interactions global attention may serve better.

The three items are different costs of the same choice: memory grows with channels × height × width × steps, sequentiality prevents hiding that cost with parallelism, and the receptive field limits what the memory can link. If your problem has small images and slow movements the bill is favourable; if it has large images and interactions at a distance, it is worth measuring before choosing.

Real traps

The first trap is violating the speed limit without noticing: the model trains well on slow movements and fails on fast ones, and the validation sample looks good if it contains mostly slow movements, so measure the distribution of speeds. The second concerns the borders: zero padding creates an "edge" that the model can learn and use, and when the shift brings an object close to the border the behaviour changes. The third is forgetting the forget-gate bias: as for the LSTM, an initial value of 0 makes the memory fragile at the start.

The fourth is getting the shapes of the initial states wrong, which must be (B, hidden, H, W) and not a simple vector, a typical mistake when adapting the LSTM code. The fifth is the order of the concatenated channels: the convolution on the concatenation depends on the order [x, h], and if you swap it between training and inference the weights no longer match. The last is thinking that the ConvLSTM "sees" motion without being trained to: it learns what the task requires, and a model trained to minimise the pixel-by-pixel error on slow images often learns to copy the last frame, an effect we will see in the next lesson.

Questions to check yourself (with answers)

A ConvLSTM has 3 input channels, 16 state channels and a 5×5 kernel. How many parameters (a single bias per gate and channel)? 4 · 16 · (3 + 16) · 25 + 4 · 16 = 30400 + 64 = 30464.

After 6 steps with a 3×3 kernel, how wide is the receptive field (with respect to the first frame) on the final pixel? Radius 6, that is a window 13 wide, consistent with the measurement (9 for T = 4, that is 2T + 1).

An object moves by 4 pixels per frame. Is a 3×3 kernel enough? No: the radius is 1 pixel per step. You need a kernel with radius at least 4 (9×9), or to reduce the resolution by a factor of 4 first.

Why can two added convolutions be written as a single one on the concatenation? Because a convolution on concatenated channels is the sum of the convolutions on the individual groups of channels, each with its own weights.

A worked example by hand: the speed limit, in one dimension

To see "the memory that moves together with the phenomenon" and its limit, we use a one-dimensional version. The state is a row of 8 cells, with an object (value 1) in cell 1: [0, 1, 0, 0, 0, 0, 0, 0]. The "kernel" that updates the state looks only at the cell to the left: new[i] = old[i − 1] (it is the kernel [1, 0, 0] of a convolution of radius 1). Applying it at every step, the object slides to the right by one cell per step: after 1 step it is in 2, after 2 in 3, after 3 in 4.

Now the real object moves by two cells per step: after 3 steps it is in 7. The memory, with a kernel of radius 1, can follow it by only one cell per step and after 3 steps it is in 4: the lag grows by one cell at every step, 3 cells after three steps. With a kernel of radius 2 ([0, 0, 0, 0, 1] over 5 cells: it looks two cells to the left) the memory would arrive exactly at 7. This is the speed limit: the memory can travel at most one kernel radius per step. For an object that moves v pixels per frame you need a radius ≥ v, or a state at lower resolution in which the same movement corresponds to fewer pixels, or more layers.

This is the simplest possible example of what in the lesson we measured with gradients: with T = 4 steps and a 3×3 kernel, the central pixel of the last output depended on a window 9 wide of the first frame (radius 4 = steps × kernel radius).

Frequently asked questions

Can a ConvLSTM learn a "global" operation (for example counting how many objects there are)? With small kernels and few steps each pixel sees only a neighbourhood; global information must be carried by many steps or layers, or by a final aggregation operation (a mean over all the pixels). It is not its strong point.

How many state channels are needed? It depends on the variety of phenomena: 8-64 are typical. More channels = more capacity and more memory (it grows linearly for the activations, and quadratically for the kernel weights).

Can I stack layers with different kernels? Yes, and it is done: for example a first layer with a large kernel to capture fast movement and a following one with a small kernel for the details.

Is a ConvLSTM or a CNN + LSTM better? If spatial details matter in the memory (image forecasting, maps), the ConvLSTM; if you only need a decision on the whole video (classification), the CNN + LSTM is cheaper and often enough. We have not measured it on the same task.

Does the forget-gate bias apply here too? Yes: the same reason (starting from a memory that is preserved), with the same effect on the gradient flow. In the implemented cell we set it to 1.

How do I initialise the initial states? To zero with shape (B, channels, H, W); alternatively they can be learned as parameters or obtained from an encoder of the context frames.

Takeaways

The ConvLSTM is an LSTM in which every number has become a map and every matrix-vector product a convolution. With this substitution the memory preserves geometry, the parameters collapse (2,624 against 18,890,752 in our example, a factor of 7,199) and the cell acquires translation equivariance (verified: difference 0.0 in the interior). The price is a receptive field that grows by one radius per step, hence a speed limit to respect, and a compute memory proportional to channels × height × width. In the next lesson we put it to work on a real task: predicting the next frame of moving MNIST digits.

In the series, the previous lesson is Hybrid Architectures: A CNN for Space and an LSTM for Time; the next one is ConvLSTM in Practice: Predicting the Next Frame of Moving Digits.