Predicting Video with VQVAE

1 INTRODUCTION . When it comes to real-world image data , deep generative models have made substantial progress . With advances in computational efficiency and improvements in architectures , it is now feasible to generate high resolution , realistic images from vast and highly diverse datasets ( Brock et al. , 2019 ; Razavi et al. , 2019 ; Karras et al. , 2017 ) . Apart from the domain of images , deep generative models have also shown promise in other data domains such as music ( Dieleman et al. , 2018 ; Dhariwa et al. , 2020 ) , speech synthesis ( Oord et al. , 2016 ) , 3D voxels ( Liu et al. , 2018 ; Nash & Williams , 2017 ) , and text ( Radford et al. , 2019 ) . One particular fledgling domain is video . While some work in the area of video generation ( Clark et al. , 2020 ; Vondrick et al. , 2016 ; Saito & Saito , 2018 ) has explored video synthesis—generating videos with no prior frame information— many approaches actually focus on the task of video prediction conditioned on past frames ( Ranzato et al. , 2014 ; Srivastava et al. , 2015 ; Patraucean et al. , 2015 ; Mathieu et al. , 2016 ; Lee et al. , 2018 ; Babaeizadeh et al. , 2018 ; Oliu et al. , 2018 ; Xiong et al. , 2018 ; Xue et al. , 2016 ; Finn et al. , 2016 ; Luc et al. , 2020 ) . It can be argued that video synthesis is a combination of image generation and video prediction . In other words , one could decouple the problem of video synthesis into unconditional image generation and conditional video prediction from a generated image . Therefore , we specifically focus on video prediction in this paper . Potential computer vision applications of video forecasting include interpolation , anomaly detection , and activity understanding . More generally , video prediction also has more general implications for intelligent systems—the ability to anticipate the dynamics of the environment . The problem is thus also relevant for robotics and reinforcement learning ( Finn et al. , 2016 ; Ebert et al. , 2017 ; Oh et al. , 2015 ; Ha & Schmidhuber , 2018 ; Racanire et al. , 2017 ) . Approaches toward video prediction have largely skewed toward variations of generative adversarial networks ( Mathieu et al. , 2016 ; Lee et al. , 2018 ; Clark et al. , 2020 ; Vondrick et al. , 2016 ; Luc et al. , 2020 ) . In comparison , we are aware of only a relatively small number of approaches which propose variational autoencoders ( Babaeizadeh et al. , 2018 ; Xue et al. , 2016 ; Denton & Fergus , 2018 ) , autoregressive models ( Kalchbrenner et al. , 2017 ; Weissenborn et al. , 2020 ) , or flow based approaches ( Kumar et al. , 2020 ) . There may be a number of reasons for this situation . One is the explosion in the dimensionality of the input space . A generative model of video needs to model not only one image but tens of them in a coherent fashion . This makes it difficult to scale up such models to large datasets or high resolutions . In addition , previous work ( Clark et al. , 2020 ) suggests that video prediction may be fundamentally more difficult than video synthesis ; a synthesis model can generate simple samples from the dataset while prediction potentially forces the model to forecast conditioned on videos that are outliers in the distribution . Furthermore , most prior work has focused on datasets with low scene diversity such as Moving MNIST ( Srivastava et al. , 2015 ) , KTH ( Schuldt et al. , 2004 ) , or robotic arm datasets ( Finn et al. , 2016 ; Ebert et al. , 2017 ) . While there have been attempts to synthesize video at a high resolution ( Clark et al. , 2020 ) , we know of no attempt—excluding flow based approaches—to predict video beyond resolutions of 64x64 . In this paper we address the large dimensionality of video data through compression . Using Vector Quantized Variational Autoencoders ( VQ-VAE ) ( van den Oord et al. , 2017 ) , we can compress video into a space requiring only 1.3 % of the bits expressed in pixels . While this compressed encoding is lossy , we can still reconstruct the original video from the latent representation with a high degree of fidelity . Furthermore , we can leverage the modularity of VQ-VAE and decompose our latent representation into a hierarchy of encodings , separating high-level , global information from details such as fine texture or small motions . Instead of training a generative model directly on pixel space , we can instead model this much more tractable discrete representation , allowing us to train much more powerful models , use large diverse datasets , and generate at a high resolution . While most prior work has focused on GANs , this discrete representation can also be modeled by likelihood-based models . Likelihood models in concept do not suffer from mode-collapse , instability in training , and lack of diversity of samples often witnessed in GANs ( Denton & Fergus , 2018 ; Babaeizadeh et al. , 2018 ; Razavi et al. , 2019 ) . In this paper , we propose a PixelCNN augmented with causal convolutions in time and spatiotemporal self-attention to model this space of latents . In addition , because the latent representation is decomposed into a hierarchy , we can exploit this decomposition and train separate specialized models at different levels of the hierarchy . Our paper makes four contributions . First , we demonstrate the novel application of VQ-VAE to video data . Second , we propose a set of spatiotemporal PixelCNNs to predict video by utilizing the latent representation learned with VQ-VAE . Third , we explicitly predict video at a higher resolution than ever before . Finally , we demonstrate the competitive performance of our model with a crowdsourced human evaluation . 2 BACKGROUND . 2.1 VECTOR QUANTIZED AUTOENCODERS . VQ-VAEs ( van den Oord et al. , 2017 ) are autoencoders which learn a discrete latent encoding for input data x . First , the output of non-linear encoder ze ( x ) , implemented by a neural network , is passed through a discretization bottleneck . ze ( x ) is mapped via nearest-neighbor into a quantized codebook e ∈ RK×D where D is the dimensionality of each vector ej and K is the number of categories in the codebook . The discretized representation is thus given by : zq ( x ) = ek where k = argminj ||ze ( x ) − ej ||2 ( 1 ) Equation 1 is not differentiable ; however , ( van den Oord et al. , 2017 ) notes that copying the gradient of zq ( x ) to ze ( x ) is a suitable approximation similar to the straight-through estimator ( Bengio et al. , 2013 ) . A decoder D , also implemented by a neural network , then reconstructs the input from zq ( x ) . The total loss function for the VQ-VAE is thus : L = ||D ( zg ( x ) ) − x||22 + ||sg [ zg ( x ) ] − e||22 + β||zg ( x ) − sg [ e ] ||22 ( 2 ) Where sg is a stop gradient operator , and β is a parameter which regulates the rate of code change . As in previous work ( van den Oord et al. , 2017 ; Razavi et al. , 2019 ) , we replace the second term in equation 2 and learn the codebook e ∈ RK×D via an exponential moving average of previous values during training : N ( t ) i : = N ( t−1 ) i ∗ γ + n ( t ) i ( 1− γ ) , m ( t ) i : = m ( t−1 ) i ∗ γ + n ( t ) i∑ j ze ( x ) ( t ) i , j ( 1− γ ) , e ( t ) i : = m ( t ) i N ( t ) i ( 3 ) Where γ is a decay parameter and n ( t ) i is the numbers of vectors in zg ( x ) in a batch that will map to ei . 2.2 PIXELCNN MODELS . PixelCNN and related models have shown promise in modeling a wide variety of data domains ( van den Oord et al. , 2016 ; Oord et al. , 2016 ; Kalchbrenner et al. , 2017 ; Weissenborn et al. , 2020 ) . These autoregressive models are likelihood-based—they explicitly optimize negative log-likelihood . They exploit the fact that the joint probability distribution input data x can be factored into a product of conditional distributions for each dimension of the data : Pθ ( x ) = n∏ i=0 pθ ( xi|x < i ) ( 4 ) Where n is the full dimensionality of the data . This factorization is implemented by a neural network , and the exact set of conditional dependencies is determined by the data domain . Image pixels may depend on regions above and to the left of them ( van den Oord et al. , 2016 ) , while temporal dimensions may depend on past dimensions ( Oord et al. , 2016 ; Kalchbrenner et al. , 2017 ; Weissenborn et al. , 2020 ) . 3 METHOD . Our approach consists of two main components . First , we compress video segments into a discrete latent representation using a hierarchical VQ-VAE . We then propose a multi-stage autoregressive model based on the PixelCNN architecture , exploiting the low dimensionality of the compressed latent space and the hierarchy of the representation . 3.1 COMPRESSING VIDEO WITH VQ-VAE . Similar to ( Razavi et al. , 2019 ) , we use VQ-VAE to compress video in a hierarchical fashion . This multi-stage composition of the latent representation allows decomposition of global , high level information from low-level details such as edges or fine motion . For image information ( Razavi et al. , 2019 ) this approach confers a number of advantages . First , the decomposition allows latent codes to specialize at each level . High level information can be represented in an even more compressed manner , and the total reconstruction error is lower . In addition , this hierarchy leads to a naturally modular generative model . We can then develop a generative model that specializes in modeling the high-level , global information . We can then train a separate model , conditioned on global information , that fills in the details and models the low-level information further down the hierarchy . In this paper , we adopt the terminology of ( Razavi et al. , 2019 ) and call the set of high-level latents the top layer and the low-level latents the bottom layer . Consistent with the experimental setup of previous work in video prediction , we deal with 16-frame videos . Most of the videos in our training dataset are 25 frames per second . We use frames at a 256 × 256 resolution . The full video voxel is thus 256 × 256 × 16 . Using residual blocks with 3D convolutions , we downsample the video spatiotemporally . At the bottom layer , the video is downsampled to a quantized latent space of 64× 64× 8 , reducing the spatial dimension by 4 and the temporal dimension by 2 . Another stack of blocks reduces all dimensions by 2 , with a top layer of 32× 32× 4 . Each of the voxels in the layer is quantized into 512 codes with a different codebook for both layers . The decoder then concatenates the bottom layer and the top layer after upsampling using transposed convolutions . From this concatentation as input , the decoder deterministically outputs the full 256 × 256 × 16 video . Overall , we reduce a 256 × 256 × 16 × 3 × log ( 256 ) space down to a 64×64×8× log ( 512 ) +32×32×4× log ( 512 ) space , a greater than 98 % reduction in bits required .

Summary: The authors propose to use a VQVAE-2 setup for video prediction. In particular, they propose a hierarchical discrete latent variable model that compresses videos into a latent space. An autoregressive model is then used to model dynamics in this latent space, which has reduced dimensionality, and can be used together with the VQVAE decoder to predict video.  Empirical results show that this model is comparable to SOTA GAN models and a human evaluation suggests that humans have a preference for the