FitVid: High-Capacity Pixel-Level Video Prediction

1 INTRODUCTION Predicting what happens next is a cornerstone of intelligence and one of the key capabilities of humans , which we heavily rely on to make decisions in everyday life ( Bubic et al. , 2010 ) . This capability enables us to anticipate future events and plan ahead to perform temporally extended tasks . While the machine learning literature has studied a wide range of prediction problems , one of the most direct challenges is to predict raw sensory inputs . In particular , prediction of future visual inputs conditioned on a context of past observations – i.e. , pixel-level video prediction – encapsulates the challenges of visual perception , modeling of physical events , and reasoning about uncertain behaviors . Video prediction can be formulated as a self-supervised problem , enabling us to use a substantial amount of unlabeled data to provide autonomous systems with powerful predictive capabilities as well as learning rich representations for downstream tasks . Already , video models have been successfully deployed in applications such as robotics ( Finn & Levine , 2017 ; Zhang et al. , 2019 ) , simulation ( Kim et al. , 2021 ; 2020 ) , compression ( Duan et al. , 2020 ; Sulun & Tekalp , 2021 ) and video synthesis from a single frame ( Endo et al. , 2019 ; Nam et al. , 2019 ) . Despite recent advances in generative models in many domains , such as images ( Karras et al. , 2017 ; Vahdat & Kautz , 2020 ) and text ( Devlin et al. , 2018 ; Brown et al. , 2020 ) , video prediction is still considered to be challenging . Current state-of-the-art methods are limited to low-resolution videos ( typically 64×64 ( Villegas et al. , 2019 ) and a maximum of 256×256 ( Kim et al. , 2021 ) ) , usually in a narrow domain , such as a single human walking Ionescu et al . ( 2014 ) , or a robotic arm pushing objects in a stationary setting Ebert et al . ( 2017 ) . Even then the quality of predicted frames tends to drop substantially after as little as 10 seconds into the future ( Saxena et al. , 2021 ) . A growing body of evidence suggests that underfitting is one of the primary reasons for low quality predictions . For example , Villegas et al . ( 2019 ) demonstrate how scaling the model , by adding more parameters , can substantially improve the prediction quality . Similarly , Castrejon et al . ( 2019 ) argue that blurry video prediction of variational methods is a sign of underfitting , exhibiting an improved test and train performance as the network capacity increases . Wu et al . ( 2021a ) also observed monotonic improvement as the number of modules in a hierarchical architecture increases . While scaling up models is common to address underfitting , it comes at the cost of more computation , memory , and integration risks ( Wolf et al. , 2019 ) as well as more complicated training regimes ( Wu et al. , 2021a ) . In this paper we take a step back and address underfitting by instead finding an architecture that uses its parameters more efficiently . We propose FitVid , a model that – with the same parameter count as current state-of-the-art models – can fit the current video prediction benchmarks successfully . In fact , FitVid fits these benchmarks so well that in many cases we find it actually begins to suffer from overfitting . To the best of our knowledge , this is the first time a video model reports substantial overfitting on theses benchmarks . Importantly , we also find that simple existing image augmentation techniques can mitigate this overfitting , leading to models that can both fit the training set and generalize well to held-out videos . As a result , FitVid achieves state-of-the-art on four challenging video datasets across a wide range of metrics . Furthermore , while each of the components of our proposed model are not novel on their own , we find that FitVid can be trained a significantly simpler training recipe . Prior works on video prediction , particularly those that make use of variational methods to provide for stochasticity , typically require a number of delicate design decisions to train successfully : curriculum training ( Oh et al. , 2015 ; Lee et al. , 2018 ; Cenzato et al. , 2019 ) , a learned prior ( Denton & Fergus , 2018 ; Castrejon et al. , 2019 ) and annealing of the weight on the VAE KL-divergence penalty ( Babaeizadeh et al. , 2018 ) . In contrast , we show that our method actually fits the training data well without any such components , training directly via optimizing the evidence lower bound with minimal hyperparameters . 2 RELATED WORK . Video prediction ( Ranzato et al. , 2014 ; Srivastava et al. , 2015 ) has been formulated in different ways such as generating videos from a single image ( Endo et al. , 2019 ; Nam et al. , 2019 ; Dorkenwald et al. , 2021 ; Yang et al. , 2021 ; Hu et al. , 2021 ) or no image ( Vondrick et al. , 2016 ; Saito & Saito , 2018 ; Clark et al. , 2019 ) , text to video generation ( Wu et al. , 2021b ) , video-to-video translation ( Wang et al. , 2019a ; 2018 ) and data-driven simulation ( Kaiser et al. , 2019 ; Kim et al. , 2021 ; 2020 ) . In this paper , our focus is on conditional video prediction , which is to predict the future frames of a video conditioned on a few initial context frames and possibly the future actions of the agents ( Finn et al. , 2016a ; Babaeizadeh et al. , 2018 ) . Conditional video prediction has a number of applications , including model-based reinforcement learning from pixels ( Hafner et al. , 2018 ; 2019 ; Kaiser et al. , 2019 ; Rafailov et al. , 2020 ) and robotics ( Boots et al. , 2014 ; Kalchbrenner et al. , 2017a ; Ebert et al. , 2017 ; 2018a ; b ; Xie et al. , 2019 ; Paxton et al. , 2019 ; Nair & Finn , 2020 ; Nair et al. , 2020 ) . Initially , video prediction was tackled using deterministic models ( Walker et al. , 2015 ; Jia et al. , 2016 ; Xue et al. , 2016 ; Walker et al. , 2016 ; Liang et al. , 2017 ; Byravan & Fox , 2017 ; Vondrick & Torralba , 2017 ; Van Amersfoort et al. , 2017 ; Liu et al. , 2017 ; Chen et al. , 2017 ; Lu et al. , 2017 ) . Later on , given the common randomness and partial observability in the real-life situations , various stochastic models were proposed to capture the stochasticity of the future . Generative adversarial networks ( GANs ) ( Goodfellow et al. , 2014 ) are used for video predictions ( Mathieu et al. , 2015 ; Lee et al. , 2018 ; Clark et al. , 2019 ; Luc et al. , 2020 ; Hong et al. , 2021 ) as sell as autoregressive models ( Kalchbrenner et al. , 2017b ; Reed et al. , 2017 ; Weissenborn et al. , 2019 ) . Flow based generative models ( Dinh et al. , 2014 ; 2016 ) are also adopted for video prediction ( Kumar et al. , 2019 ) . In this paper we focus on Variational Auto-Encoders ( VAEs ) ( Kingma & Welling , 2014 ) which are widely used for conditional video prediction ( Shu et al. , 2016 ; Babaeizadeh et al. , 2018 ; Denton & Fergus , 2018 ; Wichers et al. , 2018 ; Villegas et al. , 2019 ; Franceschi et al. , 2020 ; Castrejon et al. , 2019 ; Yan et al. , 2021 ; Rakhimov et al. , 2020 ; Lee et al. , 2021 ; Wu et al. , 2021a ; Walker et al. , 2021 ) . Recently , it has been demonstrated that underfitting on training data plays a major role in low quality blurry predictions in VAE based models . Villegas et al . ( 2019 ) ; Castrejon et al . ( 2019 ) ; Wu et al . ( 2021a ) all reported improved prediction quality as the network capacity increases . In this paper , we are interested in addressing underfitting without increasing the capacity of the current models , and instead , find a more expressive architecture which uses its capacity more efficiently . 3 THE FITVID MODEL . Striving for simplicity , we propose the FitVid model for stochastic video prediction , a convolutional non-hierarchical variational model with a fixed prior . In this section we will describe the proposed model and arcitecture which is also visualized in Fig 2 . The appendix includes more architecture details and the code is available in Supplementary Material . 3.1 VARIATIONAL VIDEO PREDICTION WITH FIXED PRIOR . Following prior work ( Finn et al. , 2016a ; Babaeizadeh et al. , 2018 ; Denton & Fergus , 2018 ; Rubinstein , 1997 ; Oprea et al. , 2020 ; Wu et al. , 2021a ) , we define the problem of pixel-level video prediction as follows : given the first c frame of a video x < c = x0 , x1 , . . . , xc−1 , our goal is to predict the future frames by sampling from p ( xc : T |x < c ) . Optionally , the predictive model can be conditioned on additional given information at , such as the actions that the agents in the video are planning to take . This is typically called action-conditioned video prediction . Variational video prediction ( Babaeizadeh et al. , 2018 ) follows the variational auto-encoder ( Kingma & Welling , 2014 ) formalism by introducing a set of latent variables z to capture the inherent stochasticity of the problem . We factorize the likelihood model to ∏T t=c pθ ( xt|x < t , z≤t ) which is parametrized in an autoregressive manner over time ; i.e . at each timestep t the video frame xt and the latent variables zt are conditioned on the past latent samples and frames . For inference , we need to compute a marginalized distribution over the latent variables z , which is intractable . To overcome this problem , we use variation inference ( Jordan et al. , 1999 ) by defining an amortized approximate posterior q ( z|x ) = ∏ t q ( zt|z < t , x≤t ) that approximates the posterior distribution p ( z|x ) . The approximated posterior is commonly modeled with an inference network qφ ( z|x ) that outputs the parameters of a conditionally Gaussian distribution N ( µφ ( x ) , σφ ( x ) ) . This network can be trained using the reparameterization trick ( Kingma & Welling , 2014 ) by optimizing the variational lower bound ( Kingma & Welling , 2014 ; Rezende et al. , 2014 ) : L ( x ) = −Eqφ ( z|x ) [ log pθ ( xt : T |x < t , z ) ] + βDKL ( qφ ( z|x ) ||p ( z ) ) where DKL is the Kullback-Leibler divergence between the approximated posterior and the fixed prior p ( z ) = N ( 0 , I ) . The hyper-parameter β represents the trade-off between minimizing frame prediction error and fitting the prior ( Higgins et al. , 2016b ; a ; Denton & Fergus , 2018 ) . For FitVid , we set β = 1 .

Summary: The paper proposes a simple and scalable variational video prediction model FitVid, which attains a better fit to video prediction datasets even with a similar parameter count as prior models. The author has observed that previous methods suffer from underfitting on these datasets, directly applying FitVid actually results in overfitting. The FitVid uses a set of existing image augmentation techniques to prevent overfitting so that it can achieve state-of-the-art results on several prediction benchmarks. FitVid's architecture is based on SE-UNet LSTM, which seems to be the common backbone for the stochastic video prediction task.