Title: DEEP AMORTIZED CLUSTERING

Abstract: We propose a deep amortized clustering (DAC), a neural architecture which learns to cluster datasets efficiently using a few forward passes. DAC implicitly learns what makes a cluster, how to group data points into clusters, and how to count the number of clusters in datasets. DAC is meta-learned using labelled datasets for training, a process distinct from traditional clustering algorithms which usually require hand-specified prior knowledge about cluster shapes/structures. We empirically show, on both synthetic and image data, that DAC can efficiently and accurately cluster new datasets coming from the same distribution used to generate training datasets.

:

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 = [x 1 , . . . , x n ] and Y = [y 1 , . . . , y m ] 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

EQUATION

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 SAB L (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 PMA k (X) = MAB(S, X) where S = [s 1 , . . . , s k ] are trainable parameters.

Note that the time-complexity of SAB is O(n 2 ) because of pairwise computation. To reduce this, Lee et al. (2019) proposed to use Induced Self-Attention Block (ISAB) defined as

EQUATION

where I = [i 1 , . . . , i m ] 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 ISAB L (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:

EQUATION

where π j is the mixing coefficient and θ j = (µ j , σ 2 j ) are the mean and variance for the jth Gaussian component. The network is trained to maximize the expected log likelihood over datasets:

EQUATION

where n X is the number of elements in X. 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] n X where n X 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. m i represents the probability of x i 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 . Assume X has k X true clusters, and let y ∈ [1, . . . , k X ] n X 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

EQUATION

where n X|j := n X 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 k X ! 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:

H X = ISAB L (X), decode cluster: H θ = PMA 1 (H X ), θ = rFF(H θ ), decode mask: H m = ISAB L (MAB(H X , H θ )), m = sigmoid(rFF(H m )).

(6) The network first encodes X into H X and extracts cluster parameters θ. Then θ together with encoded data H X 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, . . . , n X } 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 x a . The corresponding loss function is,

EQUATION

where j a denotes the the true cluster index containing a. The architecture to be trained with this loss can be implemented as encode data:

H X = ISAB L (X), H X|a = MAB(H X , h a ),

decode cluster:

H θ = PMA 1 (H X|a ), θ = rFF(H θ ),

decode mask:

EQUATION

where h a is the row vector of H X 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).

3.2. BEYOND SIMPLE PARAMETRIC FAMILIES:

When each cluster cannot be well described by a Gaussian or other simple parametric distributions, we have several choices to learn them. The first is to estimate the densities along with the filtering using neural density estimators such as Masked Autoregressive Flow (MAF) (Papamakarios et al., 2017) . Another option is to lower-bound log p(x; θ) by introducing variational distributions, for example using Variational Autoencoder (VAE) (Kingma & Welling, 2014) . See Section 5.2 and Section 6.1 for examples. If the density estimation is not necessary, we can choose not to learn log p(x; θ). In other words, instead of ( 5) and ( 7), we train

EQUATION

for MLF and AF, respectively. The corresponding architectures are ( 6) and ( 8) with parameter estimation branches removed. The DAC trained in this way implicitly learns how to define a cluster from the given training datasets and cluster labels. See Section 6.2 and Section 6.3 where we applied this to cluster image datasets.

Recall that each step of filtering yields one cluster. To solve DAC, we iterate this procedure until all clusters are found. After each filtering step, we remove from the dataset the points that were assigned to the cluster, and perform filtering again. This recursive procedure is repeated until all datapoints have been assigned to a cluster1 . We call the resulting amortized clustering algorithm as DAC2 . DAC learns both data generating distributions and cluster assignment distributions from meta-training datasets without explicit hand-engineering.

4. RELATED WORKS:

Deep clustering methods There is a growing interest in developing clustering methods using deep networks for complex data (Yang et al., 2016a; b; Xie et al., 2015; Li et al., 2017; Ji et al., 2018) . See Aljalbout et al. (2018) for a comprehensive survey on this line of work. The main focus of these methods is to learn a representation of input data amenable to clustering via deep neural networks. Learning representations and assigning data points to the clusters are usually trained alternatively. However, like the traditional clustering algorithms, these methods aim to cluster particular datasets. Since such methods typically learn a data representation using deep neural networks, the representation is prone to overfitting when applied to small datasets.

Learning to Cluster Learning to cluster refers to the task of learning a clustering algorithm from data. Such methods are trained in a set of source datasets and tested on unseen target datasets. Constrained Clustering Networks (Hsu et al., 2017; 2019) follow a two-step process for learning to cluster: they first learn a similarity metric that predicts whether a given pair of datapoints belong to the same class, and then optimize a neural network to predict assignments that agree with the similarity metric. Centroid Networks (Huang et al., 2019) learn an embedding which is clustered with the Sinkhorn K-means algorithm. While these methods combine deep networks with an iterative clustering algorithm, our framework is much more efficient as it directly identifies each cluster after one forward pass. Our experiments in Section 6 that our model is orders of magnitude faster than previous works in learning to cluster.

Amortized clustering methods:

To the best of our knowledge, the only works that consider a similar task to ours is Lee et al. (2019) and Pakman et al. (2019) . We refer the reader back to Section 2 for an outline of the amortized clustering framework presented in Lee et al. (2019) . Pakman et al. (2019) presented an amortized clustering method called Neural Clustering Process (NCP). Given a dataset, NCP sequentially computes the conditional probability of assigning the current data point to one of already constructed clusters or a new one, similar in spirit to the popular Gibbs sampling algorithm for Dirichlet process mixture models (Neal, 2000) , but without positing particular priors on partitions, but rather letting the network learn from data. However, the sequential sampling procedure makes the algorithm not parallizable using modern GPUs, limiting its scalability. Furthermore, since the clustering results vary a lot w.r.t. the sequential processing order, the algorithm needs a sufficient number of random samples to get stable clustering results. We compared our method to NCP on small-scale MOG experiments in Appendix C, and our results support our claim.

5. EXPERIMENTS ON SYNTHETIC DATASETS:



5.1. 2D MIXTURE OF GAUSSIANS:

We first demonstrate DAC with MLF on 2D MOG datasets with arbitrary number of clusters. We considered two baselines that can handle variable number of clusters: truncated Variational Bayesian Dirichlet Process Mixture Model (VBDPM) (Blei & Jordan, 2006) and the ST architecture (Lee et al., 2019) with ACT-style decoder so that it can produce arbitrary number of clusters. We describe the latter method, ACT-ST, in detail in Appendix A. See Appendix B for detailed experimental setup including data generation process and training scheme. We trained DAC and ACT-ST using random datasets with a random number of data points n ≤ n max and clusters k ≤ k max , where we set (n max , k max ) = (1000, 4) during training. We tested the resulting model on two scenarios. The first one is to test one same configurations; testing on 1,000 random datasets with (n max , k max ) = (1000, 4).

The second one is test on 1,000 random datasets with (n max , k max ) = (3000, 12) to see whether the amortized clustering methods can generalize to an unseen number of clusters. For both scenarios, we ran VBDPM on each test dataset until convergence from scratch. We used the mean of the variational distributions as the point-estimates of the parameters to compute log-likelihoods. Table 1 summarizes the results. DAC works well for both cases, even beating the oracle log-likelihood computed from the true parameters, while ACT-ST fails to work in the more challenging (n max , k max ) = (3000, 12) case (Fig. 2 ). VBDPM works well for both, but it takes considerable time to converge whereas DAC requires no such optimization to cluster test datasets.

5.2. 2D MIXTURE OF WARPED GAUSSIANS:

When the parametric form of cluster distribution is not known, we cannot directly compute the log p(x i ; θ) term in (5). In this case, we propose to estimate the density p(x; θ) via neural density estimators along with the DAC learning framework. We construct each cluster by first sampling points from a 2D unit Gaussian distribution, and then applying a random nonlinear transformation on each point. We use MAF (Papamakarios et al., 2017) to model p(x; θ) where θ is a context vector that summarizes the information about a cluster. We trained the resulting DAC with MLF using random datasets, and compared to spectral clustering (Shi & Malik, 2000) . As shown in Fig. 3 , DAC finds and estimates the densities of these nonlinear clusters. See Appendix D for more details and results.

6 EXPERIMENTS ON REAL DATASETS

6.1. CLUSTERING EMNIST WITH MIXTURE OF NEURAL STATISTICIANS:

We may approximate the likelihood p(x; θ) via a VAE (Kingma & Welling, 2014) when the data distribution is too high-dimensional or complex. Instead of directly maximizing the log-likelihood, we maximize a lower-bound on the likelihood, log p(x; θ) ≥ E q(z|x;θ) [log p(x, z; θ) -q(z|x; θ)],

where θ encodes the context of a cluster and z is a latent variable that describes x based on θ. Neural Statistician (NS) (Edwards & Storkey, 2016) proposed the idea of approximating p(x; θ) using a context θ produced by a set network; we thus call this model a mixture of NSs. We found that the DAC implemented in this way could cluster well, generalizes to an unseen number of clusters, and generate images conditioned on the cluster context θ. See Appendix E for detailed results. (Cai et al., 2009) 0.372 0.498 LSC (Chen & Cai, 2011) 0.376 0.500 CSP (Wang et al., 2014) 0.812 0.812 MPCK-means (Bilenko et al., 2004) 0.871 0.816 KCL (Hsu et al., 2017) 0.889 0.874 MCL (Hsu et al., 2019) 0.897 0.893 DAC AF (ours) n/a 0.829

6.2. CLUSTERING EMBEDDED IMAGENET:

We applied DAC to cluster the collection of miniImageNet (Vinyals et al., 2016) and tieredIma-geNet (Ren et al., 2018) . We gathered pretrained 640 dimensional features of the images released by Rusu et al. (2018) 3 . We used the training and validation features as training set and test features as test set. The resulting training set contains 620,000 samples from 495 classes, and the test set contains 218,000 samples from 176 classes, with no overlap between training and test classes.

We trained DAC without density estimations (9), using both MLF and AF. We sampled randomly clustered datasets from the training set having (n max , k max ) = (100, 4). We then generated 1,000 randomly clustered datasets from the test set with (n max , k max ) = (300, 12). We compared DAC to basic clustering algorithms (k-means, spectral clustering), Deep Embedding Clustering (DEC) (Xie et al., 2015) , and transfer learning methods (KCL (Hsu et al., 2017) and MCL (Hsu et al., 2019) ). k-means, spectral, and DEC were trained for each test dataset from scratch. For KCL and MCL, we first trained a similarity prediction network using the training set, and used it to cluster each test dataset. For DEC, KCL and MCL, we used fully-connected layers with 256 hidden units and 3 layers for both similarity prediction and clustering network. Note that the size of the test datasets are small (n max = 300 < 640), so one can easily predict that DEC, KCL, and MCL would overfit. The algorithms other than ours was given the true number of clusters. The results are summarized in Table 2 . It turns out that the pretrained features are good enough for the basic clustering algorithms to show decent performance. The deep learning based methods failed to learn useful representations due to the small dataset sizes. Ours showed the best clustering accuracies while also consuming the shortest computation time.

6.3. CLUSTERING OMNIGLOT IMAGES:

We apply our filtering architecture to the unsupervised cross-task transfer learning benchmark of Hsu et al. (2017; 2019) . This benchmark uses the Omniglot dataset (Lake et al., 2015) to measure how well a clustering method can generalize to unseen classes. The Omniglot dataset consists of handwritten characters from 50 different alphabets. Each alphabet consists of several characters, and each character has 20 images drawn by different people. Our problem setup consists of training a clustering model using images from the 30 background alphabets, and using each of the 20 alphabets as a seperate dataset to test on. We use the same VGG network backbone as in Hsu et al. (2019) and follow their experimental setup.

We show the normalized mutual information (NMI) of DAC along with other methods in Table 3 . While previous methods were also evaluated on the easier setting where the true number of clusters is given to the network, this setting is not applicable to DAC. We see that DAC is competitive with the state-of-the-art on this challenging task despite requiring orders of magnitude less computation time.

While the metrics for DAC were computed after at most 100 forward passes, previous methods all require some sort of iterative optimization. For example, KCL and MCL (Hsu et al., 2017; 2019) required more than 100 epochs of training for each alphabet to arrive at the cluster assignments in Table 3 . This difference in computation requirements is more clearly demonstrated in Table 4 : DAC requires an average of less than 5 seconds per alphabet whereas KCL and MCL required more than 100. Table 4 additionally shows that in addition to being extremely time-efficient, DAC was more accurate in estimating k. This demonstrates the efficacy of our overall structure of identifying one cluster at a time.

7. DISCUSSION AND FUTURE WORK:

We have proposed DAC, an approach to amortized clustering using set-input neural networks. DAC learns to cluster from data, without the need for specifying the number of clusters or the data generating distribution. It clusters datasets efficiently, using a few forward passes of the dataset through the network.

There are a number of interesting directions for future research. The clustering results produced by DAC is almost deterministic because we discretise the membership probabilities in the filtering process. It would be interesting to take into account uncertainties in cluster assignments. In the imagenet experiment, we found that we needed a sufficient number of training classes to make DAC generalize to unseen test classes. Training DAC to generalize to unseen image classes with smaller numbers of training classes seem to be a challenging problem. Finally, learning DAC along with state-of-the-art density estimation techniques for images in each cluster is also a promising research direction.

A DETAILED DESCRIPTION OF ACT-ST:

We first define an adaptive-computation-time version of PMA,

EQUATION

which enables an RNN-like iterative computation by sequentially extending the parameters for PMA.

The clustering network to output variable number of parameters is then defined as

EQUATION

where [:,1] and [:,2:] are numpy-like notation indexing the columns. s k is a "stop" variable where s k > 0.5 means the iteration stops at kth step and continues otherwise.

During training, we utilize the true number of clusters as supervision for training c k , yielding the overall loss function

EQUATION

)

where k true is the true number of clusters, k max ≥ k true is maximum number of steps to run.

B DETAILS OF MOG EXPERIMENTS:

We generated dataset by the following process.

n ∼ Unif(0.3n max , n max ), k -1 ∼ Binomial(k max -1, 0.5), π ∼ Dir(α k [1, . . . , 1]), (y i ) n i=1 i.i.d. ∼ Cat(π) (µ j ) k j=1 i.i.d. ∼ Normal([0, 0] , 9I), (σ j ) k j=1 i.i.d.

∼ log Normal(log(0.25)[1, 1] , 0.01I),

EQUATION

Both ACT-ST and filtering networks were trained with n max = 1, 000 and k max = 4. For each step of training, we sampled a batch of 100 datasets (sharing the same n ∼ Unif(0.3n max , n max ) to comprise a tensor of shape 100 × n × 2), and computed the stochastic gradient to update parameters. We trained the networks for 20,000 steps using ADAM optimizer (Kingma & Ba, 2015) with initial learning rate 5 × 10 -4 . The results in Table 1 are obtained by testing the trained models on randomly generated 1,000 datasets with the same generative process. 

D DETAILS OF MIXTURE OF MAFS AND WARPED GAUSSIAN EXPERIMENTS:

We model the cluster density as MAF.

EQUATION

where

EQUATION

)

We can efficiently implement this with MADE (Germain et al., 2015) . We generated the warped Gaussian datasets by the following generative process.

EQUATION

)

EQUATION

where MoG 1 (y) denotes the sampling from 1d Mixture of Gaussians with the same parameter distributions as (13).

The filtering network is constructed as (6) where θ is a 128 dimensional vector to be fed into MAF as a context vector for a cluster. We implemented log p(x; θ) as a 4 blocks of MAF with MADE (Germain et al., 2015) . The filtering network was trained using random datasets with n max = 1, 000 and k max = 4, and trained for 20,000 steps with ADAM optimizer. We set initial learning rate as 5 • 10 -4 . Table 6 compares the resulting DAC to spectral clustering (Shi & Malik, 2000) . Spectral clustering was ran for each dataset from scratch with true number of clusters given. To give a better idea how good is the estimated log-likelihood values, we trained MAF with same structure (4 blocks of MADE) but without mixture component and cluster context vectors for 100 random test datasets. We trained for 20,000 steps for each dataset using ADAM optimizer with learning rate 5 • 10 -4 . The log-likelihood values estimated with filtering on the same datasets outperforms the one obtained by exhaustively training MAF for each dataset, and got -2.570 for MAF and -2.408 for DAC. This shows that the amortized density estimation works really well. Let X = [x 1 , . . . , x n ] be an input set. We construct a filtering network as follows.

h enc,i = Encoder(x i ) for i = 1, . . . , n, H X = ISAB L ([h enc,1 , . . . , h enc,n ] ), H θ = PMA 1 (H X ), θ = rFF(H θ ), z i ∼ q(z i |h enc,i , θ) for i = 1, . . . , n h dec,i = Decoder(z i , θ) for i = 1, . . . , n, xi ∼ p(x i |h dec,i ) for i = 1, . . . , n H m = ISAB L (MAB(H X , H θ )), m = sigmoid(rFF(H m )).

(17

)

The log-likelihood for a particular cluster log p({x i |y i = j}; θ) is then approximate by the variational lower-bound, log p({x i |y i = j}; θ) ≥ i|yi=j q(z i |h enc,i , θ) log p(x i |h dec,i )p(z i ) q(z i |h enc,i , θ) dz i = i|yi=j E q(zi|henc,i,θ)) [log p(x i |h dec,i )] -KL[q(z i |h enc,i , θ) p(z i )] . (18)

Once the clustering is done, the likelihood of the whole dataset X can be lower-bounded as log p(X; θ 1 , . . . , θ k ) ≥ k j=1 π j log p({x i |y i = j}; θ j ).

The likelihood lower-bounding for each cluster corresponds to the NS except for that the context vector is constructed with ISAB and PMA, so we call this model a mixture of NSs.

E.2 EXPERIMENTS ON EMNIST DATA:

We trained the model described in (17) to EMNIST (Cohen et al., 2017) . We picked "balanced" split, with 47 class / 112,800 training images / 18,800 test images. At each training step, we generated 10 randomly clustered dataset with (n max , k max ) = (1000, 4) and trained the network by the loss function ( 5) with the log-likelihood part replaced by (18). For Encoder(x i ) and Decoder(z i , θ), we used three layers of multilayer perceptrons. For the variational distribution q(z i |h enc,i , θ), we used Gaussian distribution with parameters constructed by a fully-connected layer taking [h enc,i , θ] as inputs. Following Chen et al. (2017) , We used an autoregressive prior distribution constructed by MAF for p(z). For the likelihood distribution, p(x|h dec,i ), we used Bernoulli distribution. Each training image was stochastically binarized.

Table 7 shows that iterative filtering can decently cluster EMNIST. Fig. 6 shows that given a set of 100 images with 4 clusters, the filtering can correctly identify a cluster and learns the generative model to describe it. Fig. 7 shows the clustering results using iterative filtering for 100 images. 

F MORE RESULTS FOR OMNIGLOT EXPERIMENTS:

We present more detailed comparion of KCL, MCL and ours for each alphabet. 

In practice, we input the entire dataset along with m and assign zero attention weight to datapoints with mi = 1. This is equivalent to the described scheme, but has the added benefit of being easy to parallelize across multiple datasets.

Note that DAC with MLF is not stochastic once we discretize the membership probability.

https:github.com/deepmind/leo

https://github.com/aripakman/neural_clustering_process

Figures and Tables:

Figure 1: Our model identifies one cluster per iteration (top), allowing it to find any number of clusters (bottom).

Figure 2: (Left) Clustering by ACT-ST, (middle) one step of filtering, (right) clustering by DAC.

Figure 3: Clustering warped Gaussian data with DAC. The model was trained with k ∈ {1, . . . , 4} clusters (left) but generalizes to datasets with 12 clusters (right).

Figure 4: More clustering results. Top two rows shows ACT-ST applied to the datasets having k ∈ {1, . . . , 4}. The middle two rows show the same ACT-ST model applied to datasets having k > 4. The bottom two rows show iterative filtering applied to the same datasets.

Figure 5: More clustering results for warped Gaussian datasets.

Figure 6: A step of filtering for 100 EMNIST images with 4 clusters. (Left) images belong to the cluster identified by the filtering step. (Middle) Images generated by decoding random latent vectors z ∼ p(z) passed through the decoder, with cluster context vector extracted from the filtering step. (Right) Images do not belong to the

Figure 7: Clustering results of 100 EMNIST images with 4 ground-truth clusters by iterative filtering. Each block corresponds to a cluster.

Results on synthetic 2D MOG. We report log-likelihood (LL), clustering accuracies (adjusted Rand index (ARI)(Hubert & Arabie, 1985) and normalized mutual information (NMI), the mean absolute error between true k and estimated k, and processing time per dataset. The numbers below (n max , k max ) are oracle LL values computed by the true parameters. We report the average on 5 runs. .719 ± 0.002 0.971 ± 0.001 0.977 ± 0.001 0.079 ± 0.003 0.037 ± 0.001 ACT-ST -0.721 ± 0.008 0.974 ± 0.002 0.974 ± 0.001 0.044 ± 0.003 0.006 ± 0.000 DAC -0.692 ± 0.002 0.983 ± 0.001 0.978 ± 0.001 0.120 ± 0.009 0.008 ± 0.000 .561 ± 0.001 0.962 ± 0.000 0.970 ± 0.000 0.435 ± 0.010 0.400 ± 0.006 ACT-ST -5.278 ± 0.573 0.781 ± 0.008 0.855 ± 0.004 1.993 ± 0.082 0.024 ± 0.001 DAC -1.544 ± 0.006 0.971 ± 0.000 0.974 ± 0.000 0.279 ± 0.012 0.021 ± 0.001

Table content: <table><tr><td>(nmax, kmax) Algorithm</td><td>LL</td><td>ARI</td><td>NMI</td><td>k-MAE</td><td>Time [sec]</td></tr><tr><td>(1000,4) -0.693 VBMOG -0(3000,12) VBMOG -1</td><td/><td/><td/><td/><td/></tr><tr><td>-1.527</td><td/><td/><td/><td/><td/></tr></table>

Results for 1,000 test datasets sampled from Embedded ImageNet.

Table content: <table><tr><td/><td>ARI</td><td>NMI</td><td>k-MAE</td><td>Time (sec)</td></tr><tr><td>k-means</td><td colspan="2">0.370 ± 0.001 0.514 ± 0.001</td><td>-</td><td>0.188 ± 0.003</td></tr><tr><td colspan="3">Spectral (Shi &amp; Malik, 2000) 0.432 ± 0.000 0.568 ± 0.000</td><td>-</td><td>0.087 ± 0.002</td></tr><tr><td>DEC (Xie et al., 2015)</td><td>0.195</td><td>0.326</td><td>-</td><td>46.098</td></tr><tr><td>KCL (Hsu et al., 2017)</td><td>0.201</td><td>0.361</td><td>-</td><td>13.401</td></tr><tr><td>MCL (Hsu et al., 2019)</td><td>0.157</td><td>0.350</td><td>-</td><td>14.646</td></tr><tr><td>DACMLF</td><td colspan="4">0.400 ± 0.012 0.527 ± 0.013 2.103 ± 0.160 0.012 ± 0.001</td></tr><tr><td>DACAF</td><td colspan="4">0.451 ± 0.014 0.579 ± 0.013 1.805 ± 0.031 0.017 ± 0.001</td></tr></table>

Unsupervised cross-task transfer learning on Omniglot. Normalized mutual information (higher is better) is averaged across 20 alphabets (datasets), each of which have between 20 and 47 letters (classes). All values beside DAC were reported inHsu et al. (2019). "k given" means that the true number of clusters was given to the model.

Table content: <table><tr><td>Method</td><td colspan="2">NMI(k given) NMI</td></tr><tr><td>K-means (MacQueen et al., 1967)</td><td>0.353</td><td>0.464</td></tr><tr><td>LPNMF</td><td/><td/></tr></table>

Mean absolute error of cluster number (k) estimate and processing time per dataset on the Omniglot benchmark.

Table content: <table><tr><td/><td>k-MAE</td><td/><td/><td>Time [sec]</td><td/></tr><tr><td>KCL</td><td>MCL</td><td>DAC AF</td><td>KCL</td><td>MCL</td><td>DAC AF</td></tr><tr><td colspan="6">6.4 ± 6.4 5.1 ± 4.6 4.6 ± 2.7 129.3 ± 18.9 124.5 ± 14.4 4.3 ± 0.6</td></tr></table>

Comparison of iterative filtering and NCP on 100 random datasets with (n max , k max ) =(3000, 12). The oracle log-likelihood is -1.5309, and DAC recorded -1.5640. S is the number of samples per dataset used for NCP. Due to the sequential nature, the training procedure of NCP does not scale to the other experiments we conducted. Instead, we trained NCP for MOG data described in Appendix B with smaller scale having n max = 100 and k max = 4. We used the code released by the authors 4 with default hyperparameters. We measured the clustering performance for 100 random datasets generated with (n max , k max ) = (3000, 12) (Table5). Both method generalized well w.r.t. the number of data points n, but filtering did much better in generalizing for the number of clusters k. The performance of NCP for clustering depends heavily on the processing order, so we conducted multiple runs with different random orders and picked the best one w.r.t. the clustering probability computed from NCP. Ours outperformed NCP with S = 50 samples per dataset with at least two orders of magnitude faster processing time.

Table content: <table><tr><td/><td>ARI</td><td colspan="2">k-MAE Time [sec]</td></tr><tr><td>DAC</td><td colspan="2">0.9616 0.2800</td><td>0.0208</td></tr><tr><td colspan="3">NCP (S = 1) 0.7947 1.6333</td><td>4.1435</td></tr><tr><td colspan="3">NCP (S = 10) 0.8955 0.8000</td><td>5.8920</td></tr><tr><td colspan="3">NCP (S = 50) 0.9098 0.6444</td><td>6.7936</td></tr><tr><td colspan="4">C COMPARISON TO NEURAL CLUSTERING PROCESS</td></tr><tr><td>We compare DAC to NCP.</td><td/><td/></tr></table>

Results on warped Gaussian datasets.

Table content: <table><tr><td colspan="2">(nmax, kmax) Algorithm</td><td>LL</td><td>ARI</td><td>NMI</td><td>k-MAE</td><td>Time [sec]</td></tr><tr><td>(1000,4)</td><td>Spectral DAC</td><td colspan="5">--1.275 ± 0.015 0.974 ± 0.001 0.970 ± 0.001 0.320 ± 0.035 0.011 ± 0.000 0.845 ± 0.003 0.889 ± 0.002 -0.103 ± 0.000</td></tr><tr><td>(3000,12)</td><td>Spectral DAC</td><td colspan="5">--2.436 ± 0.029 0.923 ± 0.002 0.936 ± 0.001 1.345 ± 0.099 0.037 ± 0.001 0.592 ± 0.001 0.766 ± 0.001 -0.572 ± 0.003</td></tr></table>

The results on EMNIST.

Table content: <table><tr><td>(nmax, kmax)</td><td>LL/pixel</td><td>ARI</td><td>k-MAE</td></tr><tr><td colspan="4">(1000,4) -0.193 ± 0.002 0.887 ± 0.005 0.607 ± 0.093</td></tr><tr><td colspan="4">(3000,12) -0.199 ± 0.003 0.728 ± 0.019 1.668 ± 0.147</td></tr></table>

Peer review (iclr2020):
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. Strengths: Overall, I think the paper is well written and the relationship to previous works is well described. The empirical results seem promising, especially in terms of computational efficiency. The authors conduct some experiments on relatively large datasets, such as miniImageNet and tiereImageNet, which is indeed crucial for the practical applications of the proposed model. Weaknesses: - I think this is a good paper, but my major concern is the limited theoretical contribution, given the fact that this work is mainly based on set-input neural networks introduced in Ref. [1]. I would like the authors to clarify a bit more the novelty of the paper. - The authors claim that DAC can process data points in parallel while Ref. [2] uses a sequential sampling procedure. However, there does not seem to be sufficient details on how to parallelize the proposed algorithm. - As shown in Figs. 1 and 4, it seems that some clusters are split into two or three fragments. I think this simply means the failure of the proposed method on synthetic data. - As also mentioned in the discussion on page 8, it would be important to consider uncertainties in cluster assignments, as already done in Ref. [2]. I would recommend the authors to provide some insight on how to take the cluster assignment uncertainty into account within current model. At the moment, I recommend a weak reject as the technical contribution of the paper seems rather limited, but I could be open to increasing my score if my concerns are addressed. References: [1] J. Lee, Y. Lee, J. Kim, A. R. Kosiorek, S. Choi, and Y. W. Teh. Set transformer: a framework for attention-based permutation-invariant neural networks. In Proceedings of International Conference on Machine Learning, 2019. [2] A. Pakman, Y. Wang, C. Mitelut, J. Lee, and L. Paninski. Discrete neural processes. ArXiv:1901.00409, 2019.