N EURAL C LUSTERING P ROCESSES
Anonymous authors
Paper under double-blind review
A BSTRACT
Mixture models, a basic building block in countless statistical models, involve latent random variables over discrete spaces, and existing posterior inference methods can be inaccurate and/or very slow.
In this work we introduce a novel deep learning architecture for efﬁcient amortized Bayesian inference over mixture models.
While previous approaches to amortized clustering assumed a ﬁxed or maximum number of mixture components and only amortized over the continuous parameters of each mixture component, our method amortizes over the local discrete labels of all the data points, and performs inference over an unbounded number of mixture components.
The latter property makes our method natural for the challenging case of nonparametric Bayesian models, where the number of mixture components grows with the dataset.
Our approach exploits the exchangeability of the generative models and is based on mapping distributed, permutation-invariant representations of discrete arrangements into varying-size multinomial conditional probabilities.
The resulting algorithm parallelizes easily, yields iid samples from the approximate posteriors along with a normalized probability estimate of each sample (a quantity generally unavailable using Markov Chain Monte Carlo) and can easily be applied to both conjugate and non-conjugate models, as training only requires samples from the generative model.
We also present an extension of the method to models of random communities (such as inﬁnite relational or stochastic block models).
As a scientiﬁc application, we present a novel approach to neural spike sorting for high-density multielectrode arrays.
I NTRODUCTION
Mixture models (or equivalently, probabilistic clustering models) are a staple of statistical modelling in which a discrete latent variable is introduced for each observation, indicating its mixture component identity.
Popular inference methods in these models fall into two main classes.
When exploring the full posterior is crucial (e.g. there is irreducible uncertainty about the latent structure or many separate local optima exist), the method of choice is Markov Chain Monte Carlo (MCMC) (Neal, 2000; Jain & Neal, 2004).
This method is asymptotically accurate but time-consuming, with convergence that is difﬁcult to assess.
Models whose likelihood and prior are non-conjugate are particularly challenging, since in general in these cases the model parameters cannot be marginalized and must be kept as part of the state of the Markov chain.
Alternatively, variational methods (Blei & Jordan, 2004; Kurihara et al., 2007; Hughes et al., 2015) are typically much faster but do not come with accuracy guarantees.
As an alternative to MCMC and variational approaches, in recent years there has been steady progress on amortized inference methods, and such is the spirit of this work.
Concretely, we propose a novel technique to perform amortized approximate posterior inference over discrete latent variables in mixture models.
The basic idea is to use neural networks to express posteriors in the form of multinomial distributions (with varying support) in terms of ﬁxed-dimensional, distributed representations that respect the permutation symmetries imposed by the discrete variables.
A major advantage of our architecture, compared to previous approaches to amortized clustering, is its ability to handle an arbitrary number of clusters.
This makes the method a natural choice for nonparametric
Bayesian models, such as Dirichlet process mixture models (DPMM), and their extensions, where the number of components, a measure of the model complexity, is inferred as a posterior random variable; see (Rodriguez & Mueller, 2013) for a recent overview.
Moreover, the method can be applied to both conjugate and non-conjugate models.
The term ‘amortization’ refers to the process of investing computational resources to train a model that is later used for very fast posterior inference (Gershman & Goodman, 2014).
Concretely, the amortized approach learns a parametrized function q θ ( z | x ) that approximates p ( z | x ) for any x ; learning the model parameters θ may be computationally challenging, but once θ is in hand then evaluating q θ ( z | x ) for new data x is fast.
The amortized inference literature can be coarsely divided into two approaches.
On one side, the variational autoencoder approach (Kingma & Welling, 2013), with roots in the wake-sleep algorithm (Hinton et al., 1995), learns q θ ( z | x ) along with the generative model p φ ( x | z ) .
While p ( z ) is usually a known simple distribution, for discrete latent variables backpropagation cannnot be performed through them, and special approaches have been developed for those cases (Mnih &
Rezende, 2016; Jang et al., 2016; Maddison et al., 2016).
Our work corresponds to the alternative case: a generative model p ( x, z ) is postulated , and posterior inference is the main focus of the learning phase.
Amortized methods in this case usually involve a degree of specialization to the particular generative model of interest.
Examples include methods developed for Bayesian networks (Stuhlmüller et al., 2013), sequential Monte Carlo (Paige & Wood, 2016), probabilistic programming (Ritchie et al., 2016; Le et al., 2016), neural decoding (Parthasarathy et al., 2017) and particle tracking (Sun & Paninski, 2018).
Our work is specialized to the case where the latent variables are discrete and their range is not ﬁxed beforehand.
In the approach we present, after training the neural architecture using labeled samples from a particular generative model, we can obtain independent, parallelizable, approximate posterior samples of the discrete variables for any new set of observations of arbitrary size, with no need for expensive
MCMC steps.
These samples can be used (i) to compute approximate expectations, (ii) as high quality importance samples, or (iii) as independent Metropolis-Hastings proposals.
In Section 2 we study amortized mixtures and in Section 3 we review related works.
In Sectionwe discuss quantitative evaluations of the new method.
In Section 5 we present an extension of the method to random community graph models.
We close in Section 6 with a neuroscientiﬁc application of this method to spike sorting for high-density multielectrode probes.
A MORTIZING M IXTURE M ODELS
We start by presenting mixture models from the perspective of probabilistic models for clustering (McLachlan & Basford, 1988).
The latter introduce random variables c i denoting the cluster number to which the data point x i is assigned, and assume a generating process of the form
[Equation]
Here α 1 , α 2 are hyperparameters.
The number of clusters K is a random variable, indicating the number of distinct values among the sampled c
i ’s, and µ k denotes a parameter vector controlling the distribution of the k -th cluster (e.g., µ k could include both the mean and covariance of a Gaussian mixture component).
We assume that the priors p ( c 1: N | α 1 ) and p ( µ 1: K | α 2 ) are exchangeable,
[Equation]
where { σ i } is an arbitrary permutation of the indices, and similarly for p ( µ 1: K | α 2 ) .
Our interest in this work is in cases where K can take any value K ≤ N , such as the Chinese Restaurant
Process (CRP), or its Pitman-Yor generalization.
Of course, our methods will also work for models with K < B with ﬁxed B , such as Mixtures of Finite Mixtures (Miller & Harrison, 2018).
Given N data points x = { x i } , we would like to draw independent samples from the posterior
[Equation]
Note that p ( c 1 = 1 | x ) = 1 , since the ﬁrst data point is always assigned to the ﬁrst cluster.
While we might also be interested in the hidden variables α 1 , α 2 , µ k , the reason to focus on the discrete variables
c i ’s is that given samples from them, it is generally relatively easy to obtain posterior samples from p ( α 1 | c 1: N ) and p ( µ k , α 2 | x , c 1: N ) .
We would like to model all the factors in (1) in a uniﬁed way, with a generic factor given by
[Equation]
Here we assumed that there are K unique values in c 1: n − 1 , and therefore c n can take K + 1 values, corresponding to x n joining any of the K existing clusters, or forming its own new cluster.
We are interested in approximating (2):
[Equation]
where q θ is parameterized by a ﬂexible model such as a neural network that takes as inputs ( c 1: n − 1 , x ) , then extracts features and combines them nonlinearly to output a probability distribution on c n .
Critically, we will design the network to enforce the highly symmetric structure of the lhs of (3).
To make this symmetric structure more transparent, and in light of the expression (2), let us consider the joint distribution of the assignments of the ﬁrst n data points,
[Equation]
A neural representation of this quantity should respect the permutation symmetries imposed on the x i ’s by the values of c 1: n .
Therefore, our ﬁrst task is to build permutation-invariant representations of the observations x .
The general problem of constructing such invariant encodings was discussed recently in (Zaheer et al., 2017); to adapt this approach to our context, we consider three distinct permutation symmetries:
• Permutations within a cluster: (4) is invariant under permutations of x i ’s in the same cluster.
For each of the K clusters that have been sampled so far, we deﬁne the encoding
[Equation]
which is clearly invariant under permutations of x i ’s in the same cluster.
In general h is an encoding function we learn from data, unless p ( x | µ ) belongs to an exponential family and the prior p ( c 1: N ) is constant, as shown in Appendix A. • Permutations between clusters: (4) is invariant under permutations of the cluster labels.
In terms of the within-cluster invariants H k , this symmetry can be captured by
[Equation]
[Figure]
Permutations of the unassigned data points: (4) is also invariant under permutations of the N − n unassigned data points.
This can be captured by
[Equation]
Note that G and U provide ﬁxed-dimensional, symmetry-invariant representations of the assigned and non-assigned data points, respectively, for any values of N and K .
Encodings of this form were shown in (Zaheer et al., 2017) to lead to arbitrarily accurate approximations of symmetric functions.
xh c 1 = xh c 3 = xh c
= xh c
= x n h x n+ u x n+ u
Hg
Hg
G
H K+ g
U x 1 x 3 x n c n = c n = c n =K+ c n = c n = c n =K+ x 2 xx 1 x 3 x 2 x 4 x n
Neural Clustering Process x 1 x 3 x 2 x 4 x n
G
G
G K+ v a r i a b l e - i n p u t s o f t m a x 0.
U f f f
2. T HE V ARIABLE - INPUT S OFTMAX
After assigning values to c 1: n − 1 , each of the K + 1 possible values for c n corresponds to h ( x n ) appearing in one particular H k in (5), and yields a separate vector G k in (6).
See Figure 1 for an example.
In terms of the G k ’s and U , we propose to model (2) as
[Equation]
where we have introduced a new real-valued function f .
In other words, each value of c n corresponds to a different channel through which the encoding h ( x n ) ﬂows to the logit value f , as shown in Figure 2.
Note that
k
= K + 1 corresponds to c n forming its own new cluster with H k = h ( x n ) .
The softmax (8) differs from its usual form in, e.g., classiﬁcation networks, where a ﬁxed number of categories receive their logit values f from the ﬁxed-size ﬁnal layer of an MLP.
In our case, the discrete identity of each logit is determined by the neural path that the input h ( x n ) takes to G , thus allowing a ﬂexible number of categories.
In eq.
(8), θ denotes the parameters in the functions h, g, u and f , which we represent with neural networks.
By storing and updating G and U for successive values of n , the computational cost of a full i.i.d.
sample of c 1: N is O ( NK ) , the same as a single Gibbs sweep.
See Algorithm 1 for details; we term this approach the Neural Clustering Process (NCP).
It is relatively easy to run hundreds of copies of Algorithm 1 in parallel on a GPU, with each copy yielding a different set of samples c 1: N .
2. T HE O BJECTIVE F UNCTION
In order to learn the parameters θ of the neural networks, we use stochastic gradient descent to minimize the expected KL divergence,
[Equation]
Samples from p ( c 1: N , x ) are obtained from the generative model, irrespective of the model being conjugate.
If we can take an unlimited number of samples from the generative model, we can potentially train a neural network to approximate p ( c n | c 1: n − 1 , x ) arbitrarily accurately.
Note that the gradient here acts only on the variable-input softmax term q θ , not p ( c, x ) , so there is no problem of
Algorithm
1: h i ← h ( x i ) , u i ← u ( x i ) i = 1 . . .
N {Notation} 2: U ← (cid:80) Ni =2 u i {Initialize unassigned set} 3: H 1 ← h 1 , G ← g ( H 1 ) , K ← 1 , c
1 ←{Create ﬁrst cluster with x 1 } 4: for n ← 2 . . .
N do 5:
U ← U − u n {Remove x n from unassigned set} 6:
H K +1 ←{We deﬁne g (0) = 0 } 7: for k ← 1 . . .
K + 1 do 8:
G k ← G + g
( H k + h n ) − g
( H k ) {Add x n to cluster k } 9: q k ← e f ( G k ,U ) 10: end for 11: q k ← q k / (cid:80) K
+1 k (cid:48)
=1 q k (cid:48) , c n ∼ q k {Normalize probabilities and sample assignment} 12: if c n = K
+ 1 then 13:
K ← K +14: end if 15:
G ← G − g ( H c n ) + g ( H c n + h n ) {Add point x n to sampled cluster c n } 16:
H c n ← H c n + h n 17: end for 18: Return c 1 . . .
c N
2. T WO E XAMPLES
Clustering in 2D Gaussian models: The generative model is
[Equation]
where CRP stands for the Chinese Restaurant Process, with concentration parameter α , σ µ = 10 , and σ = 1 .
Figure 3 shows that the NCP captures the posterior uncertainty inherent in clustering this data.
Note that since the generative model is an analytical known distribution, there is no distinction here between training and test sets.
Clustering of MNIST digits: We consider next a DPMM over MNIST digits, with generative model α ∼ Exp (1) c 1: N ∼ CRP 10 ( α )
N ∼ Uniform [5 , 100] l k ∼ Unif [0 , 9] − without replacement.
k
= 1 . . .
K x i ∼
Unif [ MNIST digits with label l c
i ] i = 1 . . .
N where CRP 10 is a Chinese Restaurant Process truncated to up to 10 clusters, and d x = 28 × 28 .
Training was performed by sampling x i from the MNIST training set.
Figure 4 shows posterior samples for a set of digits from the MNIST test set, illustrating how the estimated model correctly captures the shape ambiguity of some of the digits.
Note that in this case the generative model has no analytical expression (and therefore is non-conjugate), but this presents no problem; a set of labelled samples is all we need for training.
See Appendix F for details of all the network architectures used.
[Figure]
assignments which
GPU implementation gives thousands of such samples in a fraction
Observations
Cluster Samples
Observations
Cluster Samples
Observations
Cluster Samples
Observations
Cluster Samples
Observations
Cluster Samples
Observations
Cluster Samples
Observations
Cluster Samples
Observations
Cluster Samples
Observations
Cluster Samples
Observations
Cluster Samples
Observations
Cluster Samples
Observations
Cluster Samples
Observations
Cluster Samples
Observations
Cluster Samples
Observations
Cluster Samples
Observations
Cluster Samples
Observations
Cluster Samples
Observations
Cluster Samples
Observations
Cluster Samples
Observations
Cluster Samples
K = 5 Pr: 0.
K = 4 Pr: 0.
K = 4 Pr: 0.
K = 6 Pr: 0.
K = 5 Pr: 0.
R ELATED W ORKS
Permutation-invariant neural architectures have been explored recently in (Ravanbakhsh et al., 2017;
Korshunova et al., 2018; Lee et al., 2018; Bloem-Reddy & Teh, 2019; Wagstaff et al., 2019).
The representation of a set via a sum (or mean) of encoding vectors was also used in (Guttenberg et al., 2016;
Ravanbakhsh et al., 2016; Edwards & Storkey, 2017; Zaheer et al., 2017; Garnelo et al., 2018a;b).
Most works on neural network-based clustering focus on learning features as inputs to traditional clustering algorithms, as reviewed in (Du, 2010;
Aljalbout et al., 2018; Min et al., 2018).
The works closest to ours are (Le et al., 2016) and (Lee et al., 2018).
Both present techniques for amortized inference of mixtures of Gaussians, so it is instructive to compare them in detail to our approach.
The work (Le et al., 2016) studies amortized inference of a variable number of latent variables generated during the trace of a general sequential probabilistic program.
For the case of a mixture of 2D Gaussians with a latent random number of components, a 2D histogram image of binned observations is fed to a convolutional network whose output enters into a recurrent neural network with a ﬁxed-sized softmax output layer to estimate the number of clusters.
The network also outputs the means and covariances of each cluster.
The work (Lee et al., 2018) presents Set Transformer, an attention-based architecture that improves the simple sum-based set encoding that we used above.
In their 2D Gaussian clustering application, the number of components is ﬁxed beforehand, and inference is made only on the cluster parameters.
These approaches have several limitations compared to ours.
First, the number of clusters is upper bounded by the size of the softmax layer (Le et al., 2016) or ﬁxed (Lee et al., 2018).
Second, the models perform inference on the continuous parameters µ k , but not on the discrete labels of each data point.
Finally, in (Le et al., 2016), the use of a convnet on a 2D histogram to determine the number of clusters does not scale to higher dimensional data due to the curse of dimensionality.
In Table 1 we summarize the comparison between the three approaches.
[Table]
[Figure]
t
[Figure]
E XPECTATIONS , E VALUATIONS AND DIAGNOSTICS
Samples from the NCP can be used to compute approximate expectations.
If the interest is in asymptotically exactness, the samples can be used as self-normalized importance samples, E [ f ( c )]
= (cid:80)
Mi =
[Equation]
( c
[Equation]
i ) )
w i / (cid:80) Mi
=1 w i where w i = p ( x , c
( i ) )
/q
θ ( c ( i ) | x ) .
Alternatively, the samples can be used as proposals in Metropolized independent sampling (Liu, 1996).
Of course, in both cases the variance of the estimated expectations will be lower when the NCP posterior is closer to the true posterior.
The examples presented in Sec.
2.3 provide strong qualitative evidence that our approximations to the true posterior distributions in these models are capturing the uncertainty inherent in the observed data.
But we would like to go further and ask quantitatively how well our approximations match the exact posterior.
Unfortunately, for sample sizes much larger than N = O (10) it is impossible to compute the exact posterior in these models.
Nonetheless, there are several quantitative metrics we can examine to check the accuracy of the model output.
Global symmetry from exchangeability:
Our results relied on p ( c 1: N | α 1 ) being exchangeable, which in turn implies exchangeability of the joint posterior (1).
But this is not explicit in the rhs of (1), where a particular order is chosen for the expansion.
If our model learns the conditional probabilities correctly, this symmetry should be (approximately) satisﬁed, and this can be monitored during training, as we show in Appendix C.
Estimated vs. Analytical Probabilities: Some conditional probabilities can be computed analytically and compared with the estimates output by the network; in the example shown in Figure 5, upper-right, the estimated probabilities are in close agreement with their exact values.
Geweke’s Tests: A popular family of tests that check the correctness of MCMC implementations (Geweke, 2004) can also be applied in our case: verify the (approximate) identity between the
[Figure]
[Equation]
where p ( x ) is the marginal from the generative model.
Figure 5 shows such a comparison for the 2D
Gaussian DPMM from Section 2.3, showing excellent agreement.
NCP vs. MCMC: NCP has some advantages over MCMC approaches.
First, it gives a probability estimate for each sample, in general unavailable in MCMC.
Secondly, NCP enjoys higher efﬁciency, due to parallelization of iid samples.
For example, in the Gaussian 2D example in eq.(10), in the time a collapsed Gibbs sampler produces one (correlated) sample, our GPU-based method produces more than 100 iid approximate samples.
Finally, NCP does not need a burn-in period.
NCP vs. Variational Inference: In Section 6, we compare NCP with a variational approach on clustering neural spikes. For 2000 spikes, the variational approach returned one clustering estimate in 0.76 secs., but does not properly handle the uncertainty about the number of clusters.
NCP produced 150 clustering conﬁgurations in 10 secs., efﬁciently capturing clustering uncertainty.
In addition, the variational approach requires a preprocessing step that projects the samples to lower dimensions, whereas NCP directly consumes the high-dimensional data by learning an encoder function h .
C OMMUNITIES
As an extension, we consider now a similar prior as above over cluster labels, but the observation model is more challenging:
[Equation]
where k 1 , k 2 = 1 . . .
K .
Here p ( c 1: n | α ) can be any exchangeable prior, and the binary observations x i,j represent edges in a graph of N vertices.
We focus on the symmetric graph case here, so φ k 1 ,k 2 = φ k 2 ,k 1 and x i,j ≡ x j,i .
We use a Beta model for p ( φ | β ) , but other choices are possible.
These models include stochastic block models (Holland et al., 1983; Nowicki & Snijders, 2001) and the single-type Inﬁnite Relational Model (Kemp et al., 2006; Xu et al., 2006; Schmidt & Morup, 2013).
Neural architectures for communities in graphs have been studied in (Chen et al., 2019) as a classiﬁcation problem for every node over a ﬁxed predetermined number of clusters.
We could proceed similarly to the clustering case, considering N particles, each given by a row of the adjacency matrix x i =
( x i, 1 . . .
x i,N ) .
But we should be careful when encoding these particles.
When values of c 1: n are assigned, a generic encoding h
( x i ) would ignore the permutation symmetries present among the components of x i , i.e., the columns of x i,j , as a result of the c 1: n assignments (the same three permutation symmetries discussed above for clustering models).
Moreover, a ﬁxed encoding h ( x i ) cannot accommodate the arbitrary length N of x i .
In Appendix B we present an invariant encoding that respects all these requirements.
We call our approach Neural Block Process (NBP).
See Figure 6 for an example.
A PPLICATION : SPIKE SORTING WITH NCP
Large-scale neural population recordings using multi-electrode arrays (MEA) are crucial for understanding neural circuit dynamics.
Each MEA electrode reads the signals from many neurons, and each neuron is recorded by multiple nearby electrodes.
As a key analysis step, spike sorting converts the raw signal into a set of spike trains belonging to individual neurons (Pachitariu et al., 2016; Chung et al., 2017; Jun et al., 2017; Lee et al., 2017; Chaure et al., 2018; Carlson & Carin, 2019).
At the core of many spike sorting pipelines is a clustering algorithm that groups the detected spikes into clusters, each representing a putative neuron (Figure 7).
However, clustering spikes can be challenging: (1)
Spike waveforms form highly non-Gaussian clusters in spatial and temporal dimensions, and it is unclear what are the optimal features for clustering.
(2) It is unknown a priori how many clusters there are.
(3) Existing methods do not perform well on spikes with low signal-to-noise ratios (SNR) due to increased clustering uncertainty, and fully-Bayesian approaches proposed to handle this uncertainty (Wood & Black, 2008; Carlson et al., 2013) do not scale to large datasets.
To address these challenges, we propose a novel approach to spike clustering using NCP.
We consider the spike waveforms as generated from a Mixture of Finite Mixtures (MFM) distribution (Miller &
Harrison, 2018), which can be effectively modeled by NCP.
(1) Rather than selecting arbitrary features for clustering, the spike waveforms are encoded with a convolutional neural network (ConvNet), which is learned end-to-end jointly with the NCP network to ensure optimal feature encoding.
Using a variable-input softmax function, NCP is able to perform inference on cluster labels without assuming a ﬁxed or maximum number of clusters.
(3) NCP allows for efﬁcient probablistic clustering by GPU-parallelized posterior sampling, which is particularly useful for handling the clustering uncertainty of ambiguous small spikes.
The computational cost of NCP training can be highly amortized, since neuroscientists often sort spikes form many statistically similar datasets.
We trained NCP for spike clustering using synthetic spikes from a simple yet effective generative model that mimics the distribution of real spikes, and evaluated the spike sorting performance on labeled synthetic data, unlabeled real data and hybrid test data by comparing NCP against two other methods: (1) vGMFM , variational inference on Gaussian MFM (Hughes & Sudderth, 2013).
Kilosort , a state-of-the-art spike sorting pipeline described in Pachitariu et al. (2016).
In Appendix D, we describe the dataset, neural architecture, and the training/inference pipeline of NCP spike sorting.
[Figure]
Synthetic Data.
We run NCP and vGMFM on 20 sets of synthetic test data each with 500, 1000, and 2000 spikes.
As the ground-truth cluster labels are known, we compared the clustering quality using
Adjusted Mutual Information (AMI) (Vinh et al., 2010).
The AMI of NCP is on average 11% higher than vGMFM (Figure 13), showing better performance of NCP on synthetic data.
Real Data.
We run NCP, vGMFM and Kilosort on a retina recording with white noise stimulus as described in Appendix D, and extracted the averaged spike template of each cluster (i.e. putative neuron).
Example clustering results in Figure 8 (top) shows that NCP produces clean clusters with visually more distinct spike waveforms compared to vGMFM.
As real data do not come with groundtruth cluster labels, we compared the spike templates extracted from NCP and Kilosort using retinal receptive ﬁeld (RF), which is computed for each cluster as the mean of the stimulus present at each spike.
A clearly demarcated RF provides encouraging evidence that the spike template corresponds to a real neuron.
Side-by-side comparisons of matched RF pairs are shown in Figure 8 (bottom-left) and Figure 14.
Overall, NCP found 103 templates with clear RFs, among which 48 were not found by
Kilosort.
Kilosort found 72 and 17 of them were not found by NCP (Figure 8 bottom-right), showing that NCP performs at least as well as Kilosort
[Figure]
and ﬁnds many additional templates with clear RFs.
[Figure]
[Figure]
P T P ( S U )
( D o t s )
[Figure]
F i r i n g r a t e ( H z ) ( g r e y b a r s )
[Figure]
produced
Hybrid Data.
We compared NCP against vGMFM and Kilosort on a hybrid recording with partial ground truth as in Pachitariu et al.
Spikes from 20 ground-truth templates were inserted into a real recording to test the spike sorting performance on realistic recordings with complex background noise and colliding spikes.
As shown in Figure 9, NCP recovered 13 of the 20 injected ground-truth templates, outperforming both Kilosort and vGMFM, which recovered 8 and 6, respectively.
Probabilistic clustering of ambiguous small spikes.
Sorting small spikes has been challenging due to the low SNR and increased uncertainty of cluster assignment.
By efﬁcient GPU-parallelized posterior sampling of cluster labels, NCP is able to handle the clustering uncertainty by producing multiple plausible clustering conﬁgurations.
Figure 10 shows examples where NCP separates spike clusters with amplitude as low as 3-4 × the standard deviation of the noise into plausible units that are not mere scaled version of each other but have distinct shapes on different channels.
Overall, our results show that using NCP for spike sorting provides high clustering quality, matches or outperforms a state-of-the-art method, and handles clustering uncertainty by efﬁcient posterior sampling, demonstrating substantial promise for incorporating NCP into production-scale pipelines.
A N EURAL C LUSTERING P ROCESS FOR E XPONENTIAL F AMILIES
The likelihood for an exponential family is given by
[Equation]
where t ( x ) is a vector of sufﬁcient statistics, and we deﬁned
[Equation] [Equation]
Let us denote by K and K (cid:48)
≥ K the total number of distinct values in c 1: n and c 1: N , respectively.
Consider the joint distribution
[Equation]
from which we obtain the marginal distributions
[Equation] [Equation]
(19) where we deﬁned
[Equation]
(20) and H k = 0 for k > K .
Note now that if p ( c 1: N ) is constant, all the dependence of F on c 1: n , x 1: n is encoded in the H k ’s, and F is symmetric under separate permutations of the H k ’s and the h ( x i ) ’s for i > n .
Based on these symmetries we can approximate F as
[Equation]
modulo adding to f any function symmetric on all x
i ’s, where
[Equation]
In the conditional probability we are interested in,
[Equation]
the product of the m ( x i ) ’s in (19) cancels.
Similarly, adding to f a function symmetric on all x i ’s leaves invariant our proposed approximation
[Equation]
B D ETAILS OF THE N EURAL B LOCK P ROCESS
Let us recall the generative model in this case,
[Equation]
where k 1 , k 2 = 1 . . .
K .
The prior p ( c 1: n | α ) can be any exchangeable priors for clustering, and the observations x i,j represent the presence or absence of an edge in a graph of N vertices.
We set φ k 1 ,k 2 = φ k 2 ,k 1 and x i,j ≡ x j,i , and assume for notational convenience that x ij ∈ { +1 , − 1 } .
B. E NCODING EACH ROW OF THE ADJACENCY MATRIX
In principle posterior inference in this case can proceed similarly to the clustering case, by considering
N particles, each given by a row of the adjacency matrix x i =
( x i, 1 , . . .
, x i,N ) .
But we should be careful when encoding these particles.
Consider the situation when the values of c 1: n have been assigned.
Encoding with a generic function h ( x i ) would ignore the permutation symmetries present among the components of x i , i.e., the columns of the matrix x i,j , as a result of the c 1: n assignments.
These symmetries are the same three symmetries discussed above for clustering models.
Moreover, a ﬁxed function h ( x i ) would not be able to accommodate the fact that the length of x i changes with the size N of the dataset.
Suppose that there are K clusters among the c 1: n , each with s k elements.
In order to simplify the notation, let us assume that the N − n unassigned points all belong to an additional ( K + 1) -th cluster with s K +
=
N − n , so we assume c n +1: N = K + 1 , and we have (cid:80) K
+1 k
=1 s k =
N
and s k = (cid:80)
Nj =1 δ ( c j = k ) .
Now, in each row x i , the number s k of elements in the k -th cluster can be split as
[Equation]
and note that both s − i,k and s + i,k are invariant under the symmetry of permuting the indices within cluster k .
Example: N = 5 and x 1 =
(+1 , +1 , − 1 , +1 , +1) .
If four assignments were made c 1 = c 2 = 1 , c 3 = c 4 = 2 , then K = 2 and c 5 = 3 , and from x 1 we get s +1 , 1 = 2 , s − 1 , 1 = 0 , s +1 , 2 = 1 , s − 1 , 2 = 1 , s +1 , 3 = 1 , s − 1 , 3 = 0 .
If we permute the columns 3 and 4, both from cluster k = 2 , we get x 1 =
(+1 , +1 , +1 , − 1 , +1) , but all the s ± 1 ,j ’s stay invariant.
Additional invariants can be obtained combining s + j,k and s − j,k across all rows x j ’s with c j = c i , such as
[Equation]
and similarly m − c i ,k and v − c i ,k .
Note that these invariants are the same for all rows x j with c j = k .
The motivation to consider them is that, if the partition corresponding to c 1: n is correct, then for i ≤ n and k ≤ K we have n + i,k (cid:39) m
+ c i ,k since they are both estimators of the latent Bernoulli parameter φ c i ,k .
For the same reason, if the partition is correct and those two estimators of φ
c i ,k are exact, then v
+ c i ,k (cid:39) 0 .
Similarly for m − c i ,k and v − c i ,k .
Then these values provide learning signals to the network that estimates the probability of the assignments c 1: n being correct.
Therefore we propose to encode the components of x i belonging to cluster k as
[Equation]
In order to preserve the symmetry of the ﬁrst K labels under permutations, we combine them as
[Equation]
where the encoding function is t : R 6 → R d t .
The encoding (29) of the unassigned components x i,n +1: N is kept separate and denoted as q i =
r i,K +1 .
In summary, each row x i of the adjacency matrix is represented by the ﬁxed-dimensional pair ( t i , q i ) ∈ R d t +6 in a way that respects the symmetries of the assignments c 1: n : permutations between members of a cluster, permutations of cluster labels and permutations among unassigned columns.
B. C LUSTERING THE ROWS OF THE ADJACENCY MATRIX
We can proceed now as in regular clustering, encoding each cluster of x i ’s within c 1: n as
[Equation]
and deﬁning the permutation invariant, ﬁxed-dimensional vectors
[Equation]
In terms of these quantities, the conditional probabilities are deﬁned as usual as
[Equation]
for k = 1 . . .
K + 1 , with h n = h ( t n , q n ) and with G k being the value of G for the different conﬁgurations.
Compared to the regular clustering case, here we need to learn the additional function t .
We call our approach Neural Block Process (NBP).
C M ONITORING GLOBAL PERMUTATION INVARIANCE
As mentioned in Section [Figure] we must verify the symmetry of the posterior likelihood under global permutations of all the data points.
We show such a check in Figure 11.
[Figure]
D D ETAILS OF SPIKE SORTING USING NCP
Data preprocessing.
Training and test data come from the retinal recordings in Chichilnisky &
Kalmar (2002) using a 512-channel 2D hexagonal MEA with 20 kHz sampling rate.
After spike detection (Lee et al., 2017), each multi-channel spike waveform was assigned to the channel where the waveform has the maximum peak-to-peak (PTP) amplitude (i.e. the center channel, ch0).
This partitioned the recording data by channel such that each center-channel-based partition only contains multi-channel spike waveforms centered at that channel.
Each spike waveform is represented as a× 32 array containing the 32 time steps surrounding the peak from the center channel and the same time window from the 6 immediate neighbor channels (Figure 7).
These 7 × 32 arrays are the spikes on which clustering was performed.
Neural architecture for NCP spike sorting.
The overall architecture is the same as the one described in Section 2 and Figure 2.
To extract useful features from the spatial-temporal patterns of spike waveforms, we use a 1D ConvNet as the h and u encoder functions.
The convolution is applied along the time axis, with each electrode channel treated as a feature dimension.
The ConvNet uses a
ResNet architecture (He et al., 2016) with 4 residual blocks, each having 32, 64, 128, 256 feature maps (kernel size = 3, stride =
The last block is followed by an averaged pooling layer and a ﬁnal linear layer.
The outputs of the ResNet encoder are the h
i and u i vectors of NCP, i.e. h i = ResNetEncoder ( x i ) .
We used d h = d u = 256 .
The other two functions, g and f , are identical to those in the 2D Gaussian example.
Training NCP using synthetic data.
To train NCP for spike clustering, we created synthetic labeled training data using a MFM generative model (Miller & Harrison, 2018) of noisy spike waveforms that mimic the distribution of real spikes:
N ∼ Uniform [ N min , N max ] (
[Equation]
Here, N is the number of spikes between [200 , 500] .
The number of clusters K is sampled from a shifted Poisson distribution with λ = 2 so that each channel has on average 3 clusters.
π 1: K represents the proportion of each cluster and is sampled from a Dirichlet distribution with α 1: K = 1 .
The training spike templates µ k ∈ R 7 × 32 are sampled from a reservoir of 957 ground-truth templates not present in any test data, with the temporal axis slightly jittered by random resampling.
Finally, each waveform x i is obtained by adding to µ c i Gaussian noise with covariance given by the Kronecker product of spatial and temporal correlation matrices estimated from the training data.
This method creates spatially and temporally correlated noise patterns similar to real data (Figure 12).
We trained
NCP for 20000 iterations on a GPU with a batch size of 32 to optimize the NLL loss by the Adam optimizer (Kingma & Ba, 2015).
A learning rate of 0.0001 was used (reduced by half at 10k and 17k iterations).
Probabilistic spike clustering using NCP.
At inference time, we fed the 7 x 32 arrays of spike waveforms to NCP, and performed GPU-parallelized posterior sampling of cluster labels (Figureand Figure 7).
Using beam search (Graves, 2012; Sutskever et al., 2014) with a beam size of 150, we were able to efﬁciently sample 150 high-likelihood clustering conﬁgurations for 2000 spikes in less
[Figure]
n
[Figure] [Figure]
than 10 seconds on a single GPU.
After clustering, we obtained a spike template for each cluster as the average shape of the spike waveforms.
The clustering conﬁguration with the highest probability was used for most experiments.
The spike sorting pipelines for real and hybrid data.
The real data is a 49-channel, 20-minute retina recording with white noise stimulus.
To create the hybrid test data, 20 ground-truth spike templates were manually selected from a 49-channel test recording and inserted into another test dataset according to the original spike times.
For NCP and vGMFM, we performed clustering on 2000 randomly sampled spikes from each channel (clusters containing less than 20 spikes were discarded), and assigned all remaining spikes to a cluster based on the L2 distance to the cluster centers.
Then, a ﬁnal set of unique spike templates were computed, and each detected spike was assigned to one of the templates.
The clustering step of vGMFM uses the ﬁrst 5 PCA components of the spike waveforms as input features.
For Kilosort, we run the entire pipeline using the Kilosort2 package (Pachitariu, 2019).
After extracting spike templates and RFs from each pipeline, we matched pairs of templates from different methods by
L-inﬁnity distance and pairs of RFs by cosine distance.
Electrode drift in real MEA data.
The NCP spike sorting pipeline described above does not take into consideration electrode drift over time, which is present in some real recording data.
As a step towards addressing the problem of spike sorting in the presence of electrode drift (Calabrese &
Paninski, 2011; Shan et al., 2017), we describe in Sup.
Material E a generalization of NCP to handle data in which the per-cluster parameters (e.g. the cluster means) are nonstationary in time.
[Table]
[Figure]
E P ARTICLE TRACKING
Inspired by the problem of electrode drift (Calabrese & Paninski, 2011; Pachitariu, 2019; Shan et al., 2017), let us consider now a generative model given by
[Equation]
In this model, a cluster corresponds to the points along the time trajectory of a particle, and (42) represents the time evolution of the cluster parameters.
The cluster labels c t indicate which particle is observed at time t , and note that particles can in principle appear or disappear at any time.
To take the time evolution into account, we let particles inﬂuence one another with a weight that depends on their time distance.
For this, let us introduce a time-decay constant b > 0 , and generalize the NCP equations to
[Equation]
[Equation] [Equation]
The conditional assignment probability for c t is now
[Equation] [Equation]
for k = 1 . . .
K + 1 .
The time-decay constant b is learnt along with all the other parameters.
We can also consider replacing e − b | t − t (cid:48) | with a general distance function e − d ( | t − t (cid:48) | ) .
Figure 15 illustrates this model in a simple 2D example.
We call this approach Neural Particle Tracking.
F N EURAL ARCHITECTURES IN THE EXAMPLES
To train the networks in the examples, we used stochastic gradient descent with Adam (Kingma &
Ba, 2015), with learning rate 10 − 4 .
The number of samples in each mini-batch were: 1 for p ( N ) ,for p ( c 1: N ) , 64 for p ( x | c 1: N ) .
The architecture of the functions in each case were:
[Table]
• h : MLP [2-256-256-256-128] with ReLUs • u : MLP [2-256-256-256-128] with ReLUs • g : MLP [128-256-256-256-256] with ReLUs • f : MLP [384-256-256-256-1] with ReLUs
[Table]
• h : 2 layers of [convolutional + maxpool + ReLU] + MLP [320-256-128] with ReLUs • u : same as h • g : MLP [256-128-128-128-128-256] with ReLUs • f : MLP [384-256-256-256-1] with ReLUs
[Table]
• t : MLP [6-64-64-64-256] with ReLUs • h : MLP [262-64-64-64-256] with ReLUs • u : MLP [262-64-64-64-256] with ReLUs • g : MLP [256-64-64-64-256] with ReLUs • f : MLP [512-64-64-64-64-1] with ReLUs