Multi-Facet Clustering Variational Autoencoders

1 Introduction . Clustering is the task of finding structure by partitioning samples in a finite , unlabeled dataset according to statistical or geometric notions of similarity [ 1 , 2 , 3 ] .For example , we might group items along axes of empirical variation in the data , or maximise internal homogeneity and external separation of items within and between clusters with respect to a specified distance metric . The choice of similarity measure and how one consequently validates clustering quality is fundamentally a subjective one : it depends on what is useful for a particular task [ 2 , 4 ] . In this work , we are interested in uncovering abstract , latent characteristics/facets/aspects/levels of the data to understand and characterise the data-generative process . We further assume a fully exploratory , unsupervised setting without prior knowledge on the data , which could be exploited while fitting the clustering algorithm , and in particular without given ground-truth partitions at training time . When being faced with high-dimensional data such as images , speech or electronic health records , items typically have more than one abstract characteristic . Consider the example of the MNIST dataset [ 5 ] : MNIST images possess at least two such characteristics : The digit class , which might impose the largest amount of statistical variation , and the style of the digit ( e.g . stroke width ) . This naturally raises a question : By which characteristic is a clustering algorithm supposed to partition the data ? In MNIST , both digit class and ( the sub-categories of ) style would be perfectly reasonable candidates to answer this question . In our exploratory setting described above , there is not one “ correct ” partition of the data . ∗Equal contribution . 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . Deep learning based clustering algorithms , so-called deep clustering , were particularly successful in recent years in dealing with high-dimensional data by compressing the inputs into a lower-dimensional latent space in which clustering is computationally tractable [ 6 , 7 ] . However , almost all of these deep clustering algorithms find only a single partition of the data , typically the one corresponding to the given class label in a supervised dataset [ 8 , 9 , 10 , 11 , 12 , 13 , 14 , 15 ] . When evaluating their model , said approaches validate clustering performance by treating the one supervision label ( e.g . digit class in the case of MNIST ) as the de-facto “ ground truth clustering ” . We argue that restricting our view to a single facet C1 rather than all or at least multiple facets ( C1 , C2 , . . . , CJ ) is an arbitrary , incomplete choice of formulating the problem of clustering a high-dimensional dataset . To this end , we propose Multi-Facet Clustering Variational Autoencoders ( MFCVAE ) , a principled , probabilistic model which finds multiple characteristics of the data simultaneously through its multiple Mixtures-of-Gaussians ( MoG ) prior structure . Our contributions are as follows : ( a ) Multi-Facet Clustering Variational Autoencoders ( MFCVAE ) , a novel class of probabilistic deep learning models for unsupervised , multi-facet clustering in high-dimensional data that can be optimised end-to-end . ( b ) Novel theoretical results for the optimisation of the corresponding ELBO , correcting and extending an influential , related paper for the single-facet case . ( c ) Demonstrating MFCVAE ’ s stable empirical performance in terms of multi-facet clustering of various levels of abstraction , compositionality of facets , generative , unsupervised classification , and diversity of generation . 2 Multi-facet clustering . High-dimensional data are inherently structured according to a number of abstract characteristics , and in an exploratory setting , it is clear that arbitrarily clustering by one of them is insufficient . However , the question remains whether these multiple facets should also be explicitly represented by the model . In particular , one might argue that a single partition could be used to represent all cross-combinations2 of facets C = C1 × C2 × · · · × CJ where Cj = { 1 , 2 , . . . , Kj } , as in Fig . 1 ( a ) . In this work , we explain that explicitly representing and clustering by multiple facets , as we do in MFCVAE and illustrated in Fig . 1 ( b ) , has the following four properties that are especially desirable in an unsupervised learning setting : ( a ) Discovering a multi-facet structure . We adopt a probabilistically principled , unsupervised approach , specifying an independent , multiple Mixtures of Gaussians ( MoG ) prior on the latent space . This induces a disentangled representation across facets , meaning that in addition to examples assigned to certain clusters being homogeneous , the facets ( and their corresponding clusters ) represent different abstract characteristics of the data ( such as digit class or digit style ) . Because of this multifacet structure , the total number of clusters required to represent a given multi-partition structure of the data scales linearly w.r.t . the number of data characteristics . In comparison , the number of clusters required in a single-facet model scales exponentially ( see Fig . 1 , and Appendix A for details ) . 2Note that in practice , not all cross-combinations of facets might be present . For example , in a dataset like MNIST , one might not observe ‘ right-tilted zeros ’ , even though we observe ‘ right-tilted ’ digits and ‘ zeros ’ . ( b ) Compositionality of facets . A multi-facet model has a compositional advantage : different levels of abstraction of the data are represented in separate latent variables . As we will show , this allows qualitatively diverse characteristics to be meaningfully combined . ( c ) Generative , unsupervised classification . Our method joins a myriad of single-facet clustering models in being able to accurately identify known class structures given by the label in standard supervised image datasets . However , in contrast to previous work , we are also able find interesting characteristics in other facets with homogeneous clusters . We stress that while we compare generative classification performance against other models to demonstrate statistical competitiveness , this task is not the main motivation for our fully unsupervised model . ( d ) Diversity of generated samples . In a generative sense , the structure of the latent space allows us to compose new , synthetic examples by a set of J pairs of ( continuous , discrete ) latent variables . We can in particular intervene on each facet separately . This yields a rich set of options and fine-grained control for interventions and the diversity of generated examples . We illustrate these four properties in our experiments in Section 4 . 3 Multi-Facet Clustering Variational Autoencoders Our model comprises J latent facets , each learning its own unique clustering of samples via a Mixtureof-Gaussians ( MoG ) distribution : cj ∼ Cat ( πj ) , zj | cj ∼ N ( µcj , Σcj ) ( 1 ) where πj is the jth facet ’ s Kj-dimensional vector of mixing weights , and ( µcj , Σcj ) are the mean and covariance of the cj th mixture component in facet j ( Σcj can be either diagonal or full ) . The multi-facet generative model ( Fig . 2 [ Right ] ) is thus structured as pθ ( x , ~z , c ) = pθ ( x|~z ) pθ ( ~z|c ) pθ ( c ) = pθ ( x|~z ) J∏ j=1 pθ ( zj |cj ) pθ ( cj ) , ( 2 ) where c = { c1 , c2 , ... , cJ } , ~z = { z1 , z2 , ... , zJ } , and pθ ( x|~z ) is a user-defined likelihood model , for example a product of Bernoulli or Gaussian distributions , which is parameterised with a deep neural network f ( ~z ; θ ) . Importantly , this structure in Eq . ( 2 ) encodes prior independence across facets , i.e . pθ ( ~z , c ) = ∏ j pθ ( zj , cj ) , thereby encouraging facets to learn clusterings that span distinct subspaces of ~z . The overall marginal prior pθ ( ~z ) can be interpreted as a product of independent MoGs . 3.1 VaDE tricks . To train this model , we wish to optimise the evidence lower bound ( ELBO ) of the data marginal likelihood using an amortised variational posterior qφ ( ~z , c|x ) ( Fig . 2 [ Left ] ) , parameterised by a neural network g ( x ; φ ) , within which we will perform Monte Carlo ( MC ) estimation where necessary to approximate expectations log p ( D ) ≥ L ( D ; θ , φ ) = Ex∼D [ Eqφ ( ~z , c|x ) [ log pθ ( x , ~z , c ) qφ ( ~z , c|x ) ] ] . ( 3 ) What should we choose for qφ ( ~z , c|x ) ? Training deep generative models with discrete latent variables can be challenging , as reparameterisation tricks so far developed , such as the Gumbel-Softmax trick [ 16 , 17 ] , necessarily introduce bias into the optimisation , and become unstable when a discrete latent variable has a high cardinality . Our setting where we have multiple discrete latent variables is even more challenging . First , the bias from using the Gumbel-Softmax trick compounds when there is a hierarchy of dependent latent variables , leading to poor optimisation [ 18 ] . Second , we can not necessarily avail ourselves of advances in obtaining good estimators for discrete latent variables as either they do not carry over to the hierarchical case [ 19 ] , or are restricted to binary latent variables [ 20 ] . Third , we wish for light-weight optimisation , avoiding the introduction of additional neural networks whenever possible as this simplifies both training and neural specification . Thus , we sidestep these problems , bias from relaxations of discrete variables and the downsides of additional amortised-posterior neural networks for the discrete latent variables , by developing the hierarchical version of the VaDE trick . This trick was first developed for clustering VAEs with a single Gaussian mixture in the generative model [ 10 ] . Informally , the idea ( for a single-facet model ) is to define a Bayes-optimal posterior for the discrete latent variable using the responsibilities of the constituent components of the mixture model ; these responsibilities are calculated using samples taken from the amortised posterior for the continuous latent variable . Estimating the ELBO for models of this form does not require us to take MC samples from discrete distributions—the data likelihood is conditioned only on the continuous latent variable ~z , which we sample using the reparameterization trick [ 21 ] , and the posterior for ~z is conditioned only on x . Thus , when calculating the ELBO , we can cheaply marginalise out discrete latent variables where needed . In other words , we do not have to perform multiple forward passes through the decoder as neither it nor the ~z samples we feed it depend on c. As it is fundamental to our method , we now briefly recapitulate the original VaDE trick for VAEs with a single latent mixture ( correcting a misapprehension in the original form of this idea ) and will then cover our hierarchical extension3 . Single-Facet VaDE Trick : Consider a single facet model , so the generative model is pθ ( x , z , c ) = pθ ( x|z ) pθ ( z|c ) pθ ( c ) . Introduce a posterior qφ ( z , c|x ) = qφ ( z|x ) qφ ( c|x ) where qφ ( z|x ) is a multivariate Gaussian with diagonal covariance . The ELBO for this model for one datapoint is L ( x ; θ , φ ) = Eqφ ( z , c|x ) [ log pθ ( x|z ) pθ ( z|c ) pθ ( c ) qφ ( z|x ) qφ ( c|x ) ] = Eqφ ( z , c|x ) [ log pθ ( x|z ) pθ ( z ) pθ ( c|z ) qφ ( z|x ) qφ ( c|x ) ] , ( 4 ) where we have chosen to rewrite the generative model factorisation , pθ ( z ) = ∑ c pθ ( z|c ) pθ ( c ) is the marginal mixture of Gaussians , and pθ ( c|z ) = pθ ( z|c ) pθ ( c ) /pθ ( z ) is the Bayesian posterior for c. Expanding out the ELBO , we get L ( x ; θ , φ ) = Eqφ ( z|x ) log pθ ( x|z ) −KL [ qφ ( z|x ) ||pθ ( z ) ] − Eqφ ( z|x ) KL [ qφ ( c|x ) ||pθ ( c|z ) ] . ( 5 ) We can define qφ ( c|x ) such that Eqφ ( z|x ) KL [ qφ ( c|x ) ||pθ ( c|z ) ] is minimal , by construction , which is the case if we choose qφ ( c|x ) ∝ exp ( Eqφ ( z|x ) log pθ ( c|z ) ) as we will show in Theorem 1 . This means that we can simply use samples from the posterior for z to define the posterior for c , using Bayes ’ rule within the latent mixture model . Remark : We note , however , that in the original description of this idea in [ 10 ] , it was claimed that Eqφ ( z|x ) KL [ qφ ( c|x ) ||pθ ( c|z ) ] could , in general , be set to zero , which is not the case . Rather , this KL can be minimised , in general , to a non-zero value . We discuss this misapprehension in more detail and why the empirical results in [ 10 ] are still valid in Appendix B.1.1 . Theorem 1 . ( Single-Facet VaDE Trick ) For any probability distribution qφ ( z|x ) , the distribution qφ ( c|x ) that minimises Eqφ ( z|x ) KL [ qφ ( c|x ) ||pθ ( c|z ) ] in ( 5 ) is argmin qφ ( c|x ) Eqφ ( z|x ) KL [ qφ ( c|x ) ||pθ ( c|z ) ] = π ( c|qφ ( z|x ) ) ( 6 ) with the minimum value attained being min qφ ( c|x ) Eqφ ( z|x ) KL [ qφ ( c|x ) ||pθ ( c|z ) ] = − logZ ( qφ ( z|x ) ) ( 7 ) where π ( c|qφ ( z|x ) ) : = exp ( Eqφ ( z|x ) log p ( c|z ) ) Z ( qφ ( z|x ) ) for c = 1 , . . . , K ( 8 ) Z ( qφ ( z|x ) ) : = K∑ c=1 exp ( Eqφ ( z|x ) log p ( c|z ) ) . ( 9 ) Proof : See Appendix B.1 . 3We note that the original VaDE paper , besides the misapprehension discussed in Section 3.1 and Appendix B.1.1 , proposed a highly complex training algorithm with various pre-training heuristics which we significantly simplified while maintaining or increasing performance ( details in Appendix D.5 ) . Multi-facet VaDE Trick : In this work , we consider the case of having J facets , each with its own pair of variables ( zj , cj ) . Perhaps surprisingly , we do not have to make a mean-field assumption between the J facets for c once we have made one for ~z . In other words , once we have chosen that qφ ( ~z , c|x ) = qφ ( c|x ) ∏J j=1 qφ ( zj |x ) , where qφ ( zj |x ) is defined to be a multivariate Gaussian with diagonal covariance for each j , the optimal qφ ( c|x ) similarly factorises4 . We formalise this : Theorem 2 . ( Multi-Facet VaDE Trick for factorized qφ ( ~z|x ) , p ( ~z , c ) ) For any factorized probability distribution qφ ( ~z|x ) = ∏ j qφ ( zj |x ) , the distribution qφ ( c|x ) that minimises Eqφ ( ~z|x ) KL [ qφ ( c|x ) ||pθ ( c|~z ) ] under factorized prior p ( ~z , c ) = ∏ j p ( zj , cj ) of ( 2 ) is argmin qφ ( c|x ) Eqφ ( ~z|x ) KL [ qφ ( c|x ) ||pθ ( c|~z ) ] = ∏ j πj ( cj |qφ ( zj |x ) ) ( 10 ) where the minimum value is attained at min qφ ( c|x ) Eqφ ( ~z|x ) KL [ qφ ( c|x ) ||pθ ( c|~z ) ] = − ∑ j logZj ( qφ ( zj |x ) ) ( 11 ) where πj ( cj |qφ ( zj |x ) ) : = exp ( Eqφ ( zj |x ) log pθ ( cj |zj ) ) Zj ( qφ ( zj |x ) ) , for cj = 1 , . . . , Kj ( 12 ) Zj ( qφ ( zj |x ) ) : = Kj∑ cj=1 exp ( Eqφ ( zj |x ) log pθ ( cj |zj ) ) . ( 13 ) Proof : See Appendix B.2 . Note that we use Eq . ( 12 ) as the probability distribution of assigning input x to clusters of facet j . Armed with these theoretical results , we can now write the ELBO for our model , with the optimal posterior for c , in a form that trivially admits stochastic estimation and does not necessitate extra recognition networks for c , LMFCVAE ( D ; θ , φ ) =Ex∼D [ Eqφ ( ~z|x ) log pθ ( x|~z ) − J∑ j=1 [ Eqφ ( cj |x ) KL ( qφ ( zj |x ) ||pθ ( zj |cj ) ) + KL ( qφ ( cj |x ) ||p ( cj ) ) ] ] ( 14 ) where the optimal qφ ( cj |x ) is given by Eq . ( 12 ) for each j . To obtain the posterior distributions for c , we take MC samples from qφ ( ~z|x ) and use these to construct the posterior as in Eq . ( 12 ) . We found one MC sample ( L = 1 ; for each facet and for each x ) to be sufficient . We derive the complete MC estimator which we use as the loss function of our model and ablations on two alternative forms in Appendix C .

Summary: The authors extend the VaDE model (Jiang et al.) to allow for multiple, independent facets i.e. Mixtures of Gaussians. The authors first point out and resolve a technical oversight in derivations related to estimation of $q(c \mid x)$ of Jiang et al. (VaDE). The authors implement their formulation in an architecture based on ladder VAE to extract multiple clusterings of the data simultaneously. The authors demonstrate their method on 3 different benchmark datasets.