Clockwork Variational Autoencoders

1 Introduction Video prediction involves generating highdimensional sequences for many steps into the future . A simple approach is to predict one image after another , based on the previously observed or generated images ( Oh et al. , 2015 ; Kalchbrenner et al. , 2016 ; Babaeizadeh et al. , 2017 ; Denton and Fergus , 2018 ; Weissenborn et al. , 2020 ) . However , such temporally autoregressive models can be computationally expensive because they operate in the highdimensional image space and at the frame rate of the dataset . In contrast , humans have the ability to reason using abstract concepts and predict their changes far into the future , without having to visualizing the details in the input space or at the input frequency . Latent dynamics models predict a sequence of learned latent states forward that is then decoded into the video , without feeding generated frames back into the model ( Kalman , 1960 ; Krishnan et al. , 2015 ; Karl et al. , 2016 ) . These models are typically trained using variational inference , similar to VRNN ( Chung et al. , 2015 ) , except that the generated images are not fed back into the model . Latent dynamics models have recently achieved success in reinforcement learning for learning world models from pixels ( Ha and Schmidhuber , 2018 ; Zhang et al. , 2019 ; Buesing et al. , 2018 ; Mirchev et al. , 2018 ; Hafner et al. , 2019b ; a ; 2020 ) . 1All code is publicly available at https : //github.com/vaibhavsaxena11/cwvae . 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . To better represent complex data , hierarchical latent variable models learn multiple levels of features . Ladder VAE ( Sønderby et al. , 2016 ) , VLAE ( Zhao et al. , 2017 ) , NVAE ( Vahdat and Kautz , 2020 ) , and very deep VAEs ( Child , 2020 ) have demonstrated the success of this approach for generating static images . Hierarchical latents have also been incorporated into deep video prediction models ( Serban et al. , 2016 ; Kumar et al. , 2019 ; Castrejón et al. , 2019 ) . These models can learn to separate high-level details , such as textures , from low-level details , such as object positions . However , they operate at the frequency of the input sequence , making it challenging to predict far into the future . Temporally abstract latent dynamics models predict learned features at a slower frequency than the input sequence . This encourages learning long-term dependencies that can result in more accurate predictions and enable computationally efficient long-horizon planning . Temporal abstraction has been studied for low-dimensional sequences ( Koutník et al. , 2014 ; Chung et al. , 2016 ; Mujika et al. , 2017 ) . VTA ( Kim et al. , 2019 ) models videos using two levels of latent variables , where the fast states decide when the slow states should tick . TD-VAE ( Gregor and Besse , 2018 ) models high-dimensional input sequences using jumpy predictions , without a hierarchy . Refer to Appendix B for further related work . Despite this progress , scaling temporally abstract latent dynamics to complex datasets and understanding how these models organize the information about their inputs remain open challenges . In this paper , we introduce the Clockwork Variational Autoencoder ( CW-VAE ) , a simple hierarchical latent dynamics model where all levels tick at different fixed clock speeds . We conduct an extensive empirical evaluation and find that CW-VAE outperforms existing video prediction models . Moreover , we conduct several experiments to gain insights into the inner workings of the learned hierarchy . Our key contributions are summarized as follows : • Clockwork Variational Autoencoder ( CW-VAE ) We introduce a simple hierarchical video prediction model that leverages different clock speeds per level to learn long-term dependencies in video . A comprehensive empirical evaluation shows that on average , CW-VAE outperforms strong baselines , such as SVG-LP , VTA , and RSSM across several metrics . • Long-term video benchmark In the past , the video prediction literature has mainly focused on short-term video prediction of under 100 frames . To evaluate the ability to capture long-term dependencies , we propose using the Minecraft Navigate dataset as a challenging benchmark for video prediction of 500 frames . • Accurate long-term predictions Despite the simplicity of fixed clock speeds , CW-VAE improves over the distance of accurate video prediction of prior work . On the Minecraft Navigate dataset , CW-VAE accurately predicts for over 400 frames , whereas prior work fails before or around 150 frames . • Adaptation to sequence speed We demonstrate that CW-VAE automatically adapts to the frame rate of the dataset . Varying the frame rate of a synthetic dataset confirms that the slower latents are used more when objects in the video are slower , and faster latents are used more when the objects are faster . • Separation of information We visualize the content represented at higher levels of the temporal hierarchy . This experiment shows that the slower levels represent content that changes more slowly in the input , such as the wall colors of a maze , while faster levels represent faster changing concepts , such as the camera position . 2 Clockwork Variational Autoencoder . Long video sequences contain both , information that is local to a few frames , as well as global information that is shared among many frames . Traditional video prediction models that predict ahead at the frame rate of the video can struggle to retain information long enough to learn such long-term dependencies . We introduce the Clockwork Variational Autoencoder ( CW-VAE ) to learn long-term correlations of videos . Our model predicts ahead on multiple time scales , as visualized in Figure 3 . We build our work upon the Recurrent State-Space Model ( RSSM ; Hafner et al. , 2019b ) , the details of which can be found in Appendix A. CW-VAE consists of a hierarchy of recurrent latent variables , where each level transitions at a different clock speed . We slow down the transitions exponentially as we go up in the hierarchy , i.e . each level is slower than the level below by a factor k. We denote the latent state at timestep t and level l by slt and the video frames by xt . We define the set of active timesteps Tl for each level l ∈ [ 1 , L ] as those instances in time where the state transition generates a new latent state , Active steps : Tl . = { t ∈ [ 1 , T ] | tmod kl−1 = 1 } . ( 1 ) At each level , we condition k consecutive latent states on a single latent variable in the level above . For example , in the model shown in Figure 3 with k = 2 , T1 = { 1 , 2 , 3 , . . . } , T2 = { 1 , 3 , 5 , . . . } , and both s11 and s 1 2 are conditioned on the same s 2 1 from the second level . The latent chains can also be thought of as a hierarchy of latent variables where each level has a separate state variable per timestep , but only updates the previous state every kl−1 timesteps and otherwise copies the previous state , so that ∀t /∈ Tl : Copied states : slt . = slmax τ { τ∈Tl|τ≤t } . ( 2 ) Joint distribution We can factorize the joint distribution of a sequence of images and active latents at every level into two terms : ( 1 ) the reconstruction terms of the images given their lowest level latents , and ( 2 ) state transitions at all levels conditioned on the previous latent and the latent above , p ( x1 : T , s 1 : L 1 : T ) . = ( ∏T t=1 p ( xt | s1t ) ) ( ∏L l=1 ∏ t∈Tl p ( s l t | slt−1 , sl+1t ) ) . ( 3 ) To implement this distribution and its inference model , CW-VAE utilizes the following components , ∀ l ∈ [ 1 , L ] , t ∈ Tl , Encoder : elt = e ( xt : t+kl−1−1 ) Posterior transition qlt : q ( s l t | slt−1 , sl+1t , elt ) Prior transition plt : p ( s l t | slt−1 , sl+1t ) Decoder : p ( xt | s1t ) . ( 4 ) Inference CW-VAE embeds the observed frames using a CNN . Each active latent state at a level l receives the image embeddings of its corresponding kl−1 observation frames ( dashed lines in Figure 3 ) . The diagonal Gaussian belief qlt is then computed as a function of the input features , the posterior sample at the previous step , and the posterior sample above ( solid lines in Figure 3 ) . We reuse all weights of the generative model for inference except for the output layer that predicts the mean and variance . Generation The diagonal Gaussian prior plt is computed by applying the transition function from the latent state at the previous timestep in the current level , as well as the state belief at the level above ( solid lines in Figure 3 ) . Finally , the posterior samples at the lowest level are decoded into images using a transposed CNN . Training objective Because we can not compute the likelihood of the training data under the model in closed form , we use the ELBO as our training objective . This training objective optimizes a reconstruction loss at the lowest level , and a KL regularizer at every level in the hierarchy summed across active timesteps , max e , q , p ∑T t=1 Eq1t [ ln p ( xt | s 1 t ) ] − ∑L l=1 ∑ t∈Tl Eqlt−1q l+1 t [ KL [ qlt ∥ plt ] ] . ( 5 ) The KL regularizers limit the amount of information about the images that enters via the encoder . This encourages the model to utilize the “ free ” information from the previous and above latent and only attend to the input image to the extent necessary . Since the number of active timesteps decreases as we go higher in the hierarchy , the number of KL terms per level decreases as well . Hence it is easier for the model to store slowly changing information high up in the hierarchy than to pay a KL penalty to repeatedly extracting the information from the images at the lower level or trying to remember it by passing it along for many steps at the lower level without accidental forgetting . Stochastic and Deterministic Path As shown in Figure 3 ( right ) , we split the state slt into stochastic ( zlt ) and deterministic ( h l t ) parts ( Hafner et al. , 2019b ) . The deterministic state is computed using the top-down and temporal context , which then conditions the stochastic state at that level . The stochastic variables follow diagonal Gaussians with predicted means and variances . We use one GRU ( Cho et al. , 2014 ) per level to update the deterministic variable at every active step . All components of Equation 4 jointly optimize Equation 5 by stochastic backprop with reparameterized sampling ( Kingma and Welling , 2013 ; Rezende et al. , 2014 ) . Refer to Appendix C for architecture details . 3 Experiments . We compare the Clockwork Variational Autoencoder ( CW-VAE ) on 4 diverse video prediction datasets to state-of-the-art video prediction models in Sections 3.2 to 3.5 . We then conduct an extensive analysis to gain insights into the behavior of CW-VAE . We study the content stored at each level of the hierarchy in Section 3.6 , the effect of different clock speeds in Section 3.7 , the amount of information at each level in Section 3.8 , and the change in information as a function of the dataset frame rate Section 3.9 . The source code and video predictions are available on the project website.2 Datasets We choose 4 diverse video datasets for the benchmark . The MineRL Navigate dataset ( available under the CC Attribution-NonCommercial-ShareAlike 4.0 license ) was crowd sourced by Guss et al . ( 2019 ) for reinforcement learning applications . We process this data to create a longhorizon video prediction dataset that contains ∼750k frames . The sequences show players traveling to goal locations in procedurally generated 3D worlds of the video game Minecraft ( traversing forests , mountains , villages , oceans ) . The KTH Action video prediction dataset ( Schuldt et al. , 2004 ) ( available under the CC Attribution-NonCommercial license ) contains 290k frames . The 600 videos show humans walking , jogging , running , boxing , hand-waving , and clapping . GQN Mazes ( Eslami 2https : //danijar.com/cwvae et al. , 2018 ) ( available under the Apache License 2.0 ) contains 9M frames of videos of a scripted policy that traverses procedurally generated mazes with randomized wall and floor textures . For Moving MINST ( Srivastava et al. , 2015 ) ( available under the CC Attribution-ShareAlike 3.0 license ) we generate 2M frames where two digits move with velocities sampled uniformly in the range of 2 to 6 pixels per frame and bounce within the edges . Baselines We compare CW-VAE to 3 well-established video prediction models , and an ablation of our method where all levels tick at the fastest scale , which we call NoTmpAbs . VTA ( Kim et al. , 2019 ) is the state-of-the-art for video prediction using temporal abstraction . It consists of two levels , with the lower level predicting when the higher level should step . The lower level is reset when the higher level steps . RSSM ( Hafner et al. , 2019b ) is commonly used as a world model in reinforcement learning . It predicts forward using a sequence of compact latents without temporal abstraction . SVG-LP ( Denton and Fergus , 2018 ) , or SVG for short , has been shown to generate sharp predictions on visually complex datasets . It autogressively feeds generated images back into the model while also using a latent variable at every time step . The parameter counts are shown in Table D.2 . NoTmpAbs simply sets the temporal abstraction factor to 1 and thus uses the same number of parameters as its temporally-abstract counterparts . Training details We train all models on all datasets for 300 epochs on training sequences of 100 frames of size 64× 64 pixels . For the baselines , we tune the learning rate in the range [ 10−4 , 10−3 ] and the decoder stddev in the range [ 0.1 , 1 ] . We use a temporal abstraction factor of 6 per level for CW-VAE , unless stated otherwise . Refer to Appendix D for hyperparameters and training durations . A 3-level CW-VAE with abstraction factor 6 takes 2.5 days to train on one Nvidia Titan Xp GPU . Higher temporal abstraction factors train faster because fewer state transitions need to be computed . Evaluation We evaluate the open-loop video predictions under 4 metrics : Structural Similarity index ( SSIM , higher is better ) , Peak Signal-to-Noise Ratio ( PSNR , higher is better ) , Learned Perceptual Image Patch Similarity ( LPIPS , lower is better ; Zhang et al. , 2018 ) , and Frechet Video Distance ( FVD , lower is better ; Unterthiner et al. , 2018 ) . All video predictions are open-loop , meaning that the models only receive the first 36 frames as context input and then predict forward without access to intermediate frames . The number of input frames equals one step of the slowest level of CW-VAE for simplicity ; see Appendix D .

Summary: This paper presents a video prediction model named Clockwork VAE, which leverages hierarchical latent variables and temporal abstraction to learn long-term dependency. It also proposes the new Minecraft benchmark for long-term video prediction. The model is compared with existing video prediction methods on four datasets, and achieves remarkable long-term prediction results.