A NOVEL VARIATIONAL FAMILY FOR HIDDEN NON-LINEAR MARKOV MODELS

1 INTRODUCTION . In recent years , advances on neural data acquisition have made it possible to record the simultaneous sequential activity of up to thousands of neurons ( Paninski & Cunningham ( 2017 ) ) . The analysis of these datasets often focuses on dimensionality reduction techniques that encode the activity of the population in a lower dimensional latent trajectory ( Cunningham & Yu ( 2014 ) ) . At the other extreme , there is a big body of detailed electrophysiological data coming from voltage measurements in single cells ( Jones et al . ( 2009 ) ) . In this setting it is understood that the underlying dynamics is in fact highly nonlinear and multidimensional , though the experimenter only has access to a one-dimensional ( 1D ) observation . From such 1D recordings , the task is to approximately recover the complete hidden latent space paths and dynamics . A host of sophisticated techniques has been proposed for the analysis of complex sequential data that is not well described by linear transitions and observations ( Archer et al . ( 2015 ) ; Chung et al . ( 2015 ) ; Gao et al . ( 2016 ; 2015 ) ; Hernandez et al . ( 2017 ) ; Johnson et al . ( 2016 ) ; Krishnan et al . ( 2015 ; 2016 ) ; Linderman et al . ( 2017 ) ; Pandarinath et al . ( 2018 ) ; Sussillo et al . ( 2016 ) ; Zhao & Memming Park ( 2017 ) ; Wu et al . ( 2017 ; 2018 ) ) . In that context , we here present Variational Inference for Nonlinear Dynamics ( VIND ) . The main contribution of VIND is an algorithm that allows variational inference ( VI ) from structured , intractable approximations to the posterior distribution . In particular , VIND can handle variational posteriors that ( i ) represent nonlinear evolution in the latent space , and ( ii ) disentangle the latent dynamics ( transition ) from the data encoding ( recognition ) . Crucially , the VIND approximate posterior shares the exact nonlinear structure of latent dynamics evolution with the model for data generation . This makes the VIND approximation potentially more powerful than models in which the choice of approximate posterior is made solely on grounds of tractability . VIND relies on two key ideas . Firstly , it makes use of the fact that given an intractable posterior Q ( Z|X ) , it is always possible to compute a tractable Gaussian approximation to it . This Gaussian approximation inherits its parameters from Q ( Z|X ) ( Chung et al . ( 2015 ) ) , so optimizing for it can be interpreted as indirectly optimizing Q ( Z|X ) . The second novel aspect of VIND is the use of the fixed-point iteration ( FPI ) method to significantly speed up the computation of the aforementioned Gaussian approximation . In this work we focus on a VIND variant in which the latent dynamics is represented as a Locally Linear Dynamical System ( LLDS ) . The running time of LLDS/VIND is linear in the number of time points in a trial . We are especially interested in determining LLDS/VIND ’ s ability to infer the hidden dynamics , as demonstrated by its generative / predictive capabilities . After training , can the VIND-trained model generate data that is indistinguishable from the original observations , if provided with a suitable starting point ? In the second half of this work we apply VIND to four datasets , one synthetic and three using experimental data , and show that VIND excels in this task , in some cases outperforming established methods by orders of magnitude in the predictive mean squared error ( MSE ) . 2 VARIATIONAL INFERENCE WITH NONLINEAR DYNAMICS ( VIND ) . For a set of temporally ordered , correlated , noisy observations X ≡ { x1 , . . .xT } , xt ∈ RdX , a latent variable model proposes an additional , time-ordered set of random variables Z ≡ { z1 , . . . zT } , zT ∈ RdZ that is hidden from view . The hidden state zt is endowed with a stochastic dynamics : zt+1 ∼ p ( zt+1|z1 : t ) by which it evolves . The observations xt are generated by drawing samples from a zt-dependent probability distribution . A naive objective for such a model is the marginal log-likelihood log p ( X ) , with the latent variables integrated out of the joint . However , it is well known that for anything other than the simplest distributions , marginalization with respect to Z is intractable ( Bishop ( 2006 ) ) . VI overcomes this problem by approximating the posterior p ( Z|X ) with a distribution q ( Z|X ) , the Recognition Model ( RM ) , from a tractable class . The objective becomes the celebrated ELBO , a lower bound to log p ( X ) ( Jordan et al . ( 1999 ) ) : log p ( X ) ≥ LELBO ( X ) = E q [ log p ( X , Z ) ] − E q [ log q ( Z|X ) ] . ( 1 ) Successful VI relies on the choice of the approximation q ( Z|X ) . This choice is constrained by two goals that stand in tension : expressiveness and tractability . We approach the search for a suitable class of variational posterior by considering the joint p ( X , Z ) . Our focus is on factorizations of the form : p ( X , Z ) ≡ pφ , θ ( X , Z ) = cφ , θ ·Hφ ( Z ) T∏ t=0 fθ ( xt|zt ) , ( 2 ) where the distribution parameters have been written explicitly . The unnormalized densities fθ stand for an observation model that , for the purposes of this work , can be either Gaussian , xt|zt ∼ N ( mθ ( zt ) , Σ ) , or Poisson , xt|zt ∼ Poisson ( λθ ( zt ) ) . The respective mean , mθ ( zt ) , and rate , λθ ( zt ) , are arbitrary nonlinear functions of the latent state zt , that we represent as neural networks . The standard deviation Σ of the Gaussian observation model is taken to be zt-independent . cφ , θ is a normalization constant . Hφ is the latent evolution term in Z-space with a Markov Chain structure Johnson et al . ( 2016 ) ; Krishnan et al . ( 2015 ; 2016 ) ; Archer et al . ( 2015 ) ; Gao et al . ( 2016 ) : Hφ ( Z ) = h0 ( z0 ) T∏ t=1 hφ ( zt|zt−1 ) , ( 3 ) z0 ∼ N ( a0 , Γ0 ) , ( 4 ) zt|zt−1 ∼ N ( aφ ( zt−1 ) , Γ ) , ( 5 ) where aφ ( z ) is an arbitrary nonlinearity . From Eq . ( 2 ) , the posterior distribution of the Generative Model ( GM ) can be factorized as pφ , θ ( Z|X ) = cφ , θ ∏ fθ ( xt|zt ) pφ , θ ( X ) ·Hφ ( Z ) . ( 6 ) We used Eq . ( 6 ) as a guidance to propose the inclusion of the GM evolution term Hφ ( Z ) into the variational posterior . We start then with an approximation that factorizes as : Qφ , ϕ ( Z|X ) = αφ , ϕ ( X ) Gϕ ( X , Z ) Hφ ( Z ) . ( 7 ) We refer to Qφ , ϕ ( Z|X ) as the parent distribution . By design , the factor Gϕ in the parent contains all the dependence on X . For definiteness , the case Gϕ ( X , Z ) = T∏ t=0 gϕ ( zt|xt ) , zt|xt ∼ N ( µϕ ( xt ) , Λ−1ϕ ( xt ) ) , ( 8 ) is considered in this work , where µ ( x ) and Λ ( x ) are nonlinear maps . In Eq . ( 7 ) , αφ , ϕ is a normalization constant that , regardless of the specific form of Gϕ , can not be computed in closed form . This is most easily seen by noticing that due to the nonlinearity h ( zT |zT−1 ) , integration with respect to zT yields an intractable zT−1-dependent factor . Hence , VI can not be formulated directly in terms of the parent , since expectation values with respect to Qφ , ϕ can not be computed in closed form . VIND represents a way out of this conundrum that , effectively , allows for the use of an intractable , unnormalized distribution as the Recognition Model in VI . The idea is to define a Gaussian approximation qφ , ϕ ( Z|X ) to the parent ; this child distribution is then used as the variational posterior in Eq . ( 1 ) . The inference problem becomes tractable since the child is normal . Importantly , the parameters in qφ , ϕ ( Z|X ) , with respect to which we optimize , are inherited from the parent . After training , they can be replaced back into Qφ , ϕ ( Z|X ) obtaining , in particular , the nonlinear dynamics aφ ( z ) for the latent space . In this novel way , VIND achieves the reuse of the latent dynamics h ( zt|zt−1 ) of the GM in the RM . Concretely , let qφ , ϕ be a Laplace approximation to Qφ , ϕ , qφ , ϕ ( Z|X ) = N ( Pφ , ϕ ( X ) , C −1 φ , ϕ ( X ) ) . ( 9 ) The mean Pφ , ϕ in Eq . ( 9 ) is the solution to the equation ∂ ∂Z logQφ , ϕ ( Z|X ) ∣∣∣∣ Z=P = 0 , ( 10 ) and the precision is given by [ Cφ , ϕ ( X ) ] ij = ∂2 ∂Zi∂Zj logQφ , ϕ ( Z|X ) ∣∣∣∣ Z=Pφ , ϕ ( X ) ≡ [ sφ , ϕ ( Pφ , ϕ ( X ) , X ) ] ij , ( 11 ) where Eq . ( 11 ) defines sφ , ϕ . A closed form solution for Eq . ( 10 ) is not possible in general . However , for a large class of distributions , and in particular for any Qφ , ϕ such that logQφ , ϕ includes terms quadratic in Z , it is possible to rewrite Eq . ( 10 ) in the form P = rφ , ϕ ( P , X ) . ( 12 ) In this form , the latter can be approximately solved by making use of the FPI method . This approximate solution provides the mean for the variational approximation qφ , ϕ ( Z|X ) . Henceforth , VIND ’ s algorithm includes two steps per epoch that are carried out in alternation ( see Algorithm 1 in App . A ) . The first step is a FPI that , for the current values of the parameters φ , ϕ , determines the mean and variance of a Laplace approximation to the parent . The second is a regular ADAM gradient ascent update Kingma & Ba ( 2014 ) with respect to the ELBO objective . Upon convergence , the set of parameters φ , ϕ , θ that maximizes the ELBO can be plugged into aφ ( z ) , to obtain a dynamical rule that interpolates between the different latent trajectories inferred from the data trials . In the experiments conducted in this paper , the nonlinear dynamics is specified as aφ ( z ) = Aφ ( z ) z where Aφ ( z ) is a state-space dependent dZ × dZ matrix . We call this evolution rule , a Locally Linear Dynamical System , and the resulting inference algorithm LLDS/VIND . To derive it , consider a parent distribution distribution Qφ , ϕ , as defined in Eq . ( 7 ) . The mean µϕ and the standard deviation Λ−1ϕ in Eq . ( 8 ) are represented as deep neural networks : µϕ = NNϕµ ( xt ) , Λϕ = NNϕΛ ( xt ) . ( 13 ) The remaining ingredient of Qφ , ϕ is the shared evolution law Hφ , Eq . ( 3 ) . We write the hφ factors that determine the latent evolution model as hϕ ( zt+1|zt ) = exp { −1 2 ( zt+1 −Aϕ ( zt ) zt ) T Γ ( zt+1 −Aϕ ( zt ) zt ) } , ( 14 ) where Γ is a constant precision matrix . Eq . ( 14 ) can be thought of as describing the stochastic evolution of a “ locally linear ” dynamical system : zt+1 ∼ A ( zt ) zt + noise . ( 15 ) By LLDS/VIND we refer to VIND with this specific parameterization of the nonlinearity aφ ( zt ) in Eq . ( 5 ) . LLDS/VIND has some desirable features : 1 . The limit of linear evolution is easily taken as Aφ ( zt ) → const .. 2. maxZ |Aφ ( zt ) − I| is a simple measure of the smoothness of the latent trajectories . The equation for the Laplace mean is found by taking derivatives of Qφ , ϕ with respect to the hidden path : ∂ logQφ , ϕ ∂Z ∣∣∣∣ Z=P = 0 . ( 16 ) Using Eqs . ( 3 ) and ( 8 ) we find logQφ , ϕ = logCφ , ϕ − 1 2 [ ( Z−Mϕ ) TΛϕ ( Z−Mϕ ) + ZTSφ ( Z ) Z ] ( 17 ) where Mϕ = { µϕ ( x1 ) , . . . , µϕ ( xT ) } and Sφ ( Z ) is a state-space-dependent covariance with tridiagonal form : Sφ ( Z ) = AT0 ΓA0 −AT0 Γ 0 . . . . . . 0 0 −ΓA0 AT1 ΓA1 AT1 Γ . . . . . . 0 0 0 −ΓA1 . . . . . . . . . 0 0 . . . . . . . . . . . . . . . . . . . . . 0 0 . . . . . . . . . −ATT−2Γ 0 0 0 . . . . . . −ΓAT−2 ATT−1ΓAT−1 −ATT−1Γ 0 0 . . . . . . 0 −ΓAT−1 ATTΓAT ( 18 ) Here Ai ≡ Aφ ( zi ) . Inserting Eq . ( 17 ) into Eq . ( 16 ) we obtain the LLDS/VIND equation for the posterior , Eq . ( 12 ) , with rφ , ϕ ( P , X ) = [ Λϕ + Sφ ( P ) ] −1 ( ΛϕMϕ − 1 2 PT ∂Sφ ( P ) ∂P P ) . ( 19 ) Note that it is not necessary to find the value of Cφ , ϕ . Moreover , we remark that the matrix Λϕ + Sφ ( P ) is block-tridiagonal and can be inverted in O ( T ) . Although Eq . ( 12 ) can not be solved analytically , a solution can be approximated by choosing an initial point P ( 0 ) and iterating P ( n ) = rφ , ϕ ( P ( n−1 ) , X ) , ( 20 ) The VIND method assumes that this FPI converges . In practice , this assumption can be guaranteed throughout training by appropriate choices of hyperparameters and network architectures ( see App . A.2 ) . During training , the mean P represents the best current estimate of the latent trajectory . Note that the FPI step in Eq . ( 20 ) mixes all the components in P. In particular , the i-th component of P ( n ) depends in general on all the components of P ( n−1 ) via the inverse covariance in Eq . ( 19 ) . Thus , at every training epoch , the best estimate for the path at a specific time point t contains information from the complete data , both to the past and to the future of t. This is analogous to the Kalman smoother .

Summary: The paper presents a variational inference approach for locally linear dynamical models. In particular,  the latent dynamics are drawn from a Gaussian approximation of the parent variational distribution,  enabled by Laplace approximations with fixed point updates, while the parameters are optimized the resulting stochastic ELBO. Experiments demonstrate the ability of the proposed approach to learning nonlinear dynamics, explaining data variability, forecasting and inferring latent dimensions.