A Log-likelihood Regularized KL Divergence for Video Prediction With a 3D Convolutional Variational Recurrent Network

Figure 1 provides a graphical illustration of the VRNN. The VRNN uses a latent variable $\textbf{z}_{t}$ at each timestep of a recurrent network to capture the variations in the observed data. It contains a VAE at every timestep whose mean $\mu_{t}$ and variance $\sigma_{t}$ are conditioned on the hidden unit $h_{t}$ of a recurrent network. These parameters are then used to sample the latent variable $\textbf{z}_{t}$ at each timestep. Concisely, the forward pass can be completely described by the following set of recurrence equations where the subscripts p and q denote the prior and posterior distributions respectively, and the components $\text{f}_{\text{p}}$, $\text{f}_{\text{q}}$, $\text{f}_{\text{enc}}$, $\text{f}_{\text{dec}}$ are functions implemented using neural networks.

	$\displaystyle\mu_{\text{p},t},\sigma_{\text{p},t}$	$\displaystyle=\text{f}_{\text{p}}(h_{t-1})$		(1)
	$\displaystyle\mu_{\text{q},t},\sigma_{\text{q},t}$	$\displaystyle=\text{f}_{\text{q}}(h_{t-1},\text{f}_{\text{enc}}(x_{t}))$		(2)
	$\displaystyle z_{\text{p},t}$	$\displaystyle\sim N(\mu_{\text{p},t},\sigma_{\text{p},t})$		(3)
	$\displaystyle z_{\text{q},t}$	$\displaystyle\sim N(\mu_{\text{q},t},\sigma_{\text{q},t})$		(4)
	$\displaystyle\hat{x}_{t}$	$\displaystyle=\text{f}_{\text{dec}}(z_{\text{p},t},h_{t-1})$		(5)
	$\displaystyle h_{t}$	$\displaystyle=\text{LSTM}(\text{f}_{\text{enc}}(x_{t}),h_{t-1},z_{\text{p},t})$		(6)

Here $N(\mu,\sigma)$ is a multivariate Gaussian distribution with mean $\mu$ and co-variance diag$(\sigma^{2})$. Note that the posterior network $\text{f}_{\text{q}}$ is used only during training and is discarded at test time. The entire model is then trained end-to-end for future frame prediction by minimizing a sum of the reconstruction loss ($L_{\text{rec}}$), and latent loss ($L_{\text{latent}}$) expressed as:

$\displaystyle L=\lambda_{\text{rec}}L_{\text{rec}}+\lambda_{\text{latent}}L_{\text{latent}}$

(7)

where $\lambda_{\text{rec}}$ and $\lambda_{\text{latent}}$ are the trade off hyper-parameters and the latent loss is the timestep-wise KL divergence ($L_{\text{KL}}$) between the prior ($p$) and posterior ($q$) distributions and is expressed as:

	$\displaystyle L_{\text{KL}}=\sum_{t=1}^{T}\text{KL}(q(z_{t}|X_{\leq t},Z_{<t})||p(z_{t}|X_{<t},Z_{<t})))=$		(8)
	$\displaystyle\sum_{t=1}^{T}\log(\sigma_{\text{q,t}})-\log(\sigma_{\text{p,t}})+\frac{\sigma_{\text{p,t}}^{2}+(\mu_{\text{p,t}}-\mu_{\text{q,t}})^{2}}{2\sigma_{\text{q,t}}^{2}}-0.5$		(9)

Typically, the KL divergence-based latent loss is used to regularize the latent space, enforcing it to be a Gaussian distribution with known parameters. We further enhance this regularization by appending the negative of the log-likelihood term to the latent loss. The objective here is to maximize the likelihood of the prior mean distribution w.r.t. the posterior. This is done by minimizing the negative likelihood as shown in Eq. 11 by assuming that the prior, posterior and the conditional prior mean given the posterior all follow a Gaussian distribution.

	$\displaystyle-L_{\text{LL}}$	$\displaystyle=-\log\prod_{t=1}^{T}p(\mu_{\text{p},t}|\mu_{\text{q},t},\sigma_{\text{q},t})$		(10)
		$\displaystyle=\sum_{t=1}^{T}\log({\sigma_{\text{q},t}})+(\frac{\mu_{\text{p},t}-\mu_{\text{q},t}}{\sigma_{\text{q},t}})^{2}$		(11)

The proposed latent loss is thus expressed together as:

	$\displaystyle L_{\text{KL}}-L_{\text{LL}}$
	$\displaystyle=\log(\sigma_{\text{q}})-\log(\sigma_{\text{p}})+\frac{\sigma_{\text{p}}^{2}+(\mu_{\text{p}}-\mu_{\text{q}})^{2}}{2\sigma_{\text{q}}^{2}}-0.5$
	$\displaystyle+\log({\sigma_{\text{q}}})+(\frac{\mu_{\text{p}}-\mu_{\text{q}}}{\sigma_{\text{q}}})^{2}$		(12)
	$\displaystyle=2\log(\sigma_{\text{q}})-\log(\sigma_{\text{p}})+\frac{\sigma_{\text{p}}^{2}+3(\mu_{\text{p}}-\mu_{\text{q}})^{2}}{2\sigma_{\text{q}}^{2}}-0.5$		(13)

Interestingly, the above equation is similar to the KL divergence (eq 9) but with some of its components weighted differently. In particular, the log posterior variance, $\log(\sigma_{q})$, has been scaled by a factor of 2 and the squared difference of mean, $(\mu_{p}-\mu_{q})^{2}$, by a factor of 3. This modification has 2 effects. First, it puts more emphasis on sample diversity since the log prior variance $\log(\sigma_{p})$ put out by the network must now be higher in order to match the scaled log posterior variance $2\log(\sigma_{q})$. Secondly, the scaled difference of mean $3(\mu_{p}-\mu_{q})^{2}$ serves to balance out the additional weight assigned to the variance term and thus encourages the model to continue generating samples that are representative of the dataset. As such, the log-likelihood regularized KL divergence should have no adverse effects on the model since it is simply the KL divergence with a reweighting of its components and would argue it to be more forceful if one needs to have a greater emphasis on sample diversity. Interestingly, the weights for each component can also be customized although their individual effects will not be investigated since it is not the purpose of this paper. All-in-all, our new loss function for training the VRNN is expressed together as: $L=\lambda_{\text{rec}}L_{\text{rec}}+\lambda_{\text{latent}}(L_{\text{KL}}-L_{\text{LL}})$.

The ConvLSTM was proposed in [26] to address the shortcomings of Fully-Connected LSTM, namely that latter always ends up decimating any spatial information contained in the input tensor. Intuitively, if the states are viewed as the hidden representations of moving objects, then a ConvLSTM with a larger kernel should be able to capture faster motions while one with a smaller kernel can capture slower motions. However, if the input at each timestep is a single image, then the hidden states are the only component that carry motion information in both Fully-Connected and ConvLSTMs.

In our work, we counteract this limitation by replacing all 2D convolutions (and de-convolutions) with 3D to enable every component to retain motion information instead of only the hidden states. The benefits of this are two-fold. First, the 3D ConvLSTM is no longer completely reliant on the hidden states for motion information since there is an additional source coming from the 3D image encoder. Specifically, the use of 3D convolutions on multiple frames result in an input tensor that carries short-term spatiotemporal information as opposed to a VRNN that run a 2D convolution on a single frame at every timestep. Second, we can now vary the window size and output horizon at each timestep without needing to redesign the architecture. For example, let us define M to be the window size, or the number of input frames to our model at each timestep and $\mathcal{H}$ the output horizon (output frames), or how far into the future should the model predict. Then, 3D convolutions (and de-convolutions) allow us to set a large M to efficiently capture large motions when dealing with datasets where the motion between frames is prevalently large and conversely, a large $\mathcal{H}$ to predict many frames into the future at once with minimal reconstruction errors if said motion between frames is small. All in all, this upgrade renders our 3D convolutional VRNN more effective and general than its 2D counterpart.

Our proposed architecture is shown in Figure 2. The encoder (3D-ENC) takes at each timestep a clip of M video frames of shape [M,H,W,C] to produce a tensor of shape [M/2, H/4, W/4, C/4] where H,W,C denote the height, width and channels respectively. This tensor is then concatenated with the latent tensor Z (a zero tensor in the 1st timestep) along the 4’th channel indicated by $[\circ,\circ]$ then passed to the 3D ConvLSTM with 2 hidden layers for motion learning. The LSTM hidden states at the second level with a shape of [M/2, H/4, W/4, 128] are then fed through a shallow 3D CNN with two heads to produce the parameters of the prior distribution with shape [M/2, H/4, W/4, 16] that are later sampled to produce the latent tensor Z. This latent tensor is then concatenated with the hidden states and finally propagated through the decoder (3D-DEC) to predict the frames H timesteps into the future. The model is applied recursively by using the newly generated frame as input if the ground truth is not available. Specifically, if the frames are observed up to time C, then the model will use the ground truth as input up till time C-M then a combination of the ground truth and predicted frames between time C-M to C, and then finally, only the predicted frames as input from time C onwards. During training, a separate 3D encoder (not shown in Figure 2) is used to generate the posterior distribution to optimize the KL divergence.

As shown in Figure 3, we use a 2-block 3D ResNet-18 for both the encoder and decoder in contrast to [3] that use a full ResNet. We find this to be sufficient especially since the 3D ConvLSTM itself serve as an extension of the 3D CNN for learning complex spatiotemporal features given a window of M frames. Furthermore, by truncating the number of blocks to 2, we reduce the total number of parameters significantly which allow us to devote additional resources to our 3D ConvLSTM with 128 hidden units that contains 7m parameters per level. However, we also want to avoid the other end of not having a feature extractor at all [37, 35] since they have been shown to extract useful features that tend to be task specific at the higher blocks and more general purpose at the lower blocks. In short, we propose to use a smaller feature extractor CNN and a larger LSTM. We show in our experiments that modelling the architecture in such a manner allows us to outperform the state-of-the-art while requiring fewer parameters.

We compare our approach to several state-of-the-art methods using publicly available source code and model where available with default parameters and using standard metrics such as frame-wise Mean Squared Error (MSE), Structural Similarity Index (SSIM) and Peak Signal to Noise Ratio (PSNR). We sample 50 predictions from the stochastic models for each ground truth test sequence and average the metrics across the test set. Note that sampling is done only for the purpose of evaluation in order to get the average performance and is not required for deployment.

Table 1 shows the performance of the models when using 5 frames to predict 10 and 15 frames into the future and when using 10 frames to predict 10 and 20 frames into the future. Our method demonstrates its promise, outperforming both the state-of-the-art deterministic (E3D-LSTM [35]) and stochastic (VRNN [4]) models, with the latter by a large margin. We also improve over the E3D-LSTM [35] despite having fewer parameters. We were only able to train the smallest variant of the models presented in [3] which nevertheless contains 62 million parameters. Interestingly, we also outperform them. These results thus indicate the impact of our new design and the novel latent loss.

We present some visual results in Figure 5 where the first row illustrates the input sequence $\textbf{x}_{1:10}$, the second the ground truth for the predicted sequence $\textbf{x}_{11:20}$ and all subsequent rows the predictions made by the various models $\hat{\textbf{x}}_{11:20}$. It can firstly be seen that the injection of stochasticity causes the Variational ConvLSTM to output predictions that are blurrier than its deterministic counterpart. This could principally be due to the fact that information coming from the latent nodes act as noise and thus interferes with reconstruction. Unlike the ConvLSTM however, the digits generated by the Variational ConvLSTM are closer to the ground truth resulting in a performance that is generally superior as evidenced by the quantitative scores in Table 1. We can then observe our model producing the best results. This signals that the blurry reconstructions manifested by the Variational ConvLSTM are counteracted by replacing all 2D convolutions with 3D. Intuitively, our novel loss also acts as a stronger regularizer for the reconstructions.

Effectiveness of each component: We quantize the effect of each contribution in Table 3 on the moving MNIST dataset. It can be seen that the introduction of stochasticity into the recurrent network allows it to better tackle uncertainties in the recurrent dynamics which results in better predictions, significantly lowering the MSE from 82.2 to 60.7. Additionally, swapping out the 2D convolutions in place for 3D brings about significant improvements to the model, lowering the MSE by approximately another 20 points. This makes sense since 3D convolutions operate on both the spatial and temporal axis, letting the architecture capture relationships in said dimensions. Finally, it is quite apparent that the introduction of the log-likelihood criterion has a noticeable effect, further bringing down the MSE by approximately 2 points. This can be interpreted in various ways: (1) that the resulting loss function empirically helps the network traverse towards a better local minima and (2) that the added regularizer helps the recurrent model strengthen its ability at expressing complex distributions. Intuitively, appending the log-likelihood criterion to the KL divergence has some conceptual similarity to the use of the L1+L2 loss functions that has been empirically shown in [37, 35] to be better than the individual counterparts. In conclusion, each component brings a definitive upgrade to the model and together, lend it the advantage it needs to outperform the deterministic and stochastic state-of-the-art models [4, 35, 37, 3].