Bayesian Clustering of Neural Spiking Activity Using a Mixture of Dynamic Poisson Factor Analyzers
Abstract
Modern neural recording techniques allow neuroscientists to observe the spiking activity of many neurons simultaneously.
Although previous work has illustrated how activity within and between known populations of neurons can be summarized by low-dimensional latent vectors, in many cases what determines a unique population may be unclear.
Neurons differ in their anatomical location, but also, in their cell types and response properties.
Moreover, multiple distinct populations may not be well described by a single low-dimensional, linear representation.
To tackle these challenges, we develop a clustering method based on a mixture of dynamic
Poisson factor analyzers (mixDPFA) model, with the number of clusters treated as an unknown parameter.
To do the analysis of DPFA model, we propose a novel
Markov chain Monte Carlo (MCMC) algorithm to efﬁciently sample its posterior distribution.
Validating our proposed MCMC algorithm with simulations, we ﬁnd that it can accurately recover the true clustering and latent states and is insensitive to the initial cluster assignments.
We then apply the proposed mixDPFA model to multi-region experimental recordings, where we ﬁnd that the proposed method can identify novel, reliable clusters of neurons based on their activity, and may, thus, be a useful tool for neural data analysis.
Introduction
With modern high-density probes [Jun et al., 2017], neuroscientists can observe the spiking activity of many neurons from many different anatomical regions simultaneously.
With these expanding capabilities, new methods to analyze neural data at the population-level and at the level of multiple populations become necessary.
Several recent models have been developed to extract shared latent structures from simultaneous neural recordings, assuming that neural activity can be described through low-dimensional latent states.
Many existing approaches are extensions of two basic models: the linear dynamical system (LDS) model [Macke et al., 2011] and a Gaussian process factor analysis (GPFA) model [Yu et al., 2009].
The LDS model is built on the state-space model and assumes latent factors evolve with linear dynamics.
On the other hand, GPFA models the latent vectors by non-parametric Gaussian processes.
However, in both cases, the observation model is generalized linear.
Several variants of these models have been implemented to analyze multiple neural populations and their interactions [Semedo et al., 2019, Glaser et al., 2020].
However, in many cases, the total number of distinct populations and which neurons belong to a population is unclear.
Neurons in different anatomical locations may interact with each other or receive common input from unobserved brain areas, sharing the same latent structure.
On the other hand, neurons of different cell-types within the same brain area may be better described by distinct latent structures.
From a functional point of view, neither the anatomical location nor cell type (Fig.
1A) indicates which neurons should be grouped into the same populations.
The incorrect population assignments can lead to biased and inconsistent inference on the latent structure [Ventura, 2009].
If we instead ignore multi-population structure and treat all neurons as a single population, then using linear model based methods may not describe their activity well, especially when the input is non-homogeneous.
Besides, nonlinear models such as deep learning [Pandarinath et al., 2018, Whiteway et al., 2019] and Gaussian processes [Wu et al., 2017] have been developed, but these models do not explicitly distinguish among distinct populations of neurons.
Motivated by the mixture of (Gaussian) factor analyzers (MFA,Arminger et al. 1999, Ghahramani and
Hinton 1996, Fokoué and Titterington 2003), which describes globally nonlinear data by combining a number of local factor analyzers, here we group neurons based on the latent factors (Fig. 1B).
A similar idea was previously implemented using a mixture of Poisson linear dynamical system (PLDS) model (mixPLDS, Buesing et al. 2014).
The mixPLDS model infers the subpopulations and latent factors using deterministic variational inference Wainwright and Jordan [2008], Jordan et al.
Emtiyaz Khan et al.
[2013] and the model parameters are estimated by Expectation Maximization (EM).
Unlike MFA, the mixPLDS can capture temporal dependencies of neural activity as well as interactions between clusters over time.
However, there are several limitations for mixPLDS: 1) it requires we predetermine the number of clusters, and 2) the clustering results are often sensitive to the initial cluster assignment.
Here we cluster the neurons by a mixture of dynamic Poisson factor analyzers (mixDPFA).
The
DPFA model takes the advantages of both Poisson factor analysis (FA) and PLDS and includes both a population baseline and baselines for individual neurons.
The number of clusters is treated as an unknown parameter in the mixDPFA, and the posteriors are sampled using Markov Chain Monte
Carlo (MCMC).
To sample high dimensional latent factors, we approximate the full conditional distribution of the latent state by a Gaussian, which is similar to results by sampling from exact full conditional distribution.
To improve mixing in the cluster assignments, we marginalize the loading out for clustering by Poisson-Gamma conjugacy.
We also discuss the constraints necessity for successful sampling of the proposed models.
After validating the proposed model with simulated data, we apply it to analyze multi-region experimental recordings from behaving mice: the Visual
Coding - Neuropixels Dataset from the Allen Institute for Brain Science.
Overall, the proposed method provides a way to efﬁciently cluster neurons into populations based on their activity.
Methods
Here we introduce a mixture of dynamic Poisson factor analyzers (mixDPFA) to cluster neurons based on multi-population latent structure.
The number of mixture components is treated as an unknown parameter and the posteriors are sampled by MCMC.
In this section, we ﬁrst provide the single population DPFA for a given cluster.
Then, we introduce a prior on the number of clusters and describe how we use the mixture of ﬁnite mixture model (MFM) to efﬁciently sample the posterior of the mixDPFA.
2.1 Dynamic Poisson Factor Analyzer
Denote the observed spike count of neuron i ∈ { 1 , . . .
, N } at time bin t ∈ { 1 , . . .
, T } as y it (a non-negative integer), and let y i =
( y i 1 , . . .
, y iT ) (cid:48) .
Further, let z i be the cluster indicator of neuron i .
Motivated by the nature of neural activity and the former PLDS model [Macke et al., 2011], we propose a new Poisson FA model by adding individual baselines δ i .
The proposed model is a combination of PLDS and Poisson FA, which includes both population baseline and individual baseline.
Assume neuron i belongs to the j -th cluster (i.e., z i = j ), and its spiking activity is independently Poisson distributed, conditional on the low-dimensional latent state x ( j ) t ∈ R p j and population baseline µ ( j ) t as follows:
[Equation]
with c i ∼ N p j ( 0 , I p j ) .
The neuron-speciﬁc baseline δ i is a constant across time for the i th neuron and unrelated to the cluster assignment.
For simplicity, we assume the dimension of latent factors (states) is the same for all clusters, s.t. p
j = p , however, our method easily extends to the situation when p j s differs across clusters (see Discussion).
Further, we assume the population baseline µ ( j ) t and the latent state x ( j ) t evolve linearly over time with Gaussian noise as following
[Equation]
where (cid:15) ( j ) t ∼ N (0 , σ 2( j ) ) and η ( j ) t ∼ N p ( 0 , Q ( j ) ) .
If we denote λ i = ( λ i 1 , . . .
, λ iT ) (cid:48) , µ ( j ) =
( µ ( j ) 1 , . . .
, µ ( j ) T ) (cid:48) and X ( j ) =
( x ( j ) 1 , . .
, x ( j ) T ) (cid:48) , the proposed model can be rewritten as
[Equation]
Generally, a factor model is consistent only when T/N → 0 [Johnstone and Lu, 2009], but this is often not the case for most neural spike data.
However, when we assume linear dynamics on µ ( j ) and X ( j ) , it resolves the consistency issue.
As known in a FA model, when p > 1 , the model is only identiﬁable up to orthogonal rotation on X ( j ) , with c i ∼ N ( 0 , I p ) .
With including an individual baseline δ i 1 T in our proposed DPFA model (1), it further makes the model invariant to translation of µ ( j ) and X ( j ) .
That means if µ ( j ) and X ( j ) is a set solution, then µ ( j ) + a 1 T and X ( j ) U + 1 T ⊗ m (cid:48) also satisfy the model, for any a , m and orthogonal matrix U .
Thus, to make the model identiﬁable, we need to add several constraints.
Although the clustering is invariant to orthogonal rotation, how we put constraints on translation will inﬂuence the cluster assignments.
Here we assume A ( j ) and
Q ( j ) are diagonal [Peña and Poncela, 2004, Lopes et al., 2008], and, to encourage clustering based on the trajectories of latent factors, we set (cid:80) Tt =1 µ ( j ) t = 0 and (cid:80) Tt
=1 x ( j ) t = 0 .
See Section 5 for more discussions about the choice of these constraints.
Given the parameters of the j -th cluster θ ( j ) = (cid:8) µ ( j ) , X ( j ) , h ( j ) , g ( j ) , σ 2( j ) , A ( j ) , b ( j ) , Q ( j ) (cid:9) , the spike counts of neuron i are generated by the dynamic Poisson factor analyzer (DPFA) model as [ y i | z i =
j ] ∼ DPFA ( δ i , c i , θ ( z i ) ) .
To faciliate the Bayesian computation, we have to impose priors H on θ ( j ) , see more details of prior settings in Appendix A.1.
2.2 Clustering by Mixture of Finite Mixtures Model
When the population labels z
i s are unknown, we cluster the neurons by a mixture of DPFA (mixDPFA).
Since the number of neural populations is ﬁnite but unknown, we need to put priors on it.
To make the
Bayesian computation more efﬁcient, we utilize the idea from the mixture of ﬁnite mixtures (MFM,
Miller and Harrison 2018) model, by assigning the priors for the clusters in the following way: k ∼ f k , f k is a p.m.f.
on { 1 , 2 , . . . } , π = ( π 1 , . . .
, π k ) ∼ Dir k ( γ, . .
. , γ ) given k, z 1 , . .
z N i.i.d.
∼ π given π , θ (1) , . .
, θ ( k ) i.i.d.
∼ H given k, y i =
( y i 1 , . . .
, y iT ) (cid:48)
∼ DPFA( δ i , c i , θ ( z i ) ) given δ i , c i , θ ( z i ) , z i , ∀ i = 1 , . .
, N, (2) where p .
m .
f denotes the probability mass function.
By using the MFM, we can integrate the ﬁeld knowledge about the number of neural populations into our analysis.
In the analysis of this paper, we assume k follows a geometric distribution, i.e., k ∼ Geometric ( α ) with its density deﬁned as f k ( k | α ) =
(1 − α ) k − 1 α for k = 1 , 2 , . . .
, and let γ = 1 .
The complete generative model is summarized in a graphical form shown in Fig.
1C.
2.3 Inference
Here the posteriors of the proposed mixDPFA model are sampled by an MCMC algorithm (see
Appendix A.1).
In each iteration, we sample the model parameters assuming the known cluster indices at ﬁrst, and then sample the cluster indices given the model parameters.
When sampling the (labeled) model parameters, the latent state X ( j ) and population baseline µ ( j ) have no closed-form full conditional distributions.
Although we could sample the posterior by particle MCMC directly, convergence may be too slow for clustering.
Here, we approximate the full conditional distribution for
X ( j ) and µ ( j ) by a Gaussian distribution (a Laplace approximation) and generate samples according to this approximation.
This Laplace approximation is widely used for EM [Macke et al., 2011] and variational inference [Glaser et al., 2020] with PLDS models and their variants, and Gaussian approximation of the intractable full conditional distribution of latent effects has also been used in
Bayesian mixed effects binomial regression, where Berman et al.
2022 found that the approximation provided reasonable estimation accuracy with substantial computational speedups.
We, thus, use a global Laplace approximation that can be efﬁciently computed in O ( T )
[Paninski et al., 2010].
To help convergence, sampling on (labeled) model parameters is repeated several times before updating the cluster indices.
Approximating the intractable full conditional with a Laplace approximation also makes computation of the proposed mixDPFA more efﬁcient.
However, to assess the accuracy of the approximation we also compare our approach to directly sampling from the exact posterior of the model.
We develop a Pólya-Gamma (PG) data augmentation approach [Windle et al., 2013,
Linderman et al., 2017, 2016, Polson et al., 2013] with an additional Metropolis-Hastings (MH) step [Metropolis et al., 1953, Hastings, 1970] to sample exactly from the full conditional of X ( j ) and µ ( j ) .
We ﬁnd that the proposed method using a Laplace apprixmation is faster but performs similarly as sampling from the exact joint posterior (Fig. 2, Fig.
5, Appendix A.2, and Table 1)).
Once we update the latent state X ( j ) and population baseline µ ( j ) , the cluster index is then sampled by the analogy of partition-based algorithm in Dirichlet process mixtures (DPM, Neal 2000).
See details in Miller and Harrison 2018 and the Appendix (A.1).
When doing the clustering, we need to evaluate the likelihood for neurons under each cluster.
Although we can sample c i directly and evaluate the full likelihood as in MCMC for Gaussian MFA (data-augmentation/ imputation-posterior algorithm, Fokoué and Titterington 2003), the chain has poor mixing and stops after a few iterations, because of the high dimensionality.
The heavy dependency on the starting point when ﬁtting the mixture of PLDS (mixPLDS, Buesing et al. 2014) model may suggest a similar problem.
To resolve this, we evaluate the marginal likelihood by integrating out the neuron-speciﬁc c i , i.e., the marginal likelihood of neuron i in cluster j is computed by
[Equation]
However, this marginal likelihood has no closed form.
Though we may evaluate it by a Laplace approximation, but iterating over all potential clusters for each neuron is computationally intensive.
To make faster clustering, we approximate the marginal likelihood by utilizing a Poisson-Gamma conjugacy.
This approach has been previously utilized to approximate posteriors [El-Sayyad, 1973] and predictive distributions [Chan and Vasconcelos, 2009].
In our situation, since c i ∼ N ( 0 , I p ) , we have λ it =
exp( δ i + µ ( j ) t + c (cid:48)
i x ( j ) t ) ∼ lognormal ( δ i + µ ( j ) t , x (cid:48)
( j ) t x ( j ) t ) , and then we can approximate this lognormal distribution by a gamma distribution, i.e., assume λ
it follows Gamma ( a it , b it ) with a it =
( x (cid:48) ( j ) t x ( j ) t ) − 1 and b it =
x (cid:48)
( j ) t x ( j ) t · e δ i + µ ( j ) t .
Then, by the conjugate property with Poisson and
Gamma random variables, we have
[Equation]
with ν it = a it and p it = 1 / (1 + b it ) .
Further, noticing that we have the conditional independence assumption for P ( y i | θ ( j ) , δ i ) , that is P ( y i
| θ ( j ) , δ i ) =
(cid:81)
Tt =1 P ( y it | θ ( j ) , δ i ) , we then have a closedform for Equation (3).
Another possible idea is to approximate the log-likelihood by second-order polynomials, with coefﬁcients determined by Chebyshev polynomial approximation [Keeley et al., 2019].
However, we ﬁnd that this approximation doesn’t work well in practice when spike counts have a wide range.
The model is implemented in MATLAB and the code is available at https://github.com/weigcdsb/
MFM_DPFA_clean .
Simulations
To validate and illustrate the proposed clustering method, we simulate neural data directly from the generative model (1) .
The labels for each neuron are assumed known and ﬁxed at ﬁrst to check convergence and model identiﬁability.
We then infer the labels to evaluate clustering performance.
All experiments in this paper were performed using a 3.40 GHz processor with 16 GB of RAM.
3.1 Labeled data
We ﬁrst simulate 10 clusters with 5 neurons in each, with recording length T = 1000 and p =dimensional latent factors for each cluster.
Individual baselines are generated by δ i ∼ N (0 , 0 .
5 2 ) , and the loading for the latent states are generated by c i ∼ N ( 0 , I 2 ) .
The population baseline µ ( j ) and latent vector X ( j ) are generated by the spline interpolation on 10 to 30 evenly spaced knots.
The simulations are conducted 50 times with different seeds.
Here we show results for one simulation, and the performance for the rest is similar.
Since the labels are known, each whole simulation is equivalent to 10 independent simulations, with 5 neurons in each.
Running MCMC for 10,000 iterations, we ﬁnd that the log-likelihood per spike converges rapidly for individual clusters and overall (Fig.2A).
Trace plots (Fig.2B) of the (Frobenius) norms for linear dynamics samples ( h ( j ) , g ( j ) , σ 2( j ) , A ( j ) , b ( j ) , Q ( j ) ) show rapid mixing and convergence for each DPFA.
The ﬁtted mean ﬁring rate (mean response) (Fig.
2C), µ ( j ) and
X ( j ) (Fig.
2D) match the ground truth well.
The convergence is fast, especially in terms of mean response and population baseline µ ( j ) .
Together, these results demonstrate the identiﬁability of the
DPFA with appropriate constraints.
We then compare the 10-cluster model with the simpliﬁed model ignoring the clustering structure (1-cluster model).
The p for 1-cluster model is 14, chosen by 5-fold speckled cross-validation described in [Williams et al., 2020].
To evaluate the ﬁtting performance of these two models, we hold out 1/2 of the data and compare the held-out log-likelihood per spike.
The distribution training log-likelihood is evaluated by averaging over the samples in short chains (e.g. iteration 50 to 100), which is justiﬁed by the observed fast convergence.
The same procedure is replicated for 50 times.
In this case, ignoring the clustering structure leads to a worse performance (Fig 2E), and the single population analysis cannot describe the data as accurately, since the input is non-homogeneous and the data is global nonlinear.
3.2 Clustering
Using the same simulation we now infer the cluster labels.
The latent factor dimension was ﬁrst optimized with p = 2 selected by 5-fold speckled CV on small chains (100 iterations).
We then compare three chains ﬁtting with all data: two unique chains initialized using a single cluster and one chain initialized with N = 50 clusters (i.e. all clusters are singletons).
Traceplots (Fig.
3A) of training log-likelihood and number of clusters show that all chains converge.
When ﬁt to the full data or only half, the number of clusters converges to 10, although the prior over the number of clusters is
K ∼ Geometric (0 . 2) .
When the recording length is sufﬁciently long, the likelihood will dominate, and the number of clusters will not be much affected by the prior setting.
The true mean ﬁring rate for each neuron (Fig.
3B) can be well recovered, even with half data held out.
To evaluate cluster membership, here we show a similarity matrix where the entry ( i, l ) is the posterior probability that data points i and l belong to the same cluster.
The clustering results for all 4 chains recover the true clusters, no matter what the the starting assignment is (Fig.
3C), which suggest the convergence of
MCMC.
The overall performance of the full model, where cluster membership is inferred alongside the latent states, is similar to the case when cluster labels are known and substantially higher than the 1-cluster model (Fig.
3D and E).
Multi-region neural spike recordings
We then apply the proposed clustering method to the Allen Institute Visual Coding Neuropixels dataset.
The dataset contains spiking activity from hundreds of neurons from multiple brain regions of an awake mouse.
See detailed data description in [Siegle et al., 2021].
Here we investigate the clustering structure of neurons from four anatomical sites (83 neurons): 1) hippocampal CA1 ( neurons), 2) dorsal part of the lateral geniculate complex (LGd, 36 neurons), 3) lateral posterior nucleus of the thalamus (LP, 12 neurons) and 4) primary visual cortex (VISp, 11 neurons).
And we analyze responses to 20s epochs during three visual stimuli: drifting gratings, spontaneous activity, and natural movies.
Only neurons with rates > 1Hz within the selected epochs are included (72% of 115 neurons) and we analyze data with 40ms bins.
We use a Geometric (0 . 33) prior over the number of clusters, such that p ( k ≤ 4) = 0 . 8 .
In responses to drifting gratings the 5-fold speckled CV log-likelihood is optimized with p = 2 , and, as in the simulations, the log-likelihood and number of clusters show rapid convergence and mixing (Fig 4A and B).
Low ﬁring rates and short recording lengths tend to cause confusions in clustering, reﬂecting uncertainty in cluster membership for neurons with little information.
Here the average number of clusters is 16.
To summarize the clustering results stored as posterior samples in MCMC, we give the single estimate for cluster indices
ˆ z i by maximizing the posterior expected adjusted Rand index (maxPEAR, Fritsch and Ickstadt 2009).
The maxPEAR-sorted neural activity and posterior similarity matrix are shown in Fig.
4C and D. Results sorted by Maximum a posteriori (MAP) estimate are similar and are shown in the Appendix (Fig. 5).
To examine the relationship between the clustering results and anatomy, we additionally sort the neurons according anatomical labels (upper left panel in Fig.
4E).
Although many identiﬁed clusters are neurons from the same anatomical area, clusters also include neurons from different regions and neurons within a region are often clustered into separate populations ( ˆ P ( neuron i, l in the same region | z i =
z l , { y i } Ni =1 ) = 0 . 57 ).
Together, these results suggest that a simple assignment of populations based on anatomy many not accurately represent the latent structure.
We then evaluate the clustering patterns for different visual stimuli.
We run 2 independent chains for each epoch (results from the second chain in Fig.
6D).
The similarity matrices show that the pattern is consistent for the same epoch, but will change along the time even under the same experimental settings (D1 vs. D2 and S1 vs. S2).
The changes in the clustering patterns may suggest long-term drift for neuron interactions.
To quantify the observations, we evaluate the adjusted Rand index (ARI) of maxPEAR estimates (bottom right panel in Fig.
4E).
Between-epoch comparisons tend to have lower similarity (average ARI from comparing 2 chains for each epoch) than within-epoch comparisons (different chains) for both maxPEAR and MAP (Fig.
6C).
The MAP number of clusters is largest (18) for the natural movie, suggesting this epoch has the most severe global non-linearity issue.
Here we compare three models: 1) clustering model with p = 2 , 2) single cluster model, with p = 8 selected by 5-fold speckled CV and 3) anatomical cluster model.
The anatomical cluster model ﬁts a single DPFA for each region, using the anatomical labels to deﬁne the clusters explicitly.
Using cross-validation, we ﬁnd that the optimized dimension for each region is p = (1 , 11 , 5 , 3) , respectively.
We ﬁnd that the single cluster model tends to underﬁt the data, while the anatomical cluster model tends to overﬁt the data (Fig. 4F).
Discussion
Here we introduce a Bayesian approach to cluster neural spike trains by MCMC.
Previous approaches to multi-population latent variable modeling have used anatomical information to label distinct groups of neurons, but this choice is somewhat arbitrary.
Brain region and cell-type, for instance, can give contradictory population labels.
The proposed method groups neurons by common latent factors, which may be useful for identifying "functional populations" of neurons.
Here we use a mixDPFA model and infer the number of clusters by MFM with a partition-based algorithm similar to DPM.
MFM may be more conceptually appropriate than DPM, since the number of neural "populations" is unknown but ﬁnite.
Additionally, MFM produces more concentrated, evenly dispersed clusters (see Miller and Harrison 2018 for detailed discussion).
The mixture modeling approach may also be appropriate in cases where neurons share non-homogeneous inputs, since it can approximate global nonlinearity with a mixture of locally linear models.
Here we ﬁnd that the mixture model outperforms globally linear (1-cluster) models in simulations and with experimental data.
Although the proposed method can describe data and cluster neural spiking activity successfully, there are some potential improvements.
Firstly, as mentioned above, the unconstrained model does not have unique solutions.
To ensure model identiﬁable, we put diagonal constraints on A ( j ) and Q ( j ) and constrain µ ( j ) and X ( j ) to have mean zero.
The assumption that A ( j ) and Q ( j ) are diagonal does not allow interaction between latent factors.
However, these interactions could be allowed by instead constraining X (cid:48)
( j ) X ( j ) to be diagonal [Krzanowski and Marriott, 1994a,b, Fokoué and Titterington, 2003].
Such a constraint could allow unique solutions for the (P)LDS and GPFA.
A second potential improvement would be to automatically infer the dimension of the latent factors (states).
In this paper, we assume p j is the same for all clusters, for convenience.
p is a pre-selected value or can be selected by cross-validation (CV).
This may limit the accuracy of the model, since populations of neurons in experimental data are likely to have different latent dimensionalities.
In future work, it would also be possible to treat p j as a parameter and sample the posterior by a reversible-jump (RJ)MCMC [Lopes and West, 2004], birth-death (BD)MCMC [Stephens, 2000, Fokoué and Titterington, 2003], or adaptive Gibbs sampling with shrinkage prior on X ( j ) [Bhattacharya and Dunson, 2011].
Although
RJMCMC and BDMCMC can be easily implemented, they perform poorly for high dimensional data and may be sensitive to priors.
Adaptive Gibbs sampling with shrinkage, on the other hand, has been implemented with the inﬁnite mixture of inﬁnite factor analyzers (IMIFA, Murphy et al. 2020).
The same idea may be useful here with an additional prior on linear dynamics ( A ( j ) , b ( j ) and Q ( j ) ) to encourage shrinkage in X ( j ) .
Finally, a deterministic approximation of MCMC, such as variational inference may be more computationally efﬁcient.
Standard methods for ﬁtting the PLDS could be used directly in the VI updates, and if we further use a stick-breaking representation for the MFM model, it would be straightforward to use VI for clustering as well, similar to [Blei and Jordan, 2006].
As the number of neurons and brain regions that neuroscientists are able to record simultaneously continues to grow, understanding the latent structure of multiple populations will be a major statistical challenge.
The Bayesian approach to clustering neural spike trains introduced here converges fast and is insensitive to the initial cluster assignments, and may, thus, be a useful tool for identifying "functional populations" of neurons. Acknowledgments
This material is based upon work supported by the National Science Foundation under Grant No.
Checklist
For all authors... (a) Do the main claims made in the abstract and introduction accurately reﬂect the paper’s contributions and scope?
[Yes] (b) Did you describe the limitations of your work?
[Yes]
(c) Did you discuss any potential negative societal impacts of your work?
[No] (d) Have you read the ethics review guidelines and ensured that your paper conforms to them?
[Yes] 2.
If you are including theoretical results... (a) Did you state the full set of assumptions of all theoretical results?
[N/A] (b) Did you include complete proofs of all theoretical results?
[N/A] 3.
If you ran experiments... (a) Did you include the code, data, and instructions needed to reproduce the main experimental results (either in the supplemental material or as a URL)?
[Yes] (b) Did you specify all the training details (e.g., data splits, hyperparameters, how they were chosen)?
[Yes] (c) Did you report error bars (e.g., with respect to the random seed after running experiments multiple times)?
[Yes] (d) Did you include the total amount of compute and the type of resources used (e.g., type of GPUs, internal cluster, or cloud provider)?
[Yes] 4.
If you are using existing assets (e.g., code, data, models) or curating/releasing new assets... (a) If your work uses existing assets, did you cite the creators?
[Yes] (b) Did you mention the license of the assets?
[Yes] (c) Did you include any new assets either in the supplemental material or as a URL?
[No] (d) Did you discuss whether and how consent was obtained from people whose data you’re using/curating?
[Yes] (e) Did you discuss whether the data you are using/curating contains personally identiﬁable information or offensive content?
[N/A] 5.
If you used crowdsourcing or conducted research with human subjects... (a) Did you include the full text of instructions given to participants and screenshots, if applicable?
[N/A] (b) Did you describe any potential participant risks, with links to Institutional Review Board (IRB) approvals, if applicable?
[N/A] (c) Did you include the estimated hourly wage paid to
participants and the total amount spent on participant compensation?
[N/A]
A Appendix A.1 MCMC updates
The posteriors are sampled by a Gibbs sampler.
In each iteration, the sampling scheme has two main stages: 1) to sample the model parameters assuming known labels and 2) to sample the cluster indices given model parameters.
Before moving into the second stage, the ﬁrst stage is repeated several times.
The sampling in the ﬁrst stage is conducted without considering constraints for µ ( j ) t and x ( j ) t at ﬁrst, and then we project the samples onto the constraint space for (cid:80)
Tt =1 µ ( j ) t = 0 and (cid:80)
Tt
=1 x ( j ) t = 0 as used in [Sen et al., 2018].
Update x ( j ) t and µ ( j ) t
The priors for initial population baseline and latent factor are:
[Equation]
Assume there are n j = # { i : z i = j } neurons belong to the j -th cluster, and denote the spike counts and ﬁring rates of these neurons at time t as (cid:101)
y
( j ) t =
vec ( { y it | z i = j } ) ∈
Z n j ≥ 0 and (cid:101) λ ( j ) t = vec ( { λ it | z i = j } ) ∈ R n j , where Z n j ≥ 0 denotes a n j -dimensional vector with each element being non-negative integers.
The corresponding loading and baseline for these n j neurons is
C ( j ) ∈ R n j × p and ∆ j = vec ( { δ i
| z i =
j } ) ∈ R n j , such that log (cid:101) λ ( j ) t [Equation] j
[Equation]
=
[Equation]
and
[Equation]
σ
[Equation]
The full conditional distribution P ( (cid:101) x ( j
[Equation]
| . . . )
[Equation] [Equation]
b
[Equation]
Q ( j ) ) is approximated by a global Laplace approximation, i [Equation] e.,
[Equation]
Then, taking the logarithm of the full conditional distribution, i.e., (cid:96) =
(cid:96) ( (cid:101) x ( j ) )
=
log P ( (cid:101) x
( j ) | . . . )
, we have [Equation]
const +
[Equation]
Here, (cid:96) ( (cid:101) x ( j ) ) is concave and unimodal.
By the Markovian assumption for latent state vectors, the
Hessian matrix is tri-block diagonal.
We can thus compute Newton updates to get µ (cid:101) x
( j ) for the
Laplace approximation in O ( T )
[Paninski et al., 2010], similar to the E-step for the PLDS [Macke et al., 2011].
The gradient ∇ (cid:96) and the Hessian ∇∇ (cid:96) are provided as follows.
For t = 2 , . . .
, T − 1 , the gradient of (cid:96) ( (cid:101) x ( j ) ) is:
[Equation]
And the Hessian matrix is:
[Equation]

[Equation]
[Equation]
[Equation]
[Equation]
[Equation]
[Equation]
 
[Equation]
Although we can conduct the Newton update efﬁciently in O ( T ) , bad initial values may slow down the convergence.
To facilitate convergence, we initialize the Newton update with a smoothing estimate by local Gaussian approximation.
The forward ﬁltering for a dynamic Poisson model has been previously described [Eden et al., 2004], and we use an additional backward pass to smooth [Rauch et al., 1965].
Let (cid:101) x
( j ) t | t − 1 = E ( (cid:101)
x ( j ) t | (cid:101) y
( j ) 1 , . .
, (cid:101) y
( j ) t − 1 ) and V ( j ) t | t − 1 =
V ar ( (cid:101) x
( j ) t | (cid:101)
y
( j ) 1 , . .
, (cid:101) y
( j ) t − 1 ) be the mean and variance for the one-step prediction density, where (cid:101) x
( j ) t | t − 1 = (cid:101)
A ( j ) (cid:101) x ( j ) t − 1 | t − 1 + (cid:101) b ( j ) and V ( j ) t | t − 1 = (cid:101)
A ( j ) V ( j ) t − 1 | t − 1 (cid:101) A (cid:48)
( j ) + (cid:101) Q ( j ) .
Then, after we observe the data at time t , we can do a forward ﬁltering step for the mean (cid:101) x ( j ) t | t =
E ( (cid:101) x ( j ) t | (cid:101) y
( j ) 1 , . .
, (cid:101) y
( j ) t ) and the variance V ( j ) t | t =
V ar ( (cid:101) x
( j ) t | (cid:101)
y
( j ) 1 , . .
, (cid:101) y
( j ) t ) , which are given by
[Equation]
Derivation of the ﬁltering estimates can be found in [Eden et al., 2004], and we can further get the smoothing estimates directly by standard Rauch-Tung-Striebel smoother [Rauch et al., 1965].
The smoother estimates (cid:101) x ( j ) t | T and V ( j ) t | T are updated as follows:
[Equation]
where J t − 1 = V ( j ) t − 1 | t − 1 (cid:101) A (cid:48)
( j ) [ V ( j ) t | t − 1 ] − 1 .
Update δ i and c i
We specify the prior for neuron-speciﬁc baseline δ i as δ i ∼ N (0 , 1) and we have assumed the loading c i ∼ N ( 0 , I p ) .
Then, from the matrix representation of DPFA in (1), i.e., log λ i = δ i 1 T + µ ( j ) + X ( j ) c
i , it is easy to see that given µ ( j ) and X ( j ) are known, the update of δ i and c i is just a regular Bayesian Poisson regression problem.
Thus, we can sample the full conditional distribution of δ i and c i by a Hamiltonian Monte Carlo (HMC, Duane et al. 1987) within the Gibbs sampler.
Update parameters of latent state The parameters for linear dynamics are h ( j ) , g ( j ) , σ 2( j ) , A ( j ) , b ( j ) and Q ( j ) .
To make the model identiﬁable, we simply assume A ( j ) =
diag ( a ( j ) 1 , . .
, a ( j ) p ) and Q ( j ) =
diag ( q ( j ) 1 , . . .
, q ( j ) p ) .
Therefore, we can update A ( j ) , b ( j ) and Q ( j ) for each diagonal element separately, as the update in h ( j ) , g ( j ) and σ 2( j ) .
Here, we update h ( j ) , g ( j ) and σ 2( j ) as follows.
First, we specify the priors for σ 2( j ) following IG (cid:0) ν 0 / 2 , ν 0 σ 2 0 / 2 (cid:1) and (cid:0) g
( j ) , h ( j ) (cid:1) (cid:48)
∼
N ( τ 0 , σ 2( j ) Λ − 1 0 ) , with ν 0 = 1 , σ 0 = 0 .
01 , τ 0 = (0 , 1) (cid:48) and Λ
=
I 2 .
Here, the "IG" denotes the inverse-gamma distribution.
Denote µ ( j ) 2: T = (cid:16) µ ( j ) 2 , . .
. , µ ( j ) T (cid:17) (cid:48) and (cid:101) µ ( j ) 1:( T − 1) = (cid:16) 1 T − 1 , µ ( j ) 1:( T − 1) (cid:17) , with µ ( j ) 1:( T − 1) = (cid:16) µ ( j ) 1 , . .
. , µ ( j ) T − 1 (cid:17) (cid:48) .
The full conditional distributions for σ 2( j ) and (cid:0) g ( j ) , h ( j ) (cid:1) (cid:48) are: σ 2( j
[-ν 
-][Equation]
|{ µ ( j ) [Equation] Tt =
[Equation] ∼ IG [Equation] −
[Equation]
[-τ (cid:48) 
 Λ 0 τ 0 − τ (cid:48) n Λ n τ-] n (cid:33) , (cid:16) g ( j ) , h ( j ) (cid:17) (cid:48) |{ µ ( j ) t } T t =1 ∼ N ( τ n , σ 2( j ) Λ − 1 n ) , with Λ n = (cid:101) µ (cid:48)
( j ) 1:( T − 1) (cid:101) µ ( j ) 1:( T − 1) + Λ 0 , and τ n = Λ − 1 n (cid:16)(cid:101) µ (cid:48)
( j ) 1:( T − 1) µ ( j ) 2: T
+ Λ 0 τ 0 (cid:17) .
For completeness, we also provide the update of latent state parameters when using the more parsimonious constraint, i.e. diagonal X (cid:48)
( j ) X ( j ) .
According to results from previous research
[Krzanowski and Marriott, 1994a,b, Fokoué and Titterington, 2003], this constraint is one of the most parsimonious one, which only put constraints on p ( p − 1) / 2 parameters.
The constraint can also be implemented in MCMC by "unconstrained sampling-projection" procedure [Sen et al., 2018].
Without constraints, the sampling of h ( j ) , g ( j ) and σ 2( j ) is the same as shown previously.
The update of A ( j ) , b ( j ) and Q ( j ) is the standard multivariate Bayesian linear regression.
Denote
X ( j ) 2: T = ( x ( j ) 2 , . . .
, x ( j ) T ) (cid:48) and (cid:102)
X ( j ) 1:( T − 1) =
( 1 T − 1 , X ( j ) 1:( T − 1) ) , with X ( j ) 1:( T − 1)
=
( x ( j ) 1 , . .
, x ( j ) T − 1 ) (cid:48) .
Let us use the conjugate priors as following for Q ( j ) , b ( j ) and A ( j ) :
[Equation]
Here, the W − 1 denotes the inverse-Wishart distribution.
We can set the priors as Ψ 0 = 0 .
I p , γ
=
p
+2 and T
= ( 0 p , I p ) (cid:48) .
Then, the full conditional distributions for Q ( j ) and vec (( b ( j ) , A ( j ) ) (cid:48) ) are:
[Equation]
with
[Equation]
Before updating the cluster indices z i , the sampling of { x ( j ) t , µ ( j ) t , δ i , c i , θ ( j ) } is repeated several times (5 times in both simulation and application for this paper).
To save time, we can further allow the repetitions to be pre-stopped, when the training log-likelihood converges roughly.
Update z i
To update the cluster assignments for each neuron i , we use a partition based algorithm for MFM, similarly as described in Miller and Harrison 2018.
Let C denote a partition of neurons, and C\ i denote the partition obtained by removing neuron i from
C .
Initialize C and { θ ( c ) : c ∈ C} (e.g. one cluster).
Repeat the following steps G times to obtain G samples.
For i = 1 , . . .
, N : remove neuron i from C and place it:
(a) in c ∈ C\ i with probability ∝ ( | c
| + γ )
M θ ( c ) ( y i ) , where γ is deﬁned in the MFM model in the main text (
Equation
and
M θ ( c y i denotes the marginal likelihood of neuron i
in cluster c , when integrating the loading c i out (Equation (3)).
The marginal likelihood is approximated by using Poisson-Gamma conjugacy. (b) in a new cluster c ∗ with probability ∝ γ V n ( t +
1) V n ( t ) M θ ( c ∗ ) ( y i ) , where t is the number of partitions obtained by removing the neuron i and V n ( t ) = (cid:80) ∞
j =1 j ( t ) ( γj ) ( n ) f k ( j ) , with x ( m ) =
x ( x + 1) · · · ( x + m − 1) , x ( m ) = x ( x − 1) · · · ( x − m + 1) , x (0) =
1 and x (0) = 1 .
The update is an adaptation of partition-based algorithm for DPM [Neal, 2000], but with two substitutions: 1) replace | c i | by | c i | + γ and 2) replace α by γV n ( t + 1) /V n ( t ) .
See more details and discussions in [Miller and Harrison, 2018].
When evaluating the likelihood, we marginalize the cluster-independent loading c i out.
This is necessary for the high dimensional situation, otherwise the chain will stop moving.
One issue with incremental Gibbs samplers such as Algorithm 3 and 8 in Neal [2000], when applied to DPM, is that mixing can be somewhat slow.
To further improve the mixing, we may intersperse the "split-merge" Metropolis-Hasting updates [Jain and Neal, 2007, 2004] between Gibbs sweeps, as in [Miller and Harrison, 2018].
A.2 Sample Latent Vectors Using Pólya-Gamma Augmentation and Metropolis-Hastings Algorithm
In this section, we provide details of sampling algorithms to draw the latent states X ( j ) and population baseline µ ( j ) from the exact posterior of the model.
In order to explain the algorithm much clearer, we just focus on a given cluster index, and thus in this section we suspend the superscript ( j ) as the cluster index in our notations.
We present the sampling idea with two major parts as described below.
Pólya-Gamma Augmentation
Although the mixDPFA model does not directly follow the PG augmentation scheme [Polson et al., 2013], we can approximate the Poisson distribution by a negative binomial (NB) distribution, i.e., lim r →∞ NB ( r, σ ( ψ − log r )) =
Poisson ( e ψ ) , where σ ( ψ ) =
e ψ / (1 + e ψ ) and NB ( r, p ) denotes the NB distribution with rp/ (1 − p ) as its expectation.
We approximate the proposed DPFA model using a NB, then we can instead sample a mixture of NB factor analyzers using the PG scheme [Windle et al., 2013].
Further, we use the forward-ﬁlteringbackward-sampling (FFBS) algorithm
[Carter and Kohn, 1994, Frühwirth-Schnatter, 1994] to update the latent states X and population baseline µ .
These two steps are summarized in the step 1 and step 2 of Algorithm 1 for updating (cid:102) X
= ( µ , X ) .
Metropolis-Hastings Step Next, we use the samples of (cid:102) X yielded from FFBS algorithm as a proposal, where we employ a Metropolis-Hastings (MH) step to reject or accept the proposal.
In this step, the dispersion parameter r in NB distribution becomes a tuning parameter, to balance acceptance rate and autocorrelation in MH.
When r is large, the approximation to Poisson observation is accurate and the MH performs similar to the Gibbs sampler.
Here, we allow neurons at different time points to have unique tuning parameters r it .
The MH step is summarized in step 3 of Algorithm 1.
We have further implemented Algorithm 1 to ﬁt DPFA model for cluster 1 data in Fig.
and the results have been presented in Fig.
We set the tuning parameters for each neuron at all time points as r it = 10 , to tune the acceptance rate around 0.4.
When implementing the PG-MH algorithm in DPFA, the ω it (i.e., PG random variable) was sampled using pgdraw function in R package pgdraw , which is called from our MATLAB code.
Although we can program Algorithm 1 without calling pgdraw function from R to save some time, to sample the PG random variable is often much computationally expensive in comparison to sample from the Gaussian random variable.
When ﬁtting the simulated data with N = 5 , T = 1000 and p = 2 (cluster 1 data in Fig. 2) using a 3.40 GHz processor with 16 GB of RAM, it takes 131.64s for PG-MH to draw 1000 posterior samples, while the proposed approximation method in the main context of the paper takes 83.86s.
Although the proposed method takes the approximation to the full conditional of the latent state X and population baseline µ , the results are similar as to sample directly from the exact full conditional.
Table 1 shows the similarities, deﬁned by cosine function, of ﬁtted µ and X =
( X ∗ 1 , X ∗ 2 ) between the PG-MH and the Gaussian approximation for full conditional, where X ∗ l denotes the l -th latent vectors (i.e. the l -th column of X ).
Here, ˆ ξ PGMH denotes the average from iteration 5000 to 10,000 for chain 1 of
PG-MH (Fig. 5), ˆ ξ N denotes the average from iteration 5000 to 10,000 for the chain in Fig.
2 of the Gaussian approximation, and ξ true denotes the ground truth.
In each column, ξ is replaced by µ ,
X ∗ 1 and X ∗ 2 respectively.
The posterior means for these two method are close to each other, i.e. cos (ˆ ξ PGMH , ˆ ξ N ) ≈ 1 .
A.3 Supplementary Results for Neuropixels Application
This section shows supplementary results when applying the proposed model to the Neuropixels dataset (4 Multi-region neural spike recordings).
Here, we show 1) results sorted by maximum a posteriori probability (MAP) estimates and 2) clustering results in another independent chain (Fig. 6).
ModelRecall the DPFA model using the notations deﬁned in A.1: y it ∼ Poi ( λ it ) , log λ it
= δ i + (cid:101) c (cid:48)
i (cid:101) x t , for i = 1 , . . . , n and t = 1 , . .
, T .
Here, (cid:101) c
i = (1 , c (cid:48) i ) (cid:48) and (cid:101) x t = ( µ t , x (cid:48) t ) (cid:48) .
(cid:101) x t follows linear dynamics (cid:101) x t +1 | (cid:101) x t ∼ N ( (cid:101)
A (cid:101) x t + (cid:101) b , (cid:101) Q ) .
Denote the prior as (cid:101) x 1 ∼ N ( m 0 , V 0 ) .
Given the sample from the ( G − 1) -th iteration (cid:101) x
( G − 1) t and U = { (cid:101)
c
i , δ i , (cid:101) A , (cid:101) b , (cid:101) Q } .
sample ω it from PG distribution and calculate ˆ y it , which follows N ( (cid:101) c (cid:48)
i (cid:101) x
( G − 1) t , ω − 1 it ) for t = 1 , . .
, T do for i = 1 , . . . , n do sample ω
it ∼ P PG
( r it + y it , δ i + (cid:101) c (cid:48)
i (cid:101) x
( G − 1) t − log r it ) κ it =
( y it − r it ) /
+ ω it (log r it − δ i )
ˆ y
it =
ω
−
it κ it end end 2.
Forward-ﬁltering-
-sampling (
) for (cid:102) X
Denote ˆ y t =
(ˆ y 1 t , . .
, ˆ y Nt ) (cid:48) , Ω t =
Diag ([ ω 1 t , . . .
, ω Nt ]) and (cid:101) C
=
( (cid:101) c 1 , . . .
, (cid:101) c N ) (cid:48) for t = 1 , . .
, T do m t | t − 1 = (cid:101)
Am t − 1 + (cid:101)
b
V t | t − 1 = (cid:101)
AV t − 1 (cid:101) A (cid:48)
+ (cid:101) Q
K t =
V t | t − 1 (cid:101) C (cid:48)
( (cid:101) CV t | t − 1 (cid:101) C (cid:48)
+ Ω − 1 t ) −m t = m t | t − 1 + K t
(ˆ y t − (cid:101) Cm t | t − 1 )
V t =
( I − K t (cid:101) C ) V t | t −end sample (cid:101) x ∗ T ∼ N ( m T , V T ) for t =
T − 1 , . .
. , 1 do
J t =
V t (cid:101) A (cid:48) ( (cid:101) AV t (cid:101) A (cid:48)
+ (cid:101) Q ) −m ∗ t = m
t + J t
( x ∗ t +1 − (cid:101) Am t − (cid:101) b )
V ∗ t =
( I − J t
(cid:101) A ) V t sample x ∗ t ∼ N ( m ∗ t , V ∗ t ) end or proposal
X ∗ compute the acceptance ratio ζ = π ( (cid:102)
X ∗ |{ y i } ni =1 , U ) π
( (cid:102) X ( G − 1)
|{ y i } ni =1 , U ) q
( (cid:102) X ( G − 1)
| (cid:102) X ∗ , U ) q
( (cid:102) X ∗ | (cid:102)
X ( G − 1) , U ) =
P ( { y i } ni =
[Equation]
where R =
{ r it } is matrix for dispersion parameters for each neuron at all time points.
P ( · ) denotes the Poisson likelihood, and NB ( · ) denotes the negative binomial likelihood.
Accept the proposal (cid:102) X ∗ with probability min (1 , ζ ) .
[Table]