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 (DPFA) 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 mixture of DPFAmodel to multi-region experimental recordings, where we ﬁnd that the proposedmethod can identify novel, reliable clusters of neurons based on their activity, andmay, 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 (DPFA).
The DPFAmodel takes the advantages of both Poisson factor analysis (FA) and PLDS and includes both apopulation baseline and baselines for individual neurons.
The number of clusters is treated as anunknown parameter in the mixture of DPFA, 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 (DPFA) to cluster neurons basedon multi-population latent structure.
The number of mixture components is treated as an unknownparameter and the posteriors are sampled by MCMC.
In this section, we ﬁrst provide the singlepopulation 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 mixture of DPFA.
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 ) except µ ( j ) and X ( 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.
Sincethe 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 , . . . } ,
[Equation]
[Equation]
γ ) given k, z 1 , . .
z N i.i.d.
∼ π given π , θ (1) , . .
, θ ( k ) i.i.d.
∼ H given k,
[Equation]
, y
[Equation]
, 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 ﬁ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 mixture of DPFA 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 distributionfor X ( j ) and µ ( j ) by a Gaussian distribution (a Laplace approximation), which is widely used for
EM [Macke et al., 2011] and variational inference [Glaser et al., 2020] with PLDS models and theirvariants. Gaussian approximation of the intractable full conditional distribution of latent effects hasalso been used in Bayesian mixed effects binomial regression, where Berman et al.
2022 found thatthe 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 )
[Paninskiet al., 2010].
To help convergence, sampling on (labeled) model parameters is repeated several timesbefore updating the cluster indices.
Approximating the intractable full conditional with a Laplaceapproximation also makes computation of the proposed mixture of DPFA more efﬁcient.
However, toassess the accuracy of the approximation we also compare our approach to directly sampling from theexact posterior of the model.
We develop a Pólya-Gamma (PG) data augmentation approach [Windleet 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 performssimilarly 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 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 marginallikelihood 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 then compare 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) .
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 wherethe entry ( i, l ) is the posterior probability that data points i and l belong to the same cluster.
Theclustering results for all 4 chains recover the true clusters, no matter what the the starting assignmentis (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 clusterlabels 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 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 posteriornucleus 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 and weanalyze data with 40ms bins.
We use a Geometric (0 . 33) prior over the number of clusters, such thatp ( 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 tend to cause confusions in clustering, and the average number ofclusters is 16.
To summarize the clustering results stored as posterior samples in MCMC, we givethe 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 similaritymatrix are shown in Fig.
4C and D. Results sorted by Maximum a posteriori (MAP) estimate aresimilar and are shown in the Appendix (Fig. 5).
To examine the relationship between the clusteringresults 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 alsoinclude neurons from different regions and neurons within a region are often clustered into separatepopulations ( ˆ P ( neuron i, l in the same region | z i =
z l , { y i } Ni =1 ) = 0 . 57 ).
Together, these resultssuggest that a simple assignment of populations based on anatomy many not accurately represent thelatent 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 experimentalsettings (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 single clusters for each region, and the optimized dimension foreach is p
= (1 , 11 , 5 , 3) .
We ﬁnd that the single cluster model tends to underﬁt the data, while theanatomical 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 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 mixture of DPFA 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 globalnonlinearity 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 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 potentialimprovement would be to automatically infer the dimension of the latent factors (states).
In thispaper, we assume p j is the same for all clusters, for convenience.
p is a pre-selected value or can beselected by cross-validation (CV).
However, it is also possible to treat p j as a parameter and samplethe 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 prioron X ( j ) [Bhattacharya and Dunson, 2011].
Although RJMCMC and BDMCMC can be easilyimplemented, 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 ofinﬁnite factor analyzers (IMIFA, Murphy et al. 2020).
The same idea may be useful here with anadditional prior on linear dynamics ( A ( j ) , b ( j ) and Q ( j ) ) to encourage shrinkage in X ( j ) .
Finally, adeterministic approximation of MCMC, such as variational inference may be more computationallyefﬁcient.
Standard methods for ﬁtting the PLDS could be used directly in the VI updates, and if wefurther use a stick-breaking representation for the MFM model, it would be straightforward to use VIfor 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
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 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]3.
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 they were chosen)?
[Yes]
(c) Did you report error bars (e g., with respect to the random seed after running experi- ments 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, ifapplicable?
[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 amountspent on participant compensation?
[N/A]
A AppendixA.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]
The full conditional distribution P (
[Equation]
Q
[Equation]
j ) ) isapproximated 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]
−
[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]
[Figure]
[Equation]
[Figure]
[Equation]
[Figure]
[Equation]
[Figure]
[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]
− 1 + J t − 1 ( V ( j ) t | T − V ( j ) t | t − 1 ) J (cid:48)
t − 1 , 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 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 diagonal element 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 σ 20 / 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: [Equation] |{ µ ( j )
[-ν 
-][Equation]
} Tt = [Equation]
∼ N ( τ n , σ 2( j ) Λ
[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]
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.
Because of the non-conjugacy, themarginal likelihood in terms of all cluster parameters θ ( z i ) can not be easily computed.
This issue canbe solved by introducing auxiliary variables, as in "Algorithm 8" in Neal 2000 for DPM.
A similaralgorithm can be used to infer the MFM.
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 2) and M θ ( c ) ( y i ) denotes the marginal likelihood ofneuron i in cluster c , when integrating the loading c i out.
(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 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 mixture of the DPFA model does not directly followthe PG augmentation scheme [Polson et al., 2013], we can approximate the Poisson distributionby 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 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 1150.17s for PG-MH to draw 1000 posterior samples, whilethe proposed approximation method in the main context of the paper takes only 135.30s.
Although the
ModelRecall the DPFA model using the notations deﬁned in A.1:
[Equation]
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 and calculate y 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 − [Equation] ) t − [Equation]
r it ) κ it =
( y it − r it ) /
+ ω it (log r it − δ i )
ˆ y
it =
ω
−
it κ it end end
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
∗ compute the acceptance ratio ζ = π ( (cid:102)
X ∗ |{ y i } ni =1 , U ) π
( (cid:102) X ( G − 1)
[Equation] { y i } ni =1 ,
[Equation] ) q ( (cid:102) X ( G − 1) | (cid:102) X ∗ ,
[Equation]
) 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 , ζ ) .
proposed 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 shows
[Table]
the 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.
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
A.3 Supplementary Results for Neuropixels Application
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).