Variational Dynamic Mixtures

1 INTRODUCTION . Making sense of time series data is an important challenge in various domains , including ML for climate change . One important milestone to reach the climate goals is to significantly reduce the CO2 emissions from mobility ( Rogelj et al. , 2016 ) . Accurate forecasting models of typical driving behavior and of typical pollution levels over time can help both lawmakers and automotive engineers to develop solutions for cleaner mobility . In these applications , no accurate physical model of the entire dynamic system is known or available . Instead , data-driven models , specifically deep probabilistic time series models , can be used to solve the necessary tasks including forecasting . The dynamics in such data can be highly multi-modal . At any given part of the observed sequence , there might be multiple distinct continuations of the data that are plausible , but the average of these behaviors is unlikely , or even physically impossible . Consider for example a dataset of taxi trajectories1 . In each row of Fig . 1a , we have selected 50 routes from the dataset with similar starting behavior ( blue ) . Even though these routes are quite similar to each other in the first 10 way points , the continuations of the trajectories ( red ) can exhibit quite distinct behaviors and lead to points on any far edge of the map . The trajectories follow a few main traffic arteries , these could be considered the main modes of the data distribution . We would like to learn a generative model of the data , that based on some initial way points , can forecast plausible continuations for the trajectories . Many existing methods make restricting modeling assumptions such as Gaussianity to make learning tractable and efficient . But trying to capture the dynamics through unimodal distributions can lead either to “ over-generalization ” , ( i.e . putting probability mass in spurious regions ) or on focusing only on the dominant mode and thereby neglecting important structure of the data . Even neural approaches , with very flexible generative models can fail to fully capture this multi-modality because their capacity is often limited through the assumptions of their inference model . To address this , we develop variational dynamic mixtures ( VDM ) . Its generative process is a sequential latent variable model . The main novelty is a new multi-modal variational family which makes learning and inference multi-modal yet tractable . In summary , our contributions are • A new inference model . We establish a new type of variational family for variational inference of sequential latent variables . By successively marginalizing over previous latent states , the procedure can be efficiently carried-out in a single forward pass and induces a multi-modal posterior 1https : //www.kaggle.com/crailtap/taxi-trajectory approximation . We can see in Fig . 1b , that VDM trained on a dataset of taxi trajectories produces forecasts with the desired multi-modality while other methods overgeneralize . • An evaluation metric for multi-modal tasks . The negative log-likelihood measures predictive accuracy but neglects an important aspect of multi-modal forecasts – sample diversity . In Section 4 , we derive a score based on the Wasserstein distance ( Villani , 2008 ) which evaluates both sample quality and diversity . This metric complements our evaluation based on log-likelihoods . • An extensive empirical study . in Section 4 , we use VDM to study various datasets , including a synthetic data with four modes , a stochastic Lorenz attractor , the taxi trajectories , and a U.S. pollution dataset with the measurements of various pollutants over time . We illustrate VDM ’ s ability in modeling multi-modal dynamics , and provide quantitative comparisons to other methods showing that VDM compares favorably to previous work . 2 RELATED WORK . Neural recurrent models . Recurrent neural networks ( RNNs ) such as LSTMs ( Hochreiter & Schmidhuber , 1997 ) and GRUs ( Chung et al. , 2014 ) have proven successful on many time series modeling tasks . However , as deterministic models they can not capture uncertainties in their dynamic predictions . Stochastic RNNs make these sequence models non-deterministic ( Chung et al. , 2015 ; Fraccaro et al. , 2016 ; Gemici et al. , 2017 ; Li & Mandt , 2018 ) . For example , the variational recurrent neural network ( VRNN ) ( Chung et al. , 2015 ) enables multiple stochastic forecasts due to its stochastic transition dynamics . An extension of VRNN ( Goyal et al. , 2017 ) uses an auxiliary cost to alleviate the KL-vanishing problem . It improves on VRNN inference by forcing the latent variables to also be predictive of future observations . Another line of related methods rely on particle filtering ( Naesseth et al. , 2018 ; Le et al. , 2018 ; Hirt & Dellaportas , 2019 ) and in particular sequential Monte Carlo ( SMC ) to improve the evidence lower bound . In contrast , VDM adopts an explicitly multi-modal posterior approximation . Another SMC-based work ( Saeedi et al. , 2017 ) employs search-based techniques for multi-modality but is limited to models with finite discrete states . Recent works ( Schmidt & Hofmann , 2018 ; Schmidt et al. , 2019 ; Ziegler & Rush , 2019 ) use normalizing flows in the latent space to model the transition dynamics . A normalizing flow requires many layers to transform its base distribution into a truly multi-modal distribution in practice . In contrast , mixture density networks ( as used by VDM ) achieve multi-modality by mixing only one layer of neural networks . A task orthogonal to multi-modal inference is learning disentangled representations . Here too , mixture models are used ( Chen et al. , 2016 ; Li et al. , 2017 ) . These papers use discrete variables and a mutual information based term to disentangle different aspects of the data . VAE-like models ( Bhattacharyya et al. , 2018 ; 2019 ) and GAN-like models ( Sadeghian et al. , 2019 ; Kosaraju et al. , 2019 ) only have global , time independent latent variables . Yet , they show good results on various tasks , including forecasting . With a deterministic decoder , these models focus on average dynamics and don ’ t capture local details ( including multi-modal transitions ) very well . Sequential latent variable models are described next . Deep state-space models . Classical State-space models ( SSMs ) are popular due to their tractable inference and interpretable predictions . Similarly , deep SSMs with locally linear transition dynamics enjoy tractable inference ( Karl et al. , 2017 ; Fraccaro et al. , 2017 ; Rangapuram et al. , 2018 ; Becker et al. , 2019 ) . However , these models are often not expressive enough to capture complex ( or highly multi-modal ) dynamics . Nonlinear deep SSMs ( Krishnan et al. , 2017 ; Zheng et al. , 2017 ; Doerr et al. , 2018 ; De Brouwer et al. , 2019 ; Gedon et al. , 2020 ) are more flexible . Their inference is often no longer tractable and requires variational approximations . Unfortunately , in order for the inference model to be tractable , the variational approximations are often simplistic and don ’ t approximate multi-modal posteriors well with negative effects on the trained models . Multi-modality can be incorporated via additional discrete switching latent variables , such as recurrent switching linear dynamical systems ( Linderman et al. , 2017 ; Nassar et al. , 2018 ; Becker-Ehmck et al. , 2019 ) . However , these discrete states make inference more involved . 3 VARIATIONAL DYNAMIC MIXTURES . We develop VDM , a new sequential latent variable model for multi-modal dynamics . Given sequential observations x1 : T = ( x1 , . . . , xT ) , VDM assumes that the underlying dynamics are governed by latent states z1 : T = ( z1 , . . . , zT ) . We first present the generative process and the multi-modal inference model of VDM . We then derive a new variational objective that encourages multi-modal posterior approximations and we explain how it is regularized via hybrid-training . Finally , we introduce a new sampling method used in the inference procedure . Generative model . The generative process consists of a transition model and an emission model . The transition model p ( zt | z < t ) describes the temporal evolution of the latent states and the emission model p ( xt | z≤t ) maps the states to observations . We assume they are parameterized by two separate neural networks , the transition network φtra and the emission network φdec.To give the model the capacity to capture longer range temporal correlations we parametrize the transition model with a recurrent architecture φGRU ( Auger-Méthé et al. , 2016 ; Zheng et al. , 2017 ) such as a GRU ( Chung et al. , 2014 ) . The latent states zt are sampled recursively from zt | z < t ∼ N ( µ0 , t , σ20 , tI ) , where [ µ0 , t , σ20 , t ] = φtra ( ht−1 ) , ht−1 = φGRU ( zt−1 , ht−2 ) , ( 1 ) and are then decoded such that the observations can be sampled from the emission model , xt | z≤t ∼ N ( µx , t , σ2x , tI ) , where [ µx , t , σ2x , t ] = φdec ( zt , ht−1 ) . ( 2 ) This generative process is similar to ( Chung et al. , 2015 ) , though we did not incorporate autoregressive feedback due to its negative impact on long-term generation ( Ranzato et al. , 2016 ; Lamb et al. , 2016 ) . The competitive advantage of VDM comes from a more expressive inference model . Inference model . VDM is based on a new procedure for multi-modal inference . The main idea is that to approximate the posterior at time t , we can use the posterior approximation of the previous time step and exploit the generative model ’ s transition model φGRU . This leads to a sequential inference procedure . We first use the forward model to transform the approximate posterior at time t − 1 into a distribution at time t. In a second step , we use samples from the resulting transformed distribution and combine each sample with data evidence xt , where every sample parameterizes a Gaussian mixture component . As a result , we obtain a multi-modal posterior distribution that depends on data evidence , but also on the previous time step ’ s posterior . In more detail , for every zt , we define its corresponding recurrent state as the transformed random variable st = φGRU ( zt , ht−1 ) , using a deterministic hidden state ht−1 = E [ st−1 ] . The variational family of VDM is defined as follows : q ( z1 : T | x1 : T ) = T∏ t=1 q ( zt | x≤t ) = T∏ t=1 ∫ q ( zt | st−1 , xt ) q ( st−1 | x≤t ) dst−1 . ( 3 ) Chung et al . ( 2015 ) also use a sequential inference procedure , but without considering the distribution of st. Only a single sample is propagated through the recurrent network and all other information about the distribution of previous latent states z < t is lost . In contrast , VDM explicitly maintains st as part of the inference model . Through marginalization , the entire distribution is taken into account for inferring the next state zt . Beyond the factorization assumption and the marginal consistency constraint of Eq . ( 3 ) , the variational family of VDM needs two more choices to be fully specified ; First , one has to choose the parametrizations of q ( zt | st−1 , xt ) and q ( st−1 | x≤t ) and second , one has to choose a sampling method to approximate the marginalization in Eq . ( 3 ) . These choices determine the resulting factors q ( zt | x≤t ) of the variational family . We assume that the variational distribution of the recurrent state factorizes as q ( st−1 | x≤t ) = ω ( st−1 , xt ) q̃ ( st−1 | x < t ) , i.e . it is the distribution of the recurrent state given the past observation2 , re-weighted by a weighting function ω ( st−1 , xt ) which involves only the current observations . For VDM , we only need samples from q̃ ( st−1 | x < t ) , which are obtained by sampling from the previous posterior approximation q ( zt−1 | x < t ) and transforming the sample with the RNN , s ( i ) t−1 ∼ q̃ ( st−1 | x < t ) equiv . to s ( i ) t−1 = φ GRU ( z ( i ) t−1 , ht−2 ) , z ( i ) t−1 ∼ q ( zt−1 | x < t ) , ( 4 ) where i indexes the samples . The RNN φGRU has the same parameters as in the generative model . Augmenting the variational model with the recurrent state has another advantage ; approximating the marginalization in Eq . ( 3 ) with k samples from q ( st−1 | x≤t ) and choosing a Gaussian parametrization for q ( zt | st−1 , xt ) results in a q-distribution q ( zt | x≤t ) that resembles a mixture density network ( Bishop , 2006 ) , which is a convenient choice to model multi-modal distributions . q ( zt | x≤t ) = k∑ i ω ( i ) t N ( µ ( i ) z , t , σ ( i ) 2 z , t I ) , [ µ ( i ) z , t , σ ( i ) 2 z , t ] = φ inf ( s ( i ) t−1 , xt ) . ( 5 ) We assume q ( zt | st−1 , xt ) to be Gaussian and use an inference network φinf to model the effect of the observation xt and recurrent state st−1 on the mean and variance of the mixture components . The mixture weights ω ( i ) t : = ω ( s ( i ) t−1 , xt ) /k come from the variational distribution q ( st−1 | x≤t ) = ω ( st−1 , xt ) q̃ ( st−1 | x < t ) and importance sampling3 . We are free to choose how to parametrize the weights , as long as all variational distributions are properly normalized . Setting ω ( i ) t = ω ( s ( i ) t−1 , xt ) /k : = 1 ( i = argmax j p ( xt | ht−1 = s ( j ) t−1 ) ) , ( 6 ) achieves this . In Appendix A , we explain this choice with importance sampling and in Appendix H , we compare the performance of VDM under alternative variational choices for the weights . In the next time-step , plugging the variational distribution q ( zt | x≤t ) into Eq . ( 4 ) yields the next distribution over recurrent states q̃ ( st | x≤t ) . For this , the expected recurrent state ht−1 is required . 2q̃ ( st−1 | x < t ) is the distribution obtained by transforming the previous zt−1 ∼ q ( zt−1|x < t ) through the RNN . It can be expressed analytically using the Kronecker δ to compare whether the stochastic variable st−1 equals the output of the RNN : q̃ ( st−1 | x < t ) ∝ ∫ δ ( st−1 − φGRU ( zt−1 , ht−2 ) ) q ( zt−1 | xt−1 , λt−1 ) dzt−1 . 3the ω adjusts for using samples from q̃ ( st−1 | x < t ) when marginalizing over ω ( st−1 , xt ) q̃ ( st−1 | x < t ) We approximate the update using the same k samples ( and therefore the same weights ) as in Eq . ( 5 ) . ht−1 = E [ st−1 ] = ∫ st−1 q ( st−1 | x≤t ) dst−1 ≈ k∑ i ω ( i ) t s ( i ) t−1 . ( 7 ) A schematic view of the generative and inference model of VDM is shown in Fig . 2 . In summary , the inference model of VDM alternates between Eqs . ( 4 ) to ( 7 ) . Latent states are sampled from the posterior approximation of the previous time-step and transformed by Eq . ( 4 ) into samples of the recurrent state of the RNN . These are then combined with the new observation xt to produce the next variational posterior Eq . ( 5 ) and the expected recurrent state is updated ( Eq . ( 7 ) ) . These are then used in Eq . ( 4 ) again . Approximating the marginalization in Eq . ( 3 ) with a single sample , recovers the inference model of VRNN ( Chung et al. , 2015 ) , and fails in modeling multi-modal dynamics as shown in Fig . 3 . In comparison , VDM ’ s approximate marginalization over the recurrent states with multiple samples succeeds in modeling multi-modal dynamics . Variational objective . We develop an objective to optimize the variational parameters of VDM φ = [ φtra , φdec , φGRU , φinf ] . The evidence lower bound ( ELBO ) at each time step is LELBO ( x≤t , φ ) : = 1 k k∑ i ω ( s ( i ) t−1 , xt ) Eq ( zt|s ( i ) t−1 , xt ) [ log p ( xt | zt , ht−1 = s ( i ) t−1 ) ] + 1 k k∑ i ω ( s ( i ) t−1 , xt ) Eq ( zt|s ( i ) t−1 , xt ) [ log p ( zt | ht−1 = s ( i ) t−1 ) q ( zt | s ( i ) t−1 , xt ) ] − 1 k k∑ i ω ( s ( i ) t−1 , xt ) [ logω ( s ( i ) t−1 , xt ) +C ] ( 8 ) Claim 1 . The ELBO in Eq . ( 8 ) is a lower bound on the log evidence log p ( xt | x < t ) , log p ( xt | x < t ) ≥ LELBO ( x≤t , φ ) , ( see proof in Appendix B ) . ( 9 ) In addition to the ELBO , the objective of VDM has two regularization terms , LVDM ( φ ) = T∑ t=1 Epdata [ −LELBO ( x≤t , φ ) − ω1Lpred ( x≤t , φ ) ] + ω2Ladv ( x≤t , φ ) . ( 10 ) In an ablation study in Appendix E , we compare the effect of including and excluding the regularization terms in the objective . VDM is competitive without these terms , but we got the strongest results by setting ω1,2 = 1 ( this is the only nonzero value we tried . This hyperparameter could be tuned even further . ) The first regularization term Lpred , encourages the variational posterior ( from the previous time step ) to produce samples that maximize the predictive likelihood , Lpred ( x≤t , φ ) = logEq ( st−1|x < t ) [ p ( xt | st−1 , x < t ) ] ≈ log 1 k k∑ i p ( xt | s ( i ) t−1 ) . ( 11 ) This regularization term is helpful to improve the prediction performance , since it depends on the predictive likelihood of samples , which isn ’ t involved in the ELBO . The second optional regularization term Ladv ( Eq . ( 12 ) ) is based on ideas from hybrid adversarial-likelihood training ( Grover et al. , 2018 ; Lucas et al. , 2019 ) . These training strategies have been developed for generative models of images to generate sharper samples while avoiding “ mode collapse ” . We adapt these ideas to generative models of dynamics . The adversarial term Ladv uses a forward KL-divergence , which enables “ quality-driven training ” to discourage probability mass in spurious areas . Ladv ( x≤t , φ ) = DKL ( p ( xt | x < t ) ‖pD ( xt | x < t ) ) = E [ log p ( xt | x < t ) − log pD ( xt | x < t ) ] ( 12 ) The expectation is taken w.r.t . p ( xt | x < t ) . The true predictive distribution pD ( xt | x < t ) is unknown . Instead , we can train the generator of a conditional GAN ( Mirza & Osindero , 2014 ) , while assuming an optimal discriminator . As a result , we optimize Eq . ( 12 ) in an adversarial manner , conditioning on x < t at each time step . Details about the discriminator are in Appendix G. Stochastic cubature approximation ( SCA ) . The variational family of VDM is defined by a number of modeling choices , including the factorization and marginal consistency assumptions of Eq . ( 3 ) , the parametrization of the transition and inference networks Eqs . ( 4 ) and ( 5 ) , and the choice of weighting function ω ( · ) . It is also sensitive to the choice of sampling method which we discuss here . In principle , we could use Monte-Carlo methods . However , for a relatively small number of samples k , Monte-Carlo methods don ’ t have a mechanism to control the quality of samples . We instead develop a semi-stochastic approach based on the cubature approximation ( Wan & Van Der Merwe , 2000 ; Wu et al. , 2006 ; Arasaratnam & Haykin , 2009 ) , which chooses samples more carefully . The cubature approximation proceeds by constructing k = 2d+1 so-called sigma points , which are optimally spread out on the d-dimensional Gaussian with the same mean and covariance as the distribution we need samples from . In SCA , the deterministic sigma points are infused with Gaussian noise to obtain stochastic sigma variables . A detailed derivation of SCA is in Appendix D. We use SCA for various reasons : First , it typically requires fewer samples than Monte-Carlo methods because the sigma points are carefully chosen to capture the first two moments of the underlying distribution . Second , it ensures a persistence of the mixture components ; when we resample , we sample another nearby point from the mixture component and not an entirely new location .

Summary: This paper introduces variational dynamic mixtures (VDM), a new variational family, and demonstrates that using VDM to model the approximate posterior in sequential latent variable models can better capture multi-modality in data. VDM includes a distribution over recurrent states in the inference model, such that a sampling-based marginalization of this distribution reduces the approximate posterior to a mixture model. Setting the weights such that only the most probable mixture component is selected allows other mixture components to capture other modes. The authors validate VDM on both synthetic and real multimodal datasets, which outperform baselines with respect to negative log-likelihood and a new empirical Wasserstein distance.