M I CE: M IXTURE OF C ONTRASTIVE E XPERTS FOR U NSUPERVISED I MAGE C LUSTERING
Anonymous authors
Paper under double-blind review
A BSTRACT
We present Mixture of Contrastive Experts (MiCE), a uniﬁed 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.
Empirically, we evaluate the clustering performance of MiCE on four widely adopted natural image datasets.
MiCE achieves signiﬁcantly better results than various previous methods and a strong contrastive learning baseline.
I NTRODUCTION
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., 2019) in contrastive learning has shown promise in pre-training representations transferable to downstream tasks through ﬁne-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 uniﬁed 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., 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.
Moreover, we carefully discuss the algorithmic relation between MiCE and the two-stage baseline and show that the latter is actually a special instance of the former in a certain extreme case.
Compared with existing clustering methods, MiCE has the following advantages.
(i) Methodologically uniﬁed :
MiCE conjoins the beneﬁts of both the discriminative representations learned by contrastive learning and the semantic structures captured by a latent mixture model within a uniﬁed 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 signiﬁcantly better results than a strong contrastive baseline and extensive prior clustering methods on several benchmarks without any form of pre-training.
R ELATED W ORK
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 difﬁcult in real-world clustering settings (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-speciﬁc details that can beneﬁt 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., 2019; 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 classiﬁcation (Kolesnikov et al., 2019), after ﬁne-tuning with human labels.
In particular, InstDisc (Wu et al., 2018) learns from instance-level discrimination using NCE (Gutmann & Hyv¨arinen, 2010), and maintains a memory bank to compute the loss function efﬁciently.
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 signiﬁcantly.
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.
P RELIMINARY
We introduce the contrastive learning methods based on the instance discrimination task (Wu et al., 2018;
Ye et al., 2019;
He et al., 2019; Chen et al., 2020), with a particular focus on the recent state-of-the-art method, MoCo (He et al., 2019).
Let X = { x n } Nn =1 be a set of images without the ground-truth labels, and each of the datapoint x n is assigned with a unique surrogate label y n ∈ { 1 , 2 , ..., N } such that y n (cid:54) =
y j , ∀ j (cid:54) = n 1 .
To learn representations in an unsupervised manner, instance discrimination considers a discriminative classiﬁer that maps the given image to its surrogate label.
Suppose that we have two encoder networks f θ and f θ (cid:48) that generate (cid:96) 2 -normalized embeddings v y n ∈ R d and f n ∈ R d , respectively, given the image x n with the surrogate label y n .
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 classiﬁer with:
[Equation]
where τ is the temperature hyper-parameter controlling the concentration level (Hinton et al., 2015) 2 .
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).
The choice of the two encoder networks based on deep neural networks (DNNs) in which θ (cid:48) 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., 2019) version of θ .
In particular, MoCo (He et al., 2019) learns by minimizing the InfoNCE loss:
[Equation]
where q ∈ R ν × d is a queue of size ν ≤ N storing previous embeddings from f θ (cid:48) .
While it adopts the EMA approach to avoid rapidly changing embeddings in the queue that adversely impacts the performance (He et al., 2019).
For convenience, we refer f θ and f θ (cid:48) as the student and teacher network respectively (Tarvainen & Valpola, 2017; Tsai et al., 2019) in the following sections.
In the next section, we propose a uniﬁed latent mixture model based on contrastive learning to deal with the clustering task.
M IXTURE OF C ONTRASTIVE E XPERTS
Unsupervised clustering aims to partition a dataset X with N observations into K clusters.
We introduce the latent variable z n ∈ { 1 , 2 , ..., K } to be the cluster label of the image x n and naturally extend Eq.
(1) to Mixture of Contrastive Experts (MiCE):
[Equation]
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 ( y n | x n , z n ) is one of the experts that learn to discriminate a subset of instances and p ( z n | x n ) 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 efﬁciency.
Notably, MiCE is generic to the choice of the underlying contrastive methods (Wu et al., 2018;
He et al., 2019; 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 ﬁt the high-dimensional data.
In the following, we will elaborate on how we parameterize the gating function and the experts to ﬁt the clustering task.
For simplicity, we omit the parameters in all probability distributions in this section.
[Figure]
Expert Expert
Gating function | , | , | ,
Prediction | | | |
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 deﬁne g ψ as an encoder network that outputs an embedding for each input image.
We denote the output vector for image x n as g n ∈ R d .
The gating function is then parameterized as:
[Equation]
where κ is the temperature, and ω = { ω k } Kk =1 represent the gating prototypes.
All prototypes and image embeddings are (cid:96) 2 -normalized in the R d 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 deﬁne the expert in terms of the unnormalized model Φ( · ) following Wu et al.
; He et al.
Therefore, the probability of the image x n being recognized as the y n -th one by the z n -th expert is formulated as follows:
[Equation]
where Z ( x n , z n ) = (cid:80) N i =1 Φ( x n , y i , z n ) is a normalization constant that is often computationally intractable.
Similar to MoCo, we have the student network f θ that maps the image x n into K continuous embeddings f n = { f n,k } Kk =1 ∈ R K × d .
Likewise, the teacher network f θ (cid:48) outputs v y n = { v y n ,k } Kk =1 ∈ R K × d given x n .
To be speciﬁc, f n,z n ∈ R d and v y n ,z n ∈ R
d are the student embedding and the teacher embedding for images x n under the z n -th expert, respectively.
We then parameterize the unnormalized model as:
[Equation]
where τ is the temperature and µ = { µ k } Kk =1 represent the cluster prototypes for the experts.
In
Eq.
(6), the ﬁrst 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 deﬁne 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.
I NFERENCE AND L EARNING
We ﬁrst 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.
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.
5. E VIDENCE LOWER BOUND (ELBO)
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:
[Equation]
where q ( z n | x n , y n ) is a variational distribution to infer the latent cluster label given the observed data.
The ﬁrst term in Eq.
(7) encourages q ( z n | x n , y n ) 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 deﬁned by the gating function.
With this term, the gating function is reﬁned 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.
5. EM ALGORITHM
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:
[Equation]
Comparing with the gating function p ( z n | x n ; ψ ) , 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 ( x n , z n ; θ , µ ) .
In fact, given the image x n and the cluster label z n , Z ( x n , z n ; θ , µ ) sums over the entire dataset, which is prohibitive for large-scale image dataset.
We present a simple and analytically tractable estimator to approximate them.
Speciﬁcally, we maintain a queue q ∈ R ν
× K × d that stores ν previous outputs of the teacher network, following MoCo closely.
Formally, the estimator ˆ Z ( · ) is:
[Equation]
The estimator is biased, while its bias decreases as ν increases and we can get a sufﬁcient amount of embeddings from the queue q efﬁciently 3 .
With Eq.
(9), we approximate the posterior as:
[Equation]
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 3 Even though the bias does not vanish due to the use of queue, we ﬁnd that the approximation works well empirically.
ELBO in analogy to the E-step, formulated as follows:
[Equation]
Sampling a mini-batch B of datapoints, we can construct an efﬁcient stochastic estimator of ELBO over the full dataset to learn θ , ψ and µ :
[Equation]
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:
[Equation]
where ∀ n, ˆ z n = arg max k q ( k | x n , y n ; θ , ψ , µ ) 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 ﬁx the gating prototypes ω to a set of pre-deﬁned embeddings to stabilize the training process.
However, using randomly initialized prototypes may cause unnecessary difﬁculties 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., 2019) 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 µ .
5. R ELATIONS TO A TWO - STAGE BASELINE
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., 2019) 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 f n,k and teacher embeddings v y n ,k are identical for different k
(Assumption A4), and the class-level information in Eq.
(6) is omitted (Assumption A5), we arrive at the same Softmax classiﬁer (Eq.
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 (cid:96) 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 uniﬁed framework to align the representation learning and clustering objectives in a principled manner.
See a comprehensive comparison in Tab. 1.
[Table]
E XPERIMENTS
[Table]
In this section
[Table]
we present experimental results to demonstrate the effectiveness of MiCE.
We compare MiCE with extensive prior clustering methods and the contrastive learning based two-stage baseline on four widely adopted benchmarking datasets for clustering, including STL-10 (Coates et al., 2011), CIFAR10 (Krizhevsky & Hinton, 2009), CIFAR100 (Krizhevsky & Hinton, 2009)
[Table]
and
ImageNet-Dog (Chang et al., 2017).
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 ﬁne labels.
Theclasses 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.
; Ji et al.
; 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, we use a 34-layer ResNet (ResNet-34)
(He et al., 2016) backbone following the recent methods (Ji et al., 2019; Shiran & Weinshall, 2019) for fair comparisons.
For the student network, the ﬁnal output layer is replaced by K fully connected layers.
Each of them
[Figure]
generates a 128-dimensional embedding with a unit norm.
The gating network is a ResNet-34 model with a single output layer.
The parameters of the student and gating networks are shared except for the output layers.
The architecture of the teacher network follows the student network, while the parameters are updated using the exponential moving average (EMA) of the parameters in the student network.
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.
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.
The network, hyper-parameters, and training settings are discussed detailedly in Appendix D.
6. M AIN CLUSTERING RESULTS
Comparison with existing deep clustering methods.
As shown in Tab. 1, MiCE outperforms the previous clustering approaches by a signiﬁcant 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.
Speciﬁcally, 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 signiﬁcance of aligning the representation learning and clustering objectives in a uniﬁed framework.
Together with the strong empirical evidence, 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.
Different colors denote the different ground-truth class labels.
At the beginning of training, the embeddings from distinct classes are indistinguishable.
MiCE progressively reﬁnes its estimates and ends up with embeddings that show a clear cluster structure.
The learned clusters align with the ground-truth semantics well, which veriﬁes the effectiveness of our method.
Please refer to the Appendix E for additional visualizations and the comparisons with MoCo. [Table]
6. A BLATION S TUDIES
Simpliﬁed 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.
The teacher embedding v y n is pushed to be close to all expert prototypes at the same time.
It may be difﬁcult for the simpliﬁed expert to encode the latent semantics while being discriminative.
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
[Table] 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).
C ONCLUSION
We present a principled probabilistic clustering method that conjoins the beneﬁts of the discriminative representations learned by contrastive learning and the semantic structures introduced by the latent mixture model in a uniﬁed 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 signiﬁcant margin on several benchmarking datasets.
For future work, one may explore different learning pretext tasks that potentially ﬁt 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 beneﬁcial to weakly supervised settings (Zhou, 2018) where quality labels are scarce.
A I NFERENCE AND L EARNING A. D ERIVATION OF ELBO
log p
[Equation] [Equation]
, µ ) =
[Equation] [Equation]
|
[Equation]
Y ) (cid:20)
[Equation] [Equation]
[Equation] [Equation]
[Equation] [Equation]
X
[Equation]
Y
[Equation]
, µ ) ≥ L ( θ , ψ , µ
[Equation]
X , Y ) :=
E q
( Z | X , Y ) (cid:20)
[Equation]
=
[Equation]
X
[Equation] [Equation]
=
[Equation]
(cid:88) n =
K (cid:88) k
=
[Equation]
n
[Equation]
z n
[Equation]
k
[Equation]
θ
[Equation]
µ
[Equation]
log
[Equation]
z n
[Equation]
k | x n , y n )] .
A. D ERIVATION OF E Q . (13)
Assumption A1.
Under the hard assignment, the variational distribution is given by:
[Equation]
otherwise
[Equation]
Assumption A2.
For all possible inputs, the unnormalized model Φ( x n , y n , z n ) is self-normalized similar to (Wu et al., 2018), i.e., Z ( x n , z n ) equals to a constant c , such that there is no need to approximate the normalization constant as well.
Here, we derive the analytical update of the expert prototypes based on ELBO.
Under the Assumptions
A1 and A2, we can update µ k by solving the following problem:
[Equation]
A
[Equation]
subjected to the constraint that (cid:107) µ k (cid:107) = 1 .
Introducing a Lagrange multiplier λ , we then solve the
Lagrangian of the objective function:
[Equation]
We can get the estimates of µ k and λ by differentiating the Lagrangian with regards to them and setting the derivatives to zero.
To be speciﬁc, for µ k , we have:
[Equation]
such that
[Equation]
where we deﬁne ˆ µ k :=
(cid:80) Nn =
[Equation]
for simplicity
[Equation]
Similarly, by differentiating with respect to λ , we get:
[Equation]
[Equation]
Substituting Eq.
(15) in Eq.
(16), we have:
[Equation]
Therefore, we obtain the analytical solutions:
[Equation]
The last expression is essentially the update for the prototypes of the experts, as presented in Eq.
(13) in the main text.
A. P SEUDOCODE OF M I CE
# encoder_s, encoder_t, encoder_g: student, teacher, and gating network respectively # K: number of clusters; D: embedding size # omega: gating prototypes that are fixed to the centers of MMD (KxD)
# mu: expert prototypes (KxD)
# queue: dictionary as a queue of V embeddings (VxKxD) # m: momentum # tau, kappa: temperatures for the experts and gating function respectively teacher.params =
student.params # initialize mu_hat =
zeros((K, D)) for x in loader: # load a mini-batch x with B samples x_f
= aug(x) # a randomly augmented version x_v = aug(x)
# another randomly augmented version x_g = aug(x)
# the other randomly augmented version f = encoder_s.forward(x_f) # student embeddings: BxKxD v = encoder_t.forward(x_v) # teacher embeddings: BxKxD v = v.detach() # no gradient to teacher embeddings g = encoder_g.forward(x_g)
# gating embeddings: BxD v_f = einsum("BKD,BKD->BK", [v, f]) # instance-level information v_mu = einsum("BKD,KD->BK", [v, l2normalize(mu)])
# class-level information l_pos = (v_f + v_mu)
# positive logits: BxK queue_f =
einsum("VKD,BKD->BKV", [queue.detach().clone(), f]) queue_mu = einsum("VKD,KD->KV", [queue, l2normalize(mu)]
l_neg = queue_f
+ queue_mu.view(1, K, V)
# BxKxV experts_logits = cat([l_pos.view(B, K, 1), l_neg], dim=2) /
tau
# BxKx(1+V) experts = Softmax(expert_logits, dim=2)
[:, :, 0] # BxK gating =
Softmax(einsum("KD,BD->BK", [omega.detach().clone(), g]) / kappa) # BxK variational_q = einsum("BK,BK->BK", [gating, experts]) variational_q /=
variational_q.sum(dim=1)
# BxK
ELBO = sum(variational_q * (log(gating) + log(experts) - log(variational_q)) /
B loss = -ELBO
# SGD update: student and gating networks loss.backward() update(encoder_s.params)
update(encoder_g.params) #momentumupdate:teacher network encoder_t.params =
m*encoder_t.params+(1-m)*encoder_s.params # update the queue enqueue(queue, v) # enqueue the current mini-batch dequeue(queue)
# dequeue the earliest mini-batch
aggregate the teacher embeddings hard_assignments =
argmax(variational_q, dim=1)
# Bx mu_hat += einsum("BKD,BK->KD", [v, oneHot(hard_assignments)]
# update the expert prototypes mu = mu_hat einsum : Einstein sum; cat : concatenation; oneHot : One
[Table]
hot
[Table]
[Table] [Table]
; l2normalized :
[Table]
A. C ONVERGENCE OF THE PROPOSED EM ALGORITHM
Theorem 1 (Convergence of MiCE).
Assume that the log conditional likelihood of the observed data log p ( Y | X ; θ , ψ , µ ) has an upper bound, the approximated ELBO (cid:101)
L ( θ , ψ , µ ; X , Y ) of MiCE is then upper bounded and it will converge to a certain value L ∗ with the proposed EM algorithm.
Proof.
For any given variational distribution q ( z n | x n , y n )
[Equation]
the approximated ELBO we optimize in
Eq.
(11) can be written as the original ELBO in Eq.
(7) plus an additional term:
[Equation]
Recall that the normalization constant is
[Equation]
and the approximated normalization constant is
[Equation]
As in Eq
[Equation]
where q is a queue that stores the outputs of the teacher network as described in the main text.
If we use the teacher embeddings of the entire dataset as the queue, then the two normalization constants cancel each other and we obtain the original ELBO.
In such cases, the convergence can be proved following the standard EM algorithm
(Dempster et al., 1977; Wu, 1983).
Otherwise, we can also bound the approximated ELBO.
Since that all student and teacher embeddings and expert prototypes are (cid:96) 2 -normalized, we can bound the log ratio term independent of the choice of the queue:
[Equation]
which implies
[Equation]
and
[Equation]
Given that they are both positive, we can get
[Equation]
Therefore, the approximated ELBO is bounded by:
[Equation]
Similar to the proof of original EM (Wu, 1983), since the approximated ELBO will not decrease in expectation during training, the approximated ELBO of MiCE will converge to a certain value L ∗ .
There are two remarks for Theorem 1: (1) The convergence of the proposed EM relies on similar assumptions of the standard EM algorithm.
Due to the extra log ratio term in the approximated
ELBO, it will need further analysis to know the exact convergent point comparing to MiCE learning with the standard EM algorithm.
B A LGORITHM FOR GENERATING THE CENTERS OF MMD
We provide the algorithm on generating the centers of the Max-Mahalanobis distribution (MMD) in
Algorithm 2.
The gating prototypes ω are ﬁx to these centers during training.
The algorithm closely follows the one proposed by Pang et al.
For a dataset with K ground-truth classes, we will generate K centers which are all (cid:96) 2 -normalized in the R d space.
Please kindly note that the algorithm requires K ≤ ( d + 1) (Pang et al., 2019).
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 ﬁrst prototype ω 1 with the ﬁrst unit basis vector e 1 ∈ R d .
Rest of the prototypes ω
i , i (cid:54) = 1 , are initialized with the zero vector 0 d ∈ R d 1 for i = 2 to K do for j = 1 to i − 1 do ω ij = − [1 / ( K − 1) + ω (cid:62) i ω j
] / ω jj end ω ii = (cid:112) 1 − (cid:107) ω
i (cid:107)6 end
Return: The gating prototypes ω = { ω i } Ki =1 .
C R ELATIONS TO THE TWO - STAGE BASELINE C. C ONTRASTIVE LEARNING
Assumption A3.
The gating temperature κ → ∞ such that for all k , the prior p ( z n | x n ) = 1 K , ∀ n.
Assumption A4.
There is only a single output layer for both student network f θ and teacher network f θ (cid:48) respectively.
The resulting embeddings are used across K expert, to be speciﬁc, for all possible cases,
[Equation]
where f n = f θ ( x n ) ∈ R d , v (cid:62) y n = f θ (cid:48) ( x n ) ∈ R d .
Assumption A5.
The unnormalized model Φ( · ) considers only the instance-level information such that for all possible cases,
[Equation]
Theorem 2.
If Assumptions A3-A5 holds, the overall output of MiCE become
[Equation]
Proof.
Since we the expert
[Equation]
The overall output of MiCE is then
[Equation]
The simpliﬁed model shown in Theorem 2 is essentially the non-parametric Softmax classiﬁer 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 & Hyv ¨ arinen, 2010).
For detailed comparisons, please refer to He et al.
[Equation]
Lemma 1.
If Assumptions A3-A5 hold, the posterior is uniformly distributed.
Proof.
[Equation]
Assumption A6.
The normalization constant is computationally tractable such that we can have the variational distribution being identical to the posterior distribution
[Equation]
Theorem 3.
Given that A3-A6 hold
[Equation]
p ( y n | x n ) = exp( v (cid:62) y n f n /τ ) / (cid:80) Ni =1 exp( v (cid:62) i f n /τ )
[Equation]
and the tractable version of ELBO is identical to the form of InfoNCE (Oord et al., 2018) loss used by MoCo.
Proof.
[Equation] [Equation]
i f n /τ (cid:1) .
(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. S PHERICAL k - MEANS
Theorem 4.
If Assumptions A1-A4 hold
[Equation]
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
[Equation]
the variational distribution reduces to
[Equation]
such that the hard cluster assignment is determined bt the cosine similarity between the teacher embedding and the expert prototypes:
[Equation]
With the simpliﬁed expert model, we follow the previous discussion on Eq.
(14) and get the Lagrangian of the objective function:
[Equation]
The analytical solution is therefore:
[Equation]
where the cluster assignment step and prototype update rule are the same as the spherical k -means.
D E XPERIEMENT SETTINGS
Regarding the network architecture, MiCE mainly use a ResNet-34 (He et al., 2016) as the backbone.
For the gating network, the output layer is replaced by a single fully connected layer that generates a (cid:96) 2 -normalized embeddings in R 128 .
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 θ (cid:48) is the exponential moving average (EMA) version of the
[Figure]
   3UREDELOLW\LQWHUYDO       & RXQ W VWDJH )LQDO6WDJH,QLWLDO6WDJH           &OXVWHULQGH[     & RXQ W student network f θ , which stabilizes the learning process (He et al., 2019; Tarvainen & Valpola, 2017).
The update rule follows θ (cid:48) ← m θ (cid:48) + (1 − m ) θ with m ∈ [0 , 1) being the smoothing coefﬁcient.
In practice, we let m = 0 .
999 following MoCo.
Since the images of CIFAR-10 and CIFARare smaller than ImageNet images, following (Chen et al., 2020), we replace the ﬁrst 7x7 Conv of stride 2 with a 3x3 Conv of stride 1 for all experiments on CIFAR-10 and CIFAR-100.
The ﬁrst 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 ﬁrst 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.
Besides, 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.
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 ﬂip.
Given that the validation dataset is often unavailable in the real-world scenarios, both temperatures τ and κ are set as 1 . 0 .
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., 2019).
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 areand 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.
E A DDITIONAL EXPERIMENTS AND VISUALIZATIONS
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 conﬁdent 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.
[Figure]
[Figure]
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 ﬁgures 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 uniﬁed 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 ﬁlled with images from the same classes while having some difﬁcult ones lying mostly on the boundaries.
This veriﬁes 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 speciﬁcation of the expert prototypes can lead to a bad result if it is not properly handled.
Speciﬁcally, 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.
The ﬁrst method is used as follows.
At the end of each epoch, we update µ while ﬁxing the network parameters until the pre-deﬁned 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 datasettimes.
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 beneﬁts 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., 2019) 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 ﬁne-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 speciﬁc, in the ﬁne-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
[Table]
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.
, 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 .
F A DDITIONAL DISCUSSIONS
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 ﬁle 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 ﬂexibility.