Year: 2021

Abstract:
We present Mixture of Contrastive Experts (MiCE), a unified probabilistic clustering framework that simultaneously exploits the discriminative representations learned by contrastive learning and the semantic structures captured by a latent mixture model. Motivated by the mixture of experts, MiCE employs a gating function to partition an unlabeled dataset into subsets according to the latent semantics and multiple experts to discriminate distinct subsets of instances assigned to them in a contrastive learning manner. To solve the nontrivial inference and learning problems caused by the latent variables, we further develop a scalable variant of the Expectation-Maximization (EM) algorithm for MiCE and provide proof of the convergence. Empirically, we evaluate the clustering performance of MiCE on four widely adopted natural image datasets. MiCE achieves significantly better results 1 than various previous methods and a strong contrastive learning baseline.

heading: 1 INTRODUCTION
text: 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.

heading: 2 RELATED WORK
text: 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.

heading: 3 PRELIMINARY
text: 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.

heading: 4 MIXTURE OF CONTRASTIVE EXPERTS
text: 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.

heading: 5 INFERENCE AND LEARNING
text: We first discuss the evidence lower bound (ELBO), the single objective used in MiCE, in Sec. 5.1. Then, we present a scalable variant of the Expectation-Maximization (EM) algorithm (Dempster et al., 1977) to deal with the non-trivial inference and learning of MiCE in Sec. 5.2. Lastly, in Sec. 5.3, we show that a naı̈ve two-stage baseline, in which we run a spherical k-means algorithm on the embeddings learned by MoCo, is a special case of MiCE.

heading: 5.1 EVIDENCE LOWER BOUND (ELBO)
text: The parameters to update include the parameters θ,ψ of the student and gating network respectively, and the expert prototypes µ = {µ}Kk=1. The learning objective of MiCE is to maximize the evidence lower bound (ELBO) of the log conditional likelihood of the entire dataset. The ELBO of the datapoint n is given by:
log p(yn|xn) ≥ L(θ,ψ,µ;xn, yn) := Eq(zn|xn,yn)[log p(yn|xn, zn;θ,µ)]−DKL(q(zn|xn, yn)‖p(zn|xn;ψ)), (7)
where q(zn|xn, yn) is a variational distribution to infer the latent cluster label given the observed data. The first term in Eq. (7) encourages q(zn|xn, yn) to be high for the experts that are good at discriminating the input images. Intuitively, it can relief the potential degeneracy issue (Caron et al., 2018; Ji et al., 2019), where all images are assigned to the same cluster. This is because a degenerated posterior puts the pressure of discriminating all images on a single expert, which may result in a looser ELBO. The second term in Eq. (7) is the Kullback–Leibler divergence between the variational distribution and the distribution defined by the gating function. With this term, the gating function is refined during training and considers the capability of the experts when partitioning data. Notably, MiCE does not rely on auxiliary loss or regularization terms as many prior methods (Haeusser et al., 2018; Shiran & Weinshall, 2019; Wu et al., 2019; Niu et al., 2020) do.

heading: 5.2 EM ALGORITHM
text: E-step. Inferring the posterior distribution of latent variables given the observations is an important step to apply MiCE to clustering. According to Bayes’ theorem, the posterior distribution given the current estimate of the model parameters is:
p(zn|xn, yn;θ,ψ,µ) = p(zn|xn;ψ)p(yn|xn, zn;θ,µ)∑K k=1 p(k|xn;ψ)p(yn|xn, k;θ,µ) . (8)
Comparing with the gating function p(zn|xn;ψ), the posterior provides better estimates of the latent variables by incorporating the supplementary information of the experts. However, we cannot tractably compute the posterior distribution because of the normalization constants Z(xn, zn;θ,µ). In fact, given the image xn and the cluster label zn, Z(xn, zn;θ,µ) sums over the entire dataset, which is prohibitive for large-scale image dataset. We present a simple and analytically tractable estimator to approximate them. Specifically, we maintain a queue q ∈ Rν×K×d that stores ν previous outputs of the teacher network, following MoCo closely. Formally, the estimator Ẑ(·) is:
Ẑ(xn, zn;θ,µ) = exp ( v>yn,zn (fn,zn + µzn) /τ ) + ν∑ i=1 exp ( q>i,zn (fn,zn + µzn) /τ ) . (9)
The estimator is biased, while its bias decreases as ν increases and we can get a sufficient amount of embeddings from the queue q efficiently4. With Eq. (9), we approximate the posterior as:
q(zn|xn, yn;θ,ψ,µ) = p(zn|xn;ψ)Φ (xn, yn, zn;θ,µ)/Ẑ(xn, zn;θ,µ)∑K k=1 p(k|xn;ψ)Φ (xn, yn, k;θ,µ)/Ẑ(xn, k;θ,µ) . (10)
M-step. We leverage the stochastic gradient ascent to optimize ELBO with respect to the network parameters θ,ψ and the expert prototypes µ. We approximate the normalization constants appear in ELBO in analogy to the E-step, formulated as follows:
L̃(θ,ψ,µ;xn, yn) = Eq(zn|xn,yn;θ,ψ,µ)
[ log Φ (xn, yn, zn;θ,µ)
Ẑ(xn, zn;θ,µ) ] −DKL(q(zn|xn, yn;θ,ψ,µ)‖p(zn|xn;ψ)). (11)
Sampling a mini-batch B of datapoints, we can construct an efficient stochastic estimator of ELBO over the full dataset to learn θ,ψ and µ:
L(θ,ψ,µ;X,Y) ≈ N |B| ∑ n∈B L̃(θ,ψ,µ;xn, yn). (12)
It requires additional care on the update of the prototypes, as discussed in many clustering methods (Sculley, 2010; Xie et al., 2016; Yang et al., 2017; Shiran & Weinshall, 2019). Some of them carefully adjust the learning rate of each prototype separately (Sculley, 2010; Yang et al., 2017), which can be very different from the one used for the network parameters. Since evaluating different learning rate schemes on the validation dataset is often infeasible in unsupervised clustering, we employ alternative strategies which are free from using per-prototype learning rates in MiCE.
As for the expert prototypes, we observe that using only the stochastic update can lead to bad local optima. Therefore, at the end of each training epoch, we apply an additional analytical update derived from the ELBO as follows:
µ̂k = ∑
n:ẑn=k
vyn,k, µk = µ̂k ‖µ̂k‖2 , ∀k, (13)
where ∀n, ẑn = arg maxk q(k|xn, yn;θ,ψ,µ) is the hard assignment of the cluster label. Please refer to Appendix A.2 for the detailed derivation. Intuitively, the analytical update in Eq. (13) considers all the teacher embeddings assigned to the k-th cluster, instead of only the ones in a mini-batch, to avoid bad local optima.
Beside, we fix the gating prototypes ω to a set of pre-defined embeddings to stabilize the training process. However, using randomly initialized prototypes may cause unnecessary difficulties in partitioning the dataset if some of them are crowded together. We address the potential issue by using the means of a Max-Mahalanobis distribution (MMD) (Pang et al., 2018) which is a special case of the mixture of Gaussian distribution. The untrainable means in MMD provide the optimal inter-cluster dispersion (Pang et al., 2020) that stabilizes the gating outputs. We provide the algorithm of MMD in Appendix B and a systematical ablation study in Tab. 3 to investigate the effect of the updates on ω and µ. Lastly, we provide the formal proof on the convergence of the EM algorithm in Appendix A.4.

heading: 5.3 RELATIONS TO A TWO-STAGE BASELINE
text: The combination of a contrastive learning method and a clustering method is a natural baseline of MiCE. Our analysis reveals that MiCE is the general form of the two-stage baseline in which we learn the image embeddings with MoCo (He et al., 2020) and subsequently run a spherical k-means algorithm (Dhillon & Modha, 2001) to obtain the cluster labels.
On one hand, in the extreme case where κ → ∞ (Assumption A3), the student embeddings fn,k and teacher embeddings vyn,k are identical for different k (Assumption A4), and the class-level
4Even though the bias does not vanish due to the use of queue, we find that the approximation works well empirically.
information in Eq. (6) is omitted (Assumption A5), we arrive at the same Softmax classifier (Eq. (1)) and the InfoNCE loss (Eq. (2)) used by MoCo as a special case of our method. On the other hand, of particular relevance to the analytical update on expert prototypes (Eq. (13)) is the spherical k-means algorithm (Dhillon & Modha, 2001) that leverages the cosine similarity to cluster `2-normalized data (Hornik et al., 2012). In addition to Assumptions A3 and A4, if we assume the unnormalized model is perfectly self-normalized (Assumption A2), using the hard assignment to get the cluster labels together with the analytical update turns out to be a single-iteration spherical k-means algorithm on the teacher embeddings. Please refer to the Appendix C for a detailed derivation.
The performance of the baseline is limited by the independence of the representation learning stage and the clustering stage. In contrast, MiCE provides a unified framework to align the representation learning and clustering objectives in a principled manner. See a comprehensive comparison in Tab. 1.
6 EXPERIMENTS
In this section, we present experimental results to demonstrate the effectiveness of MiCE. We compare MiCE with extensive prior clustering methods and the contrastive learning based twostage baseline on four widely adopted benchmarking datasets for clustering, including STL-10 (Coates et al., 2011), CIFAR-10 (Krizhevsky et al., 2009), CIFAR-100 (Krizhevsky et al., 2009),
and ImageNet-Dog (Chang et al., 2017). The experiment settings follow the literature closely (Chang et al., 2017; Wu et al., 2019; Ji et al., 2019; Shiran & Weinshall, 2019; Darlow & Storkey, 2020) and the numbers of the clusters are known in advance. The statistics of the datasets are summarized in
Tab. 2. We adopt three common metrics to evaluate the clustering performance, namely normalized mutual information (NMI), cluster accuracy (ACC), and adjusted rand index (ARI). All the metrics are presented in percentage (%). We use a 34-layer ResNet (ResNet-34) (He et al., 2016) as the backbone for MiCE and MoCo following the recent methods (Ji et al., 2019; Shiran & Weinshall, 2019) for fair comparisons. We set both temperatures τ and κ as 1.0, and the batch size as 256. The datasets, network, hyper-parameters, and training settings are discussed detailedly in Appendix D.

heading: 6.1 MAIN CLUSTERING RESULTS
text: Comparison with existing deep clustering methods. As shown in Tab. 1, MiCE outperforms the previous clustering approaches by a significant margin on all datasets. The comparison highlights the importance of exploring the discriminative representations and the semantic structures of the dataset.
Comparison with the two-stage baseline. Compared to the straightforward combination of MoCo and spherical k-means, MiCE explores the semantic structures of the dataset that improve the clustering performance. From Tab. 1, we can see that MiCE consistently outperforms the baseline in terms of the mean performance, which agrees with the analysis in Sec. 5.3. Specifically, regarding ACC, we improve upon the strong baseline by 8.7%, 2.7%, and 8.2% on CIFAR-10, CIFAR-100, and ImageNet-Dog, respectively. Taking the measurement variance into consideration, our performance overlaps with MoCo only on STL-10. We conjecture that the small data size may limit the performance as each expert learns from a subset of data. Nevertheless, the comparison manifests the significance of aligning the representation learning and clustering objectives in a unified framework, and we believe that MiCE points out a promising direction for future studies in clustering.
Visualization of the learned embeddings. We visualize the image embeddings produced by the gating network using t-SNE (Maaten & Hinton, 2008) in Fig. 2. Different colors denote the different ground-truth class labels. At the beginning, the embeddings from distinct classes are indistinguishable. MiCE progressively refines its estimates and ends up with embeddings that show a clear cluster structure. The learned clusters align with the ground-truth semantics well, which verifies the effectiveness of our method. Additional visualizations and the comparisons with MoCo are in Appendix E.

heading: 6.2 ABLATION STUDIES
text: Simplified model (Tab. 3 (left)). We investigate the gating function and the unnormalized model to understand the contributions of different components. Using a simpler latent variable model often deteriorates the performance. (1) With a uniform prior, the experts would take extra efforts to become specialized in a set of images with shared semantics. (2 & 3) The teacher embedding vyn is pushed to be close to all expert prototypes at the same time. It may be difficult for the simplified expert to encode the latent semantics while being discriminative. (4) The performance drop shows that the class-level information is essential for the image embeddings to capture the semantic structures of the dataset, despite the learned representations are still discriminative between instances. Without the
term, the learned embeddings are mixed up and scattered over the embedding space without a clear cluster structure.
Prototypes update rules (Tab. 3 (right)). We also conduct ablation studies to gain insights into the different ways of handling the gating and expert prototypes. We see that (a) without Eq. (13), we may be stuck in bad local optima. As mentioned in Sec. 5.2, a possible reason is that we are using the same learning rate for all network parameters and prototypes (Sculley, 2010; Yang et al., 2017), but tuning separate learning rates for each prototype is impractical for unsupervised clustering. Hence, we derive the analytical update to tackle the issue. As for (b), it shows that the current gradient update rule avoids the potential inconsistency between the expert prototypes and the teacher embeddings during the mini-batch training. Lastly, as discussed in Sec. 5.2, comparing to using (c) uniformly initiated gating prototypes projected onto the unit sphere, utilizing the means of MMD slightly improves performance. This also bypasses the potential learning rate issue that may appear in (d).

heading: 7 CONCLUSION
text: We present a principled probabilistic clustering method that conjoins the benefits of the discriminative representations learned by contrastive learning and the semantic structures introduced by the latent mixture model in a unified framework. With a divide-and-conquer principle, MiCE comprises an input-dependent gating function that distributes subtasks to one or a few specialized experts, and K experts that discriminate the subset of images based on instance-level and class-level information. To address the challenging inference and learning problems, we present a scalable variant of Expectation-Maximization (EM) algorithm, which maximizes the ELBO and is free from any other loss or regularization terms. Moreover, we show that MoCo with spherical k-means, one of the two-stage baselines, is a special case of MiCE under various assumptions. Empirically, MiCE outperforms extensive prior methods and the strong two-stage baseline by a significant margin on several benchmarking datasets.
For future work, one may explore different learning pretext tasks that potentially fit the clustering task, other than the instance discrimination one. Also, it would be an interesting and important future work to include dataset with a larger amount clusters, such as ImageNet. Besides, being able to obtain semantically meaningful clusters could be beneficial to weakly supervised settings (Zhou, 2018) where quality labels are scarce.

heading: 8 ACKNOWLEDGEMENTS
text: The authors would like to thank Tianyu Pang and Zihao Wang for the discussion and the reviewers for the valuable suggestions. This work was supported by NSFC Projects (Nos. 62061136001, 61620106010, 62076145), Beijing NSF Project (No. JQ19016), Beijing Academy of Artificial Intelligence (BAAI), Tsinghua-Huawei Joint Research Program, a grant from Tsinghua Institute for Guo Qiang, Tiangong Institute for Intelligent Computing, and the NVIDIA NVAIL Program with GPU/DGX Acceleration. C. Li was supported by the fellowship of China postdoctoral Science Foundation (2020M680572), and the fellowship of China national postdoctoral program for innovative talents (BX20190172) and Shuimu Tsinghua Scholar.

heading: B ALGORITHM FOR GENERATING THE CENTERS OF MMD
text: We provide the algorithm on generating the centers of the Max-Mahalanobis distribution (MMD) in Algorithm 2. The gating prototypes ω are fix to these centers during training. The algorithm closely follows the one proposed by Pang et al. (2020). For a dataset with K ground-truth classes, we will generate K centers which are all `2-normalized in the Rd space. Please kindly note that the algorithm requires K ≤ (d+ 1) (Pang et al., 2020).
Algorithm 2: Algorithm to craft the gating prototypes ω as the centers of MMD Input: The dimension of each gating prototypes d and the number of the clusters K. Initialization: We initialize the first prototype ω1 with the first unit basis vector e1 ∈ Rd . Rest
of the prototypes ωi, i 6= 1, are initialized with the zero vector 0d ∈ Rd 1 for i = 2 to K do 2 for j = 1 to i− 1 do 3 ωij = −[1/(K − 1) + ω>i ωj ]/ωjj 4 end 5 ωii = √ 1− ‖ωi‖22 6 end Return: The gating prototypes ω = {ωi}Ki=1.

heading: C RELATIONS TO THE TWO-STAGE BASELINE
text: C.1 CONTRASTIVE LEARNING
Assumption A3. The gating temperature κ→∞ such that for all k, the prior
p(zn|xn) = 1
K , ∀n.
Assumption A4. There is only a single output layer for both student network fθ and teacher network fθ′ respectively. The resulting embeddings are used across K expert, to be specific, for all possible cases,
Φ(xn, yn, k) = exp ( v>yn(fn + µk)/τ ) ,
where fn = fθ(xn) ∈ Rd, v>yn = fθ′(xn) ∈ R d.
Assumption A5. The unnormalized model Φ(·) considers only the instance-level information such that for all possible cases,
Φ(xn, yn, zn) = exp ( v>yn,znfn,zn/τ ) .
Theorem 2. If Assumptions A3-A5 holds, the overall output of MiCE become
p(yn|xn) = exp(v>ynfn/τ)∑N i=1 exp(v > i fn/τ) .
Proof. Since we the expert
p(yn|xn, zn) = Φ(xn, yn, zn)
Z(xn, zn) =
exp ( v>yn(fn + µzn)/τ )∑N i=1 exp ( v>i (fn + µzn)/τ
) (Assumption A4) =
exp(v>ynfn/τ)∑N i=1 exp(v > i fn/τ) . (Assumption A5)
The overall output of MiCE is then
p(yn|xn) = K∑ k=1 p(zn = k|xn)p(yn|xn, zn = k)
= K∑ k=1 1 K p(yn|xn, zn = k) (Assumption A3) = exp(v>ynfn/τ)∑N i=1 exp(v > i fn/τ) .
The simplified model shown in Theorem 2 is essentially the non-parametric Softmax classifier used by MoCo and InstDisc, which is also related to some recent contrastive learning methods (Ye et al., 2019; Bachman et al., 2019). Note that InstDisc adopts a slightly different implementation in which the teacher network is identical to the student network and the loss function is based on NCE (Gutmann & Hyvärinen, 2010). For detailed comparisons, please refer to He et al. (2020).
Lemma 1. If Assumptions A3-A5 hold, the posterior is uniformly distributed.
Proof.
p(zn|xn, yn) = p(zn|xn)Φ (xn, yn, zn)/Z(xn, zn)∑K k=1 p(k|xn)Φ (xn, yn, k)/Z(xn, k)
= Φ (xn, yn, zn)/Z(xn, zn)∑K k=1 Φ (xn, yn, k)/Z(xn, k)
(Assumption A3)
=
exp(v>yn fn/τ)∑N i=1 exp(v
> i fn/τ)∑K
k=1 exp(v>yn fn/τ)∑N i=1 exp(v > i fn/τ)
(Theorem 2)
= 1
K .
Assumption A6. The normalization constant is computationally tractable such that we can have the variational distribution being identical to the posterior distribution.
Theorem 3. Given that A3-A6 hold, p(yn|xn) = exp(v>ynfn/τ)/ ∑N i=1 exp(v > i fn/τ), and the tractable version of ELBO is identical to the form of InfoNCE (Oord et al., 2018) loss used by MoCo.
Proof.
L̃(θ,ψ,µ;xn, yn) = Eq(zn|xn,yn)[log Φ (xn, yn, zn)
Ẑ(xn, zn) ]−DKL(q(zn|xn, yn)‖p(zn|xn)
= log Φ (xn, yn, zn)
Ẑ(xn, zn) (Lemma 1)
= log exp
( v>ynfn/τ ) exp ( v>ynfn/τ ) + ∑ν i=1 exp ( q>i fn/τ ) . (Theorem 2)
According to Theorem 2 and 3, we see that MoCo can be viewed as a special case of the proposed MiCE. In other words, they are able to learn the same representations under the same experimental setting if the above assumptions are made.
C.2 SPHERICAL k-MEANS
Theorem 4. If Assumptions A1-A4 hold, the analytical update on the expert prototypes is equivalent to a single-iteration spherical k-means algorithm on the teacher embeddings.
Proof. Under the assumptions, the variational distribution reduces to
q(zn|xn, yn) = Φ (xn, yn, zn)∑K k=1 Φ (xn, yn, k)
(Assumptions A2 and A3)
= exp
( v>yn,zn fn,zn/τ + v > yn,zn µzn/τ )
∑K k=1 exp ( v>yn,k fn,k/τ + v > yn,k µk/τ )
= exp
( v>ynfn/τ + v > ynµzn/τ )∑K k=1 exp ( v>ynfn/τ + v > ynµk/τ
) (Assumption A4) = exp ( v>ynµzn/τ
)∑K k=1 exp ( v>ynµk/τ
) , such that the hard cluster assignment is determined bt the cosine similarity between the teacher embedding and the expert prototypes:
q̂(zn|xn, yn) =
{ 1, if zn = argmax
k q(k|xn, yn) = argmax k v>ynµk,
0, otherwise.
With the simplified expert model, we follow the previous discussion on Eq. (14) and get the Lagrangian of the objective function:
argmax µk λ(1− µ>k µk) + N∑ n=1 q̂(zn = k|xn, yn)v>ynµk/τ.
The analytical solution is therefore:
µ̂k := N∑ n=1 q̂(zn = k|xn, yn)vyn ,
µk = µ̂k ‖µ̂k‖ ,
where the cluster assignment step and prototype update rule are the same as the spherical k-means.

heading: D EXPERIEMENT SETTINGS
text: We mainly compare with the methods that are trained from scratch without using the pre-training model and the experiment settings follow the literature closely (Chang et al., 2017; Wu et al., 2019; Ji et al., 2019; Shiran & Weinshall, 2019; Darlow & Storkey, 2020). For CIFAR-10, CIFAR-100, and STL-10, all the training and test images are jointly utilized, and the 20 superclasses of CIFAR-100 are used instead of the fine labels. The 15 classes of dog images are selected from the ILSVRC2012 1K (Deng et al., 2009) dataset and resized to 96× 96× 3 to form the ImageNet-Dog dataset (Chang
et al., 2017; Wu et al., 2019). Note that the numbers of the clusters are known in advance as in Chang et al. (2017); Ji et al. (2019); Wu et al. (2019); Shiran & Weinshall (2019). The statistics of the datasets are summarized in Tab. 2. We adopt three common metrics to evaluate the clustering performance, namely normalized mutual information (NMI), cluster accuracy (ACC), and adjusted rand index (ARI). All the metrics are presented in percentage (%).
Regarding the network architecture, MiCE mainly use a ResNet-34 (He et al., 2016) as the backbone following the recent methods (Ji et al., 2019; Shiran & Weinshall, 2019) for fair comparisons. For the gating network, the output layer is replaced by a single fully connected layer that generates a `2-normalized embeddings in R128. As for the student network, it uses the same ResNet-34 backbone as the gating network and includes K more fully connected layers which map the images to K embeddings. Therefore, the parameters of student and gating networks are shared except for the output layers. The teacher network fθ′ is the exponential moving average (EMA) version of the student network fθ , which stabilizes the learning process (He et al., 2020; Tarvainen & Valpola, 2017). The update rule follows θ′ ← mθ′ + (1 −m)θ with m ∈ [0, 1) being the smoothing coefficient. In practice, we let m = 0.999 following MoCo. Since the images of CIFAR-10 and CIFAR-100 are smaller than ImageNet images, following (Chen et al., 2020), we replace the first 7x7 Conv of stride 2 with a 3x3 Conv of stride 1 for all experiments on CIFAR-10 and CIFAR-100. The first max-pooling operation is removed as well (Wu et al., 2018; Chen et al., 2020; Ye et al., 2019). Please kindly note that if the first max-pooling operation is not removed, MiCE can still achieve 83.6% ACC on CIFAR-10. For fair comparisons, MoCo also uses a ResNet-34 with a 128-dimensional output and follows the same hyper-parameter settings as MiCE.
As it is often infeasible to tune the hyper-parameters with a validation dataset in real-world clustering tasks (Ghasedi Dizaji et al., 2017), we set both temperatures τ and κ as 1.0. The queue size ν is set to 12800 for STL-10 because of the smaller data size and 16384 for the other three datasets. The data augmentation follows MoCo closely. Specifically, before passing into any of the embedding networks, images are randomly resized and cropped to the same size, followed by random gray-scale, random color jittering, and random horizontal flip. For a fair comparison, MoCo in the two-stage baseline also uses a ResNet-34 backbone. For all datasets, we use a batch size equals to 256. Note that the data augmentation strategy is critical to contrastive learning methods and MiCE, and we follow the one used by MoCo for fairness.
In terms of the optimization details, we use stochastic gradient descent (SGD) as our optimizer on the negative ELBO. We set the SGD weight decay as 0.0001 and the SGD momentum as 0.9 (He et al., 2020). The learning rate is initiated as 1.0 and is multiplied by 0.1 at three different epochs. For different datasets, the number of training epochs is different to accommodate the data size to have a similar and reasonable training time. For CIFAR-10/100, we train for 1000 epochs in total and multiply the learning rate by 0.1 at 480, 640, and 800 epochs. For STL-10, the total epochs are 6000 and the learning rate is multiplied by 0.1 at 3000, 4000, and 5000 epochs. Lastly, for ImageNet-Dog, the total epochs are 3000 and the learning rate is multiplied by 0.1 at 1500, 2000, and 2500 epochs. Also, the learning rate for expert prototypes is the same as the one for network parameters. All the experiments are trained on a single GPU.
For all experiment settings in the main text, we follow the recent methods (Wu et al., 2019; Ji et al., 2019; Shiran & Weinshall, 2019) closely where the models are trained from scratch. The setting is different from some of the methods including VaDE (Jiang et al., 2016), DGG (Yang et al., 2019), and LTVAE (Li et al., 2019) from two aspects. Firstly, we do not use any form of the pre-trained model. Secondly, we focus on a purely unsupervised setting. In contrast, VaDE (Jiang et al., 2016) and DGG (Yang et al., 2019) use a supervised pre-trained model on ImageNet for STL-10. For fairness, in the original submission, we compare to many previous methods that use the same settings.

heading: E ADDITIONAL EXPERIMENTS AND VISUALIZATIONS
text: Posterior distribution and cluster predictions. From the left side of Fig. 3, we can see that in the initial stage (epoch 1), MiCE is yet certain about the cluster labels of any give images. At the end of the training (epoch 1000), the major values of the approximated posterior distribution fall in the [0, 0.1) interval, indicating that the model is confident that images do not belong to those clusters.
After training, the learned model is able to generate sparse posterior distributions. The predicted cluster labels are also balanced across different clusters, which is shown on the right side of Fig. 3.
Visualization of embeddings of MiCE and MoCo. We present the t-SNE visualization of the embeddings learned by MiCE and MoCo in Fig. 4 and Fig. 5 to investigate whether the cluster structure and the latent semantics are captured by the models. The two figures differ in the way we color the datapoints.
In Fig. 4, different colors represent different cluster labels predicted by the clustering methods. We get the predicted cluster labels of MiCE based on the hard assignments using the posterior distributions. The cluster labels for MoCo are the outputs of the spherical k-means algorithm. MiCE can learn a distinct structure at the end of the training where each expert would mainly be responsible for one cluster. By comparing Fig. 4 (c) and (f), we can see that the boundaries between the clusters learned by spherical k-means do not match the structure learned by MoCo well. The divergence is caused by the independence of representation learning and clustering. MiCE solves the issue with a unified framework.
In Fig. 5, the datapoints are colored according to the ground-truth class labels that are unknown to the models. In Fig. 5(c), the cluster structure is highly aligned with the underlying semantics. Most of the clusters are filled with images from the same classes while having some difficult ones lying mostly on the boundaries. This verifies that the gating network learns to divide the dataset based on the latent semantics and allocates each of the images to one or a few experts. Please kindly note that we use the embeddings from the gating network instead of the student network for simplicity since the embeddings are from the same output head. In contrast, the cluster structure learned by MoCo does not align the semantics well. For all the above t-SNE visualizations, we set the perplexity hyper-parameter to 200.
Training time. We present the training time of MiCE and MoCO on CIFAR-10. It takes around 17 and 30 hours to train MoCo and MiCE for 1000 epochs, respectively. For all four datasets, experiments are conducted on a single GPU (NVIDIA GeForce GTX 1080 Ti).
Extra ablation studies on updating expert prototypes. A bad specification of the expert prototypes can lead to a bad result if it is not properly handled. Specifically, in Tab. (3) (a), we see that the ACC on CIFAR-10 is only 21.3%. We empirically verify two principled methods that solve the issue: (1) extra end-of-epoch training only on µ with stochastic gradient ascent or (2) using Eq. (13).
The first method is used as follows. At the end of each epoch, we update µ while fixing the network parameters until the pre-defined convergence criteria are met. The convergence can be determined based on either the norm of the prototypes or the change of the objective function. To control the training time, we stop the update at the current epoch once we iterate through the entire dataset 20 times. We discover that with additional training, we can achieve 42.3% ACC. The result shows that with a proper update on µ, the proposed model can identify semantic clusters with stochastic gradient ascent alone. However, it requires more than 10 times of training time (around 394 hours).
The above discussions manifest the benefits of using the Eq. (13) which is derived based on the same objective function. With the prototypical update, we can achieve similar results with minimal computational overhead. On average, we can achieve 42.2% ACC on CIFAR-100 as shown in Sec. 6.
Comparing MiCE to SCAN (Van Gansbeke et al., 2020). We provide additional experiment results to compare with SCAN (Van Gansbeke et al., 2020) that uses unsupervised pre-training under a comparable experiment setting. As SCAN adopts a three-step training procedure, we think that it will provide additional insights by comparing MiCE to SCAN at the steps after the pre-training step. The detail results on CIFAR-10 are presented in Tab. 4.
Firstly, we focus on SCAN with two steps of training. MiCE outperforms SCAN with two steps of training. SCAN obtains 78.7% and 81.8% on CIFAR-10 with the SimCLR augmentation and RandAugment, respectively. In contrast, MiCE can get 83.4% despite we are using a weaker augmentation strategy following MoCo (He et al., 2020) and InstDisc (Wu et al., 2018).
Since SCAN involves pre-training using SimCLR (Chen et al., 2020), it takes additional advantages when directly comparing to other methods without pre-training. Thus, we fine-tune the MiCE model with the same training protocol described in Appendix D (except the learning rate can be smaller). We discover that MiCE can obtain higher results and outperforms SCAN, as shown in the second sector in Tab. 4. To be specific, in the fine-tuning stage, we load the network parameters from the pre-trained model and a smaller initial learning rate of 0.1 with all the other settings remain the same. We show the augmentation strategy in the parenthesis if we use a different one from the pre-training stage. For the RandAugment (Cubuk et al., 2020), we follow the implementation (and hyper-parameters) based on the public code provided by SCAN. As mentioned in Van Gansbeke et al. (2020), the self-labeling stage requires a shift in the augmentation, otherwise, it will lead to a degenerated solution. In contrast, MiCE with pre-training is not prone to degeneracy can get comparable or better performance than SCAN regardless of the choice of the augmentation strategy.

heading: F ADDITIONAL DISCUSSIONS
text: The number of the experts. In the cases where the number of the experts L differs from the number of ground-truth clusters K, the model will partition the datasets into L subsets instead of K. Even though the number of experts is currently tied with K, it is not a drawback and does not prevent us from applying to the common clustering settings. Also, MiCE does not use additional knowledge comparing to the baseline methods. If the ground-truth K is not known, we may treat K as a hyperparameter and decide K following the methods described in Smyth (2000); McLachlan & Peel (2004), which is worth investigating in the future.
Overclustering. The overclustering technique (Ji et al., 2019) is orthogonal to our methods and can be applied with minor adaptations. However, it may require additional hyper-parameter tuning to ensure overclustering improves the results. From the supplementary file provided by IIC (Ji et al., 2019), we see that the numbers of clusters (for overclustering) are set differently for different datasets.
Despite overclustering is an interesting technique, we would like to highlight the simplicity of the current version of MiCE.
Similarity between the two set of prototypes. We do not expect the two prototypes to be similar even though their dimensions are the same. In fact, they have different aims: the gating ones aim to divide the datasets into simpler subtasks for the experts, while the expert ones help the expert to solve the instance discrimination subtasks by introducing the class-level information. Therefore, the derived ELBO objective does not encourage the two sets of prototypes to be similar or maintaining a clear correspondence between them.
Empirically, we calculate the cosine similarity between any pair of µ and ω. The absolute values we get are less than 0.25, which showed that they are not similar. If we force them to be similar during training, the performance may be negatively affected due to the lack of flexibility.

Review summaries:
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.