Year: 2019

Abstract:
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 efficient amortized Bayesian inference over mixture models. While previous approaches to amortized clustering assumed a fixed 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 infinite relational or stochastic block models). As a scientific application, we present a novel approach to neural spike sorting for high-density multielectrode arrays.

heading: 1 INTRODUCTION
text: 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 difficult 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 fixed-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 fixed 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 Section 4 we 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 neuroscientific application of this method to spike sorting for high-density multielectrode probes.

heading: 2 AMORTIZING MIXTURE MODELS
text: We start by presenting mixture models from the perspective of probabilistic models for clustering (McLachlan & Basford, 1988). The latter introduce random variables ci denoting the cluster number to which the data point xi is assigned, and assume a generating process of the form
α1, α2 ∼ p(α) N ∼ p(N)
c1 . . . cN ∼ p(c1, . . . , cN |α1) µ1 . . . µK |c1:N ∼ p(µ1, . . . µK |α2)
xi ∼ p(xi|µci) i = 1 . . . N
Here α1, α2 are hyperparameters. The number of clusters K is a random variable, indicating the number of distinct values among the sampled ci’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(c1:N |α1) and p(µ1:K |α2) are exchangeable,
p(c1, . . . , cN |α1) = p(cσ1 , . . . , cσN |α1) ,
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 fixed B, such as Mixtures of Finite Mixtures (Miller & Harrison, 2018).
Given N data points x = {xi}, we would like to draw independent samples from the posterior
p(c1:N |x) = p(c1|x)p(c2|c1,x) . . . p(cN |c1:N−1,x). (1)
Note that p(c1 = 1|x) = 1, since the first data point is always assigned to the first cluster. While we might also be interested in the hidden variables α1, α2, µk, the reason to focus on the discrete
variables ci’s is that given samples from them, it is generally relatively easy to obtain posterior samples from p(α1|c1:N ) and p(µk, α2|x, c1:N ). We would like to model all the factors in (1) in a unified way, with a generic factor given by
p(cn|c1:n−1,x) = p(c1 . . . cn,x) K+1∑ c′n=1 p(c1 . . . c ′ n,x) . (2)
Here we assumed that there are K unique values in c1:n−1, and therefore cn can take K + 1 values, corresponding to xn joining any of the K existing clusters, or forming its own new cluster.
We are interested in approximating (2):
p(cn|c1:n−1,x) ≈ qθ(cn|c1:n−1,x), (3) where qθ is parameterized by a flexible model such as a neural network that takes as inputs (c1:n−1,x), then extracts features and combines them nonlinearly to output a probability distribution on cn. 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 first n data points,
p(c1, . . . , cn,x) . (4)
A neural representation of this quantity should respect the permutation symmetries imposed on the xi’s by the values of c1:n. Therefore, our first 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 xi’s in the same cluster. For each of the K clusters that have been sampled so far, we define the encoding
Hk = ∑ i:ci=k h(xi) h : Rdx → Rdh (5)
which is clearly invariant under permutations of xi’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(c1: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 Hk, this symmetry can be captured by
G = K∑ k=1 g(Hk), g : Rdh → Rdg . (6)
(G1, U) (G2, U)
(G3, U) (G4, U)
Figure 1: Encoding cluster labels. The colored points have fixed labels c1:n−1, forming K = 3 clusters. The four possible labels for the circled point give four encoding vectors Gk, while the vector U encodes the 3 gray unlabeled points (Best seen in color).
• Permutations of the unassigned data points: (4) is also invariant under permutations of the N − n unassigned data points. This can be captured by
U = N∑ i=n+1 u(xi) , u : Rdx → Rdu . (7)
Note that G and U provide fixed-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.

heading: 2.1 THE VARIABLE-INPUT SOFTMAX
text: After assigning values to c1:n−1, each of the K + 1 possible values for cn corresponds to h(xn) appearing in one particular Hk in (5), and yields a separate vector Gk in (6). See Figure 1 for an example. In terms of the Gk’s and U , we propose to model (2) as
qθ(cn = k|c1:n−1,x) = ef(Gk,U)∑K+1
k′=1 e f(Gk′ ,U)
k = 1 . . .K + 1 , (8)
where we have introduced a new real-valued function f . In other words, each value of cn corresponds to a different channel through which the encoding h(xn) flows to the logit value f , as shown in Figure 2. Note that k = K + 1 corresponds to cn forming its own new cluster with Hk = h(xn).
The softmax (8) differs from its usual form in, e.g., classification networks, where a fixed number of categories receive their logit values f from the fixed-size final layer of an MLP. In our case, the discrete identity of each logit is determined by the neural path that the input h(xn) takes to G, thus allowing a flexible 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 c1: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 c1:N .1

heading: 2.2 THE OBJECTIVE FUNCTION
text: In order to learn the parameters θ of the neural networks, we use stochastic gradient descent to minimize the expected KL divergence,
Ep(N)p(x)DKL(p(c|x)‖qθ(c|x)) = −Ep(N)Ep(c1:N ,x) [ N∑ n=2 log qθ(cn|c1:n−1,x) ] + const. (9)
Samples from p(c1: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(cn|c1: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 backpropagating through discrete variables (Jang et al., 2016; Maddison et al., 2016).
1A Pytorch implementation of the algorithm is available at https://bit.ly/2lkGJ1b
Algorithm 1 O(NK) Neural Clustering Process Sampling 1: hi ← h(xi), ui ← u(xi) i = 1 . . . N {Notation} 2: U ← ∑N i=2 ui {Initialize unassigned set}
3: H1 ← h1, G← g(H1), K ← 1, c1 ← 1 {Create first cluster with x1} 4: for n← 2 . . . N do 5: U ← U − un {Remove xn from unassigned set} 6: HK+1 ← 0 {We define g(0) = 0} 7: for k ← 1 . . .K + 1 do 8: Gk ← G+ g(Hk + hn)− g(Hk) {Add xn to cluster k} 9: qk ← ef(Gk,U)
10: end for 11: qk ← qk/ ∑K+1 k′=1 qk′ , cn ∼ qk {Normalize probabilities and sample assignment} 12: if cn = K + 1 then 13: K ← K + 1 14: end if 15: G← G− g(Hcn) + g(Hcn + hn) {Add point xn to sampled cluster cn} 16: Hcn ← Hcn + hn 17: end for 18: Return c1 . . . cN

heading: 2.3 TWO EXAMPLES
text: Clustering in 2D Gaussian models: The generative model is
α ∼ Exp(1) c1:N ∼ CRP(α) N ∼ Uniform[5, 100]
µk ∼ N(0, σ2µ12) k = 1 . . .K xi ∼ N(µci , σ212) i = 1 . . . N
(10)
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) c1:N ∼ CRP10(α) N ∼ Uniform[5, 100]
lk ∼ Unif[0, 9]− without replacement. k = 1 . . .K xi ∼ Unif[MNIST digits with label lci ] i = 1 . . . N
where CRP10 is a Chinese Restaurant Process truncated to up to 10 clusters, and dx = 28 × 28. Training was performed by sampling xi 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.

heading: 3 RELATED WORKS
text: 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 fixed-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 fixed 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 fixed (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.

heading: 4 EXPECTATIONS, EVALUATIONS AND DIAGNOSTICS
text: 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)] =∑M i=1 f(c (i))wi/ ∑M i=1 wi where wi = 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(c1: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) satisfied, 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
prior p(c1:N ) and
qθ(c1:N ) ≡ ∫ dx qθ(c1:N |x) p(x) , (11)
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 efficiency, 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 configurations in 10 secs., efficiently 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.

heading: 5 COMMUNITIES
text: As an extension, we consider now a similar prior as above over cluster labels, but the observation model is more challenging:
α,N ∼ p(α), p(N) c1 . . . cN ∼ p(c1, . . . , cN |α)
φk1,k2 ∼ p(φ|β) k1 ≤ k2 xi,j ∼ Bernoulli(φci,cj ) , i ≤ j , i, j = 1 . . . N
where k1, k2 = 1 . . .K. Here p(c1:n|α) can be any exchangeable prior, and the binary observations xi,j represent edges in a graph of N vertices. We focus on the symmetric graph case here, so φk1,k2 = φk2,k1 and xi,j ≡ xj,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 Infinite 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 classification problem for every node over a fixed predetermined number of clusters.
We could proceed similarly to the clustering case, considering N particles, each given by a row of the adjacency matrix xi = (xi,1 . . . xi,N ). But we should be careful when encoding these particles. When values of c1:n are assigned, a generic encoding h(xi) would ignore the permutation symmetries present among the components of xi, i.e., the columns of xi,j , as a result of the c1:n assignments (the same three permutation symmetries discussed above for clustering models). Moreover, a fixed encoding h(xi) cannot accommodate the arbitrary length N of xi. 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.

heading: 6 APPLICATION: SPIKE SORTING WITH NCP
text: 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. (2) Using a variable-input softmax function, NCP is able to perform inference on cluster labels without assuming a fixed or maximum number of clusters. (3) NCP allows for efficient probablistic clustering by GPU-parallelized posterior sampling, which is particularly useful for handling the clustering uncertainty of ambiguous small spikes. (4) 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). (2) 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.
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 field (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, and finds many additional templates with clear RFs.
Hybrid Data. We compared NCP against vGMFM and Kilosort on a hybrid recording with partial ground truth as in Pachitariu et al. (2016). 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 efficient GPU-parallelized posterior sampling of cluster labels, NCP is able to handle the clustering uncertainty by producing multiple plausible clustering configurations. 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 efficient posterior sampling, demonstrating substantial promise for incorporating NCP into production-scale pipelines.

heading: A NEURAL CLUSTERING PROCESS FOR EXPONENTIAL FAMILIES
text: The likelihood for an exponential family is given by
p(x|µ) = eµ·t(x)−ψ(µ)m(xi) (12) = eλ·h(x)m(xi) (13)
where t(x) is a vector of sufficient statistics, and we defined
h(x) = (1, t(x)) (14) λ = (−ψ(µ), µ) (15)
Let us denote by K and K ′ ≥ K the total number of distinct values in c1:n and c1:N , respectively. Consider the joint distribution
p(c1:N ,x, µ) = p(c1:N )p(µ) K′∏ k=1 eλk· ∑ i:ci=k h(xi) N∏ i=1 m(xi) (16)
from which we obtain the marginal distributions p(c1:n,x) = ∑
cn+1...cN
p(c1:N ,x) (17)
= ∑
cn+1...cN
∫ dµp(c1:N )p(µ) K′∏ k=1 eλk·(Hk+ ∑ i>n:ci=k h(xi)) N∏ i=1 m(xi) (18)
= F (H1, . . . ,HK , h(xn+1), . . . , h(xN )) N∏ i=1 m(xi) (19)
where we defined
Hk = ∑
i≤n,ci=k
h(xi) k = 1 . . .K (20)
and Hk = 0 for k > K.
Note now that if p(c1:N ) is constant, all the dependence of F on c1:n, x1:n is encoded in the Hk’s, and F is symmetric under separate permutations of the Hk’s and the h(xi)’s for i > n. Based on these symmetries we can approximate F as
F ' ef(G,U) (21)
modulo adding to f any function symmetric on all xi’s, where
G = K∑ k=1 g(Hk) (22)
U = N∑ i=n+1 u(xi) (23)
In the conditional probability we are interested in,
p(cn|c1:n−1,x) = p(c1:n,x)∑ cn p(c1:n,x) , (24)
the product of the m(xi)’s in (19) cancels. Similarly, adding to f a function symmetric on all xi’s leaves invariant our proposed approximation
qθ(cn = k|c1:n−1,x) = ef(Gk,U)∑K+1
k′=1 e f(Gk′ ,U)
k = 1 . . .K + 1 . (25)

heading: B DETAILS OF THE NEURAL BLOCK PROCESS
text: Let us recall the generative model in this case,
α,N ∼ p(α), p(N) c1 . . . cN ∼ p(c1, . . . , cN |α) φk1,k2 ∼ Beta(α, β) k1 ≤ k2 (26) xi,j ∼ Bernoulli(φci,cj ) , i ≤ j , i, j = 1 . . . N
where k1, k2 = 1 . . .K. The prior p(c1:n|α) can be any exchangeable priors for clustering, and the observations xi,j represent the presence or absence of an edge in a graph of N vertices. We set φk1,k2 = φk2,k1 and xi,j ≡ xj,i, and assume for notational convenience that xij ∈ {+1,−1}.
B.1 ENCODING 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 xi = (xi,1, . . . , xi,N ). But we should be careful when encoding these particles. Consider the situation when the values of c1:n have been assigned. Encoding with a generic function h(xi) would ignore the permutation symmetries present among the components of xi, i.e., the columns of the matrix xi,j , as a result of the c1:n assignments. These symmetries are the same three symmetries discussed above for clustering models. Moreover, a fixed function h(xi) would not be able to accommodate the fact that the length of xi changes with the size N of the dataset.
Suppose that there are K clusters among the c1:n, each with sk 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 sK+1 = N − n, so we assume cn+1:N = K + 1, and we have ∑K+1 k=1 sk = N and
sk = ∑N j=1 δ(cj = k).
Now, in each row xi, the number sk of elements in the k-th cluster can be split as
sk = s − i,k + s + i,k
s+i,k = N∑ j=1 δ(cj = k)δ(xi,j = +1)
s−i,k = N∑ j=1 δ(cj = k)δ(xi,j = −1)
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 x1 = (+1,+1,−1,+1,+1). If four assignments were made c1 = c2 = 1, c3 = c4 = 2, then K = 2 and c5 = 3, and from x1 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 x1 = (+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 xj’s with cj = ci, such as
m+ci,k = 1
sci ∑ j:cj=ci s+j,k (27)
v+ci,k = 1
sci ∑ j:cj=ci (s+j,k −m + ci,k )2 (28)
and similarly m−ci,k and v − ci,k
. Note that these invariants are the same for all rows xj with cj = k. The motivation to consider them is that, if the partition corresponding to c1:n is correct, then for i ≤ n and k ≤ K we have n+i,k ' m + ci,k since they are both estimators of the latent Bernoulli
parameter φci,k. For the same reason, if the partition is correct and those two estimators of φci,k are exact, then v+ci,k ' 0. Similarly for m − ci,k
and v−ci,k. Then these values provide learning signals to the network that estimates the probability of the assignments c1:n being correct.
Therefore we propose to encode the components of xi belonging to cluster k as
ri,k = (s + i,k,m + ci,k , v+ci,k, s − i,k,m − ci,k , v−ci,k) ∈ R 6. (29)
In order to preserve the symmetry of the first K labels under permutations, we combine them as
ti ≡ K∑ k=1 t(ri,k) ∈ Rdt (30)
where the encoding function is t : R6 → Rdt . The encoding (29) of the unassigned components xi,n+1:N is kept separate and denoted as qi = ri,K+1.
In summary, each row xi of the adjacency matrix is represented by the fixed-dimensional pair (ti, qi) ∈ Rdt+6 in a way that respects the symmetries of the assignments c1:n: permutations between members of a cluster, permutations of cluster labels and permutations among unassigned columns.
B.2 CLUSTERING THE ROWS OF THE ADJACENCY MATRIX
We can proceed now as in regular clustering, encoding each cluster of xi’s within c1:n as Hk = ∑ i:ci=k h(ti, qi) ∈ Rdh k = 1 . . .K, (31)
and defining the permutation invariant, fixed-dimensional vectors
G = K∑ k=1 g(Hk), (32)
U = N∑ i=n+1 u(ti, qi). (33)
In terms of these quantities, the conditional probabilities are defined as usual as
qθ(cn = k|c1:n−1,x) = ef(Gk,U)∑K+1
k′=1 e f(Gk′ ,U)
(34)
for k = 1 . . .K + 1, with hn = h(tn, qn) and with Gk being the value of G for the different configurations. Compared to the regular clustering case, here we need to learn the additional function t. We call our approach Neural Block Process (NBP).

heading: C MONITORING GLOBAL PERMUTATION INVARIANCE
text: As mentioned in Section 5, 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.

heading: D DETAILS OF SPIKE SORTING USING NCP
text: 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 7 × 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 = [1, 2, 2, 2]). The last block is followed by an averaged pooling layer and a final linear layer. The outputs of the ResNet encoder are the hi and ui vectors of NCP, i.e. hi = ResNetEncoder(xi). We used dh = du = 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[Nmin, Nmax] (35) K ∼ 1 + Poisson(λ) (36)
π1 . . . πK ∼ Dirichlet(α1, . . . , αK) (37) c1 . . . cN ∼ Cat(π1, . . . , πK) (38) µk ∼ p(µ) k = 1 . . .K (39) xi ∼ p(xi|µci ,Σs ⊗ Σt) i = 1 . . . N (40)
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 ∈ R7×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 xi is obtained by adding to µci 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 (Figure 2 and Figure 7). Using beam search (Graves, 2012; Sutskever et al., 2014) with a beam size of 150, we were able to efficiently sample 150 high-likelihood clustering configurations for 2000 spikes in less
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 configuration 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 final 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 first 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-infinity 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.
KS; PTP 63.5 duplicates:2NCP; PTP 65.4 KS; PTP 43.9NCP; PTP 44.9 KS; PTP 42.5 NCP; PTP 44.5 duplicates: 2 KS; PTP 38.1NCP; PTP 39.0 KS; PTP 21.9 duplicates:2NCP; PTP 22.4
NCP; PTP 19.2 duplicates: 2 NCP; PTP 16.5 NCP; PTP 15.7 NCP; PTP 15.0 NCP; PTP 14.4
NCP; PTP 13.7 NCP; PTP 13.7 NCP; PTP 12.6 NCP; PTP 12.4 NCP; PTP 12.2
NCP; PTP 11.3 NCP; PTP 10.9 NCP; PTP 10.7 NCP; PTP 10.6 KS; PTP 10.3
KS; PTP 10.3 KS; PTP 10.0 NCP; PTP 9.8 NCP; PTP 9.3 KS; PTP 9.2 duplicates:2
NCP; PTP 9.2 duplicates: 2 NCP; PTP 9.1 NCP; PTP 9.1
NCP; PTP 9.0 duplicates: 3 NCP; PTP 8.9
KS; PTP 8.5 NCP; PTP 8.6 KS; PTP 8.5 NCP; PTP 8.6 KS; PTP 8.4 NCP; PTP 8.4 KS; PTP 8.4
KS; PTP 7.6 NCP; PTP 8.8 KS; PTP 8.0 KS; PTP 8.0 KS; PTP 7.5 NCP; PTP 8.4 duplicates: 2 KS; PTP 7.5 NCP; PTP 8.4 duplicates: 2

heading: E PARTICLE TRACKING
text: 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
ct ∼ p(ct|c1, . . . , ct−1) t = 1, . . . , T (41) µk,t ∼ p(µk,t|µk,t−1) k = 1 . . .K t = 1, . . . , T (42) xt ∼ p(xt|µct,t) t = 1, . . . , T (43)
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 ct 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 influence 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
Hk,t = t∑ t′=1:ct′=k e−b|t−t ′|h(xt′) k = 1 . . .K , (44)
Gt = K∑ k=1 g(Hk,t) , (45)
Ut = T∑ t′=t+1 e−b|t−t ′|u(xt′) . (46)
The conditional assignment probability for ct is now
qθ(ct = k|c1:t−1,x) = ef(Gk,t,Ut)∑K+1
k′=1 e f(Gk′,t,Ut)
(47)
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 ′| with a general distance function e−d(|t−t ′|). Figure 15 illustrates this model in a simple 2D example. We call this approach Neural Particle Tracking.

heading: F NEURAL ARCHITECTURES IN THE EXAMPLES
text: 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), 1 for p(c1:N ), 64 for p(x|c1:N ). The architecture of the functions in each case were:
CLUSTERS: 2D GAUSSIANS
• 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
CLUSTERS: MNIST
• 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
COMMUNITIES: IRL
• 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

Review summaries:
In this paper, the authors consider the neural amortized inference for clustering processes, in which the number of cluster can be automatically adapted based on the observed samples. The proposed algorithm largely follows the standard variational auto-encoder. The major contribution of the paper is the design of the posterior parametrization so that the posterior satisfies the permutation invariant within a cluster, between clusters, and unassigned data, based on the DeepSet method. The model can be incorporated into random communities models. Finally, the authors apply the algorithm for neural spike sorting problem. 

This paper introduces a novel deep learning architecture for efficient amortized Bayesian inference over mixture models. Unlike previous approaches to amortized clustering, the proposed method allows us to treat local discrete labels of data points and infer the unbounded number of mixture components, making it more flexible as in the case of Bayesian nonparametrics. It is shown that the resulting algorithm can be parallelized and applied to both conjugate and non-conjugate models. The authors also suggest an extension to models of random communities and a novel approach to neural spike sorting for high-density multielectrode arrays based on the proposed method.