Variational Latent Gaussian Process for Recovering Single-Trial Dynamics from Population Spike Trains
Yuan Zhao, Il Memming Park
Introduction
Neural populations implement dynamics that produce robust behavior; however, our current experimental observations of these dynamics are invariably indirect and partial. In classical analyses of neural spike trains, noisy responses are averaged over repeated trials that are presumably time-locked to a stereotypical computation process. However, neural dynamics are not necessarily time-locked nor precisely repeated from trial to trial; rather, many cognitive processes generate observable variations in the internal processes that sometimes manifest in behavior such as error trials, broad reaction time distributions, and change of mind . In addition, it is difficult to disambiguate different possible neural implementations of computation from the average trajectory since they may only differ in their trial-to-trial variability . Therefore, if we wish to understand how neural computation is implemented in neural populations, it is imperative that we recover these hidden dynamics from individual trials .
Advances in techniques for recording from larger subpopulations facilitate single-trial analysis, especially the inference of single-trial latent dynamical trajectories. Several statistical approaches have been developed for extracting latent trajectories that describe the activity observed populations . For example, latent trajectories recovered from motor cortex suggest that these methods can provide insight to the coding and preparation of planned reaching behavior . Latent trajectories also elucidate the low-dimensional noise structure of neural codes and computations .
Inference of latent dynamical trajectories is a dimensionality-reduction method for multi-variate time series, akin to Kalman smoothing or factor analysis . Given a high-dimensional observation sequence, we aim to infer a shared, low-dimensional latent process that explains the much of the variation in high-dimensional observations. A large class of methods assume an autoregressive linear dynamics model in the latent process due to its computational tractability , we refer to these as PLDS (Poisson Linear Dynamical System). Although the assumption of linear dynamics can help in smoothing, it can also be overly simplistic: interesting neural computations are naturally implemented as nonlinear dynamics, and evidence points to nonlinear dynamics in the brain in general. Therefore, we propose to relax this modeling assumption and impose a general Gaussian process prior to nonparametrically infer the latent dynamics, similar to the Gaussian process factor analysis (GPFA) method . However, we differ from GPFA in that we use a point process observation model with self-history dependence rather than an instantaneous Gaussian observation model. A Gaussian observation model is inappropriate for inference in the millisecond-range time scale since it cannot generate spike counts. The price we pay is a non-conjugate prior and, consequently, an approximate posterior inference . We use a variational approximation where we assume a Gaussian process posterior over the latents, and optimize a lower bound of the marginal likelihood for the inference. Our algorithm, we call variational latent Gaussian process (vLGP), is fast and has better predictability compared to both GPFA and PLDS at a fine timescale (1 ms bins). We compare these algorithms on simulated systems with known latent processes. We apply it to high-dimensional V1 data from anesthetized monkey to recover both the noise correlation structure and topological structure of population encoding of drifting orientation grating stimuli.
Generative model
Suppose we simultaneously observe spike trains from neurons. Let denote the spike count time-series from the -th neuron for a small time bin. We model noisy neural spike trains mathematically as a simple point process which is fully described by its conditional intensity function . We assume the following parametric form of the conditional intensity function for the point process log-likelihood :
where is a latent process and denotes the spike history vector . Each neuron is directly influenced by the observed self-historyIt is straightforward to add external covariates similar to the self-history in this point process regression (e.g., see Park et al., ). with weight and also driven by the common latent process with weight (Fig. 1). Neurons are conditionally independent otherwise: all trial-to-trial variability is attributed either to the latent process or individual point process noise (c.f., Goris et al., , Ecker et al., , Lin et al., ).
The vector denotes the -dimensional latent process at time . We assume that , since we are looking for a small number of latent processes that explain the structure of a large number of observed neurons. The vector consists of the weights of the spike history and a time-independent bias term of the log firing rate for each neuron, and is a vector of length containing the dummy value for the bias and time-step spike self-history. This parametrization assumes that at most bins in the past influence the current intensity.
Under conditional independence, the joint distribution (data likelihood) of spike trains is given by,
Note that this model is not identifiable since where is an arbitrary invertible matrix (see later sections for further discussions). Also, the mean of latent process can be traded off with the bias term in .
Our assumptions about the latent process—namely the smoothness over time in this paper—are encoded in the prior distribution over the latent process. We use the Gaussian process (GP) framework for flexible prior design of each dimension independently:
where , and are mean and covariance functions, respectively. When time is discretized, the GP prior reduces to a multi-variate Gaussian distribution over the latent time series. We use the following form:
For the analyses in this manuscript, we choose the squared exponential covariance function for general smoothness over time,
where and are hyperparameters corresponding to the magnitude and inverse time scale of the latent process, respectively.
Variational inference
Our goal is to infer the posterior distribution over the latent process and fit the model parameters given the observed data. By Bayes’ theorem, the posterior distribution of the latent process is,
However, unlike in GPFA, the posterior under a point process likelihood and Gaussian process prior does not have an analytical form . Consequently, we must turn to an approximate inference technique. We employ variational inference, which aims to find an approximate distribution of the intractable true posterior . We can introduce this approximate posterior into the likelihood by re-writing it as,
We further assume that the distribution factorizes into Gaussian distributions with respect to each dimension of the latent process, such that
where is the number of total time steps, and each temporal slice is a vector of posterior means of the latent variables at time . Each temporal slice is a diagonal matrix whose diagonal contains the variances of the latent variables at time .
Variational inference for the entire posterior over latents, parameters, and hyperparameters can all be formulated in terms of maximizing (10). We sequentially update all parameters coordinate-wise; each conditional update turns out to be a convex-optimization problem except for the hyperparameters as explained below. We derive the inference algorithm (vLGP) in the following sections, and it is summarized in Algorithm 1.
Our algorithm scales linearly in space and time per iteration (for a fixed hyperparameter) where thanks to the rank- incomplete Cholesky factorization of the prior covariance matrix. For comparison, time complexity of GPFA is , and that of PLDS is .
The variational distribution is assumed to be Gaussian and thus determined only by its mean and covariance . The optimal solution is therefore obtained by
while holding other parameters and hyperparameters fixed.
Denote the expected firing rate of neuron at time by ,
The optimal can be obtained by the Newton-Raphson method. The gradient and Hessian are given as
where is a vector of length with value 1 at and zero elsewhere. Note that the Hessian is negative definite, and hence this is a convex optimization given the other arguments and . In each iteration, the update is
If we set the derivative w.r.t. to 0,
where is a diagonal matrix. Therefore, there is no need for optimization of the covariance. This simple form of variational posterior covariance has been noted before . Also note that .
There is a redundancy between the bias term in and the mean . During optimization, we constrain the latent mean by zero-centering, and normalize the loading by its max-norm latent-wise.
The prior covariance matrix is large () and is often severely ill-conditioned. We only keep a truncated incomplete Cholesky factor of size where is the rank of the resulting approximation,
where . We obtain two useful identities as a result:
With (22) and (23), we can avoid large matrices in above equations such as,
where is the all-ones vector, and . In addition, by the one-to-one correspondence between and , we use the diagonal of as a representation of in the algorithm.
2 Weights
Denote the temporal slices of ’s by matrix . The optimal weights and given the posterior over the latents can be obtained by the Newton-Raphson method with the following derivatives and Hessians,
Once again, both Hessians are negative definite, and hence in the territory of convex optimization.
3 Hyperparameters
One way to choose hyperparameters is to maximize the marginal likelihood w.r.t. the hyperparameters. Since the marginal likelihood is intractable in the vLGP model, we instead maximize (10) once again given the parameters and posterior. Interestingly, this objective function takes the same form as the one that is maximized in the GPFA’s hyperparameters updates.
We write the squared-exponential covariance kernel as,
where is the matrix of squared distances of each time pair. Hyperparameters and corresponds to prior variance and inverse (squared) time scale. We optimize the log-transformed hyperparameter for those are positive. To the -th transformed hyperparameter of the -th latent dimension, , the derivative is given as
The optimal value can be found by common gradient algorithms for each latent dimension independently.
The above derivation of the hyperparameter optimization technique assumes a fixed posterior and parameters. Thus it requires complete prior covariance matrices and explicit posterior covariance matrices rather than low-rank decompositions. In order to avoid numerical singularity, we add a small quantity to the diagonal of prior covariance matrices. It would be extremely costly to use these complete covariance matrices for long, consecutive time series. Therefore, we randomly take many shorter temporal subsamples of the posterior for fast computation . One hyperparameter iteration is performed every fixed number of iterations of posterior and parameter optimization.
Results
We verified our inference algorithm recover the true parameters and latent variables when there is no model mismatch and then apply it to two simulated systems and one real dataset. We compare our method (vLGP) against GPFA and PLDS.
First of all, we demonstrate that vLGP converges to the correct parameters and latent variables under the assumed generative model. We applied our method to simulated spike trains driven by 2-dimensional Gaussian process. We fixed the number of time bins of a trial, GP variance and timescale (). The values of parameters were randomly drawn from standard normal distribution. There are two limits that we need to consider for the convergence; increasing the duration of observations (more trials), and increasing the number of observed neurons. The parameters and latent variables are initialized by factor analysis (FA). To identify the property of the global optima, we also initialize the parameters and latent variables at the values near the true ones (by adding zero mean and standard deviation Gaussian noises).
We calculated the mean squared error (MSE) of posterior mean and weights on a grid of different numbers of trials and neurons. Figure 2(a) shows the convergence in MSE trend. The posterior mean of the latent distribution converges to the true latent as the number of neurons grows, and the parameters converge to the true weights as the number of time bins grows. The difference between FA and near truth initialization shows relative error in FA initialization combined with the non-convex vLGP inference which is small for latent process estimation.
We also fit the vLGP model to simulated spike trains driven by 1-dimensional Gaussian process latent at four different timescales while the rest setting was the same as above simulation. For each timescale, we simulated 10 datasets. We initialized the timescale at very smooth value (), and parameters and latent variables by factor analysis. Figure 2(b) shows the learned values scattered around the ground truth for a wide range of true time scales. We note that learning the timescale in general is challenging especially in high-dimensional latent spaces (data not shown).
2 Evaluation
We use a leave-one-neuron-out prediction likelihood to compare models. For each dataset comprising of several trials, we choose one of the trials as test trial and the others as training trials. First, the weights and posterior are inferred from the training trials. Next, we leave one neuron out of the test trial and make inference on the posterior using the remaining neurons with the weights estimated from the training trials. Then the spike train of the left-out neuron is predicted by the model given the weights estimated from the training trials and the posterior inferred from the test trial. We repeat this procedure on each neuron of the chosen test trial, and choose each trial of one dataset as test trial. Finally we obtain the prediction of all spike trains in the dataset.
For simulated datasets, we know the true latent process that generates observations. Since latent space is only identifiable up to affine transformation, we can quantify using the angle between subspaces . However, due to possible mismatch in the point nonlinearity, the subspace can be distorted. To account for this mismatch, we use the mean Spearman’s rank correlation that allows invertible monotonic mapping in each direction. The Spearman’s rank correlation between the posterior and the true latent trajectory gives a measure of the goodness of the posterior. If the correlation is large, the posterior recovers more information about the underlying trajectory.
3 Simulation
We simulate two datasets: one with deterministic nonlinear dynamics, and one with linear dynamics and model-mismatched nonlinear observation. Each dataset consists of samples (simulated datasets) and each sample contains trials from neurons which last for sec. We choose a bin size of ms.
In the first dataset, the latent trajectories are sampled from the Lorenz dynamical system with the time step of . This 3-dimensional system is defined by the following set of equations,
Spike trains are simulated by (1) with -step suppressive history filter (from most recent: $$) given the latent trajectory.
In the second dataset, Poisson spike trains are simulated from a -dimensional linear dynamical system (LDS) defined as
Figure 3 shows one trial from each dataset and corresponding inferred posterior mean latent process. The posterior means are rotated toward the true latent subspace. The PLDS inference (blue) looks the farthest away from the true Lorenz latent relatively but much closer to the LDS latent because the true latent meets its assumption. However, PLDS inferred latents lack smoothness. The GPFA inference (green) is better than PLDS for Lorenz latent but shows deviations from the true LDS latent. The smoothness is kept in the inference. The inference of our method (red) are very close to the true latent in both cases along the time while being smooth at the same time.
Figure 4 shows the Spearman’s rank correlation between the posterior mean and true latent versus running time (log scale). The figures show our method (vLGP) resulted in overall larger correlation than the PLDS and GPFA after the algorithms terminated. PLDS uses nuclear norm penalized rate estimation as initialization . The rank correlation from PLDS inference only slowly improved from the initial value through the optimization. Both the GPFA and vLGP use factor analysis as initialization . Note that the GPFA divides each trial into small time segments for estimating the loading matrix and bias which it breaks the continuity within each trial. Only the final iteration infers each trial as whole. Thus the correlations of the final iterations jumps up in the figures. It is obvious that vLGP makes much improvement to the result of factor analysis in terms of the rank correlation.
To quantify predictive performance on the spike trains, we use the log-likelihood on the leave-one-neuron-out as described in the evaluations section. We normalize the test point process likelihood with respect to that of a baseline model that assumes a homogeneous Poisson process to obtain, the predictive log-likelihood (PLL), given as,
where is the leave-neuron-out prediction to the firing rate of neuron at time , and is the population mean firing rate. Positive PLL implies the model predicts better than mean firing rate, and higher PLL implies better prediction. PLL has a unit of bits per spike, and is widely used to quantify spike train prediction .
In Figure 5, we compare the three models for each dataset. Since GPFA assumes a Gaussian likelihood, it is incompatible to compare directly using a point process likelihood. We use linear rectifier to convert the GPFA predictions to non-negative rates, then compute PLL (Fig. 5(a)). Let us denote the linear predictor by omitting the neuron, time and model. Specifically, The rate prediction is given by,
4 V1 population recording
We apply our method to a large scale recording to validate that vLGP picks up meaningful known signals, and investigate the population-wide trial-to-trial variability structure. We use the dataset where 72 different equally spaced directional drifting gratings were presented to an anesthetized monkey for 50 trials each (array-5, 148 simultaneously recorded single units). We use 63 V1 neurons by only considering neurons with tuning curves that could be well approximated () by bimodal circular Gaussian functions (the sum of two von Mises functions with different preferred orientations, amplitudes and bandwidths) according to . We do not include the stimulus drive in the model, in hopes that the inferred latent processes would encode the stimulus. We used bin size of 1 ms.
We use 4-fold cross-validation to determine the number of latents. A 15–dimensional model is fitted to a subsample composed of the first trial of each direction at first. In each fold, we use its estimated parameter to infer the latent process from another subsample composed of the second trial of each direction. The inference is made by leaving a quarter of neurons out, and we predict the spike trains of the left-out neurons given the first () orthogonalized latent process corresponding to -dimension. This procedure led us to choose 5 as the dimension since the predictive log-likelihood reached its maximum.
We re-fit a 5-dimensional vLGP model using the subsample of the first trials. To quantify how much the model explains, we report pseudo- defined as
Figure 5 shows the predictive performance based on two subsets. The first one is 4 trials (0∘, 90∘, 180∘, 270∘) of the subset of 63 neurons with 5-dimensional latent process. The second one is 10 trials (5 trials of 0∘ and 5 trials of 90∘) of all 148 neurons with 4-dimensional latent process.
To evaluate the predictive performance under a longer time scale, we also cross-validated on the first subset with 20 ms time bins with GPFA and vLGP. GPFA models were fitted for both raw spike count and its square root. The mean of square root spike count by vLGP was obtained from simulation using the predicted firing rate. We report the normalized MSEs (MSE / variance of observation) of both methods for spike count and its square root respectively, GPFA: 0.713 (spike count) and 0.699 (square root), and vLGP: 0.709 (spike count) and 0.708 (square root). We use F-test to compare the MSEs and see if any one of the two methods results in a significantly larger error. The corresponding p-values are 0.588 (spike count) and 0.140 (square root). It shows that the MSEs are not significantly different between GPFA and vLGP with 20 ms time bin.
Although the parameters are estimated from a subsample, we can use them to infer the latent process of all trials of all 72 directions. Figure 6 shows inferred latent processes for two trials for two directions. We rotate the inferred latent process by the singular value decomposition (SVD; details will be given later.) Variational posterior distribution over the latents are shown for each trial. During the second half of the trial when the stimulus was off, and the firing rate was lower, the uncertainty in the latent processes increases. There are visible trial-to-trial variability in the spike trains which are reflected in the variations of latents.
First we investigate how the “signal”—defined as visual stimuli—is captured by the latent processes. We average the inferred latent processes over 50 trials with identical spatio-temporal stimuli (Fig. 7). Since the stimuli are time-locked, corresponding average latent trajectory should reveal the time-locked population fluctuations driven by the visual input. We concatenate the average latent processes along the dimension of time. Then we orthogonalize it by SVD. The dimensions of orthogonalized one are ordered by the respective singular values. The latent process of a single trial is also rotated to the same subspace.
Furthermore, we visualized the trajectories in 3D (see supplementary online videohttps://www.youtube.com/watch?v=CrY5AfNH1ik) that show how signal and noise are dynamically encoded in the state space. Figure 8 shows the projection of average latent process corresponding to each orientation to the first 3 principal components. The projection topologically preserves the orientation tuning in the V1 population. There are two continuous circular variables in the stimuli space to be encoded: orientation and temporal phase of oscillation. The simplest topological structure of neural encoding is a torus, and we observe a toroidal topology (highlighted as rings of cycle averages).
These results show that vLGP is capable of capturing both the signal—repeated over multiple trials—and noise—population fluctuation not time locked to other task variables–present in the cortical spike trains.
Discussion
We propose vLGP, a method that recovers low-dimensional latent dynamics from high-dimensional time series. Latent state-space inference methods are different from methods that only recover the average neural response time-locked to an external observation . By inferring latent trajectories on each trial, they provide a flexible framework for studying the internal neural processes that are not time-locked. Higher-order processes such as decision-making, attention, and memory recall are well suited for latent trajectory analysis due to their intrinsic low-dimensionality of computation. Our method performs dimensionality reduction on a single trial basis and allows decomposition of neural signals into a small number of smooth temporal signals and their relative contribution to the population signal.
We compare our method to two widely used latent state space modeling tools in neuroscience: GPFA and PLDS . Unlike GPFA, vLGP allows a generalized linear model observation which is more suitable for a broad range of spike train observations. Moreover, vLGP is significantly faster than PLDS, yet it shows superior performance in capturing the spatio-temporal structures in the neural data to both PLDS and GPFA at a fine timescale (1 ms bins).
To test its validity in real electrophysiological recordings, we used V1 population recording driven by fixed stimulus as a litmus test. We showed that our inferred latents contain meaningful information about the external stimuli, encoding both orientation and temporal modulation on a continuous manifold.
We only considered smoothness encoded in the GP prior in this manuscript, but a plethora of GP kernels are available . For example, to capture prior assumptions about periodicity in some of the latent processes, we can use spectral kernels . This can be particularly useful for capturing internal neural oscillations . In addition, it is straightforward to incorporate additional covariates such as external stimuli or local field potential to vLGP.
We have not found systematic issues with vLGP, but potential weaknesses could stem from the variational approximation, inappropriate assumptions on the latent processes, and particular numerical shortcuts used in the implementation. It could be challenging for our method to recover very fast-changing latent trajectories, especially when the overall firing rate is very low. Subsampling used while optimizing the hyperparameters may miss a rare but key spike when the spike trains are sparse.
The proposed method has potential application in many areas, and it will be particularly useful in discovering how specific neural computations are implemented as neural dynamics. We are working on applying this method and its extensions to sensorimotor decision-making process where the normative model guides what is being computed, but it is unclear as to how the neural system implements it.
An open-source python implementation of vLGP is available online (https://github.com/catniplab/vLGP) under MIT license.
Acknowledgment
We thank the reviewers for their constructive feedback. We are grateful to Arnulf Graf, Adam Kohn, Tony Movshon, and Mehrdad Jazayeri for providing the V1 dataset. We also thank Evan Archer, Jakob Macke, Yuanjun Gao, Chethan Pandarinath, and David Sussillo for helpful feedback. This work was partially supported by the Thomas Hartman Foundation for Parkinson’s Research.