Linear dynamical neural population models through nonlinear embeddings

Introduction

Until recently, neural data analysis techniques focused primarily upon the analysis of single neurons and small populations. However, new experimental techniques enable the simultaneous recording of ever-larger neural populations (at present, hundreds to tens of thousands of neurons). Access to these high-dimensional data has spurred a search for new statistical methods. One recent approach has focused on extracting latent, low-dimensional dynamical trajectories that describe the activity of an entire population . The resulting models and techniques permit tractable analysis and visualization of high-dimensional neural data. Further, applications to motor cortex and visual cortex suggest that the latent trajectories recovered by these methods can provide insight into underlying neural computations.

Previous work for inferring latent trajectories has considered models with a latent linear dynamics that couple to observations either linearly, or through a restricted nonlinearity . When the true data generating process is nonlinear (for example, when neurons respond nonlinearly to a common, low-dimensional unobserved stimulus), the observation may lie in a low-dimensional nonlinear subspace that can not be captured using a mismatched observation model, hampering the ability of latent linear models to recover the low-dimensional structure from the data. Here, we propose fLDS, a new approach to inferring latent neural trajectories that generalizes several previously proposed methods. As in previous methods, we model a latent dynamical state with a linear dynamical system (LDS) prior. But, under our model, each neuron’s spike rate is permitted to vary as an arbitrary smooth nonlinear function of the latent state. By permitting each cell to express its own, private non-linear response properties, our approach seeks to find a nonlinear embedding of a neural time series into a linear-dynamical state space.

To perform inference in this nonlinear model we adapt recent advances in variational inference . Using a novel approximate posterior that is capable of capturing rich correlation structure in time, our techniques can be applied to a large class of latent-LDS models. We show that our variational inference approach, when applied to learn generative models that predominate in the neural data analysis literature, perform comparably to inference techniques designed for a specific model. More interestingly, we show in both simulation and application to two neural datasets that our fLDS modeling framework yields higher prediction performance with a more compact and informative latent representation, as compared to state-of-the-art neural population models.

Notation and overview of neural data

Neuronal signals take the form of temporally fast ($\sim 1$ ms) spikes that are typically modeled as discrete events. Although the spiking response of individual neurons has been the focus of intense research, modern experimental techniques make it possible to study the simultaneous activity of large numbers of neurons. In real data analysis, we usually discretize time into small bins of duration $\Delta t$ and represent the response of a population of $n$ neurons at time $t$ by a vector $\vx_{t}$ of length $n$, whose $i^{th}$ entry represents number of spikes recorded from neuron $i$ in time bin $t$, where $i\in\left\{1,\dots,n\right\}$, $t\in\left\{1,\dots,T\right\}$. Additionally, because spike responses are variable even under identical experimental conditions, it is commonplace to record many repeated trials, $r\in\left\{1,\dots,R\right\}$, of the same experiment.

Review of latent LDS neural population models

We focus upon one thread of this literature that takes its inspiration directly from the classical Kalman filter. Under this approach, the dynamics of a population of $n$ neurons are modulated by an unobserved, linear dynamical system (LDS) with an $m$-dimensional latent state $\vz_{rt}$ that evolves according to,

where $\mA$ is an $m\times m$ linear dynamics matrix, and the matrices $\mQ_{1}$ and $\mQ$ are the covariances of the initial states and Gaussian innovation noise, respectively. The spike count observation is then related to the latent state via an observation model,

When $\mathcal{P}_{\lambda}$ is chosen to be Poisson with $f(\vz_{rt})$ to be the (element-wise) exponential of a linear transformation of $\vz_{rt}$, we recover the Poisson linear dynamical system model (PLDS),

where $M(\lambda,g(\cdot))=\sum_{k=0}^{\infty}\frac{\exp(\lambda k+g(k))}{k!}$ is the normalizing constant. The GC model can flexibly capture under- and over-dispersed count distributions.

Nonlinear latent variable models for neural populations

We relax the linear assumptions of the previous LDS-based neural population models by incorporating a per-neuron rate function. We retain the latent LDS of eq. 1 and eq. 2, but select an observation model such that each neuron has a separate nonlinear dependence upon the latent variable,

To refer to an fLDS with a given noise model $\mathcal{P}_{\lambda}$, we prepend the noise model to the acronym. In the experiments, we will consider both PfLDS (taking $\mathcal{P}_{\lambda}$ to be Poisson) and GCfLDS (taking $\mathcal{P}_{\lambda}$ to be a generalized count distribution).

2 Model Fitting: Auto-encoding variational Bayes (AEVB)

Our goal is to learn the model parameters $\theta$ and to infer the posterior distribution over the latent variables $\vz$. Ideally, we would perform maximum likelihood estimation on the parameters, $\hat{\theta}=\arg\max_{\theta}\log p_{\theta}(\vx)=\arg\max_{\theta}\sum_{r=1}^{R}\int p_{\theta}(\vx_{r},\vz_{r})\text{d}\vz_{r}$, and compute the posterior $p_{\hat{\theta}}(\vz|\vx)$. However, under a fLDS neither the $p_{\theta}(\vz|\vx)$ nor $\text{p}_{\theta}(\vx)$ are computationally tractable (both due to the noise model $\mathcal{P}_{\lambda}$ and the nonlinear observation model $f_{\psi}(\cdot)$). As a result, we pursue a stochastic variational inference approach to simultaneously learn parameters $\theta$ and infer the distribution of $\vz$.

The strategy of variational inference is to approximate the intractable posterior distribution $p_{\theta}(\vz|\vx)$ by a tractable distribution $q_{\phi}(\vz|\vx)$, which carries its own parameters $\phi$.Here, we consider a posterior $q_{\phi}(\vz|\vx)$ that is conditioned explicitly upon $\vx$. However, this is not necessary for variational inference. With an approximate posteriorThe approximate posterior is also sometimes called a “recognition model”. in hand, we learn both $p_{\theta}(\vz,\vx)$ and $q_{\phi}(\vz|\vx)$ simultanously by maximizing the evidence lower bound (ELBO) of the marginal log likelihood:

We optimize $\mathcal{L}(\theta,\phi;\vx)$ by stochastic gradient ascent, using a Monte Carlo estimate of the gradient $\nabla\mathcal{L}$. It is well-documented that Monte Carlo estimates of $\nabla\mathcal{L}$ are typically of very high variance, and strategies for variance reduction are an active area of research .

Here, we take an auto-encoding variational Bayes (AEVB) approach to estimate $\nabla\mathcal{L}$. In AEVB, we choose an easy-to-sample random variable $\epsilon\sim p(\epsilon)$ and sample $\vz$ through a transformation of random sample $\epsilon$ parameterized by observations $\vx$ and parameters $\phi$: $\vz=h_{\phi}(\vx,\epsilon)$ to get a rich set of variational distributions $q_{\phi}(\vz|\vx)$. We then use the unbiased gradient estimator on minibatches consisting of a randomly selected single trials $\vx_{r}$,

where $\epsilon^{l}$ are iid samples from $p(\epsilon)$. In practice, we evaluate the gradient in eq. 9 using a single sample from $p(\epsilon)$ ($L=1$) and use ADADELTA for stochastic optimization .

The AEVB approach to inference is appealing in its generality: it is well-defined for a large class of generative models $p_{\theta}(\vx,\vz)$ and approximate posteriors $q_{\phi}(\vz|\vx)$. In practice, however, the performance of the algorithm has a strong dependence upon the particular structure of these models. In our case, we use an approximate posterior that is designed explicitly to parameterize a temporally correlated approximate posterior . We use a Gaussian approximate posterior,

where $\mu_{\phi}(\vx_{r})$ is a ${mT\times 1}$ mean vector and $\Sigma_{\phi}(\vx_{r})$ is a ${mT\times mT}$ covariance matrix. Both $\mu_{\phi}(\vx_{r})$ and $\Sigma_{\phi}(\vx_{r})$ are parameterized by observations $\vx$ through a structured neural network, as described in detail in supplementary material. We can sample from this approximate by setting $p(\epsilon)\sim\mathcal{N}(0,I)$ and $h_{\phi}(\epsilon;\vx)=\mu_{\phi}(\vx)+\Sigma^{1/2}_{\phi}(\vx_{r})\epsilon$ , where $\Sigma^{1/2}_{\phi}$ is the Cholesky decomposition of $\Sigma_{\phi}$.

This approach is similar to that of , except that we impose a block-tridiagonal structure upon the precision matrix ${\Sigma_{\phi}}^{-1}$ (rather than a diagonal covariance), which can express rich temporal correlations across time (essential for the posterior to capture the smooth, correlated trajectories typical of LDS posteriors), while remaining tractable with a computational complexity that scales linearly with $T$, the length of a trial.

Experiments

Our AEVB approach in principle permits inference in any latent LDS model. To illustrate this flexibility, we simulate $3$ datasets from several previously-proposed models of neural responses. In our simulations, each data-generating model has a latent LDS state of $m=2$ dimensions, as described by eq. 1 and eq. 2. Further, in all data-generating models, spike rates depend linearly on the latent state variable through a fixed link function $f$ that is common across neurons. Each data-generating model has a distinct observation model (eq. 3): Bernoulli (logistic link), Poisson (exponential link), or negative-binomial (exponential link).

We compare PLDS and GCLDS model fits to each datasets, using both our AEVB algorithm and two EM-based inference algorithms: LapEM (which approximates $p(\vz|\vx)$ with a multivariate Gaussian by Laplace approximation in the E-step ) and VBDual (which approximates $p(\vz|\vx)$ with a multivariate Gaussian by variational inference, through optimization in the dual space ). Additionally, we fit PfLDS and GCfLDS models with AEVB algorithm. On this linear simulated data we do not expect these nonlinear techniques to outperform linear methods. In all simulation studies we generate $20$ training trials and $20$ testing trials, with $100$ simulated neurons and $200$ time bins for each trial. Results are averaged across $10$ repeats.

We compare the predictive performance and running times of the algorithms in Table 1. For both PLDS and GCLDS, our AEVB algorithm gives results comparable to, though slightly worse than, the LapEM and VBEM algorithms. Although PfLDS and GCfLDS assume a much more complicated generative model, both provide comparable predictive performance and running time. We note that while LapEM is competitive in running time in this relatively small-data setting, the AEVB algorithm may be more desirable in a large data setting, where it can learn model parameters even before seeing the full dataset. In constrast, both LapEM and VBDual require a full pass through the data in the E-step before the M-step parameter updates. The recognition model used by AEVB can also be used to initialize the LapEM and VBEM in the linear LDS cases.

Simulation with “grid cell” type response:

The log firing rate of each neuron, indexed by $i$, is coupled to the latent variable $\vz_{rt}$ through a sinusoid with a neuron-specific phase $\phi_{i}$ and frequency $\omega_{i}$

We generated $\phi_{i}$ uniformly at random in the region $[0,2\pi]$ and set $\omega_{i}=1$ for neurons with index $i\leq 50$ and $\omega_{i}=3$ for neurons with index $i>50$. We simulated $150$ training and $20$ testing trials, each with $T=120$ time bins. We repeated this simulated experiment $10$ times.

We compare performance of PLDS with PfLDS, both with 1-dimensional latent variable. As shown in Figure 1, PLDS is not able to adapt to the nonlinear and non-monotonic link function, and cannot recover the true latent variable (left panel and bottom right panel) or spike rate (upper right panel). On the other hand the PfLDS model captures the nonlinearity well, recovering the true latent trajectory. The one-step-ahead predictive log likelihood (PLL) on a held-out dataset for PLDS is -0.622 (se=0.006), for PfLDS is -0.581 (se=0.006). A paired t-test for PLL is significant ($p<10^{-6}$).

2 Applications to experimentally-recorded neural data

We analyze two multi-neuron spike-train datasets, recorded from primary visual cortex and primary motor cortex of the Macaque brain, respectively. We find that fLDS models outperform PLDS in terms of predictive performance on held out data. Further, we find that the latent trajectories uncovered by fLDS are lower-dimensional and more structured than those recovered by PLDS.

The dataset consists of $148$ neurons simultaneously recorded from the primary visual cortex (area V1) of an anesthetized macaque, as described in (array 5). Data were recorded while the monkey watched a $1280$ms movie of a sinusoidal grating drifting in one of $72$ orientations: (0°, 5°, 10°,…). Each of the $72$ orientations was repeated $R=50$ times. We analyze the spike activity from $300$ms to $1200$ms after stimulus onset. We discretize the data at $\Delta t=10$ms, resulting in $T=90$ timepoints per trial. Following , we consider the $63$ neurons with well-behaved tuning-curves. We performed both single-orientation and whole-dataset analysis.

We first use $12$ equal spaced grating orientation (0°, 30°, 60°,…) and analyze each orientation separately. To increase sample size, for each orientation we pool data from the $2$ neighboring orientations (e.g. for orientation °, we include data from orientation $5$°and $355$°), thereby getting $150$ trials for each dataset (we find similar, but more variable, results when we do not include neighboring orientations). For each orientation, we divide the data into $120$ training trials and $30$ testing trials. For PfLDS we further divide the $120$ training trials into $110$ trials for fitting and $10$ trials for validation (we use the ELBO on validation set to determine when to stop training). We do not include a stimulus model, but rather perform unsupervised learning to recover a low-dimensional representation that combines both internal and stimulus-driven dynamics.

We take orientation °as an example (the other orientations exhibit a similar pattern) and compare the fitted result of PLDS and PfLDS with a 2-dimensional latent space, which should in principle adequately capture the oscillatory pattern of the neural responses. We find that PfLDS is able to capture the nonlinear response charateristics of V1 complex cells (Fig. 2(a), black line), while PLDS can only reliably capture linear responses (Fig. 2(a), blue line). In Fig. 2(b)(c) we project all trajectories onto the 2-dimensional latent manifold described by the PfLDS. We find that both techniques recover a manifold that reveals the rotational structure of the data; however, by offsetting the nonlinear features of the data into the observation model, PfLDS recovers a much cleaner latent representation(Fig. 2(c)).

We assess the model fitting quality by one-step-ahead prediction on a held-out dataset; we compare both percentage mean squared error (MSE) reduction and negative predictive log likelihood (NLL) reduction. We find that PfLDS recovers more compact representations than the PLDS, for the same performance in MSE and NLL. We illustrate this in Fig. 2(d)(e), where PLDS requires approximately $10$ latent dimensions to obtain the same predictive performance as an PfLDS with $3$ latent dimensions. This result makes intuitive sense: during the stimulus-driven portion of the experiment, neural activity is driven primarily by a low-dimensional, oscillatory stimulus drive (the drifting grating). We find that the highly nonlinear generative models used by PfLDS lead to lower-dimensional and hence more interpretable latent-variable representations.

To compare the performance of PLDS and PfLDS on the whole dataset, we use $10$ trials from each of the $72$ grating orientations (720 trials in total) as a training set, and $1$ trial from each orientation as a test set. For PfLDS we further divide the $720$ trials into $648$ for fitting and $72$ for validation. We observe in Fig. 3(a)(b) that PfLDS again provides much better predictive performance with a small number of latent dimensions. We also find that for PfLDS with $4$ latent dimensions, when we projected the observation into the latent space and take the first $3$ principal components, the trajectory forms a torus (Fig. 3(c)). Once again, this result has an intuitive appeal: just as the sinusoidal stimuli (for a fixed orientation, across time) are naturally embedded into a 2D ring, stimulus variation in orientation (at a fixed time) also has a natural circular symmetry. Taken together, the stimulus has a natural toroidal topology. We find that fLDS is capable of uncovering this latent structure.

Macaque center-out reaching data:

We analyzed the neural population data recorded from the Macaque motor cortex(G20040123), details of which can be found in . Briefly, the data consist of simultaneous recordings of $105$ neurons for $56$ cued reaches from the center of a screen to $14$ peripheral targets. We analyze the reaching period ($50$ms before and $370$ms after movement onset) for each trial. We discretize the data at $\Delta t=20$ms, resulting in $T=21$ timepoints per trial. For each target we use $50$ training trials and $6$ testing trials and fit all the $14$ reaching targets together (making $700$ training trials and $84$ testing trials). We use both Poisson and GC noise models, as GC has the flexibility to capture the noted under-dispersion of the data . We compare both PLDS and PfLDS as well as GCLDS and GCfLDS fits. For both PfLDS and GCfLDS we further divide the training trials into $630$ for fitting and $70$ for validation.

As is shown in figure Fig. 4(d), PfLDS and GCfLDS with latent dimension $2$ or $3$ outperforms their linear counterparts with much larger latent dimensions. We also find that GCLDS and GCfLDS models give much better predictive likelihood than their Poisson counterparts.On figure Fig. 4(b)(c) we project the neural activities on the $2$ dimensional latent space. We find that PfLDS (Fig. 4(c)) clearly separates the reaching trajectories and orders them in exact correspondence with the true the spatial location of the targets.

Discussion and Conclusion

We have proposed fLDS, a modeling framework for high-dimensional neural population data that extends previous latent, low-dimensional linear dynamical system models with a flexible, nonlinear observation model. Additionally, we described an efficient variational inference algorithm suitable for fitting a broad class of LDS models – including several previously-proposed models. We illustrate in both simulation and application to real data that, even when a neural population is modulated by a low-dimensional linear dynamics, a latent variable model with a linear rate function fails to capture the true low-dimensional structure. In constrast, a fLDS can recover the low-dimensional structure, providing better predictive performance more interpretable latent-variable representations.

extends the linear Kalman filter by using neural network models to parameterize both the dynamic equation and the observation equation, they uses RNN based recognition model for inference. composes graphical models with neural network observations and proposes structured auto encoder variational inference algorithm for inference. Ours focus on modeling count observations for neural spike train data, which is orthogonal to the papers mentioned above.

Our approach is distinct from related manifold learning methods .While most manifold learning techniques rely primarily on the notion of nearest neighbors, we exploit the temporal structure of the data by imposing strong prior assumption about the dynamics of our latent space. Further, in contrast to most manifold learning approaches, our approach includes an explicit generative model that lends itself naturally to inference and prediction, and allows for count-valued observations that account for the discrete nature of neural data.

Future work includes relaxing the latent linear dynamical system assumption to incorporate more flexible latent dynamics (for example, by using a Gaussian process prior or by incorporating a nonlinear dynamical phase space ). We also anticipate our approach may be useful in applications to neural decoding and prosthetics: once trained, our approximate posterior may be evaluated in close to real-time.

A Python/Theano implementation of our algorithms is available at http://github.com/earcher/vilds.

References

Supplementary material: Linear dynamical neural population models through nonlinear embedding

Acknowledgments

Funding for this research was provided by the Sloan Foundation, the McKnight Foundation, and Simons Foundation Global Brain Research Awards 325233, 325171, 365002, ONR N00014-14-1-0243, ARO MURI W911NF-12-1-0594, DARPA N66001- 15-C-4032 (SIMPLEX), and a Google Faculty Research award; in addition, this work was supported by the Intelligence Advanced Research Projects Activity (IARPA) via Department of Interior/ Interior Business Center (DoI/IBC) contract number D16PC00003. The U.S. Government is authorized to reproduce and distribute reprints for Governmental purposes notwithstanding any copyright annotation thereon. Disclaimer: The views and conclusions contained herein are those of the authors and should not be interpreted as necessarily representing the official policies or endorsements, either expressed or implied, of IARPA, DoI/IBC, or the U.S. Government. Thanks to Arnulf Graf, Adam Kohn, Tony Movshon, and Mehrdad Jazayeri for providing the V1 data. Thanks to Krishna V. Shenoy, Byron Yu, Gopal Santhanam and Stephen Ryu for providing the motor cortical data.

Appendix A Temporally correlated approximate posterior

Below we detail the recognition model $q_{\phi}(\vz|\vx)$ we used in our auto-encoder variational inference algorithm in fitting fLDS. For further details, see .

We construct $q_{\phi}(\vz|\vx)$ as a product of factors across time,

This product of Gaussian factors also has a Gaussian functional form, with block-tridiagonal inverse covariance. Normalizing recovers the multivariate Gaussian representation of eq. 11, where

Appendix B Neural network structure for generative model and approximate posterior

Appendix C Description of the video

We include a video (https://www.dropbox.com/s/cluev4fzfsob4q9/video_fLDS.mp4?dl=0) illustrating the latent-space projection of the macaque V1 data we analyzed in the paper. We fit the data with a PfLDS with $4$ latent dimensions, and plot the first $3$ principal components of the inferred latent trajectory. We use data from $300$ms to $1200$ms after stimulus onset, and for each grating orientation we plot the mean trajectory of the $10$ training trials used to fit the model. To illustrate the relationship between the latent space projection and the observed data, we show the latent trajectories for $3$ single directions ($0\degree$, $60\degree$ and $120\degree$) alongside the spike rasters of one an associated trial. We then show the latent trajectories for all $36$ orientations, from $0\degree$ to $175\degree$ (results for directions from $180\degree$ to $355\degree$ are similar). We observe that the projection of neuron activity corresponding to each grating orientation forms a circle, agreeing well with the periodicity of the sinusoidal stimulus (temporal frequency $6.25$Hz) and also the periodicity of the neural activity (as can be seen from the spike raster). The projection of the whole dataset forms a torus, which also agrees well with the stimulus structure.