Spatio-temporal video autoencoder with differentiable memory

In plain words

A video is mostly the same picture moving a little from one frame to the next. This paper builds a small network that watches a video one frame at a time, keeps a short-term memory of what it has seen, and guesses the next frame. The memory is laid out like the image itself, so it remembers where things are rather than a flat list of numbers.

The guess is made by first predicting how each pixel moves, then sliding the current picture along those motions. The only training signal is how close the guess is to the real next frame, so no labels are needed to learn motion. The paper then uses the learned motion as a tool: if a few frames of a video have hand labels, the labels can be carried to the neighbouring frames along the motion. This is a workshop-length study on small datasets.

TL;DR

The model is a convolutional encoder and decoder around a convolutional LSTM memory, with a head that predicts a dense optical flow map. The flow warps the current frame's features to predict the next frame, and the whole system is trained end to end on next-frame error. On moving MNIST it has the lowest reported per-pixel error among four variants, and on CamVid it is used to propagate SegNet predictions across frames. All numbers are author-reported on small setups. The released code covers only the moving-MNIST variant (see Paper vs code).

Problem

• Video models need labels, and labelled real video for semantic segmentation is scarce §1 §5.2.
• Convolutional networks do not exploit the temporal redundancy between frames §1.
• Recurrent video autoencoders that flatten frames into vectors, with fully connected LSTM gates, carry many parameters and lose spatial structure §2.
• Optical flow methods such as DeepFlow and FlowNet are supervised for matching frame pairs, then smooth variationally; the authors' setup is end to end and needs no labelled flow §2.
• Global motion models (for example, ego-motion) cannot describe scenes where objects move differently; a dense flow map can §2.

Method

The network is a spatial autoencoder (encoder E, decoder D) with a temporal part nested inside it. At each time step it takes frame Y_t, produces a predicted next frame, and is trained on the error to frame Y_{t+1} §3 fig. 1.

1. Spatial encoder E. One convolution, a tanh non-linearity, then 2×2 max-pooling. It maps the frame to d feature maps of size h×w §3.1.
2. Memory (convolutional LSTM). The fully connected products of a standard LSTM are replaced by spatial convolutions, so the memory keeps the layout of the feature maps. The gates and the cell update are given in eq. 1. The recurrence runs only over time; there are no spatial recurrent links §3.2.1.
3. Optical flow regressor Θ. Two convolutions with large 15×15 kernels map memory features to a two-channel dense map (one horizontal and one vertical displacement per pixel). A large kernel is used because the flow magnitude is limited by the filter size §3.2.2.
4. Smoothness penalty H. An identity in the forward pass. In the backward pass it adds the gradient of a Huber penalty on the spatial gradient of the flow map, to push the flow towards local smoothness §3.2.2 eq. 2.
5. Grid generator GG and sampler S. Based on the Spatial Transformer Network sampler, but with one translation per pixel instead of one transform per image. The source position of each output position is its own position plus the flow §3.2.3 eq. 3.
6. Warp and decode. The sampler warps the encoder features of the current frame with the flow to predict the next feature map, and decoder D turns it back into a frame. D mirrors E, with nearest-neighbour upsampling and no non-linearity §3.1 §3.2.3.
7. Loss. The reconstruction error between the predicted and true next frame, plus w_H times the flow smoothness penalty §3.2.4 eq. 4.

Equation for the memory cell eq. 1, with ∗ convolution and ⊙ element-wise product:

\begin{aligned} i_t &= \sigma(x_t \ast w_{xi} + h_{t-1} \ast w_{hi} + w_{i\text{bias}}) \\ f_t &= \sigma(x_t \ast w_{xf} + h_{t-1} \ast w_{hf} + w_{f\text{bias}}) \\ \tilde{c}_t &= \tanh(x_t \ast w_{x\tilde{c}} + h_{t-1} \ast w_{h\tilde{c}} + w_{\tilde{c}\text{bias}}) \\ c_t &= \tilde{c}_t \odot i_t + c_{t-1} \odot f_t \\ o_t &= \sigma(x_t \ast w_{xo} + h_{t-1} \ast w_{ho} + w_{o\text{bias}}) \\ h_t &= o_t \odot \tanh(c_t) \end{aligned} eq. 1

Huber penalty used on the flow gradient, with δ = 10⁻³ in the experiments eq. 2 §3.2.2:

\mathcal{H}_\delta(a) = \begin{cases} \tfrac{1}{2} a^2 & |a| \le \delta \\ \delta\,(|a| - \tfrac{1}{2}\delta) & \text{otherwise} \end{cases} \qquad \nabla\mathcal{H}_\delta(a) = \begin{cases} a & |a| \le \delta \\ \delta\,\mathrm{sign}(a) & \text{otherwise} \end{cases} eq. 2

Sampling position from the flow map T, with T(x, y) = [[1, 0, t_x], [0, 1, t_y]] eq. 3:

\begin{pmatrix} x_s \\ y_s \end{pmatrix} = \mathcal{T}(x_o, y_o) \begin{pmatrix} x_o \\ y_o \\ 1 \end{pmatrix} eq. 3

Training objective per step, with w_H = 10⁻² eq. 4:

\mathcal{L}_t = \lVert \tilde{Y}_{t+1} - Y_{t+1} \rVert_2^2 + w_H\, \mathcal{H}(\nabla \mathcal{T}) eq. 4

Step-by-step forward pass, as described in the text. The paper gives no pseudocode, so this follows §3 and fig. 1.

# One time step of the video autoencoder  [§3] [fig. 1]
# Input:  frame_t (B, 1, H, W); memory (h_prev, c_prev) from step t-1
# Output: next_frame (B, 1, H, W); memory (h_t, c_t)

def step(frame_t, h_prev, c_prev):
    # Spatial encoder E: conv, tanh, 2x2 max-pool  [§3.1]
    x_t = max_pool(tanh(conv(frame_t)))           # (B, d, h, w)

    # Convolutional LSTM memory, same layout as input  [§3.2.1] [eq. 1]
    h_t, c_t = conv_lstm(x_t, h_prev, c_prev)     # (B, d_m, h, w)

    # Dense optical flow from memory, two large convolutions  [§3.2.2]
    flow = flow_regressor(h_t)                    # (B, 2, h, w)

    # Warp the features of frame t with the flow  [§3.2.3] [eq. 3]
    grid = base_grid(h, w) + flow                 # (B, 2, h, w)
    x_next = sample(x_t, grid)                    # (B, d, h, w)

    # Spatial decoder D: nearest upsampling, conv, no non-linearity  [§3.1]
    next_frame = decoder(x_next)                  # (B, 1, H, W)
    return next_frame, flow, h_t, c_t

Training loop for one sequence. The paper states the update rule and the loss, but not how step losses are combined or how sequences are sampled. Those choices are left open here.

# Training on one video sequence  [§4] [§3.2.4]
# Input:  frames y_1 .. y_T (values in [0, 1] for MNIST)
# Output: updated parameters

def train_sequence(frames, params):
    # Memory starts empty and is reset after each sequence  [§4]
    h, c = zeros_memory()
    loss = 0
    # Each frame predicts the next one  [§3.2.4]
    for t in range(1, T):
        pred, flow, h, c = step(frames[t], h, c)
        # Squared error to the true next frame  [eq. 4]
        loss = loss + sum((pred - frames[t + 1]) ** 2)
        # Huber penalty on the flow gradient  [eq. 2] [eq. 4]
        # (not in the released MNIST graph, see Paper vs code)
        loss = loss + w_h * sum(huber(flow_gradient(flow), delta=1e-3))
    # RMSprop with clipping; clip range not stated  [§4]
    grads = clip(gradients(loss, params), -clip_value, clip_value)
    params = rmsprop_update(params, grads, lr=1e-4)   # [§4]
    return params

Hyperparameters

Values for the paper's experiments, with the released demo (moving MNIST) noted where it differs. Learning-rate schedule and other RMSprop constants are in Paper vs code.

Results

Author-reported. Setup: moving MNIST, 64×64, four architectures; per-architecture training details are not stated, and the error is a per-pixel average. Table 1 gives the errors; the paper does not define the error metric beyond "per-pixel average error" tab. 1.

Motion check: sampler warping on Sintel, with ground-truth flow and no encoder or decoder, gives an average per-pixel error of 0.004 on values in [0, 1]. Artefacts appear at occlusion boundaries §5 fig. 2.

Weakly supervised segmentation on CamVid, 12-class labelling, 240×180 frames. Setup: SegNet as the per-frame segmenter, and SegNet-flow, which propagates SegNet predictions across frames with the learned flow and is trained jointly with the reconstruction loss. Test-set accuracy in per cent §5.2 tab. 2:

Source: tab. 2. The text says the gains are for large classes and losses for thin structures §5.2. Qualitative results: moving MNIST predictions and flow fig. 5; real-video flow fig. 6; CamVid segmentations fig. 8. The real-video results have no numbers §5.1.

Paper vs code

Official code: the paper's bibliography links https://github.com/viorik/ConvLSTM (the owner name matches the first author), read at commit 03dbdbd7af5b4f3645679e1abc042c29b632dcc1. The README says it is a demo on a modified moving-MNIST set and that its model "does not explicitly estimate the optical flow" README.md. The released code therefore covers the MNIST setting only; no real-video, CamVid or SegNet code is included. The Torch dependencies (rnn, extracunn, stn) are not in the repository README.md.

• Loss. Paper: squared error on intensities plus a Huber flow term eq. 4; and for moving MNIST, binary cross-entropy §5.1. Code: squared error between the spatial gradients of the predicted and true frames, with no intensity term model.lua main-mnist.lua. The Huber term is not in the flow graph flow.lua model.lua.
• Smoothness penalty. Paper: active with δ = 10⁻³ and w_H = 10⁻² eq. 2 eq. 4. Code: the module exists with default threshold 0.001 and L1 weight 0.01 SmoothHuberPenalty.lua, but its use is commented out flow.lua and the config value 0.001 is unused opts-mnist.lua. The sizeAverage or true default is always true, so the normalisation is fixed SmoothHuberPenalty.lua.
• Forget-gate bias. Paper: 1 §4. Code: initBias(0,0), so the forget bias is 0 main-mnist.lua ConvLSTM.lua.
• Learning-rate decay. Paper: ×0.9 every 5 epochs §4. Code: the rate is multiplied by 0.5^(epoch/50), where the epoch counter advances every 20,000 iterations main-mnist.lua.
• Filter counts. Paper text: 16 encoder and decoder filters, 64 memory filters §3.3. Code: 32 and 45 opts-mnist.lua. The code's flow output channels (2) and kernels (15×15) match flow.lua opts-mnist.lua.
• First time step. The released model uses an untied first step with no memory input, which the paper does not describe UntiedConvLSTM.lua model.lua.
• Flow scaling and clamp. Not in the paper. The flow is multiplied by 2/H and 2/W to fit the sampler's [−1, 1] coordinates, then clamped to ±0.5 flow.lua opts-mnist.lua. At 32×32 memory, ±0.5 is a quarter of the map width, i.e. 8 cells.
• Warp direction. The grid is the base grid plus the flow, so the flow gives the source offset for each output pixel, a pull convention DenseTransformer2D.lua. This matches eq. 3 and the text's "source position". Units are not stated in the paper.
• Gradient clipping. Paper: "predefined range". Code: ±50 applied to the full gradient main-mnist.lua opts-mnist.lua.
• Batch and sampling. Paper: not stated. Code: one sequence per update, with every frame but the last as input and the last as the target main-mnist.lua data-mnist.lua. The data are scaled to [0, 1] data-mnist.lua.
• Optimizer constants. Only the learning rate is passed to RMSprop; other constants are not set in the repository main-mnist.lua main-mnist.lua.
• Data. The MNIST data file is loaded from disk and is not in the repository data-mnist.lua; the README points to the author's site README.md.

No discrepancies found in the spatial encoder, decoder, the convolutional LSTM gates eq. 1, or the warp operation, beyond the items above.

Pitfalls

• Loss mismatch. The paper's equation eq. 4 has an intensity term, but the released MNIST code trains on spatial gradients only. An implementation that follows the equation will not match the code's behaviour, and vice versa.
• Smoothness penalty is missing from the code. Implement it from eq. 2 and eq. 4 only if you want the paper's version; the released model has no active flow smoothness.
• Forget bias. Paper says 1; code uses 0. Pick one and say which.
• Filter counts and parameter count. The 1,035,067 total §3.3 tab. 1 comes from the paper's configuration, and the released code's filter counts differ. Do not assume the code reproduces Table 1's count without checking.
• Undefined loss and metric. The per-pixel error in Table 1 is not defined, and the text says binary cross-entropy for that experiment. Results from a different metric cannot be compared directly.
• Undefined loss over time. The paper does not say how per-step losses are summed or averaged over a sequence, nor how sequences are cut for training.
• Undefined schedule and length. Batch size, number of updates, sequence length and early stopping are not stated for the paper's experiments. The MNIST demo uses 30,000 single-sequence updates, which may not match the paper's runs.
• Flow magnitude is bounded. The released code clamps flow to a quarter of the map width. Larger motions cannot be predicted by that model. The paper calls Sintel's frame-to-frame displacements very large and unrealistic for its setting §5.1 flow.lua.
• Flow resolution. Flow is computed at the encoder's downsampled resolution, so small structures may get poor flow. The paper cites this as one reason for weaker results on thin classes §5.2.
• Occlusion artefacts. Warping the features of the last frame produces artefacts at occlusions §5 fig. 2. The decoder is expected to smooth them, which is not measured.
• Two tasks, one encoder. The segmentation variant adds a cross-entropy term on labelled frames to the reconstruction loss §5.2. The weighting of the two terms is not stated.
• Small evaluation. CamVid training uses about 2.5k frames from 6-frame windows, and the paper attributes weak results on small classes partly to this §5.2.
• Hardware. A 12 GB GPU is cited as necessary for the full SegNet-flow setup §5.2. The code assumes cuDNN ConvLSTM.lua.
• Figures. Figure 1 (architecture) and Figure 2 (warp test) carry the clearest description of the data flow; use them with the text.

Lineage

• Builds on: the LSTM §2; the convolutional LSTM idea §3.2.1, which the paper notes was recently used for precipitation nowcasting §3.2.1; the Spatial Transformer Network sampler and grid generator §3.2.3; LSTM video autoencoders with fully connected gates, which it compares against §2 §5.1.
• Compared against: fully connected LSTM video autoencoder tab. 1; plain convolutional autoencoder tab. 1; DeepFlow and FlowNet, cited as supervised optical flow §2; ego-motion features from freely available data §1.
• Superseded by: not stated in this card's sources.