Conditional Flow Variational Autoencoders for Structured Sequence Prediction

1 INTRODUCTION . Anticipating future states of the environment is a key competence necessary for the success of autonomous agents . In complex real world environments , the future is highly uncertain . Therefore , structured predictions , one to many mappings of the likely future states of the world , are important . In many scenarios , these tasks can be cast as sequence prediction problems . Particularly , Conditional Variational Autoencoders ( CVAE ) ( Sohn et al. , 2015 ; Bayer & Osendorfer , 2014 ; Chung et al. , 2015 ) have been very successful – from prediction of pedestrians trajectories ( Lee et al. , 2017 ; Bhattacharyya et al. , 2018 ; Pajouheshgar & Lampert , 2018 ) to outcomes of robotic actions ( Babaeizadeh et al. , 2018 ) . The distribution of future sequences is diverse and highly multi-modal . CVAEs model diverse futures by factorizing the distribution of future states using a set of latent variables which are mapped to likely future states . However , CVAEs assume a standard Gaussian prior on the latent variables which induces a strong model bias ( Hoffman & Johnson , 2016 ; Tomczak & Welling , 2018 ) which makes it challenging to capture multi-modal distributions . This also leads to missing modes due to posterior collapse ( Bowman et al. , 2016 ; Razavi et al. , 2019 ) . Recent work ( Tomczak & Welling , 2018 ; Wang et al. , 2017 ; Gu et al. , 2018 ) has therefore focused on more complex Gaussian mixture based priors . Gaussian mixtures still have limited expressiveness and optimization suffers from complications e.g . determining the number of mixture components . Normalizing flows are more expressive and enable the modelling of complex multi-modal priors . Recent work on flow based priors ( Chen et al. , 2017 ; Ziegler & Rush , 2019 ) , have focused only on the unconditional ( plain VAE ) case . However , this not sufficient for CVAEs because in the conditional case the complexity of the distributions are highly dependent on the condition . In this work , 1 . We propose Conditional Flow Variational Autoencoders ( CF-VAE ) based on novel conditional normalizing flow based priors In order to model complex multi-modal conditional distributions over sequences . In Figure 1 , we show example predictions of MNIST handwriting stroke of our CF-VAE . We observe that , given a starting stroke , our CF-VAE model with data dependent normalizing flow based latent prior captures the two main modes of the conditional distribution – i.e . 1 and 8 – while CVAEs with fixed uni-modal Gaussian prior predictions have limited diversity . 2 . We propose a regularization scheme that stabilizes the optimization of the evidence lower bound and leads to better fit to the target data distribution . 3 . We leverage our conditional flow prior to deal with posterior collapse which causes standard CVAEs to ignore modes in sequence prediction tasks . 4 . Finally , our method outperforms the state of the art on three structured sequence prediction tasks – handwriting stroke prediction on MNIST , trajectory prediction on Stanford Drone and HighD . 2 RELATED WORK . Normalizing Flows . Normalizing flows are a powerful class of density estimation methods with exact inference . ( Dinh et al. , 2015 ) introduced affine normalizing flows with triangular Jacobians . ( Dinh et al. , 2017 ) extend flows with masked convolutions which allow for complex ( non-autoregessive ) dependence between the dimensions . In ( Kingma & Dhariwal , 2018 ) , 1 × 1 convolutions were proposed for improved image generation compared to ( Dinh et al. , 2017 ) . In ( Huang et al. , 2018 ) normalizing flows are auto-regressive and ( Behrmann et al. , 2019 ) extend it to ResNet . ( Lu & Huang , 2019 ) extended normalizing flows to model conditional distributions . Here , we propose conditional normalizing flows to learn conditional priors for variational latent models . Variational Autoencoders . The original variational autoencoder ( Kingma & Welling , 2014 ) used uni-modal Gaussian prior and posterior distributions . Thereafter , two lines of work have focused on developing either more expressive prior or posterior distributions . Rezende & Mohamed ( 2015 ) propose normalizing flows to model complex posterior distributions . Kingma et al . ( 2016 ) ; Tomczak & Welling ( 2016 ) ; Berg et al . ( 2018 ) present more complex inverse autoregessive flows , householder and Sylvester normalizing flow based posteriors . Here , we focus on the orthogonal direction of more expressive priors and the above approaches are compatible with our approach . Recent work which focus more expressive priors include ( Nalisnick & Smyth , 2017 ) which proposes a Dirichlet process prior and ( Goyal et al. , 2017 ) which proposes a nested Chinese restaurant process prior . However , these methods require sophisticated learning methods . In contrast , ( Tomczak & Welling , 2018 ) proposes a mixture of Gaussians based prior ( with fixed number of components ) which is easier to train and shows promising results on some image generation tasks . ( Chen et al. , 2017 ) , proposes a inverse autoregressive flow based prior which leads to improvements in complex image generation tasks like CIFAR-10 . ( Ziegler & Rush , 2019 ) proposes a prior for VAE based text generation using complex non-linear flows which allows for complex multi-modal priors . While these works focus on unconditional priors , we aim to develop more expressive conditional priors . Posterior Collapse . Posterior collapse arises when the latent posterior does not encode useful information . Most prior work ( Yang et al. , 2017 ; Dieng et al. , 2019 ; Higgins et al. , 2017 ) concentrate on unconditional VAEs and modify the training objective – the KL divergence term is annealed to prevent collapse to the prior . Liu et al . ( 2019 ) extends KL annealing to CVAEs . However , KL annealing does not optimize a true lower bound of the ELBO for most of training . Zhao et al . ( 2017 ) also modifies the objective to choose the model with the maximal rate . Razavi et al . ( 2019 ) propose anti-causal sequential priors for text modelling tasks . Bowman et al . ( 2016 ) ; Gulrajani et al . ( 2017 ) proposes to weaken the decoder so that the latent variables can not be ignored , however only unconditional VAEs are considered . Wang & Wang ( 2019 ) shows the advantage of normalizing flow based posteriors for preventing posterior collapse . In contrast , we study for the first time posterior collapse in conditional models on datasets with minor modes . Structured Sequence Prediction . Helbing & Molnar ( 1995 ) ; Robicquet et al . ( 2016 ) ; Alahi et al . ( 2016 ) ; Gupta et al . ( 2018 ) ; Zhao et al . ( 2019 ) ; Sadeghian et al . ( 2019 ) consider the problem of traffic participant trajectory prediction in a social context . Notably , ( Gupta et al. , 2018 ; Zhao et al. , 2019 ; Sadeghian et al. , 2019 ) use generative adversarial networks to generate socially compliant trajectories . However , the predictions are uni-modal . Starting from Bayer & Osendorfer ( 2014 ) ; Chung et al . ( 2015 ) , more recently Lee et al . ( 2017 ) ; Bhattacharyya et al . ( 2018 ) ; Rhinehart et al . ( 2018 ) ; Deo & Trivedi ( 2019 ) ; Pajouheshgar & Lampert ( 2018 ) considers structured ( one to many ) predictions using – a CVAE , improved CVAE training , pushforward policies for vehicle ego-motion prediction , motion planning , spatio-temporal convolutional network respectively . Kumar et al . ( 2019 ) proposes a normalizing flow based model for video sequence prediction , however the sequences considered have very limited diversity compared to the trajectory prediction tasks considered here . Here , we focus on improving structured predictions using conditional normalizing flows based priors . 3 CONDITIONAL FLOW VARIATIONAL AUTOENCODER ( CF-VAE ) . Our Conditional Flow Variational Autoencoder is based on the conditional variational autoencoder ( Sohn et al. , 2015 ) which is a deep directed graphical model for modeling conditional data distributions pθ ( y|x ) . Here , x is the sequence up to time t , x = [ x1 , · · · , xt ] and y is the sequence to be predicted up to time T , y = [ yt+1 , · · · , yT ] . CVAEs factorize the conditional distribution using latent variables z . In detail , pθ ( y|x ) = ∫ pθ ( y|z , x ) p ( z|x ) dz , where p ( z|x ) is the prior on the latent variables . During training , amortized variational inference is used and the posterior distribution qφ ( z|x , y ) is learnt using a recognition network . The ELBO is maximized , given by , log ( pθ ( y|x ) ) ≥ Eqφ ( z|x , y ) log ( pθ ( y|z , x ) ) −DKL ( qφ ( z|x , y ) ||p ( z|x ) ) . ( 1 ) In practice , to simplify learning , simple unconditional standard Gaussian priors are used ( Sohn et al. , 2015 ) . However , the complexity e.g . the number of modes of the target distributions pθ ( y|x ) , is highly dependent upon the condition x . An unconditional prior demands identical latent distributions irrespective complexity of the target conditional distribution – a very strong constraint on the recognition network . Moreover , the latent variables can not encode any conditioning information and this leaves the burden of learning the dependence on the condition completely on the decoder . Furthermore , on complex conditional multi-modal data , Gaussian priors have been shown to induce a strong model bias ( Tomczak & Welling , 2016 ; Ziegler & Rush , 2019 ) . It becomes increasingly difficult to map complex multi-modal distributions to uni-modal Gaussian distributions , further complicated by the sensitivity of the RNNs encoder/decoders to subtle variations in the hidden states ( Bowman et al. , 2016 ) . Moreover , the standard closed form estimate of the KL-divergence pushes the encoded latent distributions to the mean of the Gaussian leading to latent variable collapse ( Wang et al. , 2017 ; Gu et al. , 2018 ) while discriminator based approaches ( Tolstikhin et al. , 2017 ) lead to underestimates of the KL-divergence ( Rosca et al. , 2017 ) . Therefore , we propose conditional priors based on conditional normalizing flows to enable the latent variables to encode conditional information and allow for complex multi-modal latent representations . Next , we introduce our new conditional non-linear normalizing flows followed by our regularized Conditional Flow Variational Autoencoder ( CF-VAE ) formulation . 3.1 CONDITIONAL NORMALIZING FLOWS . Recently , normalizing flow ( Tabak et al. , 2010 ; Dinh et al. , 2015 ) based priors for VAEs have been proposed ( Chen et al. , 2017 ; Ziegler & Rush , 2019 ) . Normalizing flows allows for complex priors by transforming a simple base density e.g . standard Gaussian to a complex multi-modal density through a series of n layers of invertible transformations fi , f1←→ h1 f2←→ h2 · · · fn←→ z . ( 2 ) However , such flows can not model conditional priors . In contrast to prior work , we utilize conditional normalizing flows to model complex conditional priors . Conditional normalizing flows also consists of a series of n layers of invertible transformations fi ( with parameters ψ ) , however we modify the transformations fi such that they are dependent on the condition x , |x f1|x←→ h1|x f2|x←→ h2|x · · · fn|x←→ z|x . ( 3 ) Further , in contrast to prior work ( Lu & Huang , 2019 ; Atanov et al. , 2019 ; Ardizzone et al. , 2019 ) which use affine flows ( fi ) , we build upon ( Ziegler & Rush , 2019 ) and introduce conditional nonlinear normalizing flows with split coupling . Split couplings ensure invertibility by applying a flow layer fi on only half of the dimensions at a time . To compute ( 5 ) , we split the dimensions zD of the latent variable into halfs , zL = { 1 , · · · , D/2 } and zR = { D/2 , · · · , d } at each invertible layer fi . Our transformation takes the following form for each dimension zj alternatively from zL or zR , f−1i ( z j |zR , x ) = j = a ( zR , x ) + b ( zR , x ) × zj + c ( z R , x ) 1 + ( d ( zR , x ) × zj + g ( zR , x ) ) 2 . ( 4 ) where , zj ∈ zL . Details of the forward ( generating ) operation fi are in Appendix A . To ensure that the generated prior distribution is conditioned on x , in ( 4 ) and in the corresponding forward operation fi , the coefficients { a , b , c , d , g } ∈ R are functions of both the other half of the dimensions of z and the condition x ( unlike Ziegler & Rush ( 2019 ) ) . Finally , due to the expressive power of our conditional non-linear normalizing flows , simple spherical Gaussians base distributions were sufficient .

Summary: The work proposes a method to improve conditional VAE with a learnable prior distribution using normalizing flow. The authors also design two regularization methods for the CF-VAE to improve training stability and avoid posterior collapse. The paper is clearly motivated and easy to follow. Experiment results on MNIST, Stanford Drone and HighD datasets show the proposed that the model achieves better results than previous state-of-the-art models by significant margins.