A N OVEL V ARIATIONAL F AMILY FOR H IDDEN N ONLINEAR M ARKOV M ODELS
Anonymous authors
Paper under double-blind review
A BSTRACT
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 dynamics and latent dynamics.
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 transition functions from sequential data.
The framework includes a structured approximate posterior, and an algorithm that relies on the ﬁxed-point iteration method to ﬁnd 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.
I NTRODUCTION
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.
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 ( 1 D) observation.
From such 1 D 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.
; Chung et al.
Gao et al.
(2016; 2015); Hernandez et al.
; Johnson et al.
; Krishnan et al.
Linderman et al.
; Pandarinath et al.
; Sussillo et al.
; Zhao & Memming Park (2017);
Wu et al.
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 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.
, so optimizing for it can be interpreted as indirectly optimizing Q ( Z | X ) .
The second novel aspect of VIND is the use of the ﬁxed-point iteration (FPI) method to signiﬁcantly 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 hidden dynamics uncovered by VIND generate data that is indistinguishable from the original data, when 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).
V ARIATIONAL I NFERENCE WITH N ONLINEAR D YNAMICS (VIND)
For a set of temporally ordered, correlated, noisy observations X ≡ { x 1 , . . .
x T } , x t ∈ R d X , a latent variable model proposes an additional, time-ordered set of random variables Z ≡ { z 1 , . . .
z T } , z T ∈ R
d
Z that is hidden from view.
The hidden state z t is endowed with a stochastic dynamics:
z t +1 ∼ p ( z t +1 | z 1: t ) by which it evolves.
The observations x t are generated by drawing samples from a z t -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.
[Equation]
[Equation]
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:
[Equation]
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, x t | z t ∼
N (cid:0) m θ ( z t ) , Σ (cid:1) , or Poisson, x t | z t ∼ Poisson (cid:0)
λ θ ( z t ) (cid:1) .
The respective mean, m θ ( z t ) , and rate, λ θ ( z t ) , are arbitrary nonlinear functions of the latent state z t , that we represent as neural networks.
The standard deviation Σ of the Gaussian observation model is taken to be z t -independent.
c φ,θ is a normalization constant.
H φ is the latent evolution term in Z -space with a Markov Chain structure
Johnson et al.
; Krishnan et al.
; Archer et al.
; Gao et al.
[Equation]
(5) where a φ ( z ) is an arbitrary nonlinearity.
From Eq.
(2), the posterior distribution of the Generative Model (GM) can be factorized as
[Equation]
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:
[Equation]
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 deﬁniteness, the case
[Equation]
is considered in this work, where µ ( x ) and Λ ( x ) are nonlinear maps.
In Eq.
(7), α φ,ϕ is a normalization constant that, regardless of the speciﬁc form of G ϕ , cannot be computed in closed form.
This is most easily seen by noticing that due to the nonlinearity h ( z T | z T − 1 ) , integration with respect to z T yields an intractable z T − 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 deﬁne a Gaussian approximation q φ,ϕ ( Z | X ) to the parent; this child distribution is then used as the variational posterior in Eq.
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
( z t | z t − 1 ) of the
GM in the RM.
Concretely, let q φ,ϕ be a Laplace approximation to Q φ,ϕ ,
[Equation]
The mean P φ,ϕ in Eq.
(9) is the solution to the equation
[Equation]
and the precision is given by
[Equation]
where Eq.
(11) deﬁnes 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 log Q φ,ϕ includes terms quadratic in Z , it is possible to rewrite Eq.
(10) in the form
[Equation]
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 ﬁrst 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 speciﬁed as a φ ( z )
=
A φ ( z ) z where A φ ( z ) is a state-space dependent d Z × d Z 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 deﬁned in Eq.
The mean µ ϕ and the standard deviation Λ − 1 ϕ in Eq. (8) are represented as deep neural networks: [Equation]
The remaining ingredient of Q φ,ϕ is the shared evolution law H φ , Eq.
We write the h φ factors that determine the latent evolution model as
[Equation]
where Γ is a constant precision matrix.
Eq.
(14) can be thought of as describing the stochastic evolution of a “locally linear” dynamical system:
[Equation]
By LLDS/VIND we refer to VIND with this speciﬁc parameterization of the nonlinearity a φ ( z t ) in
Eq.
LLDS/VIND has some desirable features:
max Z | A φ ( z t ) − 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:
[Equation]
Using Eqs.
(3) and (8) we ﬁnd
[Equation]
where M ϕ = { µ ϕ ( x 1 ) , . .
. , µ ϕ ( x T ) } and S φ ( Z ) is a state-space-dependent covariance with tridiagonal form:
[Equation]

[Equation]
[Equation]
[Equation]
[Equation]
A
[Equation]
[Equation]
[Equation]
[Equation]
[Equation]
 (18)
Here A i ≡ A φ ( z i ) .
Inserting Eq.
(17) into Eq.
(16) we obtain the LLDS/VIND equation for the posterior, Eq.
(12), with
[Equation]
Note that it is not necessary to ﬁnd 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 [Equation]
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.
Thus, at every training epoch, the best estimate for the path at a speciﬁc 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.
R ELATION TO P REVIOUS W ORK
The problem of inference for sequential data has been treated extensively in the literature.
The GfLDS and PfLDS models introduced in Archer et al.
; Gao et al.
(2016) are particular cases of VIND in which the dynamics in the latent space is linear and time-invariant, i.e. z t | z t − 1 ∼ N (cid:0)
A z t − 1 , Q (cid:1) .
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. x t | z t ∼ N (cid:0)
C z t + d, R (cid:1) , 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.
, 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.
Finally, in
Pandarinath et al.
; 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 difﬁcult 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.
R ESULTS
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 ﬁnd 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 ﬁnd 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 x t , 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 ﬁt’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 (MSE k ) on test data, and its normalized version, the R 2 k , deﬁned as
[Equation]
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 z t , ii) using k times the LLDS/VIND evolution equation z t +1 =
A ϕ ( z ) z t , or z t +
=
A z t for the LDSs, to ﬁnd 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 R 2 .
The more general R 2 k 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.
for details).
4. L ORENZ SYSTEM
The Lorenz system is a classical nonlinear differential equation in 3 independent variables.
[Equation]
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 speciﬁed 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 ﬁt to this data are shown in Fig.
The left panel shows the R 2 k comparison for
VIND and GfLDS ﬁts, with d Z = 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 d Z = 3 , the true dimensionality of this system.
In the rightmost panel, all the paths inferred by VIND have been put together, showing the famous butterﬂy pattern.
4. E LECTROPHYSIOLOGY
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.
; Li et al.
In this task, the animals were trained to discriminate the location of a pole using whiskers.
The pole was presented at t = − 1 .
3 s, 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.
; Li et al.
; 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 ﬁtted for d Z = 5 , using a Poisson observation model.
The ﬁt 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 ﬁt using VIND’s Gaussian and Poisson observation models with up to 12 dimensions in the latent space, regardless of trial type.
The R 2 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.
; Gao et al.
The results are shown in the center panel in Fig.
We found that nonlinear Gaussian VIND performs the best regarding explained variance of the data.
Next we analyzed the forward interpolation capabilities of the ﬁtted models to the neural data using
R 2 k .
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.
These two results combined exemplify the VIND tradeoff between explained variance and forward interpolation capabilities.
Using Poisson observations, VIND is less able to ﬁt 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.
4. S INGLE C ELL V OLTAGE D ATA
We demonstrate VIND’s versatility to uncover underlying dynamics of a system by applying it to 1 D 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.
Intracellular voltage recordings from cells from the Primary Visual Cortex of the mouse, area layerwere 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 151 pA. 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 ﬁt to this data for d Z = 2 , . .
. , 8 , repeated across 10 runs.
The top three ﬁts were averaged and the results are summarized in Fig.
The center panel displays the R 210 values for each choice of latent dimensionality.
The ﬁts consistently improve up to d Z = 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 ﬁt for this data.
The right panel displays R 2 k with d Z = 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 =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 ﬁnd that in state-space, spikes are represented by big cycles (red), while interspiking ﬂuctuations 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.
4. S PONTANEOUS A CTIVITY IN W IDEFIELD I MAGING D ATA
The unsupervised modeling of spontaneous brain activity is inherently challenging due to the lack of task structure.
Here, we model the temporal dynamics of wideﬁeld optical mapping (WFOM) data and simultaneous behavior recorded from an awake head-ﬁxed 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 d X = 148 , and d Z = 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 ﬁne-tuned dynamics in X , thus also leading to better interpolations.
We highlight VIND’s capability to roughly capture the dynamics of the whole superﬁcial dorsal cortex using a 9 -D latent vector and the corresponding evolution and generative network.
Next, a VIND model was ﬁt to the brain dynamics of only the left hand side (LHS) of the brain, after similar preprocessing of the data.
Here, d X LHS = 60 , d Z LHS
= 9 .
A separate neural network was ﬁt from the latents learned on the left hand side ( Z LHS ) to the temporal dynamics of the right hand side (RHS) of the brain ( X RHS ; d X RHS = 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 R 2 = 0 .for entire data).
For comparison, we ran a baseline CCA analysis which yielded an R 2 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.
D ISCUSSION
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 speciﬁc 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 ﬁts 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 ﬁts 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 wideﬁeld 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 tensorﬂow and the source code is publicly available.
A LLDS/VIND
In this appendix, we provide details of the VIND framework for the LLDS parameterization of the hidden dynamics.
A. A LGORITHM
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 φ,ϕ :
[Equation]
) where (cid:15) is a standard normal, (Kingma & Welling (2013);
Jimenez Rezende et al.
A. I MPLEMENTATION 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 µ ϕ ( X i ) .
• 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 ( max P i | A φ ( z t ) − I | (cid:46) 0 .for
every trial i ).
• To encourage smoothness of the latent dynamics, A φ ( z t ) was speciﬁed as
[Equation]
where A is a state-space-independent linear transformation initialized to the identity, α is a nontrainable hyperparameter of the model, and B φ ( z t ) =
NN φ B ( z t ) .
This set up has the added beneﬁt that α = 0 is equivalent, both the statistical model and the algorithm, to GfLDS/PfLDS (Archer et al.
; Gao et al.
• The local transformation
A φ ( z t ) is redundant (it is akin to a gauge transformation in physics parlance).
To see this, note that for every z t , the image of the transformation A φ ( z t ) z t is at most R n .
On the other hand A φ ( z t ) has dimensionality R n 2 .
In other words, given z t and z t +1 , there is a continuum of matrices A φ ( z t ) that satisfy z t +
=
A φ ( z t ) z t .
It follows that A φ ( z t ) can be substantially restricted without loss of generality (“ﬁxing the gauge”).
We found that the best results were obtained when A φ ( z t ) 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
[Equation]
when the Jacobian J of the map r φ,ϕ satisﬁes 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 speciﬁc 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 φ ( z t ) , Eq. 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.
B D ETAILS OF THE ELECTROPHYSIOLOGY FIT
paths differ for posterior and anterior pole locations.
Also note visible changes in dynamics at t = − 1 .(stimulus), and t = 0 (go cue).
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 ﬁt covers time points (almost the entire recordings) from -0.5 sec from sample onset to 2.0 sec from response (wherefore the trial containstime 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, d z = 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 ﬁner timescales (dash lines vs dot lines;
Poisson vs
Gaussian VIND models).
This observation agrees with higher performance of explained variance ( R 2 ; Fig. 2).
The latent dynamics are smoother in the Poisson VIND model, (Fig.
6BC).
We have found that in
VIND ﬁts, 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 ﬁt 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.
C D ETAILS OF THE SINGLE CELL VOLTAGE DATA FITS
For the Allen data, ﬁg. 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.
D P REPROCESSING OF W IDEFIELD I MAGING D ATA
Macro-scale wide-ﬁeld 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 ﬂuorescence of genetically-encoded calcium (GCaMP6f) indicators using LED illumination and camera detection scheme.
We use methods for correcting ﬂuorescence recordings of neural activity for confounding contamination by changes in hemoglobin concentration and oxygenation as in Ma et al.
(2016b), by measuring both neural ﬂuorescence signals and hemodynamics.
This correction provides us with an accurate change in ﬂuorescence 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 1 D signal consisting of the standard deviation of the difference in value of all pixels from time t −to time t .
We use a WFOM recording of length 2 minutes, where the signals are sampled at 10 Hz, 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.
First, we ﬁt an anisotropic Wiener ﬁlter 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 deﬁned on their empirical autocorrelation function, and components that fall within a 99% conﬁdence 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 1 D 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.