Diffusion Models for Video Prediction and Infilling
Abstract
Video prediction and infilling require strong, temporally coherent generative capabilities.
Diffusion models have shown remarkable success in several generative tasks, but have not been extensively explored in the video domain.
We present Random-Mask Video Diffusion (RaMViD), which extends image diffusion models to videos using 3D convolutions, and introduces a new conditioning technique during training.
By varying the mask we condition on, the model is able to perform video prediction, infilling, and upsampling.
Due to our simple conditioning scheme, we can utilize the same architecture as used for unconditional training, which allows us to train the model in a conditional and unconditional fashion at the same time.
We evaluate the model on two benchmark datasets for video prediction, on which we achieve state-of-the-art results, and one for video generation.
High-resolution videos are provided at https://sites.google.com/view/video-diffusion-prediction
Introduction
Videos contain rich information about the world, and a vast amount of diverse video data is available.
Training models on this data for video prediction or video infilling—i e., observing a part of a video and generating missing frames —can be used in planning, estimating trajectories and video processing.
In addition, video models can be valuable for downstream tasks such as action recognition [17] and pose estimation [26].
Video prediction can be modelled in a deterministic or stochastic form.
Deterministic modelling [34, 35, 41, 42] tries to predict the most likely future but this often leads to averaging the future states [19].
Therefore, most recent methods are based on generative modeling, either using variational methods [2, 3, 8, 30, 44] or GANs [7, 20].
Diffusion models [1, 10, 13, 21, 23, 31, 32] have seen tremendous progress on static visual data, even outperforming GANs in image synthesis [9].
Only a few concurrent works have recently considered diffusion models for video generation.
[47] uses diffusion models for autoregressive video prediction, by modeling residuals for a predicted frame, [14] focuses on unconditional video generation, [12] uses diffusion models to predict long videos, and [39], the most closely related to our work, also considers video prediction and infilling.
The essence of diffusion models are two stochastic (diffusion) processes implemented by Stochastic Differential Equations (SDEs), a forward and a backward one.
Let x 0 ∈ R
d be a sample from the empirical data distribution, i.e., x 0 ∼ p data ( x 0 ) and d be the data dimension.
The forward diffusion process takes x 0 as the starting point and creates the random trajectory x [0 ,T ] from t = 0 to the final time t = T .
The forward process is designed such that p ( x T | x 0 ) has a simple unstructured distribution.
One example of such SDEs is
[Equation]
[Equation]
where w is the Brownian motion.
A desirable property of this process is the fact that the conditional distribution p ( x t | x 0 ) takes a simple analytical form:
[Equation]
Upon learning the gradient of p ( x t ) for each t , one can reverse the above process and obtain the complex data distribution from pure noise as
[Equation]
where w ′ is a Brownian motion independent of the one in the forward direction.
Hence, generating samples from the data distribution boils down to learning ∇ x log p ( x ) .
The original score matching objective [15]:
[Equation]
is the most intuitive way to learn the score function, but is unfortunately intractable.
Denoising Score
Matching (DSM) provides a tractable alternative objective function:
[Equation]
whose equivalence with the original score matching objective was shown by [38] and used to train energy models in [29].
Similar to many recent works, we use the DSM formulation of score matching in this work to learn the score function.
In this paper, we extend diffusion models to the video domain via several technical contributions.
We use 3D convolutions and a new conditioning procedure incorporating randomness.
Our model is not only able to predict future frames of a video but also fill in missing frames at arbitrary positions in the sequence.
Therefore, our Random-Mask Video Diffusion (RaMViD) can be used for several video completion tasks.
We summarize our technical contributions as follows:
• A novel diffusion-based architecture for video prediction and infilling. • Competitive performance with recent approaches across multiple datasets. • Introduce a schedule for the random masking.
Random-Mask Video Diffusion
Our method, Random-Mask Video Diffusion (RaMViD), consists of two main features.
First, the way we introduce conditional information to the network is different from what has been used so far.
Second, by randomizing the mask, we can directly use the same approach for video prediction and video completion (infilling).
In the following, we detail each of these aspects of the proposed method.
2.1 Conditional training
Let x 0 ∈ R L,W,H,C be a video with length L .
We partition the video x 0 into two parts, the unknown frames x U 0 ∈ R L − k, [Equation] ,H,C and the conditioning frames x C 0 ∈ R
k,W,H,C , where U and C are sets of indices such that U ∩ C = ∅ and U ∪ C = { 0 , 1 , . . .
, L − 1 } .
We write x 0 = x
U 0 ⊕ x C 0 with the following definition for the ⊕ operator:
[Equation]
where the superscript i indicates tensor indexing and in our case corresponds to selecting a frame from a video, and the subscript t indicates the diffusion step, with t = 0 corresponding to the data and
t = T to the prior Gaussian distribution.
If we use an unconditionally trained model and sample via the replacement method [32], we find that the predicted unknown frames x U 0 do not harmonize well with the conditioning frames x C 0 .
To mitigate this issue, we propose to train the model conditionally with randomized masking .
Conditional diffusion models usually optimize
[Equation]
where x C 0 is typically given as a separate input through an additional layer [6] or it is expanded to the dimension of x t (e.g., via padding) and concatenated with the input [4, 24, 25].
We, on the other hand, feed the entire sequence to the network s θ but only add noise to the unmasked frames: x U t ∼ N (cid:0) x U 0 , (cid:0) σ 2 ( t ) − σ 2 (0) (cid:1) I (cid:1) .
The input to the network is then a video where some frames are noisy and some are clean: x t = x
U t ⊕ x C 0 .
The loss is computed only with respect to x U t :
[Equation]
Note that the score function ∇ x U t log p ( x U t | x 0 ) has the same dimension as x U t , whereas in Eq.
(7) it had the dimension of the entire video x t .
The reversed diffusion process then becomes:
[Equation]
2.2 Randomization
As previously mentioned, the proposed model is able to perform several tasks.
We achieve this by sampling C at random.
At each training step, we first choose the number of conditioning frames |C| = k ∈ { 1 , . .
, K } , where K is a chosen hyperparameter.
Then we define C by selecting k random indices from { 0 , . . .
, L − 1 } , and we refrain from applying the diffusion process to the corresponding frames.
After training, we can use RaMViD by fixing C to the set of indices of the known frames ( C can be any arbitrary subset of { 0 , . . .
, L − 1 } ) and generating the unknown frames (those with indices in U ).
Our approach allows us to use the exact same architecture of unconditionally trained models, thus enabling mixed training , where we train the model conditionally and unconditionally at the same time.
We set C = ∅ (i.e., the model does not have any conditional information x C t ) with probability p U , which is a fixed hyperparameter.
If C = ∅ , our objective in Eq.
(8) becomes the same as the objective in Eq.
(5) used for unconditional training.
Experiments
To compare our model to prior work, we train it on three datasets: BAIR robot pushing [11] and
Kinetics-600 [5] for video prediction and completion and UCF-101 for unconditional generation [33].
We train all datasets on 64 × 64 resolution and choose K = 4 .
To quantitatively evaluate prediction, we use the Fréchet Video Distance (FVD) [37], 1 and to evaluate unconditional generation the Inception Score (IS) [28] with the implementation from [27].
3.1 BAIR
We train four models on the BAIR dataset on 20 frames with p U ∈ { 0 , 0 . 25 , 0 . 5 , 0 . 75 } respectively.
The models are trained for 250,000 iterations with a batch size of 32 on 8 GPUs.
First, we test our method with the typical evaluation protocol for BAIR (predicting 15 frames, given one conditional frame).
With all values of p U , we can achieve state-of-the-art performance, as shown in Table 1.
By using p U > 0 , we can even increase the performance of our method.
However, it seems that there is a tipping point after which the increasing unconditional rate hurts the prediction performance of the model.
Thanks to the randomized masking, RaMViD is also able to perform video infilling.
Quantitative results are shown in Appendix B.
The method works very well for prediction and infilling.
However, since the BAIR dataset is arguably rather simple and not very diverse, we will now evaluate RaMViD on the significantly more complex Kinetics-600 dataset.
3.2 Kinetics-
For the Kinetics-600 dataset, we increase the batch size to 64 and train for 500,000 iterations on[-
-]
GPUs, but train only on 16 frames.
First, we evaluate the model on prediction (predict 11 frames given 5 frames).
When comparing our models to concurrent work, we find that RaMViD achieves state-of-the-art results by a significant margin (see Table 2). Nevertheless, it struggles with fast
[Table]
[Table]
movements: objects moving quickly often get deformed.
Similar to what we have seen in Table 1, having an unconditional rate p U > 0 increases the performance up to a tipping point.
However, differently from the model trained on BAIR, the FVD score now drops significantly with p U = 0 .
We conjecture that this drop in performance is due to the complexity of the data distribution.
In
BAIR, the conditional and unconditional distributions are rather similar, while this is not true for
Kinetics-600.
Also on Kinetics-600 we can perform several video completion tasks.
For further experiments, we refer to Appendix C. Since we were able to generate videos unconditionally with the models
RaMViD ( p U = 0 . 5 ) and RaMViD ( p U = 0 . 75 ), we show quantitative comparision on
UCF-101 in Appendix D.
Conclusion
We have shown that diffusion models, which have been demonstrated to be remarkably powerful for image generation, can be extended to videos and used for several video completion tasks.
The way we introduce conditioning information is novel, simple, and does not require any major modification to the architecture of existing diffusion models, but it is nonetheless surprisingly effective.
Although the proposed method targets conditional video generation, we also introduce an alternative masking schedule in an attempt to improve the unconditional generation performance without sacrificing performance on conditional generation tasks.
Finally, the focus of this work has been on the diffusionbased algorithm for videos rather than on optimizing the quality of each frame.
It has been shown in concurrent works that including super-resolution modules helps create high-resolution videos.
Adding a super-resolution module to RaMViD would be a relevant direction for future work.
A Implementation details
Our implementation relies on the official code of [23], and we use the same U-Net architecture but adapted to video data by using 3D convolutions.
3 We do not encode the time dimension and we use two ResNet blocks per resolution for the BAIR dataset, and three blocks for Kinetics-600 and
UCF-101.
We set the learning rate for all our experiments to 2e-5, use a batch size of 32 for BAIR and 64 for Kinetics-600 and UCF-101, and fix T = 1000 .
We found, especially on the more diverse datasets like Kinetics-600 and UCF-101, that larger batch sizes produce better results.
Therefore, to increase the batch size, we use gradient accumulation by computing the gradients for micro-batches of size 2 and accumulate for several steps before doing back-propagation.
Even though most previous work uses the cosine noise schedule, we found that the linear noise schedule works better when training the model conditionally.
A detailed graphic of our model is shown in Fig.
B BAIR
Since we train with randomized masking, we can also perform video infilling with the same models, without retraining.
We condition on the first and last frame (i.e., set C = { 0 , 15 } for sampling) and compute the FVD of the 14 generated frames.
Again we find that the performance is very similar for different values of p U (see Table 3), however, similarly to Table 1, we observe the best results when using p U = 0 .
[Table]
C Kinetics-
We also evaluate RaMViD on two video completion tasks on Kinetics-600.
The first task is to fill in a video given the two first and last frames (i.e., C = { 0 , 1 , 14 , 15 } ): the challenge here is to harmonize the observed movement at the beginning with the movement observed at the end.
In the second task, the conditioning frames are distributed evenly over the sequence (i.e., C = { 0 , 5 , 10 , 15 } ), hence the model has to infer the movement from the static frames and harmonize them into one realistic video.
RaMViD excels on both tasks, as shown quantitatively in Table 4.
Especially when setting
C = { 0 , 5 , 10 , 15 } RaMViD is able to fill the missing frames with very high quality and coherence.
This setting can be easily applied to upsampling by training a model on high-FPS videos and then sampling a sequence conditioned on a low-FPS video.
[Table]
D UCF-
We have mentioned that we found that RaMViD ( p U = 0 . 5 ) and RaMViD ( p U = 0 . 75 ) can generate unconditional videos on Kinetics-600.
To quantify RaMViD’s unconditional generation, we will evaluate these models on the UCF-101 dataset and compare it to other work.
We train RaMViD with the same setting as used for Kinetics-600 but for 450,000 iterations.
Table 5 shows that our model achieves competitive performance on unconditional video generation, although it does not reach state-of-the-art.
The trained models can successfully generate scenes with a static background and a human performing an action in the foreground, consistent with the training dataset.
However, the actions are not always coherent and moving objects can deform over time.
Note that UCFis a very small dataset given its complexity.
Therefore we do observe some overfitting.
Since for each action we only have around 25 different settings, our model does not learn to combine those but generates very similar videos to the training set.
Due to the characteristics of this dataset we think with more extensive hyperparameter tuning, one can achieve better results with RaMViD in unconditional generation.
But our focus does not lie on this.
[Table] [Figure]