Year: 2018

Abstract:
Latent variable models have been widely applied for the analysis and visualization of large datasets. In the case of sequential data, closed-form inference is possible when the transition and observation functions are linear. However, approximate inference techniques are usually necessary when dealing with nonlinear evolution and observations. Here, we propose a novel variational inference framework for the explicit modeling of time series, Variational Inference for Nonlinear Dynamics (VIND), that is able to uncover nonlinear observation and latent dynamics from sequential data. The framework includes a structured approximate posterior, and an algorithm that relies on the fixed-point iteration method to find the best estimate for latent trajectories. We apply the method to several datasets and show that it is able to accurately infer the underlying dynamics of these systems, in some cases substantially outperforming state-of-the-art methods.

heading: 1 INTRODUCTION
text: 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).

heading: 2 VARIATIONAL INFERENCE WITH NONLINEAR DYNAMICS (VIND)
text: 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ϕ, cannot 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 cannot be formulated directly in terms of the parent, since expectation values with respect to Qφ,ϕ cannot 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) cannot 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.

heading: 3 RELATION TO PREVIOUS WORK
text: The problem of inference for sequential data has been treated extensively in the literature. The GfLDS and PfLDS models introduced in Archer et al. (2015); Gao et al. (2016) are particular cases of VIND in which the dynamics in the latent space is linear and time-invariant, i.e. zt|zt−1 ∼ N ( Azt−1, Q ) . In the jargon used in this paper, this corresponds to the situation in which the parent distribution is Gaussian, and therefore equal to its own Laplace approximation. Eq. equation 10 can be solved analytically in this case and no FPI step is needed. Gaussian Process Factor Analysis (GPFA) Yu et al. (2009) assumes linear, time-invariant dynamics as well as a linear observation model, i.e. xt|zt ∼ N ( Czt + d, R ) , for some C, d, and R. We will be explicitly comparing our results to results obtained by these models.
In Krishnan et al. (2015), Deep Kalman Filters (DKF) were proposed to handle variational posterior distributions that describes nonlinear evolution in the latent space. Their approximate posterior, analogous to the parent distribution in this paper, is plugged directly into the ELBO. This imposes
some restrictions in the form the posterior can take - for instance, it must be Gaussian conditioned on the observations. VIND can handle factorizations of the parent distribution that are not restricted in this way, an example being LLDS/VIND, which has the form in Eq. equation 7. VIND’s ability to handle unnormalizable parent distributions is due to the fact that VIND’s actual approximate posterior is always strictly normal. The same authors built upon their idea in Krishnan et al. (2016), where a variational posterior was proposed that partially uses the conditional structure implied by the generative model. In this paper, a similar prescription is used by assuming that Qφ,ϕ and pφ,ϕ share exactly the same factorization for the latent evolution.
The authors of Johnson et al. (2016) combine probabilistic graphical models with message passing in an approach based on conjugate priors. The approximate posterior distributions considered in that work are restricted by the conjugacy requirements, in particular, the evolution term must belong to the exponential family. VIND’s parent distribution is not subject to this requirement. However, since VIND’s actual approximate posterior is still Gaussian, it may be possible to combine the two methods into one that can handle both nonlinear evolution and discrete latent variables.
In Chung et al. (2015), Gaussian noise is added to the deterministic evolution rule of an RNN in the context of a variational autoencoder, termed VRNN. Similarly to LLDS/VIND, these authors share the evolution factorization between the generative model and the approximate posterior and, indeed, the only difference between the structure of their model and that of LLDS/VIND is that the evolution there is expressed as an RNN instead of as an LLDS. However, their inference algorithm only uses past data to estimate the hidden state at any given time. VIND’s algorithm, based on the FPI, uses information both from the past and from the future to estimate the latent paths. In Kalantari et al. (2018) a non-parametric approach was taken to determine the best latent dimension in an LDS. It would be interesting to apply those same methods to VIND. Finally, in Pandarinath et al. (2018); Sussillo et al. (2016) a sophisticated, bidirectional, Deep Learning-based RNN architecture called LFADS was proposed with neuroscience applications in mind. For both LFADS and DKF, we found difficult to modify their code to compute the quantities that are used in this paper to evaluate the quality of training. However, given the expressive power of these works, we expect them to perform comparably to VIND in the tasks considered in the next section.

heading: 4 RESULTS
text: We demonstrate the capabilities of LLDS/VIND by applying it to four characteristic datasets, each of which illustrates a VIND feature. First, we use synthetically generated 10D data with added Gaussian noise, and latent evolution determined by an Euler discretization of the Lorenz system. This dataset is the simplest and cleanly illustrates VIND’s ability to infer the underlying nonlinear dynamics. Secondly we use a multi-electrode neural recording from a mouse performing a delayeddiscrimination task. LLDS/VIND is run with both Gaussian and Poisson observation models. It is found that while a Gaussian observation model is superior for the explaining the variance in the data, the Poisson model performs better when it comes to interpolation of the dynamics. This is a common VIND tradeoff. The third dataset consists of a 1D voltage measurement from single-cell recordings. The problem in this case is not dimensionality reduction but rather to determine the underlying dynamics. We find that the minimum number of latent dimensions VIND requires to describe this dataset coincides with the naive expectation. Finally, we apply VIND to discovering dynamics in a more complex dataset coming from dorsal cortex calcium imaging, and find that it is possible to model the data using a surprisingly low number of latent dimensions. We also use VIND to reconstruct dynamics from one side of the brain to the other.
Given an inferred starting point in state-space, the quality of the dynamics uncovered by LLDS/VIND can be ascertained by evolving the system k steps into the future without any input data. To clarify terminology, this is not strict prediction in the sense of pure extrapolation, since we use information about all xt, both in the past and in the future, to infer the starting point. In order to avoid doubt, we use the term forward interpolate. Forward interpolation essentially tests the extent to which the dynamics are accurately learned. We take VIND’s capability for forward interpolation as the main measure of the fit’s success. As we will show, this task remains highly challenging for simpler smoothing priors like the latent LDS, and it is one of the key strengths of VIND.
To make this analysis quantitative, we compute the k-step mean squared error (MSEk) on test data, and its normalized version, the R2k, defined as
MSEk = T−k∑ t=0 (xt+k − x̂t+k)2 , R2k = 1− MSEk∑T−k t=0 (xt+k − x̄) 2
(21)
where x̄ is the data average for this trial and x̂t+k is the prediction at time t+k. The latter is obtained by i) using the full data X to obtain the best estimate for zt, ii) using k times the LLDS/VIND evolution equation zt+1 = Aϕ(z)zt, or zt+1 = Azt for the LDSs, to find the latent state k time steps in the future, and iii) using the generative network to compute the forward-interpolated observation. Note that in particular, k = 0 corresponds to the standard R2. The more general R2k ensures that VIND yields more than just a good autoencoder. We will be comparing results obtained with LLDS/VIND to several models, namely, GfLDS, PfLDS, and GPFA (see Sec. 3 for details).

heading: 4.1 LORENZ SYSTEM
text: The Lorenz system is a classical nonlinear differential equation in 3 independent variables.
ż1 = σ(z2 − z1) , ż2 = z1(ρ− z3)− z2 , (22) ż3 = z1z2 − βz3 .
This is a well studied system with chaotic solutions and that serves to cleanly demonstrate VIND’s capabilities for inferring nonlinear dynamics. We generated numerical solutions of the Lorenz system from randomly generated initial conditions, for σ = 10, ρ = 28, β = 8/3, and additive Gaussian noise. Gaussian observations in a 10D space were then generated with the mean specified by a z-dependent neural network. The complete synthetic data consisted of 100 trials, each comprising 250 time-steps, of which 80% was used for training and the remaining for validation.
The results of the fit to this data are shown in Fig. 1. The left panel shows the R2k comparison for VIND and GfLDS fits, with dZ = 3. Strikingly, for this dataset, VIND’s performance is 2-3 orders of magnitude better than GfLDS. The right panel illustrates VIND’s capability to infer properties of the underlying dynamics: VIND hits peak performance at dZ = 3, the true dimensionality of this system. In the rightmost panel, all the paths inferred by VIND have been put together, showing the famous butterfly pattern.

heading: 4.2 ELECTROPHYSIOLOGY
text: We analyzed data collected from mice performing a delayed discrimination task in a simultaneous recording session, using multi-unit electrophysiology (64 channel Janelia silicon probe) Guo et al. (2014); Li et al. (2015). In this task, the animals were trained to discriminate the location of a pole using whiskers. The pole was presented at t = −1.3s, and an auditory go cue at t = 0 signaled the beginning of the response epoch. During response, the mice reported the perceived pole position by licking one of two lick ports. Neurons in this task exhibit complex dynamics across behavioral
epochs; some neurons show ramping and persistent activity from sample to delay, which relates to the preparation of the choice at response Guo et al. (2014); Li et al. (2015); Wei et al. (2018), while some other neurons show the peaking activity in response to the behavioral epochs, see Fig. (2, left).
We asked whether VIND can capture the variety of neural dynamics using a few latent observations. The data was fitted for dZ = 5, using a Poisson observation model. The fit not only reproduces the neural observation, but also provides insights to the dynamics in the latent space. Details can be found in Fig. 6 in App. B. Subsequently, a 10-fold cross-validation method was used to decide the performance of fit using VIND’s Gaussian and Poisson observation models with up to 12 dimensions in the latent space, regardless of trial type. The R2 was computed to determine the performance of VIND as compared to other models. For VIND, both Poisson and Gaussian observation models were used. These are compared to a Peristumulus Time Histogram (PSTH), a GPFA model Yu et al. (2009), as well as GfLDS and PLDS Archer et al. (2015); Gao et al. (2016). The results are shown in the center panel in Fig. 2. We found that nonlinear Gaussian VIND performs the best regarding explained variance of the data.
Next we analyzed the forward interpolation capabilities of the fitted models to the neural data using R2k. In this case the Poisson observation model gives a substantially better forward interpolation, signaling a dynamical system that more accurately represents the data evolution. This can be seen in the right panel in Fig. 6. These two results combined exemplify the VIND tradeoff between explained variance and forward interpolation capabilities. Using Poisson observations, VIND is less able to fit the higher frequency components of the data. The resulting dynamical system, however, is smoother and more appropriately captures the evolution of the system, see App. B.

heading: 4.3 SINGLE CELL VOLTAGE DATA
text: We demonstrate VIND’s versatility to uncover underlying dynamics of a system by applying it to 1D voltage electrophysiology data recorded from single cells. This is not a dimensionality reduction problem but rather one of recovering the latent phase space from a single variable to identify the ‘true dimensionality’ of the system under study. The data is the publicly available Allen Brain Atlas dataset Jones et al. (2009).
Intracellular voltage recordings from cells from the Primary Visual Cortex of the mouse, area layer 4 were selected. Trials with no spikes were removed, resulting in 44 trials from 7 different cells. The input for each of the remaining trials consists of a step-function with an amplitude between 80 and 151pA. Observations were split into training (30 trials) and validation sets (14 trials). The data was then down-sampled from 50, 000 time bins (sample rate of 50 kHz) to 5, 000 in equal-time intervals, and subsequently normalized by dividing each trial by its maximal value.
LLDS/VIND was fit to this data for dZ = 2, . . . , 8, repeated across 10 runs. The top three fits were averaged and the results are summarized in Fig. 3. The center panel displays the R210 values for each choice of latent dimensionality. The fits consistently improve up to dZ = 5, after which there are diminishing returns. Single cell voltage data has traditionally been modeled using variants of the classical Hodgkin-Huxley neuron model (Hodgkin & Huxley (1952)), a set of nonlinear differential
equations in 4 independent variables, plus an optional independent input current. It is interesting that 5 is exactly the minimal number of latent dimensions that provide a good VIND fit for this data. The right panel displays R2k with dZ = 5 for VIND and for GfLDS. VIND outperforms GfLDS by an order of magnitude.
We show the forward-interpolated observations and sample paths for selected runs using VIND and GfLDS in Fig. 4 . The left panel represents the observations over a rolling window, k = 10 time-points in advance for both VIND and GfLDS. The dynamics inferred by GfLDS are unable to capture the nonlinearities for both the hyperpolarization and depolarization epochs. The latent trajectories for this data are plotted in the center panel. The dimensions inferred by VIND exhibit similar behavior to that of Hodgkin-Huxley gating variables. As shown in the right panel, we find that in state-space, spikes are represented by big cycles (red), while interspiking fluctuations correspond to separate regions of phase space (blue). For each trial, the height of the voltage spike is coded in the diameter of the cycles.

heading: 4.4 SPONTANEOUS ACTIVITY IN WIDEFIELD IMAGING DATA
text: The unsupervised modeling of spontaneous brain activity is inherently challenging due to the lack of task structure. Here, we model the temporal dynamics of widefield optical mapping (WFOM) data and simultaneous behavior recorded from an awake head-fixed mouse during spontaneous activity Ma et al. (2016a). This data was recorded and corrected for hemodynamics in the Laboratory for Functional Optical Imaging at Columbia University. The recording and preprocessing details for both cortical dynamics and the movement speed signal are provided in App. D. An example frame of the data is shown in Fig. 5 (top-left). The preprocessing of the WFOM cortical data leads to reduced-dimension, denoised cortical activity (as detailed in App. D). The temporal activity of the
cortex and the movement speed (jointly called X) are simultaneously modeled using both GfLDS and VIND, with the results for validation data on one mouse shown in Fig. 5, where dX = 148, and dZ = 9. The k−step forward interpolation is shown in Fig. 5 (bottom-left), with varying k, for both VIND and GfLDS. 4 of the 148 dimensions of X and X̂ on validation data are shown in Fig. 5 (center). VIND is seen to outperform GfLDS, capturing the fine-tuned dynamics in X , thus also leading to better interpolations. We highlight VIND’s capability to roughly capture the dynamics of the whole superficial dorsal cortex using a 9-D latent vector and the corresponding evolution and generative network.
Next, a VIND model was fit to the brain dynamics of only the left hand side (LHS) of the brain, after similar preprocessing of the data. Here, dXLHS = 60, dZLHS = 9. A separate neural network was fit from the latents learned on the left hand side (ZLHS) to the temporal dynamics of the right hand side (RHS) of the brain (XRHS ; dXRHS = 66), with an MSE loss function. The goal was to infer dynamics from one half of the brain to the other. 5 out of 66 reconstructions of the temporal dynamics of the RHS in held-out data are shown in Fig. 5 (right) (variance weighted average R2 = 0.49 for entire data). For comparison, we ran a baseline CCA analysis which yielded an R2 of 0.45. This shows that the latent variables learned by VIND on one half of the brain are useful to coarsely reconstruct the temporal dynamics of the other half.

heading: 5 DISCUSSION
text: In this work we introduced VIND, a novel variational inference framework for nonlinear latent dynamics that is able to handle intractable distributions. We successfully implemented the method for the specific case of Locally Linear Dynamical Systems, which allows for a fast inference algorithm (linear in T ). When applied to real data, VIND consistently outperforms other methods, in particular methods that rely on an approximate posterior representing linear dynamics. Furthermore, VIND’s fits yield insights about the dynamics of these systems. Highlights are the ability to identify the transition points and distinguish among trial types in the electrophysiology task, the dimensionality suggested by VIND’s fits for the single-cell voltage data, and the ability of the latents learned from one half of the brain to reconstruct activity from the other half in widefield imaging data. Moreover, VIND can be naturally extended to handle labelled data and data with inputs. This is work in progress.
LLDS/VIND is written in tensorflow and the source code is publicly available.

heading: A LLDS/VIND
text: In this appendix, we provide details of the VIND framework for the LLDS parameterization of the hidden dynamics.
A.1 ALGORITHM
As detailed in Alg. 1, each training epoch consists of two steps that are carried in alternate fashion: the FPI that updates the best estimate of the latent path and a gradient descent step that updates the model parameters. For optimization of the ELBO we used stochastic gradient descent. As it is customary, in order to estimate the gradients, the so called “reparameterization trick” is used. Samples are extracted from the variational approximation qφ,ϕ:
Zi = P (ep) φ,ϕ(Xi) +
[ C
(ep) φ,ϕ(Xi)
]−1/2 (23)
where is a standard normal, (Kingma & Welling (2013); Jimenez Rezende et al. (2014)).
A.2 IMPLEMENTATION DETAILS
We provide here extra details of the LLDS/VIND implementation that was used throughout the paper:
• The initial path estimates P(0)i are taken to be µϕ(Xi). • VIND is initialization-sensitive. Empirically, we found that it is important that the initial path
estimates fall within a region where the nonlinearity is not severe (maxPi |Aφ(zt)−I| . 0.1 for every trial i). • To encourage smoothness of the latent dynamics, Aφ(zt) was specified as
Aφ(zt) = A + α ·Bφ(zt) (24)
where A is a state-space-independent linear transformation initialized to the identity, α is a nontrainable hyperparameter of the model, and Bφ(zt) = NNφB (zt). This set up has the added benefit that α = 0 is equivalent, both the statistical model and the algorithm, to GfLDS/PfLDS (Archer et al. (2015); Gao et al. (2016)). • The local transformation Aφ(zt) is redundant (it is akin to a gauge transformation in physics
parlance). To see this, note that for every zt, the image of the transformation Aφ(zt)zt is at most Rn. On the other hand Aφ(zt) has dimensionality Rn 2 . In other words, given zt and
zt+1, there is a continuum of matrices Aφ(zt) that satisfy zt+1 = Aφ(zt)zt. It follows that Aφ(zt) can be substantially restricted without loss of generality (“fixing the gauge”). We found that the best results were obtained when Aφ(zt) was constrained to be symmetric. • Experiments showed that setting the number of FPIs at n = 2 in Eq. (20) is enough for
producing good convergence results across datasets. • In all experiments we found that there is no noticeable decrease in performance if the
gradient terms in Eq. (19), and the corresponding ones for sφ,ϕ are neglected. These terms are subleading compared to ΛϕMϕ both because they are proportional to the nonlinearity, small as required by smoothness, and because the gradient is applied on a deep neural network. • The FPI in Eq. (12) is in the contractive regime within a domain D, P(0)i ∈ D; D ⊂ RT×dZ
when the Jacobian J of the map rφ,ϕ satisfies that the absolute value of its determinant is smaller than 1. This can be guaranteed for any starting point P(0)i by invoking the Gershgorin Circle Theorem (Eisstein), which yields a bound on the absolute value of the eigenvalues of a matrix. For the specific case of LLDS/VIND, the entries of J are suppressed both by the small hyperparameter α and to the gradients of the deep neural network Bφ(zt), Eq. (24. This makes the Gershgorin bound easy to satisfy. Convergence of the algorithm is not straightforward from this argument, since it must cycle through different batches of training data, which may lie in different basins of the FPI map. Nevertheless, in practice we found the Jacobian suppression to be enough to guarantee convergence.

heading: B DETAILS OF THE ELECTROPHYSIOLOGY FIT
text: The experimental details in Sec. 4.2 are as follows. The recording session contains 18 simultaneous recorded units, with 74 lick-right trials (in blue; posterior pole location) and 100 lick-right trials (in red; anterior pole location). Spike counts were binned in a 67 ms non-overlapped time window, where the maximum spike count is 10 and minimum is 0. The fit covers time points (almost the entire recordings) from -0.5 sec from sample onset to 2.0 sec from response (wherefore the trial contains 77 time bins).
Fig. 6A shows the average neuronal activity of 3 representative cells in the recordings. Cell #1 is a typical neuron with small separation of trials, but strong peaking activity at transition from delay to response epochs. Cells #2, 3 exhibit the stereotypical ramping activity and separations of different trial types, which are assumed for preparation of the movements. Both setups of VIND models (Poisson and Gaussian, nonlinear evolutions, dz = 5) can reproduce the complex and variable neural dynamics in the held-out trials (9 lick-left trials; 9 lick-right trials). Particular, Gaussian VIND model can capture the changes of dynamics on even finer timescales (dash lines vs dot lines; Poisson vs Gaussian VIND models). This observation agrees with higher performance of explained variance (R2; Fig. 2).
The latent dynamics are smoother in the Poisson VIND model, (Fig. 6BC). We have found that in VIND fits, smoother trajectories are correlated with superior performance in the forward interpolation tasks. Intuitively, for noisier latent paths, the algorithm attempts to ascribe some of the variance to the dynamical system, which hurts the forward interpolation capabilities. In the Poisson VIND fit represented in Fig. 6, the latent dynamics in dimension 2 and 3 appears to represent the preparation of the choice where the neural dynamics for different trial types gradually diverges with time. The dynamics in dimension 1 shows rapid peaking dynamics at the transitions of the behavioral epochs. However, those two types of dynamics were mixed and separations of trial types were in Gaussian VIND model. In general, ramping and peaking dynamics is not operated by distinguishable groups of neurons, yet to our surprise they are separated in the latent space.

heading: C DETAILS OF THE SINGLE CELL VOLTAGE DATA FITS
text: For the Allen data, Fig. 7 shows simulated paths (forward interpolation with noise) versus the corresponding real data. Fig. 8 shows several views of the same two latent paths corresponding to two different input currents.

heading: D PREPROCESSING OF WIDEFIELD IMAGING DATA
text: Macro-scale wide-field optical mapping (WFOM) is an increasingly popular technique for surveying neural activity over very large areas of cortex with high temporal resolution. WFOM can image the fluorescence of genetically-encoded calcium (GCaMP6f) indicators using LED illumination and camera detection scheme. We use methods for correcting fluorescence recordings of neural activity for confounding contamination by changes in hemoglobin concentration and oxygenation as in Ma et al. (2016b), by measuring both neural fluorescence signals and hemodynamics. This correction provides us with an accurate change in fluorescence of neural regions (∆F/F ). An example frame
of the data is shown in Fig. 5A, 464-by-473 pixels. The activity of the mouse is simultaneously recorded using a webcam pointed at the mouse’s body, and the movement speed at time t is taken as a 1D signal consisting of the standard deviation of the difference in value of all pixels from time t− 1 to time t.
We use a WFOM recording of length 2 minutes, where the signals are sampled at 10Hz, thus leading to 1200 time points. We normalize ∆F/F to lie between 0 and 1 for every video, and then apply block singular value decomposition (SVD) to the videos for denoising and dimensionality reduction Buchanan et al. (2018). First, we fit an anisotropic Wiener filter in a 4 × 4 neighborhood of each pixel to reduce uncorrelated noise while preserving spatially-local, time-correlated signals. Next, the video is partitioned into 25 (5× 5) blocks, and SVD is performed on the pixels in each block. The temporal components are ranked according to a metric defined on their empirical autocorrelation function, and components that fall within a 99% confidence interval of Gaussian white noise are discarded. Moreover, those temporal components that have a signal-to-noise ratio lower than 1.6 are also discarded. The remaining temporal components from each block are concatenated, and these form the X matrix, here 147× 1200. This is augmented using a 1D behavior signal that is extracted using the standard deviation of successive frames from a webcam recording the lateral view of the mouse’s body, representing the speed of the mouse’s movements in arbitrary units. We used different sessions of recording from the same mouse, preprocessed in the same way, to obtain training and validation data.

Review summaries:
This paper discusses a algorithm for variational inference of a non-linear dynamical models. In this paper model assumption is to use single stage Markov model in latent space with every latent variable Z_t to be defined Gaussian distributed with mean depends on Z_(t-1) and time invariant variance matrix lambda. The non linearity in transition is encoded in mean of Guassian distribution. For modeling the likelihood and observation model, the Poisson or Normal distribution are used with X_t being sampled from another Gaussian or Poisson distribution with the non-linearty being encoded in the parameters of distribution with variable Z_t.  This way of modeling resembles so of many linear dynamical model with the difference of transition and observation distribution have nonlinearity term encoded in them. 

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.