Neural Clustering Processes
Authors: Ari Pakman
Yueqi Wang
Catalin Mitelut
JinHyung Lee
Liam Paninski
Venue:  (2020)
Keywords: amortized inference
probabilistic clustering
mixture models
exchangeability
spike sorting

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. 


Under review as a conference paper at ICLR 2020
NEURAL CLUSTERING PROCESSES
Anonymous authors Paper under double-blind review
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 exchangcability of the generative representations
models and is based on mapping distributed, cermutation-invaric of discrete arrangements into varying-size multinomial conditional probabilities. The resulting algorithm parallelizes easily, yields iid samples from the approximate posteriors along with normalized probability estimate of each sample quantity generally unavailable using Markov Chain Monte Carlo) and can easily be applied 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 novel approach to neural spike sorting for high-density multielectrode arrays.
INTRODUCTION
Mixture models (or equivalently, probabilistic clustering models) are 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 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 cone-consuming. with convergence that is difficult assess. Models whose likelihood and prior 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 iet 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-dimensiona, distributed representations that respect the permutation symmetries imposed by the discrete variables. A major advantage our architecture, compared to previous approaches to amortized clustering, is its ability to handle arbitrary number of clusters. This makes the method natural choice for nonparametric Bayesian models, such as Dirichlet process mixture models (DPMM), and their extensions, where the number of components, measure of the model complexity, is inferred posterior random variable; see (Rodriguez & Mueller (013)) for recent overview. Moreover, the method can be applied to both conjugate and non-conjugate models.

---

Under review as a conference paper at ICLR 2020
The term 'amortization"' refers to the process of investing computational resources to train model that is later used for very fast posterior inference (Gershman & Goodman 2014). Concretely, the amortized approach learns parametrized function Yo (2) ()) that approximates p(z[z) for any æ; learning the model parameters may be computationally challenging, but once 0 is in hand then evaluating (=) for new data 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 ()(2(z) along with the generative model pa (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: generative model p(r. 2) is postulated, and posterior inference is the main focus of the learning phase. Amortized methods in this case usually involve 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 al. 2016), neural decoding (Parthasarathy cet al., 2017) and particle tracking (Sun & Paninski 2018). Our work specialized to the case where the latent variables 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, can obtain independent, parallelizable, approximate posterior samples of the discrete variables for any new set of observations arbitrary size, with no need for expensive MCMC steps. These samples can be used (i) compute approximate expectations, (ii) as high quality importance samples, or (iii) as independent Metropolis-Hasting proposals.
In Section 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 Seccion 6 with neuroscientific application of this method to spike sorting for high-density multielectrode probes.
2 AMORTIZING MIXTURE MODELS
We start by presenting mixture models from the perspective of probabilistic models for cluster- ing (McLachlan & Basford 1988). The latter introduce random variables o denoting the cluster number to which the data point z; is assigned, and assume a generating process of the form
a1 02 ~ p(a)
N p(N)
CI...cy CN ~
M1 MI... MKl p(u1, MK a2)
I2~ p(z.le)))) i 1...N
Here æ1, 02 are hyperparameters. The number of clusters K is random variable, indicating the number of distinct values among the sampled o; ss and μı denotes a parameter vector controlling the distribution of the k-th cluster (e.g., P4 could include both the mean and covariance of a Gaussian mixture component). We assume that the priors p(CI.N (a1) and p(M1: (a:)) are exchangeable,
Ic, cN =
where {or} is an arbitrary permutation of the indices, and similarly for p(u llg)). Our interest in this work is in cases where K can take any value K < N, such as the Chinese Restaurant Process (CRP), its Pitman-Y Yor generalization. Of course, 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 { we would like to draw independent samples from the posterior
p pC xx) p(ci ())((c))1, x) h.t....ppc/c..... p(cN((IN--1.))
Note that p(c, = 1|x) = 1, since the first data point is always assigned to the first cluster. While might also be interested in the hidden variables on 02, M4c, the reason to focus the discrete
2

---

Under review as a conference paper at ICLR 2020
variables 's is that given samples from them, it is generally relatively easy to obtain posterior samples from p(an (1:N) and p(Ak, a22x, C1:N).
We would like to model all the factors in in unified way, with generic factor given by
CIn-,,) = p(c ... Gx,
(2)
Σ p(c ... Cr x)
Here we assumed that there are K unique values in C1n-11, and therefore en can take K values, corresponding to joining any of the K existing clusters, or forming its own new cluster.
We are interested approximating
(Im z0 qe(cn (1cm(11mm)))))
(3)
where de is parameterized by a flexible model such as neural network that takes as inputs (1,
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 Ihs of (3)
To make this symmetric structure more transparent, and in light of the expression the joint distribution of the assignments of the first n data points.
2p. let us consider
p(c Co x).
(4)
neural representation of this quantity should respect the permutation symmetries imposed on the x; by the values of Ci:n- Therefore, our first task is to build permutation-invarian representations of the observations The general problem of constructing such invariant encodings discussed recently in (Zaheer et al. 2017); to adapt this approach to our context, we consider three distinct permutation symmetries:
Permutations within a cluster: is invari- ant under permutations of 2 in the same clus- ter. For each of the clusters that have been sampled so far, we define the encoding
(G.,)) .i
(G,.4)
H₂ Σ h(x;)
h Fd + 1d
(5)
which is clearly invariant under permutations x;'s 's in the same cluster. In general is an encoding function
(G.,
learn from data, unless p(z|j4) belongs to an exponential family and the prior p(CI.N) is constant, as shown Ap- pendix A
Permutations between clusters: is invari- ant under permutations of the cluster labels. In terms of the within-cluster invariants Hg. this symmetry can be captured by
Encoding cluster labels. The col- ored points have fixed labels C1:n-1, forming K = 3 clusters. The four possible labels for the circled point give four encoding vectors G₁, while the vector U encodes the 3 gray unlabeled points (Best seen in color).
Figure 1:
G Σ (Hx), 1=1
g RR. Rdd
(6)
Permutations of the unassigned data points: 4 is also invariant under permutations of the N unassigned data points. This can be captured by
N
U
u(zz),
W: R Rd
1=++1
Note that G and U provide fxed-diimensionl. symmetry invariant representations of the assigned and non-assigned data points, respectively, for any values of and K. Encodings of this form were shown in (Zaheer et al. 2017) to lead arbitrarily accurate approximations of symmetric functions.
-inamntt-invvriannt
3

---

Under review as a conference paper at ICLR 2020
Neural Clustering Process
G_=1
X %,
Vo voasia imput softmax
0.08
c-=2
H
x x
K X,
0.02
C 1
x x
% x.
0.90
Figure 2: Architecture of the Neural Clustering Process. The full model composed by the deep networks h, 9, u, f. Left: After assigning the cluster labels Cia--1, each possible discrete value for en gives a different symmetry-invarian encoding of Ziaa into the vector Gg. using the functions and g. The remaining, vet-unassignes points In+1:NV are encoded by and summed into the vector U. Right: Each pair Gj, U is mapped by f into real number (logit), which in turn is mapped into the multinomial distribution (c (1 via a variable-inpu softmax.
2.1 THE VARIA SOFTMAX
After assigning values to C1:m-1+ 1. each of the K + possible values for c, corresponds h(x,) appearing in one particular H& in (5) and yields a separate vector G2 in (6 See Figure ] for an example. In terms of the G 's and U, propose model (2) as
(GG.U
go (c, = kk/1:an111
k=1..
(8)
ΣØ ((G
where we have introduced a new real-valued function f. In other words, each value of e, corresponds to different channel through which the encoding h(x,)) flows to the logit value f. shown in Figure Note that & = + 1 corresponds to c, forming its own new cluster with H₃. = h(xx).
The softmax differs from its usual form in, c.g., classification networks, where 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(,,)) takes to G, thus allowing a flexible number of categories.
In eq. 0 denotes the parameters in the functions ht g, and f. which we represent with neural networks. By storing and updating G and U for successive values of n, the computational cost of single Gibbs sweep. See Algorithm for details;
full i.i.d. sample of C1:N O(NK). the same we term this approach the Neural Clustering Process (NCP), It is relatively easy to run hundreds of copies of Algorithm parallel on GPU, with each copy yielding a different of samples C1:N
2.2 THE OBJECTIVE FUNCTION
In order to learn the parameters 0 of the neural networks, we use stochastic gradient descent to minimize the expected KI. divergence,
N
Bp(N) DieL Dkx.(p((\x)l999(((e( -PP(N) Ep(oux.x xx)
Σ
log go (Cn (etm//mmm-1- 1,, + const.
Samples from p(CI: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 approximate p(cm//imm-1, Cimm-1, x) arbitrarily accurately, Note that the gradient here acts only on the variable-input softmax term qe. not p(c, x), so there is no problem of backpropagating through discrete variables (Jang al. 2016 Maddison al. 2016).
A Pytorch implementation of the algorithm available https:///iiiii http:////ttt.t/1kkot
4

---

Under review as a conference paper at ICLR 2020
Algorithm O(NK) Neural Clustering Process Sampling
1: h < h(r;),
u(z;)
i= 1=...
[Notation |
2: U <- Σ Li=z
(Initialize unassigned set) (Create first cluster with x1
3: H1 f mn, G g(H)), K 1,
4: for + 2...N do
5:
U+U-W,
Reeoov from unassigned set)
6:
HK+1 +
{ (We define g = 0}
7:
k 1...K K + do
8:
Gk - G g H + hnn) - g Hs)
{ (Add Tya to cluster k)
9:
Ok + e=f(G)
10:
end for
11:
+ qk Σk'
k-=
Cn ~ Qk
[Normalize probabilities and sample assignment]
if Cn K then
13:
K K+ K +
end if
15:
G G - g( ((He))) g(He + ha)
[Add point T, to sampled cluster C,
He. < He, + hn
17:
end for
18: Return CI...CN
2.3 TWO EXAMPLES
Clustering in 2D Gaussian models: The generative model is
Exp(1) C1.N CRP(a) N Uniform(5, 100]
14 ~ NOO, == 32 N(pe, x212) i 1...N
(10)
where CRP stands for the Chinese Restaurant Process, with concentration parameter a, on = 10, and Figure 3/shows that the NCP captures the posterior uncertainty inherent in clustering this data. Note that since the generative model an analytical known distribution, there 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) C1NN CRP1o(a)
1. Unif 0, 9] without replacement. I: = 1 K 2e MNiS digits with label lc,] i
N Uniform 5, 100)
N
where CRP1o is a Chinese Restaurant Process truncated to up to 10 clusters, and dz = 28 x 28. Training was performed by sampling æ; from the MNIST training set. FFiure 4 shows posterior samples for 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 this case the generative model has no analytical expression (and therefore is non-conjugate), but this presents problem; a set of labelled for details of all the network architectures used.
samples all we need for training. See Appendix F
to
Chidwes Proh: 0.32
Caaten Prbc 2.16
Elluders Putt:
Custtrs Prit: 0.04
Chuders Pro: 6.00
Figure 3: Mixture of 2D Gaussians: Given the observations in the leftmost panel, we show samples from the NCP posterior. Note that less-reasonable samples are assigned lower probability by the NCP. The dotted ellipses indicate assignments which differ from the first, highest-probability sample. Our GPU implementation gives thousands of such samples in a fraction of a second. (Best seen in color.)
5

---

Under review as a conference paper at ICLR 2020
Observations
8
9
O 8
6
9
Cluster Samples
K = 5 Pr: 0.68
6
8
8
9 9
6
4
Pr: 0.11
8
8
8
8
9 9 7
6
K 4
Pr: 0.06
8
8
8 8
6 8
7
9
7
0
K 6 Pr: 0.03
8
8
9 9
9 9
7
+
0
6
K 5 Pr: 0.02
8
8
8
9
9
7
+
7
7
O
6
Figure 4: NCP trained on MNIST clusters. The top row shows 20 images from the MNIST test set. The five rows below show five samples of C1:20 from the NCP posterior. Note that each sample captures some ambiguity suggested by the form of particular digits.
RELATED WORKS
Permutation coemutation-invariann neural architectures have been explored recently in (Ravanbakhsh al. 2017
Korshunova et al. 2018 Lee et al.] 2018: Bloem-RRede Reddy & Teh [019 Wagstaff et al. The
representation of a set via sum (or mean) of encoding vectors was used in (Guttenberg et
2016; Ravanbakhsh et al., 2016 Edwards & Storkey 2017 Zaheer et al. 2017] Garnelo et al. 2018a b).
Most works 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 mixtures of Gaussians, so is instructive to compare them in detail to our approach.
The work (Le et al. 2016) studies amortized inference of variable number of latent variables generated during the trace of general sequential probabilistic program. For the case of mixture of 2D Gaussians with latent random number of components, 2D histogram image of binned observations is fed to a convolutional network whose output enters into a recurrent neural network with fixed-sized softmax output layer estimate the number of clusters. The network also outputs the means and covariances of each cluster.
The work (Lee al. 2018) presents Set Transformer, 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 al. 2016) fixed (Lee al. 2018). Second, the models perform inference on the continuous parameters point. Finally, in (Le al. 2016). the use of convnet 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.
but not on the discrete labels of each data
2D histogram determine the number of
Property
NCP
Program Compilation
Set Transformer Fixed
Number of mix. components
Arbitrary
Bounded
Amortizes discrete labels
Yes
No
No
Amortizes component parameters
No
Yes
Yes
Scales to high dimensional data
Yes
No
Yes
Table 1: Comparing amortized approaches to Gaussian mixtures. We compare NCP with Pro- gram Compilation (Le al. 2016) and Set Transformer (Lee et al. 20188), two previous approaches to amortized mixtures of Gaussians. Note however that NCP can be applied to any mixture model.

---

Under review as a conference paper at ICLR 2020
coord rate
20
Mixture of 2D Gaustians
CMP bution number of closters for N-30
NCP Samples
1GT Sumpka
Exect
0.2
0.00
Number
6.10
6.5
000
Points
Samberr
Figure 5: Quantitative Evaluations. Upper left: Two 2D clusters of 50 points each = 0, and line over possible locations of a 101st last point. Upper right: Assuming the 2D model from (TO). the posterior p(C101 |C1:100, x) can be computed exactly, and we compare to the NCP estimate as function of the horizontal coordinate of T101, as this point moves over the gray line on the upper left panel. Geweke's Tests. Lower left: The curves compare the exact mean (# one std.) of the number of clusters K for different N's from the CRP prior (with a 0.7), with sampled estimates using equation Lower right: Similar comparison for the full histogram of K for = 30 points.
/(1101/C1:100-
EXPECTATIONS, EVALUATIONS 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, B[f(c)] =
Eis ) / Zi=1 where uy = p(x, ( qo (t) 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 the true posterior distributions in these models are capturing the uncertainty inherent in the observed data. But we would like 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 these models. Nonetheless, there are several quantitative metries we can examine to check the accuracy of the model output.
Global symmetry from exchangeability: Our results relied on p(CEN)(a)) being exchangeable, which in turn implies exchangeability of the joint posterior (T). But this is not explicit in the rhs of ( 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 AAppeddix C
Some conditional probabilities be computed analyt- ically 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.
Estimated vs. Analytical Probabilities:
Geweke's Tests: A popular family of tests that check the conrectness of MCMC implementa- tions (Geweke 2004) can also be applied in our case: verify the (approximate) identity between the
7

---

Under review as a conference paper at ICLR 2020
Custer:
Custers Prob: 0.156
Quaters
Clusters
Proz:
Clustess
Figure 6: Community Detection with Neural Block Processes. The model is single-type Infinite Relational Model (Kemp et al. 2006 Xu et al. 2006), with a CRP prior with a = 0.7. The entries in each block are Bernoulli samples, with block parameter sampled from Beta(O.2, 0.2) prior. From left to right: (i) the original block structure, sampled from the generative model, (ii) the observed random permutation of rows and columns, (iii) four samples from the NBP posterior, along with their estimated probabilities. Fach sample from the posterior here corresponds to plausible partition.
prior p(C.,)) and
(e(C).)) = f (o1(111/1//)) p(x),
(11)
where p(x) is the marginal from the generative model. Figure 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, 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. (T), 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 burn-in period.
NCP vs. Variational Inference: In Section 6 we compare NCP with 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 10 secs., efficiently capturing clustering uncertainty. In addition, the variational approach requires preprocessing step that projects the samples to lower dimensions, whereas NCP directly consumes the high-dimensional data by learning an encoder function h.
5 COMMUNITIES
As extension, we consider now similar prior as above over cluster labels, but the observation model is more challenging:
a, N ~ p(a),p(N)
k p(013)
k S K2
C1 C.... CN ~ .y,l,
Fz. (emnoull(@e,.e))) isj, i.j=
where k1, 12 = 1... ...... K. Here p(c a)) can be any exchangeable prior, and the binary observations graph of vertices. We focus on the symmetric graph case here, so ΦkI.k = ŌR and Tı.j = æj.;- We use a Beta model for p(o13), but other choices are possible.
Jiy represent edges
These models include stochastic block models (Holland al.) 1983 Nowicki & Snijders 2001) and the single-type Infinite Relational Model (Kemp et al.) [2066] Xu et al. 20066; Schmiitt & Morup 2013). Neural architectures for communities in graphs have been studied (Chen al. (019) as 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 row of the adjacency matrix x; = (xt_1 ... Fi,N). But we should be careful when encoding these particles. When values of Cimm are assigned, a generic encoding (xx)) would ignore the permutation symmetries present among the components of x,, i.e., the columns of Ti,j, as a result of the Ci:n. assignments (the same three permutation symmetries discussed above for clustering models). Moreover, a fixed encoding h(x;) cannot accommodate the arbitrary length N of x,. In Appenddi B we present an invariant encoding that respects all these requirements. We call our approach Neural Block Process (NBP). See Figure 6 for example.
8

---

Under review as a conference paper at ICLR 2020
6 APPLICATION: SPIKE SORTING WITH NCP
Large-scale scale neural population recordings using multi-electrode arrays (MEA) are crucial for under- standing 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 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 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 mully-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 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 fixed maximum number of clusters. (3) NCP allows for efficient probablistic clustering by EPU-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 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). Appendix D we describe the dataset, neural architecture, and the training/inferenee pipeline of NCP spike sorting.
Observations
Clusters
cho w w mm nm
chl
mymm
Channels
ch2 Mm wym
ch3 wwm mym
ch4 v weme ma
ch5 vywammm wor mom
che
Overlay Average
Figure 7: Clustering multi-channel spike waveforms using NCP. Each row is an electrode channel.
Spikes with the same color belong to the same cluster. (Scale bar: 5x standard deviation (SD)).
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 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 ground- truth cluster labels, compared the spike templates extracted from NCP and Kilosort using retinal

---

Under review as a conference paper at ICLR 2020
receptive field (RF), which is computed for each cluster the mean of the stimulus present each spike. A clearly demarcated RF provides encouraging evidence that the spike template corresponds to real neuron. Side-by-side comparisons of matched RF pairs shown in Figure 8 (bottom-lefh 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 bottoom-rggg) showing that NCP performs at least as well as Kilosort, and finds many additional templates with clear RFs.
NCP
VGMFM
do A/
cho
dh1
ch1
th2
ch2
NCP
dM
Clusters:
vGMFM
0:955
100
m:00
NCP
Kilosorrt
CS 4TP 43 NOP 4L2 KS PIP
NCP F.
NCh: PTP 15.7
Nch PIP 6 Ks PTP 20-3
4B
55
17
Receptive
Figure 8: Spike sorting on real data. 2000 spikes from real data were clustered by NCP (top-left) and vGMFM (top-mid). Each column shows the spikes assigned one cluster (overlaying traces and their average). Each row is one electrode channel. Top-right: t-SNE visualization of the spike clusters. Bottom-left: Example pairs of matched RFs recovered by NCP (red boxes) and Kilosort (blue boxes). Blank indicates no matched counterpart. Bottom-right: Venn diagram of recovered RFs.
Example
Example
Hound
Clusters
Clusters
SOVNN
d10
ch1
LIVED 1
IHO Sptey MayS
JNT D 77 210
d16
Figure 9: Spike sorting on hybrid data. Top: NCP, Kilo- sort, vGMFM recovered 13, 8, and of the 20 injected ground-truth templates. Bottom: Peak-to-peak (PTP) size and firing rate of each injected template. (Smaller tem- plates with lower firing rates are more challenging.)
Figure 10: Clustering ambiguous small spikes. In both examples, multiple plausi- ble clustering results of small spikes were produced by sampling from the NCP pos- terior. (scale bar = 5x SD)
Hybrid Data. We compared NCP against vGMFM and Kilosort on 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 GPUU-paraleliized posterior sampling of cluster labels, NCP is able handle the clustering uncertainty by producing multiple plausible clustering configurations. Figure 10 shows examples where NCP separates spike clusters with amplitude as low 3-4x 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 method, and handles clustering uncertainty by efficient posterior sampling, demonstrating substantial promise for incorporating NCP into production-scale pipelines.
10

---

Under review as a conference paper at ICLR 2020
REFERENCES
Elie Aljalbout, Vladimir Golkov, Yawar Siddiqui, and Daniel Cremers. Clustering with Deep Learning: Taxonomy and New Methods. arXiv preprint arXiv: 10644,7 2018.
David M. Blei and Michael I. Jordan. Variational Methods for the Dirichlet Process. Proceedings of the Twenty-first International Conference on Machine Learning, ICML '04, 2004.
Benjamin Bloem- Bloem-Reddy Reddy and Yee Whye Teh. Probabilistic symmetry and invariant neural networks. arXiv preprint XXX::1901.00022 2019.
Ana Calabrese and Liam Paninski. Kalman filter mixture model for spike sorting of non-stationary data. Journal of neuroscience methods, (96(1))159-160) 2011.
David Carlson and Lawrence Carin. Continuing progress of spike sorting in the era of big data.
Current opinion in neurobiology, 55:90-96, 2019.
David Carlson, Joshua Vogelstein, Qisong Wu, Wenzhao Lian, Mingyuan Zhou, Colin R Stoetzner, Daryl Kipke, Douglas Weber, David B Dunson, and Lawrence Carin. Multichannel clectrophysiological spike sorting via joint dictionary learning and mixture modeling. IEEE Transactions Biomedical Engineering, 61(1):11-54,)) 2013.
Fernando Chaure, Hernan G Rey, and Rodrigo Quian Quiroga. novel and fully automatic spike-sorting implementation with variable number of features. Journal of neurophysiology. 120 (4):1859-1871, 2018. doi: 10. 10.11522 1155//.00333.200000
Zhengdao Chen, Lisha Li, and Joan Bruna. Supervised Community Detection with Line Graph Neural Networks. ICLR, 2019.
E. J. Chichilnisky and Rachel S. Kalmar. Functional asymmetries in on and off ganglion cells of primate retina. Journal of Neuroscience, 22(7):2737-2477, 2002. ISSN 0270-6474. doi: 10. D0.1222NNEUURSSC22200000077.77..
URL http://://wee URLppcp://www..uuussst..... http:oocc...gg/coneeett/t/
22/7/2737
Jason Chung. Jeremy F Magland, Alex H Barnett, Vanessa M Tolosa, Angela C Tooker, Kye Lee, Kedar G Shah, Sarah H Felix. Loren M Frank, and Leslie F Greengard. A fully automated approach to spike sorting. Neuron, 95(6): 2017.
K-L Du. Clustering: A neural network approach. Neural networks, 23(1):89-107, 2010.
Harrison Edwards and Amos Storkey. Towards a neural statistician. ICLR, 2017.
Marta Garnelo, Dan Rosenbaum, Chris Maddison, Tiago Ramalho, David Saxton, Murray Shanahan, Yee Whye Teh, Danilo J Rezende, and SM Eslami. Conditional neural processes. In International Conference on Machine Learning, 2018a.
Marta Garnelo, Jonathan Schwarz, Dan Rosenbaum, Fabio Viola, Danilo Rezende, SM Eslami, and Yee Whye Teh. Neural processes. ICML 2018 workshop on Theoretical Foundations and Applications of Deep Generative Models, 2018b.
Samuel Gershman and Noah Goodman. Amortized inference in probabilistic reasoning, In Proceed- ings of the annual meeting of the cognitive science society, volume 36, 2014.
John Geweke. Getting it right: Joint distribution tests of posterior simulators. Journal of the American Statistical Association, 99(447):799-880 -044 2004.
Alex Graves. Sequence transduction with recurrent neural networks. CoRR, 1211
Nicholas Guttenberg, Nathaniel Virgo, Olaf Witkowski, Hidetoshi Aoki, and Ryota Kanai. Permutationn-qquivrriat neural networks applied to dynamics prediction. arXiv preprint arXiv.1612.045300 2016.
Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recognition. The IEEE Conference on Computer Vision and Pattern Recognition (CVPR), June 2016.
11

---

Under review as a conference paper at ICLR 2020
Geoffrey E Hinton, Peter Dayan, Brendan Frey, and Radford M Neal. The" wake-sleep" algorithm for unsupervised neural networks. Science, 168(5214):1158-11 -161, 1995.
Paul W Holland, Kathryn Blackmond Laskey, and Samuel Leinhardt. Stochastic blockmodels: First steps. Social networks, 5(2):199-337, 1983.
Michael Hughes, Dae I1 Kim, and Erik Sudderth. Reliable and scalable variational inference for the hierarchical Dirichlet process. In Artificial Intelligence and Statistics, pp. 370-378, 2015.
Michael C Hughes and Erik Sudderth. Memoized online variational inference for dirichlet process mixture models. In Advances in Neural Information Processing Systems 26, pp. 1133-1144 111. 2013.
Sonia Jain and Radford M Neal. A split-merge Markov chain Monte Carlo procedure for the Dirichlet process mixture model. Journal of compulational and Graphical Statistics, 13(1):158-182, 2004.
Eric Jang, Shixiang Gu, and Ben Poole. Categorical reparameterization with Gumbel-softmax. arXiv preprint arXiv: 161111111000 1.01 144, 2016.
James J. Jun, Catalin Mitelut, Chongxi Lai, Sergey L. Gratiy, Costas A. Anastassiou, and Timothy D. Harris. Real-time spike sorting platform for high-density extracellular probes with ground-truth validation and drift correction. bioRxiv, 2017.
Charles Kemp, Joshua B Tenenbaum, Thomas L Griffiths, Takeshi Yamada, and Naonori Ueda. Learning systems of concepts with an infinite relational model. In AAAI, volume pp. 2006.
Diederik Kingma and Jimmy Ba. Adam: method for stochastic optimization. ICLR, 2015.
Diederik Kingma and Max Welling. arXiv:1312.6114, 2013.
Auto-encoding variational bayes. arXiv preprint
Iryna Korshunova, Jonas Degrave, Ferene Huszar, Yarin Gal, Arthur Gretton, and Joni Dambre. Bruno: A deep recurrent model for exchangeable data. In Advances in Neural Information Processing Systems 31, 2018.
Kenichi Kurihara, Max Welling, and Yee whye Teh. Collapsed Variational Dirichlet Process Mixture Models. In IJCAI, volume 7, pp. 2796-2801, 2007.
2007.
Tuan Anh Le, Atilim Gunes Baydin, and Frank Wood. Inference compilation and universal proba- bilistic programming, arXiv preprint 1110::6610009000 2016.
Jin Hyung Lee, David E Carlson, Hooshmand Shokri Razaghi, Weichi Yao, Georges Goetz, Espen Hagen, Eleanor Batty, EJ Chichilnisky, Gaute T Einevoll, and Liam Paninski. Yass: Yet another spike sorter. Advances in Neural Information Processing Systems, pp. 4002-4012, 2017.
Juho Lee, Yoonho Lee, Jungtack Kim, Adam Kosiorek, Seungjin Choi, and Yec Whye Teh. Set transformer. arXiv preprint 18X.v:8810.0088 2018.
Jun S Liu. Metropolized independent sampling with comparisons to rejection sampling and impor- tance sampling. Statistics and computing, 6(2):113-119) 1996.
Chris J Maddison, Andriy Mnih, and Yee Whye Teh. The concrete distribution: A continuous relaxation of discrete random variables. arXiv preprint arXiv:1611.0077 2016.
Geoffrey J McLachlan and Kaye E Basford. Mixture models: Inference and applications clustering, volume 84, Marcel Dekker, 1988.
Jeffrey W Miller and Matthew T Harrison. Mixture models with prior on the number of components. Journal of the American Statistical Association, 113(521):340-355) 2018.
Erxue Min, Xifeng Guo, Qiang Liu, Gen Zhang, Jianjing Cui, and Jun Long. A survey of clustering with deep learning: From the perspective of network architecture. IEEE Access, 6:39501-39514 2018.
Andriy Mnih and Danilo Rezende. Variational inference for monte carlo objectives. arXiv preprint arXiv: 2016.
12

---

Under review as a conference paper at ICLR 2020
Radford M Neal. Markov chain sampling methods for Dirichlet process mixture models. Journal of computational and graphical statistics, 9(2):249 225,, 2000.
Krzysztof Nowicki and Tom A B Snijders. Estimation and prediction for stochastic blockstructures. Journal of the American statistical (usoooccaiioonn64555111111701101 2001. 2001.
Marius Pachitariu. Kilosort2.
2019.
Marius Pachitariu, Nicholas Steinmetz, Shabnam Kadir, Matteo Carandini, and Kenneth D Harris. Kilosort: realtime spike-sorting for extracellular electrophysiology with hundreds of channels. BioRxiv, pp. 061481, 2016.
Brooks Paige and Frank Wood. Inference networks for sequential Monte Carlo in graphical models. In International Conference on Machine Learning, pp. 3040-3049, 2016.
Nikhil Parthasarathy, Eleanor Batty, William Falcon, Thomas Rutten, Mohit Rajpal, E.J. Chichilnisky, and Liam Paninski. Neural Networks for Efficient Bayesian Decoding of Natural Images from Retinal Neurons. Advances in Neural Information Processing Systems 30, pp. 6434-6445. 2017.
Siamak Ravanbakhsh, Jeff Schneider, and Barnabas Poczos. Deep learning with sets and point clouds.
arXiv preprint
2016.
Siamak Ravanbakhsh, Jeff Schneider, and Barnabás Póczos. Equivariance through parameter-sharing In Proceedings of the 34th International Conference on Machine Learning, 2017.
Daniel Ritchie, Paul Horsfall, and Noah D Goodman. Deep amortized inference for probabilistic programs. arXiv preprint arXiv:1610.05755 2016.
Abel Rodriguez and Peter Mueller. NONPARAMETRIS BAYESIAN INFERENCE. NSF-CBMS Regional Conference Series in Probability and Statistics,, 911-110, 2013.
Mikkel N Schmidt and Morten Morup. Nonparametric bayesian modeling of complex networks: An introduction. IEEE Signal Processing Magazine, 30(3):1 10-128, 2013.
Kevin Q Shan, Evgueniy V Lubenov, and Athanassios G Siapas. Model-based spike sorting with a mixture of drifting t-distributions. Journal of neuroscience methods, 288:82-99 -8. 2017.
Andreas Stuhlmüller, Jacob Taylor, and Noah Goodman. Learning stochastic inverses. In Advances in neural information processing systems, pp. 3048-3056, 2013.
Ruoxi Sun and Liam Paninski. Scalable approximate Bayesian inference for particle tracking data. In Proceedings of the 35th International Conference on Machine Learning, 2018.
Ilya Sutskever, Oriol Vinyals, and Quoc V. Le. Sequence to sequence learning with neural networks. In NIPS. 2014.
Nguyen Xuan Vinh, Julien Epps, and James Bailey. Information theoretic measures for clusterings Variants, properties, normalization and correction for chance. Journal of Machine Learning Research, 11(0c)):837-22854 2010.
comparison:
Edward Wagstaff, Fabian B Fuchs, Martin Engelcke, Ingmar Posner, and Michael Osborne. On the limitations of representing functions sets. arXiv preprint 1rXiv:1901.09006 2019.
Frank Wood and Michael Black. A nonparametric bayesian alternative to spike sorting. Journal of neuroscience methods, 173, 2008.
Zhao Xu, Volker Tresp, Kai Yu, and Hans-Peter Kriegel. Learning infinite hidden relational models. Uncertainity in Artificial Intelligence (UAI2006), 2006.
Manzil Zaheer, Satwik Kottur, Siamak Ravanbakhsh, Barnabás Póczos, Ruslan Salakhutdinov, and
Alexander J. Smola. Deep sets. In Advances in neural information processing systems, 2017.
13

---

Under review as a conference paper at ICLR 2020
A
NEURAL CLUSTERING PROCESS FOR EXPONENTIAL FAMILIES
The likelihood for an exponential family is given by
p(alp)) (utt())-(()))
(12)
m(
(13)
where ((x) is vector of sufficient statistics, and defined
(1,)(z)) h(z) = (-v(4),
(14)
(15)
Let us denote by and K' K the total number distinct values ciun and CIN,, respectively. Consider the joint distribution
K' II k=
N
P(CEN, x. = p(cc:)()p)))) (p(p)
e
de Euerus b(x))
II m(x:)
(16)
from which we obtain the marginal distributions
p(ei:: x)
Σ x(C.,N,))
(17)
CntI...N
/ dupleun))(w))
K'
Σ entlmen
II k=1
Av Ae Ei h(zz))
k(z,))
m(x;)
(18)
N
F(HI..... HK. h(thhh(nn........)...)))(()))))) ((x)) II m(x;)
(19)
where we defined
H Σ Max)) iSm,==
k k=1...K K
(20)
and H₁ = O for k > K.
Note that if p(CI:N) is constant, all the dependence of F on Cisn F1.n encoded in the Hg's, and F' is symmetric under separate permutations of the Hy 's and the h(z;)'s for i > n. Based on these symmetries can approximate
F ((G,U)
(21)
modulo adding to f any function symmetric on all x3 ss's, where
K
G =
Σ (H&)
(22)
k=1
U u(x.)
(23)
the conditional probability we are interested in,
p(e x)
plc//cmm-115
(24)
Ee, p(Cim, x)
the product of the m(xx)'s in 19 cancels. Similarly, adding f a function symmetric on all x;'s
leaves invariant our proposed approximation
((G...))
(c, = S/c1m--1,X)
1 K + 1.
(25)
K+1 (f(G
_K=1
14

---

Under review as a conference paper at ICLR 2020
B DETAILS OF THE NEURAL BLOCK PROCESS
Let us recall the generative model in this case,
a, ~ p(a).
C1 ... CN p(c,, , cNla)
k2 ~ Beta(a, B) 1n S l/2
(26)
Fi,
(e isj, i.j=1...N
where 1 1/2 = 1... K. The prior p(Cim/a) (a) can be any exchangeable priors for clustering, and the observations Ti.j represent the presence or absence of an edge in a graph vertices. We set A2 0b2,ds Ay and Fijj = j... and assume for notational convenience that Tij E -1],
B.1 ENCODING EACH ROW OF THE ADIACENCY MATRIX
In principle posterior inference in this case can proceed similarly to the clustering case, by considering N particles, each given by row the adjacency matrix x; "er.,,..,,,..... But we should be careful when encoding these particles. Consider the situation when the values of Ci.n have been assigned. Encoding with generic function l((;;) would ignore the permutation symmetries present among the components of x,, i.e., the columns of the matrix Ti.i. as a result of the cı.m assignments. These symmetrics are the same three symmetries discussed above for clustering models. Moreover, a fixed function h(x,) would not be able to accommodate the fact that the length of x, changes with the size of the dataset.
Suppose that there are K clusters among the Ci:ns each with Są elements. In order simplify the notation, let us assume that the n unassigned points all belong to an additional (K + 1) cluster with 8K+1 N
we assume Cr+1:N K + 1, and we have k=1 8k = N and
8% Σド= 8(c, k).
Now, in each x,, the number su of elements in the k-th cluster can be split as
8k
81
+ s
N
Σ
5(cj = =))(ais = +1)
=
N
Σ
8(cz h))((aa =
and note that both $i and s are invariant under the symmetry of permuting the indices within cluster k.
C2 = 1, C = 44 2, then = and cs = 3, and from X1 we get sT = 2, TT, = 0, $i.2 = 1, 1,81, =
Example: = 5 and x₁ (+1,+1,-1,+1,,
If four assignments were made
1. s 1, = 0. If we permute the columns and 4, both from cluster k 2, we get but all the st. 's stay invariant.
x1
Additional invariants can be obtained combining 81 and $J across all rows xj's with cj = C₁. such as
1
m
m
(27)
Set
jic
(st
m x
(28)
Sec
j:cy=C
and similarly mayk and V Note that these invariants are the same for all rows x, with = k. The motivation to consider them is that, the partition corresponding to C is correct, then for i < and l < K we have n m since they are both estimators of the latent Bernoulli
15

---

Under review as a conference paper at ICLR 2020
parameter Pc,,k- For the same reason, if the partition is correct and those two estimators of &c exact, then ce,.k 0. Similarly for m and Vo,..- Then these values provide learning signals to the network that estimates the probability of the assignments cim. being correct.
Therefore we propose to encode the components of x; belonging to cluster as
Ti (si m ky ut Gu 81,00 m 2
R.
(29)
In order to preserve the symmetry of the first K labels under permutations, combine them
4 Σ (Tri))
E R
(30)
where the encoding function is t RR He The encoding 29 of the unassigned components Tin+1:N is kept separate and denoted as = 1+1
In summary, each row of the adjacency matrix is represented by the fixed-dimensional pair (t,, 9) E R in way that respects the symmetries of the assignments Cimn' permutations between members of cluster, permutations of cluster labels and permutations among unassigned columns.
B.2 CLUSTERING ROWS THE ADIACENCY MATRIX
We can proceed now as in regular clustering, encoding each cluster of x, 's within as
H
Σ h(1..)))
RRd
k k=1...K K,
(31)
kccy=k
and defining the permutation invariant, ixxeddimeensionll vectors
G (Ha).
(32)
U
Σ u(ti. q:).
(33)
In terms of these quantities, the conditional probabilities are defined as usual as
(G2.0)
go(Cn k: k1C11m-1, <KT ef(Gu.))
(34)
for = 1... K with h, h(Lns (x) and with G5 being the value of for the different configurations. Compared the regular clustering case, here need to learn the additional function t. We call our approach Neural Block Process (NBP).
16

---

Under review as a conference paper at ICLR 2020
C MONITORING GLOBAL PERMUTATION INVARIANCE
As mentioned in Section 5, we must verify the symmetry of the posterior likelihood under global permutations all the data points. We show such a check in Figure 11
2020
2002
4200
5020
5000
Iteration
NLL std/mean
5000
Iteration
Figure 11: Global permutation invariance. Training curves for the NCP model 2D Gaussians in Section 2. Each minibatch was evaluated for 8 random permutations of the order of the points in the dataset. Above: Mean of the NLL over the permutations. Below: NLL standard deviation/NLL mean. Note that the ratio is of order 10
17

---

Under review as a conference paper at ICLR 2020
D DETAILS 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, chO). This partitioned the recording data by channel such that each center-channel-base partition only contains multi-channe spike waveforms centered at that channel. Each spike waveform is represented x 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.
7
Neural architecture for NCP spike sorting. The overall architecture is the same the one described in Section 2 and Figure 2 To extract useful features from the spatial-temporal patterns of spike 1D ConvNet the h and encoder functions. 'he convolution is applied along the time axis, with each electrode channel treated as 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 11 2, 2, 2D.. The last block is followed by an averaged pooling layer and a final linear layer. The outputs of the ResNet encoder are the h2 and u, vectors of NCP, i.e.
waveforms,
h; ResNetEnccode((zz). We used da d, = 256. The other two functions, y and f, are identical to those 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 Nmmnn Nmaz]
(35)
CN ~ Cat(Ti,..... TK)
(38)
K 1 + Poisson (X)
(36)
pi ~ p(u)
1: = 1..... K 24 p( Me E, @ Ex) == (40)
(39)
T1 TK ~ Dirichlet((1,.... .x)
(37)
Here, N is the number of spikes between [200, 500]. The number of clusters K is sampled from a shifted Poisson distribution with x= the proportion of each cluster and is sampled from a Dirichlet distribution with 0ı: = 1. The training spike templates μ4 E TR. 7x32 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; is obtained by adding to (lc. 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 T2). 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).
that each channel has average 3 clusters. Ti.K represents
Probabilistic spike clustering using NCP. At inference time, we fed the 7 x 32 arrays of spike waveforms to NCP, and performed GPU-parallelizer posterior sampling of cluster labels (Figure 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 likeihioel clustering configurations for 2000 spikes in less
Cluster 1 (n 230)
Cluster (n = 205)
Cluster 3 (n =65)
Bomples
Gealayy Awnrge
cho
1r 1* kp
1p 1/- Ap=
ch1
www. P/W
ypm Nbw ame
Channels
ch2
A
ch3
Aup
mpn- ap-
11--
ch4
1/0 m//
WA
Soy
ch5
WM
1/000
ch6
w/4. Now
A//- 1/-
1/
1p-
Figure 12: Synthetic data examples. Example of 500 synthetic spikes from 3 clusters.
18

---

Under review as a conference paper at ICLR 2020
500 Spikes
1000 Spikes
2000 Spikes
0.8
AMI
0.2
vGMFM
NCP
vGMFM
NCP
VGMFM
NCP
Figure 13: Clustering synthetic data. The AMI scores for clustering 20 sets of 500, 1000, and 2000 unseen synthetic spikes.
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 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 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 cluster based on the I.2 distance to the cluster centers. Then, final set 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 input features. For Kilosort, we run the entire pipeline using the Kilosort2 package (Pachitariu) (201)). 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 step towards addressing the problem of spike sorting in the presence electrode drift (Calabrese & Paninski] 2011: Shan al. 2017), describe in Sup. Material E generalization of NCP to handle data in which the per-cluster parameters (e.g. the cluster means) are nonstationary in time.
19

---

Under review as a conference paper at ICLR 2020
KS; PTP 63.5 auplicates NCP; PTP 65.4 KS; PTP 43,9NCP, FTP 44.9 KS; PTP 42.5 duplicates:
NCP: PTP 44.5
PTP 21.9
KS; PTP 38,1NCP; PTP 39.0
duplicates.7 PTP 22.4
NCP: PTF 13.2
duplicates:
NCP: PTP 16.5
NCP PTP 15.7
NCP: PTP 15.0
NCP. PTP 14.4
12.4
NCP: PTP 11.3
NCP: PTP 10.9
NCP. PTP 10.7
NCP: PTP 10.6
KS: PTP 10.3
KS: PTF
KS: 10.3
KS: 10.0
NCP: 9.8.
9.3
2uplicales:2
NCF: 97 duplicates:
NCP:
NCP: P1P
suplicates
KS: 8.P NCP:
NCP:PTP PTP 8.6
KS: PTP 8.5 NCP: PTP 8.6
KS: PTP 8.4
NCP: PTP 6.4
KS: PTP 8.4
FT duplicates.
NCR: PTP
PTP NCP,
KS;
KS: 7.5
KS; 7.5
duplicates:
PTP PTP
NCP.
KS:
KS:
KS: PTP
K3: PTP..8 NCP. NCP PTP duplicates. duplicates;
KS. 5.5 NCP: PIP duplicatesi2 suplicates
NCP: PIP 6.0
NCP: PTP 6.2
KS; PTP 5.8
KS; PTP 4.8 duplicates:
RS,
NCR: PTP 4.0 duplicates:
Figure 14:
Spike sorting real data. Receptive fields of 55 randomly selected pairs of units recovered from Kilosort and NCP spike sorting. (Red boxes indicate units found by NCP: blue boxes by Kilosort.) Both approaches find the spikes with the biggest peak-to-peak (PTP) size. For smaller-PTP units often one sorting method finds a cell that the other sorter misses. NCP and KS find comparable number units with receptive fields here with NCP finding few more than KS: text for details.
20

---

Under review as a conference paper at ICLR 2020
Trajectories Observations Clusters
Clusters
Prob: 0.296
Preb:
Figure 15: Neural Particle Tracking. Left: Time trajectorics of 2D particles. Note that particles can appear or disappear at arbitrary times. Middle and right: Two posterior samples. Note that since only one particle is observed at each time, particle observed for some time leads possible ambiguity the number of particles. (Best seen in color.)
E PARTICLE TRACKING
Inspired by the problem of electrode drift (Calabrese & Paninski 2011 2017/, let us consider now a generative model given by
Pachitariu 2019
Shan et al.
olc.le......4. Ct-1)
t= 1=.,..,, T h 1=1...K K t=
(41)
FRRt
1=1,...,7
(42)
E2 p(zi)Maat)
(43)
In this model, cluster corresponds to the points along the time trajectory of a particle, and represents the time evolution of the cluster parameters. The cluster labels c indicate which partiele observed at time t, and note that particles can in principle appear or disappear at any time.
(42)
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 0, and generalize the NCP equations to
H
Σ
1=1... K,
(44)
t==loe=
K
Gt
Σ (Hn.))
(45)
k=1
U. Σ
Me-
(((()))
(46)
w=
The conditional assignment probability for is now
qe (c k((c)-11
(47)
EKH (GGG,.)))))
for k: = 1....KK The time-decay constant is learnt along with all the other parameters. We can
also consider replacing e bbjt-tt with a general distance function e-d([r-rl)) Figure 15 illustratess this model in simple 2D example. We call this approach Neural Particle Tracking.
F NEURAL ARCHITECTURES IN THE EXAMPLES
To train the networks in the examples, used stochastic gradient descent with Adam (Kingma & Ba 2015). with learning rate 10 10-4 The number of samples each mini-batch were: for p(N),
for p(CI:N). 64 for p(x(C1.N). The architecture of the functions in each case were:
21

---

Under review as a conference paper at ICLR 2020
CLUSTERS: 2D GAUSSIANS
MLP [2-256-256-256-222 128] with ReLUs
MLP [2-26--56-226--12]] with ReLUs
g: MLP (128-256-256-256-2556 with ReLUs
f: MLP [384-256-256-25--] with ReLUs
CLUSTERS: MNIST
layers of [convolutional maxpool + ReLUJ MLP [320-256-12) 128] with ReLUs
same
g: MLP [256-128-12--11--1122556 with ReLUs
f: MLP [384-256-256-256-1] with ReLUs
COMMUNITIES: IRL
t: MLP (6-64-64-644256) with ReLUs
h: MLP (262-64-64-64-256) with ReLUs
MLP [262-64-64-64-256] with ReLUs
gi MLP [256-64--4-64-2566 with ReLUs
MLP (512-64-66-66--6411 with ReLUS
22


Reviews:
[{"review_id": "SJxZrE8BcS", "reviewer": "AnonReviewer1", "summary": "", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "\nIn 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. \n\n\nThe paper is well-organized and easy to follow. However, there are two issues should be addressed:\n\n1, The novelty of the proposed algorithm might not enough. The two major components in this paper, i.e., VAE and DeepSet, are all carefully investigated before. This paper applies the DeepSet parameterization in the VAE framework.\n\n2, The details of the amortized inference training is not clearly explained. It is well-known that gradient through the discrete random variable is quite difficult. How the gradient for the parameter of the proposed model is calculated should be carefully discussed. The REINFORCE gradient in this model, whose support of c can be as large as the number of samples, can be quite huge. \n\n3, In the empirical evaluation, I was curious why the mean-field and MCMC have not been considered in the spike sorting problem.\n\nI am expecting the authors can address my concerns during rebuttal. \n\n=======================================================================\n\nThanks for the responses to clarify my concerns. \n\nThe learning procedure and experiments are clear now.  Indeed, as a purely variational inference paper,  the discrete variable problem is absent, since the model is always *fixed* and not updated. \n\nThe major contribution of the paper becomes the design of the posterior parametrization by DeepSets, which I think still not enough. \n\nI will keep my score.", "extra_scores": {"experience_assessment": "I have published in this field for several years.", "review_assessment:_checking_correctness_of_experiments": "I assessed the sensibility of the experiments.", "review_assessment:_thoroughness_in_paper_reading": "I read the paper at least twice and used my best judgement in assessing the paper.", "review_assessment:_checking_correctness_of_derivations_and_theory": "I assessed the sensibility of the derivations and theory."}, "extra_text": {}}, {"review_id": "Hyemu9Byqr", "reviewer": "AnonReviewer2", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "The paper presents a neural network based clustering process where the number of clusters is not known a-priori.  Proposed approach requires conjecturing a generative process (where number of clusters/classes is a random variable) and the model learns to uncover the posterior distribution over clusters given the observed samples.\nOverall I think it is a valuable contribution, well written paper with good results.  \nSpecific comments:\na) Even though the model allows for variable number of clusters, I feel there may be a strong dependence between the number of clusters the model can hypothesize and the number of clusters in the training data.  It will be useful to give further insights into this.  For instance, if an MNIST model is trained with only digits 0-5 training data, how well would it perform in detecting all 10 clusters at test time?  Understanding model’s biases based on training data is one area I feel is important and the paper could add to.\nb) The neural clustering process could potentially be viewed as a transductive inference model for classification of test data.  Typically at test time classification is done for each test sample independently, and the clustering process allows one to bring in other similar test samples to help with classification.  Have the authors considered this and have any comments on potential value / feasibility of this?\nc) In the examples presented in Section 2.3, please clarify how training & testing was done.  Specifically what training data was used (all of MNIST training data?), and the test time clustering was done on a subset of MNIST test data?\nd) Use of ‘q’ for a neural-network and ‘q_\\theta’ for posterior distribution is a little confusing, will be better to have different notation for these.\n", "extra_scores": {"experience_assessment": "I have published in this field for several years.", "review_assessment:_thoroughness_in_paper_reading": "I read the paper thoroughly.", "review_assessment:_checking_correctness_of_experiments": "I carefully checked the experiments.", "review_assessment:_checking_correctness_of_derivations_and_theory": "I assessed the sensibility of the derivations and theory."}, "extra_text": {}}, {"review_id": "H1g5dPCpFr", "reviewer": "AnonReviewer3", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "Summary:\nThis 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.\n\nStrengths:\nThe paper is generally well written and the relationship to previous works is well described. Empirical results seem quite convincing, for example, the clustering results presented in Fig. 2 and Fig. 3 clearly show not only the inferred number of clusters, but also the posterior probability which indicates that reasonable samples are assigned higher probability.\n\nWeaknesses:\n- Overall, the idea looks very original and promising, but I find some technical details are not easy to understand under the current form, especially for non-experts in this domain. I would recommend the authors to elaborate a bit more on the proposed architecture and the variable-input soft-max function in Sect. 2.1.\n- On page 8, the authors mention that the NCP is much more efficient compared with MCMC, for example, in the Gaussian 2D example. However, regarding the DPMM clustering model, it is known that MCMC methods are generally slower compared with variational inference, which is computationally faster. I think it would be interesting to add a discussion or comparison with variational inference in terms of computational efficiency.\n- If I understand correctly, the NCP is essentially based on a sequential sampling procedure. The authors claim that the proposed method is easily parallelized using a GPU, but there does not seem to be sufficient details on the GPU-parallelization of sequential sampling.\n\nMinor comments:\nThe size of some figures appears too small, for example Fig. 6 and Fig. 10, which may hinder readability.", "extra_scores": {"experience_assessment": "I have read many papers in this area.", "review_assessment:_thoroughness_in_paper_reading": "I read the paper at least twice and used my best judgement in assessing the paper.", "review_assessment:_checking_correctness_of_experiments": "I assessed the sensibility of the experiments.", "review_assessment:_checking_correctness_of_derivations_and_theory": "I assessed the sensibility of the derivations and theory."}, "extra_text": {}}]

Author rebuttal:

Meta-review: This paper uses neural amortized inference for clustering processes to automatically tune the number of clusters based on the observed data. The main contribution of the paper is the design of the posterior parametrization based on the DeepSet method.  The reviewers feel that the paper has limited novelty since it mainly follows from existing methodologies. Also, experiments are limited and not all comparisons are made. 
Decision: Reject