Title: MIXTURE-OF-EXPERTS VARIATIONAL AUTOEN-CODER FOR CLUSTERING AND GENERATING FROM SIMILARITY-BASED REPRESENTATIONS

Abstract: Clustering high-dimensional data, such as images or biological measurements, is a long-standing problem and has been studied extensively. Recently, Deep Clustering gained popularity due to the non-linearity of neural networks, which allows for flexibility in fitting the specific peculiarities of complex data. Here we introduce the Mixture-of-Experts Similarity Variational Autoencoder (MoE-Sim-VAE), a novel generative clustering model. The model can learn multi-modal distributions of high-dimensional data and use these to generate realistic data with high efficacy and efficiency. MoE-Sim-VAE is based on a Variational Autoencoder (VAE), where the decoder consists of a Mixture-of-Experts (MoE) architecture. This specific architecture allows for various modes of the data to be automatically learned by means of the experts. Additionally, we encourage the lower dimensional latent representation of our model to follow a Gaussian mixture distribution and to accurately represent the similarities between the data points. We assess the performance of our model on synthetic data, the MNIST benchmark data set, and a challenging real-world task of defining cell subpopulations from mass cytometry (CyTOF) measurements on hundreds of different datasets. MoE-Sim-VAE exhibits superior clustering performance on all these tasks in comparison to the baselines and we show that the MoE architecture in the decoder reduces the computational cost of sampling specific data modes with high fidelity.

1. INTRODUCTION:

Clustering has been studied extensively (Aljalbout et al., 2018; Min et al., 2018) in machine learning. Recently, many Deep Clustering approaches were proposed, which modified (Variational) Autoencoder ((V)AE) architectures (Min et al., 2018; Zhang et al., 2017) or with varying regularization of the latent representation (Dizaji et al., 2017; Jiang et al., 2017; Yang et al., 2017; Fortuin et al., 2019) .

Reconstruction error usually drives the definition of the latent representation learned from an AE or VAE. The representation for AE models is unconstrained and typically places data objects close to each other according to an implicit similarity measure that also yields favorable reconstruction error. In contrast, VAE models regularize the latent representation such that the represented inputs follow a certain variational distribution. This construction enables sampling from the latent representation and data generation via the decoder of a VAE. Typically, the variational distribution is assumed standard Gaussian, but for example, Jiang et al. (2017) introduced a mixture of Gaussian variational distribution for clustering purposes.

A key component of clustering approaches is the choice of similarity metric for the considered data objects which we try to group (Irani et al., 2016) . Such similarity metrics are either defined a priori or learned from the data to specifically solve classification tasks via a Siamese network architecture (Chopra et al., 2005) . Dimensionality reduction approaches, such as UMAP (McInnes et al., 2018) or t-SNE (van der Maaten & Hinton, 2008) , allow to specify a similarity metric for projection and thereby define the data separation in the inferred latent representation.

In this work, we introduce the Mixture-of-Experts Similarity Variational Autoencoder (MoE-Sim-VAE), a new deep architecture that performs similarity-based representation learning, clustering of the data and generation of data from each specific data mode. Due to a combined loss function, Figure 1 : Overview of the proposed model MoE-Sim-VAE. Data (in panel A) gets encoded via a encoder network (B) into a latent representation (C) which is trained to be a mixture of standard Gaussians. Via a clustering network (G), which is trained to reconstruct a user-defined similarity matrix (F), the encoded samples get assigned to the data mode-specific decoder subnetworks (which we call experts) in the MoE Decoder (D). The experts reconstruct the original input data and can be used for data generation when sampling from the variational distribution (E). it can be jointly optimized. We assess the scope of the model on synthetic data and we present superior clustering performance on MNIST. Moreover, in an ablation study, we show the efficiency and precision of MoE-Sim-VAE for data generation purposes in comparison to the most related state-of-the-art method (Jiang et al., 2017) . Finally, we show an application of MoE-Sim-VAE on a real-world clustering problem in biology on multiple datasets. 

2. MIXTURE-OF-EXPERTS SIMILARITY VARIATIONAL AUTOENCODER:

Here we introduce the Mixture-of-Experts Similarity Variational Autoencoder (MoE-Sim-VAE, Figure 1 ). The model is based on the Variational Autoencoder (Kingma & Welling, 2014) . While the encoder network is shared across all data points, the decoder of the MoE-Sim-VAE consists of a number of K different subnetworks, forming a Mixture-of-Experts architecture (Shazeer et al., 2017) . Each subnetworks constitutes a generator for a specific data mode and is learned from the data.

The variational distribution over the latent representation is defined to be a mixture of multivariate Gaussians, first introduced by Jiang et al. (2017) . In our model, we aim to learn the mixture components in the latent representation to be standard Gaussians

EQUATION

where ω k are mixture coefficients, µ k are the means for each mixture component, I is the identity matrix and K is the number of mixture components. The dimension of the latent representation z needs to be defined to suit the demands of Gaussian mixtures which have limitations in higher dimensions (Bishop, 1995) . Similar to optimizing an Evidence Lower Bound (ELBO), we penalize the latent representation via the reconstruction loss of the data L reconst and by using the Kullback-Leibler (KL) divergence for multivariate Gaussians (Jiang et al., 2017) on the latent representation

L KL = D KL (N 0 , N 1 ) = 1 2 {tr(Σ -1 1 Σ 0 ) + (µ 1 -µ 0 ) T Σ -1 1 (µ 1 -µ 0 ) -k + ln |Σ 1 | |Σ 0 | } (2)

where k is a constant, N 0 ∼ N (µ 0 , Σ 0 = I), and I is the identity matrix. Further,

N 1 ∼ N (µ 1 , Σ 1 = diag(σ j ))

, where σ j for j = 1, . . . , D, for a number of dimensions D, is estimated from the samples of the latent representation. Finally, we assume µ 0 = µ 1 resulting in the following simplified objective

EQUATION

to penalize exclusively the covariance of each cluster. It remains to define the reconstruction loss L reconst , where we choose a Binary Cross-Entropy

EQUATION

between the original data x (scaled between 0 and 1) and the reconstructed data x reconst , where i iterates the batch size N and d the dimensions of the data D. Finally the loss for the VAE part is defined by

EQUATION

with a weighting coefficient π 1 which can be optimized as a hyperparameter.

SIMILARITY CLUSTERING AND GATING OF LATENT REPRESENTATION:

Training of a data mode-specific generator expert requires samples from the same data mode. This necessitates to solve a clustering problem, that is, mapping the data via the latent representation into K clusters, each corresponding to one of the K generator experts. We solve this clustering problem via a clustering network, also referred to as gating network for MoE models. It takes as input the latent representation z i of sample i and outputs probabilities p ik ∈ [0, 1] for clustering sample i into cluster k. According to this cluster assignment, sample i is then gated to expert k = argmax k p ik for each sample i. We further define the cluster centers µ k for k ∈ {1, . . . , K} similar as in the Expectation Maximization (EM) algorithm for Gaussian Mixture models (Bishop, 2006) as

EQUATION

where N k is the absolute number of data points assigned to cluster k based on highest probability p ik for each sample i = 1, . . . , N . The Gaussian mixture distributed latent representation (via KL loss in Equation 3) is motivation for the empirical computation of the cluster means and further, similar as in the EM algorithm, it allows iterative optimization of the means of the Gaussians. We train the clustering network to reconstruct a data-driven similarity matrix S, using the Binary Cross-Entropy

EQUATION

to minimize the error in P P T ≈ S, with P := {p ik } i∈{1,...,N },k∈{1,...,K} where N is the number of samples (e.g., batch size). Intuitively, P P T approximates the similarity matrix S since values in P P T are only close to 1 when similar data objects are assigned to the same cluster, similar to the entries in the adjacency similarity matrix S. This similarity matrix is derived in an unsupervised way in our experiments (e.g. UMAP projection of the data and k-nearest-neighbors or distance thresholding to define the adjacency matrix for the batch), but can also be used to include weaklysupervised information (e.g., knowledge about diseased vs. non-diseased patients). If labels are available, the model could even be used to derive a latent representation with supervision. The similarity feature in MoE-Sim-VAE thus allows to include prior knowledge about the best similarity measure on the data.

Moreover, we apply the DEPICT loss from Dizaji et al. (2017) , to improve the robustness of the clustering. For the DEPICT loss, we additionally propagate a noisy probability pik through the clustering network using dropout after each layer. The goal is to predict the same cluster for both, the noisy pik and the clean probability p ik (without applying dropout). Dizaji et al. (2017) derived as objective function a standard cross-entropy loss

EQUATION

whereby q ik is computed via the auxiliary function

q ik = p ik /( i p i k ) 1 2 k p ik /( i p i k ) 1 2 (9)

where we refer to Dizaji et al. ( 2017) for exact derivation. The DEPICT loss encourages the model to learn invariant features from the latent representation for clustering with respect to noise (Dizaji et al., 2017) . Looking at it from a different perspective, the loss helps to define a latent representation which has those invariant features to be able to reconstruct the similarity and therefore the clustering correctly. The complete clustering loss function L Clustering is then defined by

EQUATION

with a mixture coefficient π 2 which can be optimized as a hyperparameter.

MOE-SIM-VAE LOSS FUNCTION:

Finally, the MoE-Sim-VAE model loss is defined by 11) which consists of the two main loss functions L V AE , acting as a regularization for the latent representation, and L Clustering , which helps to learn the mixture components based on an a priori defined data similarity. The model objective function L M oE-Sim-V AE can then be optimized end-to-end to train all parts of the model.

EQUATION

3. RELATED WORK:

(V)AEs have been extensively used for clustering (Xie et al., 2016; Dizaji et al., 2017; Li et al., 2017; Yang et al., 2017; Saito & Tan, 2017; Chen et al., 2017; Aljalbout et al., 2018; Fortuin et al., 2019) . The most related approaches to MoE-Sim-VAE are Jiang et al. (2017) and Zhang et al. (2017) . Jiang et al. (2017) introduced the VaDE model, comprising a mixture of Gaussians as underlying distribution in the latent representation of a Variational Autoencoder. Optimizing the Evidence Lower Bound (ELBO) of the log-likelihood of the data can be rewritten to optimize the reconstruction loss of the data and KL divergence between the variational posterior and the mixture of Gaussians prior. Jiang et al. (2017) motivate the use of to two separate networks for reconstruction and the generation process of the model. Further, to effectively generate images from a specific data mode and to increase image quality, the sampled points have to surpass a certain posterior threshold and are otherwise rejected. This leads to an increased computational effort. The MoE Decoder of our model, which is used for both reconstruction and generation, does not need such a threshold, as we discuss in more detail in Section 4.2.1. Assignment Network predicts probabilities which are used in the Mixture Aggregation to form the output of the generator network. Each AE model learns the manifold of a specific cluster, similarly to our MoE Decoder. However, MIXAE does not optimize a variational distribution, such that generation of data from a distribution over the latent representation is not possible, in contrast to the MoE-Sim-VAE (Figure 2 ).

4. EXPERIMENTS:

We evaluate the MoE-Sim-VAE using synthetic data and the MNIST data set of handwritten digits (LeCun et al., 1998) for clustering and data generation. Furthermore, we performed an ablation study to demonstrate the importance of the MoE Decoder. Finally, we present experiments on a real-world application of defining cellular subpopulations from mass cytometry measurements (Bandura et al., 2009) of multiple publicly available datasets (Weber & Robinson, 2016; Bodenmiller et al., 2012) . Model implementation details are reported in the appendix in section A.1

We found that our model achieves superior clustering performance compared to other models on synthetic, MNIST and real-world datasets. Moreover, we show that MoE-Sim-VAE can more effectively and efficiently generate data from specific modes in comparison to other methods.

4.1. EVALUATION OF MOE-SIM-VAE ON SYNTHETIC DATA:

We evaluated our model using data sampled from a 100-dimensional multivariate Gaussian with equal mixture weights for each component. We tested two aspects of our model: Firstly, we evaluated up to how many clusters our model can fit well. Therefore, we sampled data from distributions with up to a hundred mixture components. For this experiment, we assume knowledge of the true number of clusters in the data for both methods, MoE-Sim-VAE and GMMs. Secondly, we tested if our model is able to identify the true number of clusters in the data. The similarity matrix S was defined as an adjacency matrix over the data items. Adjacency indicators were based on projecting the data via dimensionality reduction with UMAP (McInnes et al., 2018) and selecting neighbors according to a distance threshold. Details on model parameters can be found in Section A.1.1.

MoE-Sim-VAE performs better or comparable to the baseline for the number of clusters of up to 40 (Figure A1a ). The model predicts with a close to perfect F-measure until reaching a true number We compare the models with the Normalized Mutual Information (NMI) criterion but also classification accuracy (ACC) (Table 1 ). The MoE-Sim-VAE outperforms the other methods w.r.t. clustering performance when comparing NMI and achieves the second-best result when comparing ACC. Note that we used the number of experts k = 10 in our model to fit the existing number of digits in MNIST. Regarding the similarity measure, we decided to use as similarity a UMAP projection (McInnes et al., 2018) of MNIST and then apply k-nearest-neighbors of each sample in a batch. More details on the model are reported in Section A.1.2. In an ablation study we show the importance of the similarity matrix to create a clear separation of the different digits in the latent representation (Figure A4 )

In addition to the clustering network, we can make use of the latent representation for image generation purposes. The latent representation is trained as a mixture of standard Gaussians. The means of these Gaussians are the centers of the clusters trained via the clustering network. Therefore, the variational distribution can be sampled from and gated to the cluster-specific expert in the MoE-decoder.

The expert then generates new data points for the specific data mode. Results and the schematic are displayed in Figure 2 and in more detail and with greater sample size in the Appendix in Figure A2 .

4.2.1. WHY DOES A MOE DECODER ACTUALLY MATTER?:

In an ablation study, we compare the two models MoE-Sim-VAE and VaDE (Jiang et al., 2017) on generating MNIST images with the request for a specific digit. The goal is to show that a MoE decoder, as proposed in our model, is beneficial. We focus our comparison to VaDE since this model, as the MoE-Sim-VAE, resorts to a mixture of Gaussian latent representation but differs in generating images by means of a single decoder network instead of a Mixture-of-Expert decoder network. The rationale for our design choice is to ensure that smaller sub-networks learn to reproduce and generate specific modes of the data, in this case of specific MNIST digits.

To show that both models' latent representations are separating the different clusters well, we computed the Maximum Mean Discrepancy (MMD), defined in Section A.1.2. The MMD can be inter-preted as a distance between distributions computed based on samples drawn from these distributions. The heatmaps of the MMDs for VaDE and MoE-Sim-VAE as well as an UMAP projection of the latent representation colored with the mixture component confirm visually the separation of the clusters in the latent representations of both models (Fig. A3 ). As a result, we can conclude that both latent representations can separate the clusters of respective digits well, such that the decoder gets well-defined samples to generate the requested digit. Therefore, the main difference of generating specific digits arises in the decoder/generator networks.

We evaluated the importance of the MoE-Decoder to (1) accurately generate requested digits and

(2) be efficient in generating requested digits. Specifically, we sampled 10, 000 points from each mixture component in the latent representation, generated images, and used the model's internal clustering to assign a probability to which digits were generated. To generate correct and highquality images with VaDE, the posterior of the latent representation needs to be evaluated for each sample. This was done for the different thresholds φ ∈ [0.0, 0.1, 0.2, • • • , 0.9, 0.999]. The default threshold Jiang et al. (2017) used was φ = 0.999. Instead of thresholding the latent representation, we ran the generation process for MoE-Sim-VAE for each threshold with the same settings. To generate images from VaDE we used the Python implementation1 and model weights publicly available from Jiang et al. (2017) .

As a result of this analysis we report a confusion matrix for MoE-Sim-VAE in Figure A6 , the confusion matrices for each threshold for VaDE in Figure A7 , the accuracy of generating a requested digit and the number of runs required in Figure A5 . In summary, one can see that the MoE-Sim-VAE generates digits more accurately with fewer resources required. This can especially be seen when comparing the number of iterations required to fulfill the default posterior threshold of 0.999. VaDE needs nearly 2 million iterations to find samples that fulfill the aforementioned threshold criterion whereas the MoE-Sim-VAE only requires 10, 000 for a comparable sample accuracy. In comparison the mean accuracy over all thresholds for MoE-Sim-VAE is 0.970, whereas VaDE reaches on average 0.944. VaDE reaches a maximum accuracy of 0.995, which costs the aforementioned 2 million iterations for generating 100, 000 images, whereas MoE-Sim-VAE reaches a maximum accuracy of 0.971 with 100, 000 runs, without accounting for a systematic generating/clustering error (confusing 5 and 8) of MoE-Sim-VAE which can be seen in the confusion matrix in Figure A6 .

4.3. LEARNING CELL TYPE COMPOSITION IN PERIPHERAL BLOOD MONONUCLEAR CELLS USING CYTOF MEASUREMENTS:

In the following, we want to show representation learning performance on a real-world problem in biology. Specifically, we focus on cell type definition from single-cell measurements. Cytometry by time-of-flight mass spectrometry (CyTOF) (Bandura et al., 2009) is a state-of-the-art technique allowing measurement of up to 1, 000 cells per second and in parallel over 40 protein markers of the cells (Kay et al., 2013) . Defining biologically relevant cell subpopulations by clustering this data is a common learning task (Aghaeepour et al., 2013; Weber & Robinson, 2016) .

Many methods have been developed to tackle the problem introduced above and were compared on four publicly available datasets in Weber & Robinson (2016) . The best out of 18 methods were FlowSOM (Gassen et al., 2015) , PhenoGraph (Levine et al., 2015) and X-shift (Samusik et al., 2016) . These are based on k-nearest-neighbors heuristics, either defined from a spanning graph or from estimating the data density. In contrast to these methods, MoE-Sim-VAE can map new cells into the latent representation, assign probabilities for cell types and infer an interpretable latent representation allowing intuitive downstream analysis by domain experts.

We applied MoE-Sim-VAE to the same datasets as in Weber & Robinson (2016) and achieve superior results in classification using the F-measure (Equation 12) in three out of four datasets. Similarly as in Weber & Robinson (2016) we trained MoE-Sim-VAE 30 times and report in Table 2 (adopted from Weber & Robinson (2016) ) the means across all runs. The reproducibility of our model for each dataset can be seen in Figure A8 .

Further, we trained a MoE-Sim-VAE model on 268 datasets from Bodenmiller et al. (2012) (more details on the data in A.1.3), and achieve superior classification results of cell subpopulations in the data when comparing to state-of-the-art methods in this field (PhenoGraph, X-Shift, FlowSOM). Exact results can be seen in Table A1 or visualized in Figure 3 . More details on the MoE-Sim-VAE setting used for all results on CyTOF data are reported in the appendix (Section A.1.3).

5. CONCLUSION:

Our MoE-Sim-VAE model can infer similarity-based representations, perform clustering tasks, and efficiently as well as accurately generate high-dimensional data. The training of the model is performed by optimizing a joint objective function consisting of data reconstruction, clustering, and KL loss, where the latter regularizes the latent representation. On synthetic data, we have shown the strengths and limitations of the model. On the benchmark dataset of MNIST, we presented superior clustering performance and the efficiency and accuracy of MoE-Sim-VAE in generating high-dimensional data. On the biological real-world task of defining cell subpopulations in complex single-cell data, we show superior clustering performances compared to state-of-the-art methods on over 270 datasets and therefore demonstrate MoE-Sim-VAE's real-world usefulness.

Future work might include to add adversarial training to the MoE decoder, which could improve image generation to create even more realistic images. Also, specific applications might benefit from replacing the Gaussian with a different mixture model. So far the MoE-Sim-VAE's similarity measure has to be defined by the user. Relaxing this requirement and allowing for learning a useful similarity measure automatically for inferring latent representations will be an interesting extension to explore. This could be useful in a weakly-supervised setting, which often occurs for example in clinical data consisting of healthy and diseased patients. Minor details between a healthy and diseased patient might make a huge difference and could be learned from the data using neural networks.

Figure 3 : Comparison of MoE-Sim-VAE to the most popular competitor methods on defining cell types in peripheral blood mononuclear cell data via CyTOF measurements. On the x-axis different inhibitor treatments are listed whereas the y-axis reports the respective F-measure, defined in Equation 12, as performance measure of the methods. Each violin plot represents a run on a different inhibitor with multiple wells, whereas the line connects the means of the performance on the specific inhibitor.

Elie Aljalbout, Vladimir Golkov, Yawar Siddiqui, Maximilian Strobel, and Daniel Cremers. Clustering with deep learning: Taxonomy and new methods. arXiv, 2018.

A APPENDIX:

A.1 EXPERIMENTAL DETAILS

In the following sections we provide more details on model implementations, metrics used and additional result figures for the experiments described in the main text.

A.1.1 EVALUATION OF MOE-SIM-VAE ON SYNTHETIC DATA:

Model and training details:

• number of experts: {2, . . . , 40}

• batch size: 512

• code size: 10

• Number of iterations: 5000

• activation function; elu

• loss coefficient data reconstruction: 0.487

• loss coefficient clustering : 0.487

• loss coefficient mixture of Gaussian: 0.024

• learning rate: 0.001

• batch normalization:

• dropout rate: 0.5

• distance threshold (perplexity parameter): 2

• depth clustering network: 5

• internal size clustering network: 100

• trainable parameters: depending on number of experts

We compare results based on F-measure (Aghaeepour et al., 2013) , which is defined as follows

EQUATION

)

where N is the number of samples C{c 1 , c 2 , . . . , c n } and K{k 1 , k 2 , . . . , k m } are the cluster result and the reference cluster, respectively. Further F (c i , k j ) is the harmonic mean of precision and recall according to

EQUATION

whereby P r(c i , k j ) is the precision and Re(c i , k j ) is the recall. Results are shown in Table 2 . • number of experts: 10

• batch size: 128

• code size: 68 One estimator of the Maximum Mean Discrepancy (MMD) (Gretton et al., 2008) is defined as

EQUATION

where • trainable parameters: 37563 (Weber & Robinson, 2016) , 22228 (Bodenmiller et al., 2012) Results are computed setting the loss coefficient for the KL loss 3 equal to zero, since we do not intend to generate any data, but rather give the chance to the AE to pick up the correct subpouplations. , 2018) . Figure A3a shows the results for the clusters of VaDE at a posterior threshold of 0.8 which is the first threshold which shows total separation of all clusters. Figure A3b shows the separation of the clusters in latent space learned from MoE-Sim-VAE. For both methods, all distributions belonging to clusters of different respective digits show a larger distance compared to the diagonal of matching distributions, such that we generate images from a well-separated latent representation for both methods and therefore the main difference comes from the decoders.

X = {x 1 , • • • , xm } iid ∼ P , Y = {ŷ 1 , • • • , ŷm } iid ∼ Q are samples

Also here we use the F-measure defined in Equation 12 as metric to evaluate the models. For the data compared in Weber & Robinson (2016) we ran each model 30 times and report reproducability of our results in A8. The model was trained on all data and validated on the on with labels.

For the data from Bodenmiller et al. (2012) we run each model on one time on the each of the 268 datasets. Hereby we focused on the following surface markers: CD3(110:114)Dd, CD45(In115)Dd, CD4(Nd145)Dd, CD20(Sm147)Dd, CD33(Nd148)Dd, CD123(Eu151)Dd, CD14(Gd160)Dd, IgM(Yb171)Dd, HLA-DR(Yb174)Dd, CD7(Yb176)Dd. The subpopulations were originally defined via the SPADE algorithm (Qiu et al., 2011) , which is a visualization tool using Agglomerative Figure A4a shows the results on MoE-Sim-VAE trained with the similarity matrix. The different digits separate well which can also be seen in the heatmap showing the MMD statistics between all digits. In comparison, Figure A4b shows results of the MoE-Sim-VAE model ignoring the similarity matrix setting the loss coefficient to zero. One can observe that the MMD statistic, which can be seen as a similarity measure of two distributions, is way lower compared to the model including the similarity matrix. Further, also the UMAP projection confirms less separation in the latent representation between the different digits.

hierarchical clustering and minimum spanning trees. The gating of the cells is done manually via coloring of the tree leaves. With MoE-Sim-VAE we try to reconstruct the defined manually defined subpopulations. Bodenmiller et al. (2012) performed experiments on multiple well plates were different inhibitors and their effect was tested. We selected for each well plate row A to test our model on. We decided for all methods to discard subpopulations which are smaller then 30 cells. As a similarity measure for MoE-Sim-VAE we reduced the dimension of the data using UMAP (McInnes et al., 2018) using the Canberra distance

d(p, q) = n i=1 |p i -q i | |p i | + |q i | (17) (a) (b)

Figure A5 : Comparison of data generation process between Moe-Sim-VAE and VaDE (Jiang et al., 2017) . Figure A5a shows the accuracy of how certain a specific digit can be generated from the respective cluster in the latent representation whereas Figure A5b compares the number of runs until a sample from the latent representation satisfied the posterior criterion from VaDE. It needs to be mentioned that MoE-Sim-VAE does not require any thresholding such that we ran the data generation process multiple times with the same settings to compare with VaDE. In total 10000 samples are generated for each digit.

where p = (p 1 , p 2 , . . . , p n ) and q = (q 1 , q 2 , . . . , q n ). Cells were defined to be similar in MoE-Sim-VAE when the distance between the cells in the UMAP-projection was smaller then a threshold. We trained and tested MoE-Sim-VAE on a splitted dataset with rations 0.8/0.2 and evaluated the performance on the unseen test dataset. In comparison the compatitor methods were trained and Table A1 : Results of MoE-Sim-VAE on data published in Bodenmiller et al. (2012) . CyTOF measurements from peripheral blood mononuclear cells (PBMCs) were taken and the goal is to define the different cell types present in the data. The ground truth was definied using the SPADE algorithm (Qiu et al., 2011) , which can visualize the high dimensional data in such a way to be able to manual gate the cells. We compare to other fully unsupervised methods as FlowSOM, X-shift and PhenoGraph and achieve in most cases the best F-measure, which is defined as in Equation 12. 

:

https://github.com/slim1017/VaDE

https://github.com/dougalsutherland/opt-mmd/blob/master/two_sample/ mmd_test.py

Figures and Tables:

Our main contributions are to • Develop a novel autoencoder architecture for similarity-based representation learning unsupervised clustering accurate and efficient data generation • Embed the Mixture-of-Expert architecture into a Variational Autoencoder setup to train a separate generator for each data mode • Show superior clustering performance of the model on benchmark dataset and real-world biological data

Figure2: Generation of MNIST digit images. Data points from the latent representation were sampled from the variational distribution (A) which is learned to be a mixture of standard Gaussians and then clustered and gated (B) to the data-mode-specific experts of the MoE Decoder (C). (D) All samples from the variational distribution were correctly classified and therefore also correctly gated.

Figure A1: Testing MoE-Sim-VAE on data sampled from a Gaussian mixture model with random sampled parameters. Figure A1a: Testing with exact numbers of experts. When comparing to classification results with GMMs one can see that our model achievs better results until around 30 mixture components and is still compatitive until 40 mixture components. With more then 40 mixture components ourr MoE-Sim-VAE is not able anymore to compete with a GMM. Figure A1b: Testing for specific number of synthetic mixture components and iterating number of experts. Until a number of GMM components of 23 MoE-Sim-VAE is very precise in learning the real number of clusters even when allowing the model to have 40 experts.

from two distributions (e.g. samples from two different clusters of the latent representation, for MNIST of two different digits) and k is a kernel function, where we use the popular RBF kernel. Based on that estimator Sutherland et al. (2019) introduced the hypothesis test H0 : Y ). The distribution for P and Q is not required to be known. Sutherland et al. (2019) used MMD and this test to train a Generative Adversarial Network (GAN) and also to evaluate the generative performance of the model. In this work we use M M D 2 (X, Y ) to test if samples of different clusters of the latent representation are similar, or in other words the distance of the distributions. We used the Python implementation 2 from Sutherland et al. (2019). A.1.3 LEARNING CELL TYPE COMPOSITION IN PERIPHERAL BLOOD MONONUCLEAR CELLS USING CYTOF MEASUREMENTS Model and training details for all experiments on CyTOF data:• number of experts: 25(Weber & Robinson, 2016), 15 (Bodenmiller et al., 2012) • batch size: 128 • code size: 9 • Number of iterations: 30000 (Weber & Robinson, 2016), 20000 (Bodenmiller et al., 2012) • activation function: relu • loss coefficient data reconstruction: 1 • loss coefficient clustering : 1 • loss coefficient mixture of Gaussian: 0 • learning rate: 0.001 (Weber & Robinson, 2016), 0.005 (Bodenmiller et al., 2012) • batch normalization • dropout rate: 0.5 • distance threshold (perplexity parameter): 2 • distance metric: correlation • depth clustering network: 5 • internal size clustering network: 9

Figure A2: More detailed overview of results and generated samples of MNIST images. The plot on the left side shows the latent representation where the red crosses are the cluster centers. Those can be used as a mean to sample from a standard Gaussian for data generation via the MoE Decoder. The boxplots on the right show the clustering and gating result of each sample from the variational distribution.

Figure A3: Comparison of two sample MMD test (Sutherland et al., 2019) on the distributions from the different mixture components in the latent representation. The heatmaps on the left side show the estimation of the MMD which can be seen as the distance between pairs of distributions. The figures on the right side show the separation of the cluster in the latent representation based on a dimensionality reduction via UMAP (McInnes et al., 2018). FigureA3ashows the results for the clusters of VaDE at a posterior threshold of 0.8 which is the first threshold which shows total separation of all clusters. FigureA3bshows the separation of the clusters in latent space learned from MoE-Sim-VAE. For both methods, all distributions belonging to clusters of different respective digits show a larger distance compared to the diagonal of matching distributions, such that we generate images from a well-separated latent representation for both methods and therefore the main difference comes from the decoders.

Figure A4: Ablation study on the similarity matrix S. Both figures show the MMD statistic and UMAP (McInnes et al., 2018) projection of reconstructed MNIST digits computed on the latent representation.FigureA4ashows the results on MoE-Sim-VAE trained with the similarity matrix. The different digits separate well which can also be seen in the heatmap showing the MMD statistics between all digits. In comparison, FigureA4bshows results of the MoE-Sim-VAE model ignoring the similarity matrix setting the loss coefficient to zero. One can observe that the MMD statistic, which can be seen as a similarity measure of two distributions, is way lower compared to the model including the similarity matrix. Further, also the UMAP projection confirms less separation in the latent representation between the different digits.

Figure A6: Confusion map for data generation using MoE-Sim-VAE. Besides the systematic error of confusing digit 5 and 8, which can also depend on the clustering network, the digit generation of our model performs very precise with a high accuracy of generating the digit asked for. In comparison to VaDE (Jiang et al., 2017) our model does not need any threshold on samples from the latent representation which reduces the computational costs by far.

Figure A7: Confusion maps for data generation using VaDE. Figure A7d Posterior threshold 0.3. Figure A7e Posterior threshold 0.4. Figure A7f Posterior threshold 0.5.

FigureA7: Confusion maps for data generation using VaDE. FigureA7gPosterior threshold 0.6. FigureA7hPosterior threshold 0.7. FigureA7iPosterior threshold 0.8.

Figure A7: Confusion maps for data generation using VaDE. Figure A7j Posterior threshold 0.9. Figure A7k Posterior threshold 0.999. (default for Jiang et al. (2017))

Performance comparison of our method MoE-Sim-VAE with several published methods. The Table is mainly extracted from Aljalbout et al. (2018) and complemented with results of interest. (" -": metric not reported) Euclidean dist.), Shaham et al. (2018) 0.687±0.004 0.622±0.008 SpectralNet (input space, Siamese dist.), Shaham et al. (2018) 0.884±0.02 0.826±0.03 SpectralNet (code space, Euclidean dist.), Shaham et al. (2018) 0.814±0.008 0.800±0.003 SpectralNet (code space, Siamese dist.), Shaham et al.

Table content: <table><tr><td>METHOD</td><td>NMI</td><td>ACC</td></tr><tr><td>JULE, Yang et al. (2016b)</td><td>0.915</td><td>-10</td></tr><tr><td>CCNN, Hsu &amp; Lin (2017)</td><td>0.876</td><td>-</td></tr><tr><td>DEC, Xie et al. (2016)</td><td>0.8</td><td>0.843</td></tr><tr><td>DBC, Li et al. (2017)</td><td>0.917</td><td>0.964</td></tr><tr><td>DEPICT, Dizaji et al. (2017)</td><td>0.916</td><td>0.965</td></tr><tr><td>DCN, Yang et al. (2017)</td><td>0.81</td><td>0.83</td></tr><tr><td>Neural Clustering, Saito &amp; Tan (2017)</td><td>-</td><td>0.966</td></tr><tr><td>UMMC, Chen et al. (2017)</td><td>0.864</td><td>-</td></tr><tr><td>VaDE, Jiang et al. (2017)</td><td>-</td><td>0.945</td></tr><tr><td>TAGnet, Wang et al. (2016)</td><td>0.651</td><td>0.692</td></tr><tr><td>IMSAT, Hu et al. (2017)</td><td>-</td><td>0.984</td></tr><tr><td>Aljalbout et al. (2018)</td><td>0.923</td><td>0.961</td></tr><tr><td>MIXAE, Zhang et al. (2017)</td><td>-</td><td>0.945</td></tr><tr><td>Spectral clustering, Shaham et al. (2018)</td><td>0.754</td><td>0.717</td></tr><tr><td>SpectralNet (input space, (2018)</td><td colspan="2">0.924±0.001 0.971±0.001</td></tr><tr><td>MoE-Sim-VAE (proposed)</td><td>0.935</td><td>0.975</td></tr></table>

Comparison of MoE-Sim-VAE performance to competitor methods in defining cell type composition in CyTOF measurements. The results in the table are extracted from the review paper ofWeber & Robinson (2016), where 18 methods are compared on four different datasets. Our model outperforms the baselines on four out of five data sets.

Table content: <table><tr><td>Method</td><td colspan="4">Levine 32dim Levine 13dim Samusik 01 Samusik all</td></tr><tr><td>ACCENSE</td><td>0.494</td><td>0.358</td><td>0.517</td><td>0.502</td></tr><tr><td>ClusterX</td><td>0.682</td><td>0.474</td><td>0.571</td><td>0.603</td></tr><tr><td>DensVM</td><td>0.66</td><td>0.448</td><td>0.239</td><td>0.496</td></tr><tr><td>FLOCK</td><td>0.727</td><td>0.379</td><td>0.608</td><td>0.631</td></tr><tr><td>flowClust</td><td>NA</td><td>0.416</td><td>0.612</td><td>0.61</td></tr><tr><td>flowMeans</td><td>0.769</td><td>0.518</td><td>0.625</td><td>0.653</td></tr><tr><td>flowMerge</td><td>NA</td><td>0.247</td><td>0.452</td><td>0.341</td></tr><tr><td>flowPeaks</td><td>0.237</td><td>0.215</td><td>0.058</td><td>0.323</td></tr><tr><td>FlowSOM</td><td>0.78</td><td>0.495</td><td>0.707</td><td>0.702</td></tr><tr><td>FlowSOM pre</td><td>0.502</td><td>0.422</td><td>0.583</td><td>0.528</td></tr><tr><td>immunoClust</td><td>0.413</td><td>0.308</td><td>0.552</td><td>0.523</td></tr><tr><td>kmeans</td><td>0.42</td><td>0.435</td><td>0.65</td><td>0.59</td></tr><tr><td>PhenoGraph</td><td>0.563</td><td>0.468</td><td>0.671</td><td>0.653</td></tr><tr><td>Rclusterpp</td><td>0.605</td><td>0.465</td><td>0.637</td><td>0.613</td></tr><tr><td>SamSPECTRAL</td><td>0.512</td><td>0.253</td><td>0.263</td><td>0.138</td></tr><tr><td>SPADE</td><td>NA</td><td>0.127</td><td>0.169</td><td>0.13</td></tr><tr><td>SWIFT</td><td>0.177</td><td>0.179</td><td>0.202</td><td>0.208</td></tr><tr><td>Xshift</td><td>0.691</td><td>0.47</td><td>0.679</td><td>0.657</td></tr><tr><td>MoE-Sim-VAE (proposed)</td><td>0.70</td><td>0.68</td><td>0.76</td><td>0.74</td></tr></table>

Peer review (iclr2020):
Summary: The paper proposes to expand the VAE architecture with a mixture-of-experts latent representation, with a mixture-component-specific decoder that can specialize in a specific cluster. Importantly, the method can take advantage of a similarity matrix to help with the clustering. Overall, I recommend a weak accept. The method seems reasonable, and the paper is well-written, but the results are only marginally better than other methods, and there are several weaknesses with the proposed architecture and experimental setup. Positives: * The idea of a more expressive variational distribution seems good, although it is not novel. * The ability to have multiple decoder networks seems reasonable. * The ability to incorporate domain knowledge (in the form of a similarity matrix S) is a plus. * The experiments are thorough, although the method is generally only slightly better than competing methods. Negatives: * It*s not clear if the similarity matrix S is already solving the clustering problem - in which case, why do we need the rest of the model? For example, in your experiments you often used UMAP to cluster data. How does using UMAP by itself work? (Along these lines, it was not clear if your GMM experiments clustered data in the original space, or in the UMAP*d space - please clarify this). A good ablation would be to somehow remove the S matrix, to see if the model can accurately cluster samples. * There is little variance in the generated samples. * There is not a one-to-one mapping of clusters to labels, so it is hard to use this method to generate a specific type of data (for example, it is hard to generate a specific digit). This is a big difference from, say, a conditional sampler as learned by a GAN. This also arises in Fig. 3, where it is clear that latent cluster assignments do not match human-interpretable cluster assignments. I suppose this is to be expected, but taken with the previous point (little variance in generated samples) I think it seriously weakens the paper*s claim that this is an *accurate an efficient data generation method.* * The method does not do well when the number of clusters is large. Regular GMMs seem to outperform it. * I felt that this paper made excessive use of the appendix. The paper is not self-contained enough, effectively violating the length restrictions. Please make an effort to move key results back in to the main body of the paper. Experiments to run: An ablation regarding the similarity matrix S. Clarification of whether GMM experiments are run in data-space, or UMAP*d space. MIXAE features prominently in your related works, but is not compared to in your experiments. It sounds like a natural comparison. Please run this experiment, or explain why it is not a comparable method.