MiCE: Mixture of Contrastive Experts for Unsupervised Image Clustering

1 INTRODUCTION . Unsupervised clustering is a fundamental task that aims to partition data into distinct groups of similar ones without explicit human labels . Deep clustering methods ( Xie et al. , 2016 ; Wu et al. , 2019 ) exploit the representations learned by neural networks and have made large progress on high-dimensional data recently . Often , such methods learn the representations for clustering by reconstructing data in a deterministic ( Ghasedi Dizaji et al. , 2017 ) or probabilistic manner ( Jiang et al. , 2016 ) , or maximizing certain mutual information ( Hu et al. , 2017 ; Ji et al. , 2019 ) ( see Sec . 2 for the related work ) . Despite the recent advances , the representations learned by existing methods may not be discriminative enough to capture the semantic similarity between images . The instance discrimination task ( Wu et al. , 2018 ; He et al. , 2020 ) in contrastive learning has shown promise in pre-training representations transferable to downstream tasks through fine-tuning . Given that the literature ( Shiran & Weinshall , 2019 ; Niu et al. , 2020 ) shows improved representations can lead to better clustering results , we hypothesize that instance discrimination can improve the performance as well . A straightforward approach is to learn a classical clustering model , e.g . spherical k-means ( Dhillon & Modha , 2001 ) , directly on the representations pre-trained by the task . Such a two-stage baseline can achieve excellent clustering results ( please refer to Tab . 1 ) . However , because of the independence of the two stages , the baseline may not fully explore the semantic structures of the data when learning the representations and lead to a sub-optimal solution for clustering . To this end , we propose Mixture of Contrastive Experts ( MiCE ) , a unified probabilistic clustering method that utilizes the instance discrimination task as a stepping stone to improve clustering . In particular , to capture the semantic structure explicitly , we formulate a mixture of conditional models by introducing latent variables to represent cluster labels of the images , which is inspired by the mixture of experts ( MoE ) formulation . In MiCE , each of the conditional models , also called an expert , learns to discriminate a subset of instances , while an input-dependent gating function partitions the dataset into subsets according to the latent semantics by allocating weights among experts . Further , we develop a scalable variant of the Expectation-Maximization ( EM ) algorithm ( Dempster et al. , ∗Corresponding author . 1Code is available at : https : //github.com/TsungWeiTsai/MiCE 1977 ) for the nontrivial inference and learning problems . In the E-step , we obtain the approximate inference of the posterior distribution of the latent variables given the observed data . In the M-step , we maximize the evidence lower bound ( ELBO ) of the log conditional likelihood with respect to all parameters . Theoretically , we show that the ELBO is bounded and the proposed EM algorithm leads to the convergence of ELBO . Moreover , we carefully discuss the algorithmic relation between MiCE and the two-stage baseline and show that the latter is a special instance of the former in a certain extreme case . Compared with existing clustering methods , MiCE has the following advantages . ( i ) Methodologically unified : MiCE conjoins the benefits of both the discriminative representations learned by contrastive learning and the semantic structures captured by a latent mixture model within a unified probabilistic framework . ( ii ) Free from regularization : MiCE trained by EM optimizes a single objective function , which does not require auxiliary loss or regularization terms . ( iii ) Empirically effective : Evaluated on four widely adopted natural image datasets , MiCE achieves significantly better results than a strong contrastive baseline and extensive prior clustering methods on several benchmarks without any form of pre-training . 2 RELATED WORK . Deep clustering . Inspired by the success of deep learning , many researchers propose to learn the representations and cluster assignments simultaneously ( Xie et al. , 2016 ; Yang et al. , 2016 ; 2017 ) based on data reconstruction ( Xie et al. , 2016 ; Yang et al. , 2017 ) , pairwise relationship among instances ( Chang et al. , 2017 ; Haeusser et al. , 2018 ; Wu et al. , 2019 ) , multi-task learning ( Shiran & Weinshall , 2019 ; Niu et al. , 2020 ) , etc . The joint training framework often ends up optimizing a weighted average of multiple loss functions . However , given that the validation dataset is barely provided , tuning the weights between the losses may be impractical ( Ghasedi Dizaji et al. , 2017 ) . Recently , several methods also explore probabilistic modeling , and they introduce latent variables to represent the underlying classes . On one hand , deep generative approaches ( Jiang et al. , 2016 ; Dilokthanakul et al. , 2016 ; Chongxuan et al. , 2018 ; Mukherjee et al. , 2019 ; Yang et al. , 2019 ) attempt to capture the data generation process with a mixture of Gaussian prior on latent representations . However , the imposed assumptions can be violated in many cases , and capturing the true data distribution is challenging but may not be helpful to the clustering ( Krause et al. , 2010 ) . On the other hand , discriminative approaches ( Hu et al. , 2017 ; Ji et al. , 2019 ; Darlow & Storkey , 2020 ) directly model the mapping from the inputs to the cluster labels and maximize a form of mutual information , which often yields superior cluster accuracy . Despite the simplicity , the discriminative approaches discard the instance-specific details that can benefit clustering via improving the representations . Besides , MIXAE ( Zhang et al. , 2017 ) , DAMIC ( Chazan et al. , 2019 ) , and MoE-Sim-VAE ( Kopf et al. , 2019 ) combine the mixture of experts ( MoE ) formulation ( Jacobs et al. , 1991 ) with the data reconstruction task . However , either pre-training , regularization , or an extra clustering loss is required . Contrastive learning . To learn discriminative representations , contrastive learning ( Wu et al. , 2018 ; Oord et al. , 2018 ; He et al. , 2020 ; Tian et al. , 2019 ; Chen et al. , 2020 ) incorporates various contrastive loss functions with different pretext tasks such as colorization ( Zhang et al. , 2016 ) , context autoencoding ( Pathak et al. , 2016 ) , and instance discrimination ( Dosovitskiy et al. , 2015 ; Wu et al. , 2018 ) . The pre-trained representations often achieve promising results on downstream tasks , e.g. , depth prediction , object detection ( Ren et al. , 2015 ; He et al. , 2017 ) , and image classification ( Kolesnikov et al. , 2019 ) , after fine-tuning with human labels . In particular , InstDisc ( Wu et al. , 2018 ) learns from instance-level discrimination using NCE ( Gutmann & Hyvärinen , 2010 ) , and maintains a memory bank to compute the loss function efficiently . MoCo replaces the memory bank with a queue and maintains an EMA of the student network as the teacher network to encourage consistent representations . A concurrent work called PCL ( Li et al. , 2020 ) also explores the semantic structures in contrastive learning . They add an auxiliary cluster-style objective function on top of the MoCo ’ s original objective , which differs from our method significantly . PCL requires an auxiliary k-means ( Lloyd , 1982 ) algorithm to obtain the posterior estimates and the prototypes . Moreover , their aim of clustering is to induce transferable embeddings instead of discovering groups of data that correspond to underlying semantic classes . 3 PRELIMINARY . We introduce the contrastive learning methods based on the instance discrimination task ( Wu et al. , 2018 ; Ye et al. , 2019 ; He et al. , 2020 ; Chen et al. , 2020 ) , with a particular focus on the recent state-of-the-art method , MoCo ( He et al. , 2020 ) . Let X = { xn } Nn=1 be a set of images without the ground-truth labels , and each of the datapoint xn is assigned with a unique surrogate label yn ∈ { 1 , 2 , ... , N } such that yn 6= yj , ∀j 6= n2 . To learn representations in an unsupervised manner , instance discrimination considers a discriminative classifier that maps the given image to its surrogate label . Suppose that we have two encoder networks fθ and fθ′ that generate ` 2-normalized embeddings vyn ∈ Rd and fn ∈ Rd , respectively , given the image xn with the surrogate label yn . We show the parameters of the networks in the subscript , and images are transformed by a stochastic data augmentation module before passing to the networks ( please see Appendix D ) . We can model the probability classifier with : p ( Y|X ) = N∏ n=1 p ( yn|xn ) = N∏ n=1 exp ( v > ynfn/τ ) ∑N i=1 exp ( v > i fn/τ ) , ( 1 ) where τ is the temperature hyper-parameter controlling the concentration level ( Hinton et al. , 2015 ) 3 . The recent contrastive learning methods mainly differ in : ( 1 ) The contrastive loss used to learn the network parameters , including NCE ( Wu et al. , 2018 ) , InfoNCE ( Oord et al. , 2018 ) , and the margin loss ( Schroff et al. , 2015 ) . ( 2 ) The choice of the two encoder networks based on deep neural networks ( DNNs ) in which θ′ can be an identical ( Ye et al. , 2019 ; Chen et al. , 2020 ) , distinct ( Tian et al. , 2019 ) , or an exponential moving average ( EMA ) ( He et al. , 2020 ) version of θ . In particular , MoCo ( He et al. , 2020 ) learns by minimizing the InfoNCE loss : log exp ( v > ynfn/τ ) exp ( v > ynfn/τ ) + ∑ν i=1 exp ( q > i fn/τ ) , ( 2 ) where q ∈ Rν×d is a queue of size ν ≤ N storing previous embeddings from fθ′ . While it adopts the EMA approach to avoid rapidly changing embeddings in the queue that adversely impacts the performance ( He et al. , 2020 ) . For convenience , we refer fθ and fθ′ as the student and teacher network respectively ( Tarvainen & Valpola , 2017 ; Tsai et al. , 2019 ) . In the following , we propose a unified latent mixture model based on contrastive learning to tackle the clustering task . 4 MIXTURE OF CONTRASTIVE EXPERTS . Unsupervised clustering aims to partition a dataset X with N observations into K clusters . We introduce the latent variable zn ∈ { 1 , 2 , ... , K } to be the cluster label of the image xn and naturally extend Eq . ( 1 ) to Mixture of Contrastive Experts ( MiCE ) : p ( Y , Z|X ) = N∏ n=1 K∏ k=1 p ( yn , zn = k|xn ) 1 ( zn=k ) = N∏ n=1 K∏ k=1 p ( zn = k|xn ) 1 ( zn=k ) p ( yn|xn , zn = k ) 1 ( zn=k ) , ( 3 ) where 1 ( · ) is an indicator function . The formulation explicitly introduces a mixture model to capture the latent semantic structures , which is inspired by the mixture of experts ( MoE ) framework ( Jacobs et al. , 1991 ) . In Eq . ( 3 ) , p ( yn|xn , zn ) is one of the experts that learn to discriminate a subset of instances and p ( zn|xn ) is a gating function that partitions the dataset into subsets according to the latent semantics by routing the given input to one or a few experts . With a divide-and-conquer principle , the experts are often highly specialized in particular images that share similar semantics , which improves the learning efficiency . Notably , MiCE is generic to the choice of the underlying 2The value of the surrogate label can be regarded as the index of the image . 3Due to summation over the entire dataset in the denominator term , it can be computationally prohibitive to get Maximum likelihood estimation ( MLE ) of the parameters ( Ma & Collins , 2018 ) . contrastive methods ( Wu et al. , 2018 ; He et al. , 2020 ; Chen et al. , 2020 ) , while in this paper , we focus on an instance based on MoCo . Also , please see Fig . 1 for an illustration of MiCE with three experts . In contrast to the original MoE used in the supervised settings ( Jacobs et al. , 1991 ) , our experts learn from instance-wise discrimination instead of human labels . In addition , both gating and expert parts of MiCE are based on DNNs to fit the high-dimensional data . In the following , we will elaborate on how we parameterize the gating function and the experts to fit the clustering task . For simplicity , we omit the parameters in all probability distributions in this section . Gating function . The gating function organizes the instance discrimination task into K simpler subtasks by weighting the experts based on the semantics of the input image . We define gψ as an encoder network that outputs an embedding for each input image . We denote the output vector for image xn as gn ∈ Rd . The gating function is then parameterized as : p ( zn|xn ) = exp ( ω > zngn/κ ) ∑K k=1 exp ( ω > k gn/κ ) , ( 4 ) where κ is the temperature , and ω = { ωk } Kk=1 represent the gating prototypes . All prototypes and image embeddings are ` 2-normalized in the Rd space . Hence , the gating function performs a soft partitioning of the dataset based on the cosine similarity between the image embeddings and the gating prototypes . We can view it as a prototype-based discriminative clustering module , whereas we obtain the cluster labels using posterior inference to consider additional information in the experts . Experts . In MiCE , every expert learns to solve the instance discrimination subtask arranged by the gating function . We define the expert in terms of the unnormalized model Φ ( · ) following Wu et al . ( 2018 ) ; He et al . ( 2020 ) . Therefore , the probability of the image xn being recognized as the yn-th one by the zn-th expert is formulated as follows : p ( yn|xn , zn ) = Φ ( xn , yn , zn ) Z ( xn , zn ) , ( 5 ) where Z ( xn , zn ) = ∑N i=1 Φ ( xn , yi , zn ) is a normalization constant that is often computationally intractable . Similar to MoCo , we have the student network fθ that maps the image xn into K continuous embeddings fn = { fn , k } Kk=1 ∈ RK×d . Likewise , the teacher network fθ′ outputs vyn = { vyn , k } Kk=1 ∈ RK×d given xn . To be specific , fn , zn ∈ Rd and vyn , zn ∈ Rd are the student embedding and the teacher embedding for images xn under the zn-th expert , respectively . We then parameterize the unnormalized model as : Φ ( xn , yn , zn ) = exp ( v > yn , zn ( fn , zn + µzn ) /τ ) , ( 6 ) where τ is the temperature and µ = { µk } Kk=1 represent the cluster prototypes for the experts . In Eq . ( 6 ) , the first instance-wise dot product explores the instance-level information to induce discriminative representations within each expert . The second instance-prototype dot product incorporates the class-level information into representation learning , encouraging a clear cluster structure around the prototype . Overall , the learned embeddings are therefore encoded with semantic structures while being discriminative enough to represent the instances . Eq . ( 6 ) is built upon MoCo with the EMA approach , while in principle , many other potential solutions exist to define the experts , which are left for future studies . Besides , the parameters θ and ψ are partially shared , please refer to the Appendix D for more details on the architecture .

Summary: Authors present “mixture of experts” type of method to solve a clustering with unsupervised learning problem. Method is called as Mixture of Contrastive Experts (MiCE) which uses contrastive learning as a base module and combines it with latent mixture models. Authors develop a scalable algorithm for MiCE and empirically evaluate the proposed method for image clustering.