Neural Clustering Processes

1 INTRODUCTION . 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 can not 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 . 2 AMORTIZING MIXTURE MODELS . 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 . 2.1 THE VARIABLE-INPUT SOFTMAX . 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

Summary: 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.