Deep amortized clustering

1 INTRODUCTION Clustering is a fundamental task in unsupervised machine learning to group similar data points into multiple clusters . Aside from its usefulness in many downstream tasks , clustering is an important tool for visualising and understanding the underlying structures of datasets , as well as a model for categorisation in cognitive science . Most clustering algorithms have two basic components - how to define a cluster and how to assign data points to those clusters . The former is usually defined using metrics to measure distances between data points , or using generative models describing the shapes of clusters . The latter , how to assign data points to the clusters , is then typically optimized iteratively w.r.t . objective functions derived based on the cluster definitions . Note that cluster definitions are user defined , and are reflections of the user ’ s prior knowledge about the clustering process , with different definitions leading to different clusterings . However , cluster definitions used in practice are often quite simple , for example clusters in k-means are defined in terms of ` 2 distance to centroids , while Gaussians are a commonly used generative model for clusters in mixture models . Recently , advances in deep learning has facilitated the approximation of complex functions in a black-box fashion . One particular application of relevance to the problem of clustering in this paper is that of amortized inference ( Gershman & Goodman , 2014 ; Stuhlmüller et al. , 2013 ) , where neural networks are trained to predict the states of latent variables given observations in a generative model or probabilistic programme . In the context of learning set-input neural networks ( Zaheer et al. , 2017 ) , Lee et al . ( 2019 ) showed that it is possible to amortize the iterative clustering process for a Mixture of Gaussians ( MOG ) , while Pakman et al . ( 2019 ) demonstrated that it is possible to train a neural network to sequentially assign data points to clusters . Both approaches can be interpreted as using neural networks for amortized inference of cluster assignments and parameters given a dataset . Note that once neural networks are used for amortized clustering , we can take advantage of their flexibility in working with more complex ways to define clusters . Further , the amortization networks can be trained using generated datasets where the ground truth clusterings are known . This can be interpreted as implicitly learning the definition of clusters underlying the training datasets , such that amortized inference ( approximately ) produces the appropriate clusterings . In a sense this shares a similar philosophy as Neural Processes ( Garnelo et al. , 2018b ; a ) , which meta-learns from multiple datasets to learn a prior over functions . In this paper , we build on these prior works and propose Deep Amortized Clustering ( DAC ) . As in prior works , the amortization networks in DAC are trained using generated datasets where the ground truth clusterings are known . Like Lee et al . ( 2019 ) , DAC uses a Set Transformer , but differs from Lee et al . ( 2019 ) in that it generates clusters sequentially , which enables to produce a varying number of clusters depending on the complexity of the dataset ( Fig . 1 ) . Our approach also extends Lee et al . ( 2019 ) from MOG to problems with more complex cluster definitions , which are arguably harder to hand specify and easier to meta-learn from data . Our work also differs from Pakman et al . ( 2019 ) in that our network processes data points in parallel while Pakman et al . ( 2019 ) processes them sequentially , which is arguably less scalable and limits applicability to smaller datasets . This paper is organized as follows . We begin by describing in Section 2 the permutation-invariant set transformer modules that we use throughout the paper . In Section 3 , we describe how we implement our core idea of identifying one cluster at a time , and describe our framework for clustering , the DAC There are several challenges in solving DAC on complex datasets , and we structured our paper roughly in order of difficulty . We apply DAC to clustering synthetic data ( Section 5 ) and image data ( Section 6 ) ; some settings required additional components , which we describe when needed . 2 A PRIMER ON SET TRANSFORMER AND AMORTIZED CLUSTERING . In this section , we briefly review the set-input neural network architectures to be used in the paper , and describe how Lee et al . ( 2019 ) used them to solve amortized clustering for MOG . 2.1 SET TRANSFORMER . The Set Transformer ( ST ) is a permutation-invariant set-input neural network that uses self-attetntion operations as building blocks . It utilizes multi-head attention ( Vaswani et al. , 2017 ) for both encoding elements of a set and decoding encoded features into outputs . The fundamental building block of a ST is the Multihead Attention Block ( MAB ) , which takes two sets X = [ x1 , . . . , xn ] > and Y = [ y1 , . . . , ym ] > and outputs a set of the same size as X . Throughout this article , we represent sets as matrices where each row corresponds to an element . An MAB is defined as MAB ( X , Y ) = H + rFF ( H ) where H = X + rFF ( MultiheadAtt ( X , Y ) ) , ( 1 ) where rFF ( · ) is a feed-forward layer applied row-wise ( i.e. , for each element ) . MAB ( X , Y ) computes the pairwise interactions between the elements in X and Y with sparse weights obtained from attention . A Self-Attention Block ( SAB ) is simply MAB applied to the set itself : SAB ( X ) , MAB ( X , X ) . We can model high-order interactions among the items in a set by stacking multiple SABs ; we denote such a stack of L SABs applied to set X as SABL ( X ) . To summarize a set into a fixed-length representation , ST uses an operation called Pooling by Multihead Attention ( PMA ) . A PMA is defined as PMAk ( X ) = MAB ( S , X ) where S = [ s1 , . . . , sk ] > are trainable parameters . Note that the time-complexity of SAB is O ( n2 ) because of pairwise computation . To reduce this , Lee et al . ( 2019 ) proposed to use Induced Self-Attention Block ( ISAB ) defined as ISAB ( X ) = MAB ( X , MAB ( I , X ) ) , ( 2 ) where I = [ i1 , . . . , im ] > are trainable inducing points . ISAB indirectly compares the elements of X through the inducing points , reducing the time-complexity to O ( nm ) . Similarly to the SAB , we write ISABL ( X ) to denote a stack of L ISABs . 2.2 AMORTIZED CLUSTERING WITH SET TRANSFORMER . Lee et al . ( 2019 ) presented an example using ST for amortized inference for a MOG . A dataset X is clustered by maximizing the likelihood of a k component MOG , and a ST is used to output the parameters as : HX = ISABL ( X ) , Hθ = PMAk ( HX ) , ( logitπj , θj ) k j=1 = rFF ( SABL′ ( Hθ ) ) , ( 3 ) where πj is the mixing coefficient and θj = ( µj , σ2j ) are the mean and variance for the jth Gaussian component . The network is trained to maximize the expected log likelihood over datasets : Ep ( X ) [ nX∑ i=1 log k∑ j=1 πj logN ( xi ; µj , σ2j ) ] , ( 4 ) where nX is the number of elements inX . Clustering is then achieved by picking the highest posterior probability component for each data point under the MOG with parameters output by the ST . 3 DEEP AMORTIZED CLUSTERING . An apparent limitation of the model described in Section 2.2 is that it assumes a fixed number of clusters generated from Gaussian distributions . In this section , we describe our method to solve DAC in the more realistic scenario of having a variable number of clusters and arbitrarily complex cluster shapes . 3.1 FILTERING : INFERRING ONE CLUSTER AT A TIME . The objective ( 4 ) is not applicable when the number of clusters is not fixed nor bounded . A remedy to this is to build a set-input neural network f that identifies the clusters iteratively and make it to learn “ when to stop ” , similar to Adaptive Computation Time ( ACT ) for RNNs ( Graves , 2016 ) . One may think of several ways to implement this idea ( we present an illustrative example that simply augments ACT to ST in Section 5.1 ) . Here we propose to train f to solve a simpler task - instead of clustering the entire dataset , focus on finding one cluster at a time . The task , what we call as filtering , is defined as a forward pass through f that takes a set X and outputs a parameter θ to describe a cluster along with a membership probability vector m ∈ [ 0 , 1 ] nX where nX is the number of elements in X . The meaning of the parameter θ depends on the specific problem . For example , θ for MOG is ( µ , σ2 ) , the parameters of a Gaussian distribution . mi represents the probability of xi belonging to the cluster described by θ . To filter out the datapoints that belong to the current cluster , we use 0.5 as the threshold to discretize m to a boolean mask vector . The resulting smaller dataset is then fed back into the neural network to produce the next cluster and so on . Minimum Loss Filtering Now we describe how to train the filtering network f . AssumeX has kX true clusters , and let y ∈ [ 1 , . . . , kX ] nX be a cluster label vector corresponding to the true clustering of X . Then we define the loss function for one filtering iteration producing one θ and one m as L ( X , y , m , θ ) = min j∈ { 1 , ... , kX } ( 1 nX nX∑ i=1 BCE ( mi,1 { yi=j } ) − 1 nX|j ∑ i|yi=j log p ( xi ; θ ) ) , ( 5 ) where nX|j : = ∑nX i=1 1 { yi=j } , BCE ( · , · ) is the binary cross-entropy loss , and p ( x ; θ ) is the density of x under cluster parameterised by θ . This loss encourages θ to describe the data distribution of a cluster , and m to specify which datapoints belong to this particular cluster . The rationale to take minimum across the clusters is follows . One way to train f to pick a cluster at each iteration is to impose an ordering on the clusters ( e.g . in order of appearance in some arbitrary indexing of X , or in order of distance to origin ) , and to train f to follow this order . However , this may introduce unnecessary inductive biases that deteriorates learning . Instead , we let f find the easiest one to identify , thus promoting f to learn its own search strategy . Note that there are kX ! equally valid ways to label the clusters in X . This combinatorial explosion makes learning with standard supervised learning objectives for y tricky , but our loss ( 5 ) is inherently free from this problem while being invariant to the labelling of clusters . We use the following architecture for the filtering network f : Section 2 : encode data : HX = ISABL ( X ) , decode cluster : Hθ = PMA1 ( HX ) , θ = rFF ( Hθ ) , decode mask : Hm = ISABL′ ( MAB ( HX , Hθ ) ) , m = sigmoid ( rFF ( Hm ) ) . ( 6 ) The network first encodesX intoHX and extracts cluster parameters θ . Then θ together with encoded data HX are further processed to produce the membership probabilities m. We call the filtering network with architecture ( 6 ) and trained with objective ( 5 ) Minimum Loss Filtering ( MLF ) . Anchored Filtering An alternative strategy that we found beneficial for harder datasets is to use anchor points . Given a dataset X and labels y constructed from the true clustering , we sample an anchor point with index a ∈ { 1 , . . . , nX } uniformly from X . We parameterize a set-input network f to take both X and a is input , and to output the cluster that contains the anchor point xa . The corresponding loss function is , L ( x , y , a , m , θ ) = 1 nX nX∑ i=1 BCE ( mi,1 { yi=ja } ) − 1 nX|ja ∑ i|yi=ja log p ( xi ; θ ) , ( 7 ) where ja denotes the the true cluster index containing a . The architecture to be trained with this loss can be implemented as encode data : HX = ISABL ( X ) , HX|a = MAB ( HX , ha ) , decode cluster : Hθ = PMA1 ( HX|a ) , θ = rFF ( Hθ ) , decode mask : Hm = ISABL′ ( MAB ( HX|a , Hθ ) ) , m = sigmoid ( rFF ( Hm ) ) , ( 8 ) where ha is the row vector of HX corresponding to the index a . We train ( 8 ) by randomly sampling a for each step , and thus promoting f to find clusters by comparing each data point to the random anchor point . Note that the loss is also free from the label order ambiguity given anchor points . We call this filtering strategy Anchored Filtering ( AF ) .

Summary: The paper presents an amortized clustering method, called DAC, which is a neural architecture that allows efficient data clustering using a few forward passes. The proposed method is essentially based on the idea behind set-input neural networks [1], which consists of modeling the interaction between instances within a given dataset. Compared with the previous work [1], the main difference is that DAC does not need to specify the number of clusters, as in the case of Bayesian nonparametrics, making it more flexible for clustering complex datasets. It is empirically shown that DAC can efficiently and accurately cluster new datasets coming from the same distribution for both synthetic and image data.