Bayesian Clustering of Neural Spiking Activity Using a Mixture of Dynamic Poisson Factor Analyzers
Abstract
Modern neural recording techniques allow neuroscientists to observe the spikingactivity of many neurons simultaneously.
Although previous work has illustratedhow activity within and between known populations of neurons can be summarizedby low-dimensional latent vectors, in many cases what determines a unique population may be unclear.
Neurons differ in their anatomical location, but also, in theircell types and response properties.
Moreover, multiple distinct populations maynot be well described by a single low-dimensional, linear representation.
To tacklethese challenges, we develop a clustering method based on a mixture of dynamic
Poisson factor analyzers (mixDPFA) model, with the number of clusters treated asan unknown parameter.
To do the analysis of DPFA model, we propose a novel
Markov chain Monte Carlo (MCMC) algorithm to efﬁciently sample its posteriordistribution.
Validating our proposed MCMC algorithm with simulations, we ﬁndthat it can accurately recover the true clustering and latent states and is insensitiveto the initial cluster assignments.
We then apply the proposed mixDPFA model tomulti-region experimental recordings, where we ﬁnd that the proposed method canidentify 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 activityof many neurons from many different anatomical regions simultaneously.
With these expandingcapabilities, new methods to analyze neural data at the population-level and at the level of multiplepopulations become necessary.
Several recent models have been developed to extract shared latentstructures from simultaneous neural recordings, assuming that neural activity can be describedthrough 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 assumeslatent factors evolve with linear dynamics.
On the other hand, GPFA models the latent vectors bynon-parametric Gaussian processes.
However, in both cases, the observation model is generalizedlinear.
Several variants of these models have been implemented to analyze multiple neural populationsand their interactions [Semedo et al., 2019, Glaser et al., 2020].
However, in many cases, the totalnumber 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 fromunobserved brain areas, sharing the same latent structure.
On the other hand, neurons of differentcell-types within the same brain area may be better described by distinct latent structures.
From afunctional point of view, neither the anatomical location nor cell type (Fig.
1A) indicates whichneurons 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 ignoremulti-population structure and treat all neurons as a single population, then using linear model basedmethods 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 Gaussianprocesses [Wu et al., 2017] have been developed, but these models do not explicitly distinguishamong 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 combininga number of local factor analyzers, here we group neurons based on the latent factors (Fig. 1B).
Asimilar 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 latentfactors 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 asinteractions between clusters over time.
However, there are several limitations for mixPLDS: 1) itrequires we predetermine the number of clusters, and 2) the clustering results are often sensitive tothe 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 botha population baseline and baselines for individual neurons.
The number of clusters is treated as anunknown 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 conditionaldistribution of the latent state by a Gaussian, which is similar to results by sampling from exactfull conditional distribution.
To improve mixing in the cluster assignments, we marginalize theloading out for clustering by Poisson-Gamma conjugacy.
We also discuss the constraints necessityfor successful sampling of the proposed models.
After validating the proposed model with simulateddata, 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 proposedmethod 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 neuronsbased on multi-population latent structure.
The number of mixture components is treated as anunknown parameter and the posteriors are sampled by MCMC.
In this section, we ﬁrst provide thesingle population DPFA for a given cluster.
Then, we introduce a prior on the number of clusters anddescribe how we use the mixture of ﬁnite mixture model (MFM) to efﬁciently sample the posterior ofthe 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 (anon-negative integer), and let y i =
( y i 1 , . . .
, y iT ) (cid:48) .
Further, let z i be the cluster indicator of neuroni .
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 isa combination of PLDS and Poisson FA, which includes both population baseline and individualbaseline.
Assume neuron i belongs to the j -th cluster (i.e., z i = j ), and its spiking activity isindependently Poisson distributed, conditional on the low-dimensional latent state x ( j ) t ∈ R p
j andpopulation 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 neuronand 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 situationwhen 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
[Equation]
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) , theproposed model can be rewritten as
[Equation]
Generally, a factor model is consistent only when T/N → 0 [Johnstone and Lu, 2009], but this isoften 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 onlyidentiﬁable up to orthogonal rotation on X ( j ) , with c i ∼ N ( 0 , I p ) .
With including an individualbaseline δ 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, howwe 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 basedon 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 formore 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) , thespike 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 imposepriors 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: [Equation] f k is a p.m.f.
on { 1 , 2 , . . . } ,
[Equation]
given k,
[Equation]
z [Equation]
i.i.d.
∼ π given π , [Equation] [Equation]
i.i. [Equation]
[Equation]
H given k,
[Equation]
, y
[Equation]
[Equation] ∼ DPFA( δ i , c i , θ ( z i ) ) given δ i , c i , θ ( z i ) , z i , ∀ i = 1 , . . , N, [Equation]
[Equation] where p .
m .
f denotes the probability mass function.
By using the MFM, we can integrate the ﬁeldknowledge 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ﬁnedas f k ( k | α ) =
(1 − α ) k − 1 α for k = 1 , 2 , . . .
, and let γ = 1 .
The complete generative model issummarized 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 clusterindices 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-formfull 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 accordingto 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 Gaussianapproximation 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 approximationprovided reasonable estimation accuracy with substantial computational speedups.
We, thus, use aglobal Laplace approximation that can be efﬁciently computed in O ( T )
[Paninski et al., 2010].
Tohelp convergence, sampling on (labeled) model parameters is repeated several times before updatingthe cluster indices.
Approximating the intractable full conditional with a Laplace approximationalso makes computation of the proposed mixDPFA more efﬁcient.
However, to assess the accuracyof the approximation we also compare our approach to directly sampling from the exact posteriorof 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 similarlyas 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 sampledby the analogy of partition-based algorithm in Dirichlet process mixtures (DPM, Neal 2000).
Seedetails in Miller and Harrison 2018 and the Appendix (A.1).
When doing the clustering, we needto evaluate the likelihood for neurons under each cluster.
Although we can sample c i directly andevaluate the full likelihood as in MCMC for Gaussian MFA (data-augmentation/ imputation-posterioralgorithm, 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 themixture of PLDS (mixPLDS, Buesing et al. 2014) model may suggest a similar problem.
To resolvethis, 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 Laplaceapproximation, 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-Gammaconjugacy.
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 ) , wehave λ 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 approximatethis lognormal distribution by a gamma distribution, i.e., assume λ
it follows Gamma ( a
it , b it ) witha 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 independenceassumption 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, withcoefﬁ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.
Simulations
To validate and illustrate the proposed clustering method, we simulate neural data directly from thegenerative model (1) .
The labels for each neuron are assumed known and ﬁxed at ﬁrst to checkconvergence 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 = 2dimensional 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.
Thesimulations 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 perspike 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 mixingand 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 meanresponse 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-validationdescribed in [Williams et al., 2020].
To evaluate the ﬁtting performance of these two models, wehold out 1/2 of the data and compare the held-out log-likelihood per spike.
The distribution traininglog-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 singlepopulation analysis cannot describe the data as accurately, since the input is non-homogeneous andthe data is global nonlinear.
3.2 Clustering
Using the same simulation we now infer the cluster labels.
The latent factor dimension was ﬁrstoptimized with p = 2 selected by 5-fold speckled CV on small chains (100 iterations).
We thencompare three chains ﬁtting with all data: two unique chains initialized using a single cluster andone chain initialized with N = 50 clusters (i e. all clusters are singletons).
Traceplots (Fig.
3A) oftraining log-likelihood and number of clusters show that all chains converge.
When ﬁt to the full dataor 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 ratefor each neuron (Fig.
3B) can be well recovered, even with half data held out.
To evaluate clustermembership, here we show a similarity matrix where the entry ( i, l ) is the posterior probability thatdata points i and l belong to the same cluster.
The clustering results for all 4 chains recover the trueclusters, 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 alongsidethe latent states, is similar to the case when cluster labels are known and substantially higher than the1-cluster model (Fig.
3D and E).
Multi-region neural spike recordings
We then apply the proposed clustering method to the Allen Institute Visual Coding Neuropixelsdataset.
The dataset contains spiking activity from hundreds of neurons from multiple brain regionsof an awake mouse.
See detailed data description in [Siegle et al., 2021].
Here we investigate theclustering structure of neurons from four anatomical sites (83 neurons): 1) hippocampal CA1 (24neurons), 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 weanalyze 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.17%of 115 neurons) and we analyze data with 40ms bins.
We use a Geometric (0 . 33) prior over thenumber 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 averagenumber 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 Randindex (maxPEAR, Fritsch and Ickstadt 2009).
The maxPEAR-sorted neural activity and posteriorsimilarity 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 theclustering results and anatomy, we additionally sort the neurons according anatomical labels (upperleft 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 clusteredinto 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 accuratelyrepresent the latent structure.
We then evaluate the clustering patterns for different visual stimuli.
We run 2 independent chains foreach epoch (results from the second chain in Fig.
6D).
The similarity matrices show that the patternis 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-termdrift 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 tendto have lower similarity (average ARI from comparing 2 chains for each epoch) than within-epochcomparisons (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 mostsevere 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ﬁnethe clusters explicitly. Using cross-validation, we ﬁnd that the optimized dimension for each region isp =
(1 , 11 , 5 , 3) , respectively.
We ﬁnd that the single cluster model tends to underﬁt the data, whilethe anatomical cluster model tends to overﬁt the data (Fig. 4F).
-out log-likelihood per spike, from iteration 500 to 1000 for three models: 1) clustering model, 2) single cluster model, and 3) anatomical cluster model.
Discussion
Here we introduce a Bayesian approach to cluster neural spike trains by MCMC.
Previous approachesto multi-population latent variable modeling have used anatomical information to label distinct groupsof neurons, but this choice is somewhat arbitrary.
Brain region and cell-type, for instance, can givecontradictory 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 mixDPFAmodel 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 beappropriate 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 outperformsglobally linear (1-cluster) models in simulations and with experimental data.
Although the proposed method can describe data and cluster neural spiking activity successfully, thereare some potential improvements.
Firstly, as mentioned above, the unconstrained model does not haveunique solutions.
To ensure model identiﬁable, we put diagonal constraints on A ( j ) and Q ( j ) andconstrain µ ( j ) and X ( j ) to have mean zero.
The assumption that A ( j ) and Q ( j ) are diagonal doesnot allow interaction between latent factors.
However, these interactions could be allowed by insteadconstraining 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 potentialimprovement 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 selectedby cross-validation (CV).
This may limit the accuracy of the model, since populations of neurons inexperimental data are likely to have different latent dimensionalities.
In future work, it would alsobe 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], oradaptive 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 dataand may be sensitive to priors.
Adaptive Gibbs sampling with shrinkage, on the other hand, has beenimplemented with the inﬁnite mixture of inﬁnite factor analyzers (IMIFA, Murphy et al. 2020).
Thesame idea may be useful here with an additional prior on linear dynamics ( A ( j ) , b ( j ) and Q ( j ) ) toencourage shrinkage in X ( j ) .
Finally, a deterministic approximation of MCMC, such as variationalinference may be more computationally efﬁcient.
Standard methods for ﬁtting the PLDS could beused directly in the VI updates, and if we further use a stick-breaking representation for the MFMmodel, 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 simultaneouslycontinues to grow, understanding the latent structure of multiple populations will be a major statisticalchallenge.
The Bayesian approach to clustering neural spike trains introduced here converges fastand is insensitive to the initial cluster assignments, and may, thus, be a useful tool for identifying"functional populations" of neurons.
Checklist
(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 tothem?
[Yes]
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]
If you ran experiments...
(a) Did you include the code, data, and instructions needed to reproduce the main experi-
mental 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 theywere 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., typeof GPUs, internal cluster, or cloud provider)?
[Yes]
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]
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 mainstages: 1) to sample the model parameters assuming known labels and 2) to sample the cluster indicesgiven 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 asused 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 spikecounts 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 = ∆ j + µ ( j
[Equation]
=
[Equation]
and
[Equation]
The full conditional distribution P ( (cid:101) x (
[Equation]
) | [Equation]
[Equation]
b
[Equation]
) , (cid:101) Q ( j ) ) isapproximated by a global Laplace approximation, i.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(cid:96)
= 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 [Mackeet 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]
(cid:48) (
[Equation]
· · · · · · · · · ∂ [Equation] (cid:96) ∂
[Equation]
  , ∂ [Equation] ∂
[Equation]
=
[Equation] [Equation] ) − [Equation]
[Equation]
j
[Equation]
[Equation] [Equation] Although we can conduct the Newton update efﬁciently in O ( T ) , bad initial values may slow downthe convergence.
To facilitate convergence, we initialize the Newton update with a smoothingestimate by local Gaussian approximation.
The forward ﬁltering for a dynamic Poisson model hasbeen 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 andvariance 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 ﬁlteringstep 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 thesmoothing 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
[Equation] ( j ) [ V ( j ) t | t − 1 ] − [Equation]
Update δ i and c i
We specify the prior for neuron-speciﬁc baseline δ i as δ i ∼ N (0 , 1) and wehave 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 updateof δ i and c i is just a regular Bayesian Poisson regression problem.
Thus, we can sample the fullconditional distribution of δ i and c i by a No U-Turn Sampler (NUTS, Hoffman and Gelman 2011),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 diagonalelement separately, as the update in h ( j ) , g ( j ) and σ 2( j ) .
Here, we update h ( j ) , g ( j ) and σ 2( j ) asfollows.
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 theinverse-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 ) |{ µ ( j ) t } Tt =1 ∼ IG (cid:32)
{+ν 
+}
+ T − , ν 0 σ 20 + µ (cid:48)
( j ) 2: T µ ( j ) 2: T
[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 moreparsimonious 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 mostparsimonious one, which only put constraints on p ( p − 1) / 2 parameters.
The constraint can also beimplemented 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.
Theupdate 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]
n ) , with
[Equation]
Before updating the cluster indices z i , the sampling of { x ( j ) t , µ ( j ) t , δ i , c i , θ ( j ) } is repeated severaltimes (5 times in both simulation and application for this paper).
To save time, we can further allowthe 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 algorithmfor 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 MFMmodel in the main text (Equation 2) and M θ
( c ) ( y i ) denotes the marginal likelihood ofneuron i in cluster c , when integrating the loading c i out (Equation (3)).
The marginallikelihood 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 ofpartitions obtained by removing the neuron i and V n ( t ) = (cid:80) ∞
j =1 j ( t ) ( γj ) ( n ) f k ( j ) , withx ( m ) =
x ( x + 1) · · · ( x + m − 1) , x ( m ) = x ( x − 1) · · · ( x − m + 1) , x (0) =
1 andx (0)
= 1 .
The update is an adaptation of partition-based algorithm for DPM [Neal, 2000], but with twosubstitutions: 1) replace | c i | by | c i | + γ and 2) replace α by γV n ( t + 1) /V n ( t ) .
See more detailsand discussions in [Miller and Harrison, 2018].
When evaluating the likelihood, we marginalize thecluster-independent loading c i out.
This is necessary for the high dimensional situation, otherwisethe chain will stop moving.
One issue with incremental Gibbs samplers such as Algorithm 3 and 8 in Neal [2000], when appliedto DPM, is that mixing can be somewhat slow.
To further improve the mixing, we may interspersethe "split-merge" Metropolis-Hasting updates [Jain and Neal, 2007, 2004] between Gibbs sweeps, asin [Miller and Harrison, 2018].
A.2 Sample Latent Vectors Using Pólya-Gamma Augmentation and Metropolis-HastingsAlgorithm
In this section, we provide details of sampling algorithms to draw the latent states X ( j ) and populationbaseline µ ( 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 thecluster 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 PGaugmentation scheme [Polson et al., 2013], we can approximate the Poisson distribution by anegative 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 NBfactor 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 updatethe latent states X and population baseline µ .
These two steps are summarized in the step 1 and step2 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 aproposal, where we employ a Metropolis-Hastings (MH) step to reject or accept the proposal.
In thisstep, the dispersion parameter r in NB distribution becomes a tuning parameter, to balance acceptancerate and autocorrelation in MH.
When r is large, the approximation to Poisson observation is accurateand the MH performs similar to the Gibbs sampler.
Here, we allow neurons at different time points tohave 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.
2 and theresults have been presented in Fig.
We set the tuning parameters for each neuron at all time pointsas r it = 10 , to tune the acceptance rate around 0.4.
When implementing the PG-MH algorithmin DPFA, the ω it (i.e., PG random variable) was sampled using pgdraw function in R packagepgdraw , which is called from our MATLAB code.
Although we can program Algorithm 1 withoutcalling pgdraw function from R to save some time, to sample the PG random variable is often muchcomputationally expensive in comparison to sample from the Gaussian random variable.
When ﬁttingthe simulated data with N = 5 , T = 1000 and p = 2 (cluster 1 data in Fig. 2) using a 3.40 GHzprocessor with 16 GB of RAM, it takes 131.64s for PG-MH to draw 1000 posterior samples, whilethe proposed approximation method in the main context of the paper takes 83.86s.
Although theproposed method takes the approximation to the full conditional of the latent state X and populationbaseline µ , the results are similar as to sample directly from the exact full conditional. Table 1 showsthe similarities, deﬁned by cosine function, of ﬁtted µ and X = ( X ∗ 1 , X ∗ 2 ) between the PG-MHand 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 ofthe 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
[Figure]
This section shows supplementary results when applying the proposed model to the Neuropixelsdataset (4 Multi-region neural spike recordings).
Here, we show 1) results sorted by maximum aposteriori 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]