H ARD R EGULARIZATION TO P REVENT C OLLAPSE IN
O NLINE
D EEP
C LUSTERING
WITHOUT
D ATA
A UGMENTATION
Anonymous authors
Paper under double-blind review
A BSTRACT
Online deep clustering refers to the joint use of a feature extraction network and a clustering model to assign cluster labels to each new data point or batch as it is processed.
While faster and more versatile than offline methods, online clustering can easily reach the collapsed solution where the encoder maps all inputs to the same point and all are put into a single cluster.
Successful existing models have employed various techniques to avoid this problem, most of which require data augmentation or which aim to make the average soft assignment across the dataset the same for each cluster.
We propose a method that does not require data augmentation, and that, differently from existing methods, regularizes the hard assignments.
Using a Bayesian framework, we derive an intuitive optimization objective that can be straightforwardly included in the training of the encoder network.
Tested on four image datasets, we show that it consistently avoids collapse more robustly than other methods and that it leads to more accurate clustering.
We also conduct further experiments and analysis justifying our choice to regularize the hard cluster assignments.
I NTRODUCTION
Deep clustering refers to the combination of deep learning and clustering, where the data are first encoded with a deep neural network to a feature space, and then clustering is performed in the feature space.
Deep clustering models (and clustering models more generally) can be classified as offline or online.
Offline models process the entire dataset, and then assign all cluster labels at once.
Online models, by contrast, assign a cluster label to each data point, or each batch, as it is processed.
Offline methods tend to produce more accurate clusterings (e.g., (Mahon & Lukasiewicz, 2021; Niu et al., 2021)) reaching close to supervised performance, as they can leverage information from later data points when assigning cluster labels to earlier data points.
However, they are more expensive to train, as they must alternate between encoding the entire dataset/training the encoder for some number of epochs, and clustering the encoded data.
Online models, on the other hand, can jointly encode and cluster so are less computationally expensive.
They are also more versatile, being applicable in realworld settings where new data is constantly becoming available, as opposed to offline methods, which are limited to having a fixed, pre-defined dataset (Silva et al., 2013).
The disadvantage of online methods is that they are more difficult to train.
In particular, they run the risk of producing a degenerate solution where the large majority of data points are concentrated in a small number of clusters.
In the extreme case, all points are placed into the same cluster.
For example, an intuitive way to formulate a training procedure for online deep clustering is to update both the encoder network and the clustering parameters to make the encoding of each point close to its cluster centroid and far from other cluster centroids.
(The clustering parameters are, e.g., the centroids in K-means.)
However, this training objective is trivially minimized by the encoder network mapping all data points to the same point in feature space, where this point is equal to one of the cluster centroids.
Techniques for avoiding collapse we refer to as partition support.
Several partition support methods have been proposed, but they mostly require data augmentation (DA), and those that do not are often ad hoc and lack a rigorous technical foundation.
Additionally, they take the soft assignments as a measure of partition collapse and propose to regularize the soft assignments to make them more uniform.
This paper focuses on deep clustering without DA, which is an advantage because relying on DA limits a method to domains in which have sufficient prior knowledge to perform class-preserving augmentations.
Additionally, we argue that soft assignments are not an accurate measure of collapse, and that we should instead focus on hard assignments.
We propose a DA-free partition support by regularizing hard assignments.
Specifically, we consider the problem of how to optimally assign each data point in a batch to the most appropriate cluster.
We express this problem probabilistically in a Bayesian framework, where the regularizing element is captured by a uniform prior across clusters.
We then use this expression to derive a precise optimization objective, which we also show to be equivalent, up to a small error term, to the objective of maximizing the mutual information of the cluster assignments and the data index.
This objective itself is too slow to solve exactly, but we devise a greedy approximation algorithm that can be implemented straightforwardly and results in an intuitive method for fully online clustering, which we term combination assignment.
This method outperforms existing online DA-free clustering methods on four popular image clustering datasets.
Finally, we analyze the role of hard vs. soft cluster assignments in our partition support method, and in previous methods, and make the case that regularizing hard assignments is a more effective approach.
Note, although existing methods can easily convert soft assignments to hard assignments, this is very different from regularizing the hard assignments, as we propose.
While the relation between hard and soft clustering has been studied before (Bora et al., 2014; Kearns et al., 1998), its study in the context of regularizing online deep clustering collapse is new.
Our contributions are briefly summarized below.
• We articulate a clear Bayesian framework of the problem of hard assignments in online deep clustering models, which is more theoretically correct than requiring uniformity in each batch, as done by existing methods.
• We use this framework to derive an optimization objective and prove that it is approximately equivalent to an information-theoretic framework that maximizes the mutual information of the cluster assignments and the data index.
• We show empirically that the resulting method significantly outperforms existing partition support methods, both in avoiding partition collapse and resulting clustering accuracy.
• We conduct further analysis of
the performance of different partition support methods, which justifies our choice to focus on hard assignments.
The rest of this paper is organized as follows.
Section 2 gives an overview of related work, while
Section 3 lays the theoretical foundations of our method, describes our greedy algorithm for optimizing the resulting objective and proves the equivalence to mutual information maximization.
Section 4 reports our empirical results and analysis, and finally Section 5 summarizes our findings.
R ELATED W ORK
A key component of online deep clustering methods is how they avoid the collapsed solution, where (almost) every data point is placed in the same cluster.
Methods designed for contrastive learning, and those clustering models that employ it as a part of the training procedure, are more resistant to collapse, because the negative pairs are encouraged to be represented differently.
This can mean that they are to be placed in different clusters (Huang & Gong, 2021), or just that the encodings should be far apart (Zhang et al., 2021;
Cai et al., 2021).
Either approach helps resist all points having the same representation.
Even without negative pairs, data augmentation is essential for some methods to avoid collapse.
Grill et al.
(2020) perform representation learning without negative pairs, only positive pairs that are encouraged to be similar.
Then an “online” network is trained to predict the output of a “target” network, where the target network is a slow-moving average of the online network.
Zbontar et al.
(2021) use data-augmented pairs to reduce redundancy in the representations by minimizing off-diagonals in the cross-correlation matrix.
Among online clustering models, several partition support methods have been proposed.
Kulshreshtha & Guha (2018) freeze the encoder during clustering, but this requires pretraining it on a separate task.
Gao et al.
(2020) proposes an online autoencoder-based (AE) clustering model, where the reconstruction loss helps to avoid partition collapse, but this limits the method to simple datasets, because the AE’s reconstruction loss requires that different images from the same class have significant pixel overlap.
Zhan et al.
(2020) continually reweight the loss function to encourage the assignment to smaller clusters.
Additionally, when clusters decrease below a certain threshold, they are deleted, and the largest cluster is split in two using K-means.
Cai et al.
(2021) treat soft cluster assignments as representations, and use data-augmented pairs with the same method as Zbontar et al.
Zhong et al.
(2020) minimize the sum of squares of the probability (i.e., soft assignment) of each cluster, marginalized over each training batch.
Decomposing the expectation of the square, we see that this also involves minimizing the variance.
However, the authors also use contrastive learning and DA to avoid collapse, so they do not fully rely on sum-of-squares minimization.
A similar idea, employed by (Li et al., 2020; Hu et al., 2017; Niu et al., 2021; 2020; Van Gansbeke et al., 2020), is to maximize the entropy of soft assignments, marginalized over each training batch.
For an input distribution X with corresponding soft assigned labels Y , both approximated over a batch, an extra term − H
( Y ) is added to the loss function.
As the entropy of a multinomial is maximized at the uniform distribution, this encourages more equally sized clusters.
Entropy maximization has the advantage of a solid formal interpretation as part of the maximization of the mutual information between data and cluster assignments: by maximizing H ( Y ) but minimizing H ( Y | X ) (the latter is minimized explicitly in (Hu et al., 2017) and implicitly via contrastive learning in (Li et al., 2020)), we are maximizing
[Equation]
However, the marginal entropy term tends to only be partially successful at preventing partition collapse, often the entire dataset is still put into only a small number of clusters.
As well as being confirmed in our experiments in Section 4, this empirical weakness of entropy maximization for avoiding partition collapse is reported in (Hu et al., 2017), and better results are found by explicitly constraining the soft cluster assignments using non-linear programming.
Li et al.
(2020) do not rely entirely on entropy maximization, because they employ contrastive learning, which (as explained above) also helps to avoid partition collapse.
Thus, while theoretically sound, entropy maximization is not sufficiently effective empirically.
Another approach is to directly impose a constraint on cluster assignments.
Based on earlier work by Asano et al.
, Caron et al.
Deshmukh et al.
; Kumar et al.
(2021) proposed to constrain the soft cluster assignments to be marginally uniform across each training batch.
Though effective in preventing partition collapse, the constraint of uniformity across each batch is too strict, as noted by Kumar et al.
A given batch is very unlikely to be exactly uniform across classes, so the ground-truth distribution of classes will almost certainly violate this constraint.
M ETHOD 3. P ROBLEM F ORMULATION
We want to simultaneously (a) train the encoder and (b) make hard assignments to each batch of points at a time, based on the features extracted by that encoder.
If the encoder is f θ 1 : R n → R m , our clustering model is g θ 2 : R m → { 1 , . .
, K } , parametrized by θ 2 , and our batch size is N , then we seek a function of the form Γ : R N × m → { 1 , . .
, K } N .
(The difference between Γ and g θ 2 is that the former is used during training to assign an entire batch, while the latter is used at inference time and can assign each point individually.)
If we have a method for batchwise assignment, then the training objective can be formulated as minimizing the distance of points from the centroids of their assigned clusters.
This objective can be used to train both the encoder and the clustering model:
[Equation]
where x is the input batch, k i
=
Γ( f θ 1 ( x ))
i is the cluster label assigned to the encoding of x i under f , and µ k is the k th cluster centroid.
Finding a suitable assignment function Γ is non-trivial.
It should be more likely to assign points to the clusters whose centroids are closer.
However, using only this rule, and assigning every vector to its closest centroid, we allow a collapsed solution where equation 1 is minimized by f θ 1 mapping every point to a single centroid: ∀ i, f θ 1 ( x i ) =
c k , for some k ∈ { 1 , . . .
, K } .
The heart of our method is the choice of a suitable Γ , which avoids the collapsed solution.
We refer to it as combination assignment, because it uses a prior over the combination of such labels under a multinoulli distribution
3. C OMBINATION A SSIGNMENT
Combination assignment is based on a Bayesian formulation of the online clustering problem.
Let
D be the data distribution, X ∼ D be a sampled batch of data, Z be the random variable defined by applying the feature extraction to X (i.e., Z is the output of a deep encoder network), and let Y be the assigned cluster labels.
Here, we consider the encoder to be fixed, and we are just interested in the best hard assignment of cluster labels, given the extracted features.
That is, we want to determine the values of Y with the maximum probability under the a posteriori distribution p
( Y | Z ) .
We assume a uniform prior over K cluster labels, p ( Y = k )
= 1 /K for k ∈ { 1 , . . .
, K } .
Note how this is a weaker and more realistic assumption than each batch containing exactly uniform numbers of each class.
Then, the prior probability of a batch containing exactly n k labels for each cluster k ∈ { 1 , . .
, K } is
[Equation]
where N = (cid:80)
Kk =1 n k is the batch size.
We model each cluster as a multivariate normal distribution in feature space, so that the likelihood is given by
[Equation]
where d is the dimension of the feature space, and µ k and Σ k are the centroid and covariance matrix of the k th cluster, respectively.
If we further assume each cluster is spherical, with the same isotropic variance across all clusters, i.e., Σ k = σI, for k ∈ { 1 , . .
. , K } , and we maximize the posterior corresponding to the prior in equation 2 and the likelihood in equation 3, then we get the following optimization problem (details in appendix).
[Equation]
In matrix notation, the objective is
[Equation]
where ˜ Q is the matrix of unnormalized log probabilities, i.e., ˜ Q i,j = ||
z i − µ j || 2 , ⟨· , ·⟩ denotes the
Frobenius inner product, 1 a is an a -dimensional vector of all ones, and log and factorial are applied to each element in the K -dimensional vector 1 TN Q .
The requirement that Q be a Boolean matrix specifies that we seek hard assignments, and the constraint Q 1 K = 1 N specifies that the rows of Q are one-hot vectors, i.e., each latent vector is assigned to exactly one cluster.
3. S OLVING THE O PTIMIZATION P ROBLEM
Unfortunately, it is too slow to solve equation 5 exactly (see appendix for details).
However, to motivate an approximation, we can consider the simpler problem of assigning the last data point in a batch, given that the rest have already been assigned.
That is, we maximize the conditional probability of the N th assignment in a batch, conditioned on the N − 1 previous assignments.
This gives the following optimization problem (details in appendix).
[Equation]
(6) where n k is the number of points assigned to cluster k before the N th assignment.
To approximately solve equation 5, we employ a greedy algorithm that iteratively solves equationwith respect to the most confident assignment possible.
That is, at each iteration, select the pair ( i, k ) (corresponding to ( N, k N ) above) for which equation 6 is smallest, with i ranging over the indices of still-unassigned points, and k ranging over all clusters, and assign the i th point to cluster k .
To get an intuition on our method, recall that we generally want to assign points to the closest centroid, but that we should try to resist assigning to a cluster that already has lots of points assigned to it.
Even if such a cluster has its centroid closest to N th point, it may be better to instead assign to a smaller, further away cluster.
This suggests choosing the cluster assignment so as to minimize a combination of the distance to the centroid and some increasing function of cluster size, which is precisely what equation 6 expresses.
The first term says to pick a nearby cluster, and the second term says to penalize clusters that already have many points assigned.
It would not be obvious, a priori, what the increasing function of cluster-size should be exactly, but the derivation of equation 6 shows that log ( n + 1) is an appropriate choice.
3. I NTERPRETATION
Here, we show that the greedy algorithm that iteratively solves equation 6 can be interpreted as iteratively making whatever assignment will maximize the mutual information between the batch index i and the cluster labels.
First, we show that our method is very closely equivalent to maximizing the entropy of cluster labels in each batch.
Let H ( k ) be the marginal hard entropy of cluster labels after a new assignment to cluster k :
[Equation]
and consider the difference H ( k ) − H ( k ′ ) between the entropy after making assignment k vs. after making a different assignment k ′ .
It can be shown (see appendix for full proof) that
[Equation]
Thus, if we were to make each assignment so as to maximize log( L k ( x )) + λH ( k ) , where L ( k ) is the likelihood of the new data point under cluster k and λ is some hyperparameter, then, subject to the above approximation, for each k, k ′ ∈ { 1 , . . .
, K } , we would prefer to assign to k iff
[Equation]
Modelling clusters as multivariate normal distributions with isotropic variance σ (as above), and setting λ = N + 1 , equation 8 becomes equivalent to equation 6.
(See appendix for full proof.)
Thus, our method closely approximates a maximization of the entropy of cluster labels.
There is some similarity to those methods, discussed in Section 2, that use an additional loss term to encourage greater entropy of soft assignments in each batch, but an important difference here is that we are maximizing the entropy of hard assignments.
This means that the entropy of cluster labels given batch index is automatically zero.
Therefore, using the decomposition I ( X ; Y )
=
H ( X ) − H
( X | Y )
, the mutual information of the batch index i and the cluster labels equals the entropy of cluster labels, and so our method is a close approximation to maximizing this mutual information.
3. T RAINING P ROCEDURE
Our model comprises an encoder network f θ 1 and a set of cluster centroids θ 2 = { µ 1 , . .
. , µ K } .
To train on an input batch, we first encode the raw data using f θ 1 , then we employ the combination assignment method of Section 3.3, and use these assignments to minimize equation 4 with respect to both θ 1 and θ 2 .
As the second term in equation 4 has no gradient, the updates are made only with respect to the first term, so equivalently to equation 1.
At inference time, we do not perform combination assignment.
Instead, we simply assign each point to the cluster with the nearest centroid, so the resulting clustering model can assign points individually, and is not restricted to assigning cluster labels batchwise.
The full methods for training and inference are described by the functions
TrainOnBatch and PredictBatch, respectively, in Algorithm 1.
Algorithm 1: (combination assignment):
During training, cluster labels are assigned batchwise, with partition support provided by a uniform prior across clusters.
During inference, cluster labels are assigned pointwise, without any explicit partition support.
function A SSIGN B ATCH (Z) counts ← isAssigned ← N -dimensional Boolean array, initially all False;
D ← N
K matrix, where D ij = ||
Z i − µ j || 2 ; for r=1,...,N do ˜
D isAssigned loss ←
+ D return loss function T RAIN O N B ATCH (X)
Z ← f θ 1 ( X ) , encodings for a batch of N data points; loss ← AssignBatch ( Z ) ; take a gradient descent step on loss with respect to θ 1 , θfunction P REDICT D ATA P OINT (x)
z ← f θ ( x ) encodings for a batch of N data points;
assignment = arg min j =1 ,...,K || z − µ j || 2 ; return assignment
E XPERIMENTAL E VALUATION 4. D ATASETS AND M ETRICS
We report results on four popular image clustering datasets: CIFAR 10, CIFAR 100, FashionMNIST and STL, with image sizes 32,28 and 96 respectively.
We use the standard clustering metrics of accuracy (ACC), normalized mutual information (NMI), and adjusted Rand index (ARI), defined as, e.g., in (Sheng & Huber, 2020).
We also report variance of cluster sizes.
All datasets have uniform numbers of ground-truth classes, so ideally, this variance should be close to zero.
4. C LUSTER SIZES AND CLUSTERING ACCURACY
Table 1 compares our method with three existing methods: sum of squares minimization, denoted “SS” (Zhong et al., 2020), the Sinkhorn-Knopp algorithm for optimal transport, denoted “SK” (Caron et al., 2020; Kumar et al., 2021), and marginal entropy maximization (Li et al., 2020), denoted “Ent”.
Each is described in Section 2.
To make further explicit the phenomenon of partition collapse, we also include a model without any partition support.
Our method significantly outperforms others across all datasets and metrics.
Observe that the unregularized model exhibits total collapse in all experiments, placing all points in the same cluster and consequently achieving a cluster performance no better than random guessing.
In many of our experiments, the performance of SS is not much better.
By making a slight change to the author’s original method, we could actually significantly improve its results (see appendix), but it was still unreliable and less accurate than our method.
The other two existing partition support methods do a reasonable job of avoiding partition collapse.
However, entropy maximization occasionally also reaches the state with all points in the same cluster (this is consistent with previous literature, e.g., (Hu et al., 2017)).
The Sinkhorn-Knopp method is more reliable, but by far the most uniform cluster sizes are produced by our method.
Note that we do not employ our assignment algorithm at inference time, instead we just assign each point to the cluster with the nearest centroid.
This shows that our cluster centroids are well-distributed around the data manifold, each capturing a sizeable subset of the data even when the explicit support is removed.
Together, these figures show
[Table]
that (a) some form of partition support is necessary to learn anything meaningful, (b) our method of combination assignment is better at avoiding partition collapse than previous methods
[Table]
and (c) this leads to our model producing a better clustering performance
[Table]
To better isolate the effect of our proposed partition support method, we do not perform hyperparameter tuning
[Table]
and use a relatively simple architecture for all datasets.
This is the same for all methods being compared.
The network consists of two convolutional layers with filter sizes 6 and 16 , and
ReLU activations, followed by a fully connected layer to a latent space of dimension 128 .
Training uses Adam, with learning rate 1e-3, β 1 = 0 .
9 , β 2 = 0 .
99 , and batch size 256.
We follow previous works in setting K , the number of clusters, to the number of ground truth classes.
4. H ARD VS . S OFT A SSIGNMENT R EGULARIZATION
A key element of our method is the regularization of hard assignments, whereas previous methods regularize soft cluster assignments.
This is a fundamentally different form of regularization, and one which we argue is better able to prevent collapse.
For example, assume there are only three clusters and a batch size of four, and consider the following two matrices of assignment probabilities
[Equation]
The hard and soft entropy (i.e. the entropy of marginal hard and soft assignments respectively) for
D 1 are 1 . 5 and 1 . 1 respectively, and for D 2 they are 0 and 1 .
respectively.
That means D 2 has higher soft entropy than D 1 but a much lower hard entropy.
Indeed, despite having nearly maximum soft entropy, D 2 is collapsed with zero hard entropy.
A similar scenario is shown for an idealized batch in Figure 1, again revealing that near-uniform soft assignments does not guarantee avoiding collapse, and is in fact quite a different objective.
To investigate whether such differences between hard vs. soft regularization manifest in practice, we empirically measure the variance and entropy of hard and soft assignment probabilities, marginalized across the entire dataset, for each method tested.
The results are shown in Table 2.
The most striking difference between hard and soft entropy is in SS.
There, the soft entropy is often close to the maximum value (equal to the logarithm of the number of clusters), but the hard entropy
[Table]
is consistently close to zero.
This shows that this form of regularization produces batch assignment probabilities similar to matrix D 2 above, where the probabilities for each data point are squeezed close to one another, without much change in the order of highest to lowest, in particular the argmax.
A similar discrepancy is found in SK, which produces near-perfect uniformity in the soft assignments, with a variance of zero and the maximum possible entropy (up to rounding) on each dataset.
This is because it is a (close approximation to a) hard constraint problem.
However, SK’s hard assignments still show significant variability, and markedly lower entropy than the soft assignments.
This suggests that applying the SK algorithm also produces batch assignment probabilities somewhat similar to D 2 above.
Our method, on the other hand, explicitly forces the argmax to be more evenly distributed during training and, as Table 2 shows, this transfers to the testing setting as well.
(Recall that, during testing, we simply assign each point to the cluster with the nearest centroid.)
Our hard entropy is only slightly lower than our soft entropy, and is consistently higher than that of the other three methods.
This supports our argument that the mean soft assignments do not contain sufficient information to determine if the clustering model is learning a meaningful partition, and regularizing this quantity is not an optimal way to prevent collapse.
The pattern in the argmax is also important and is lost if we look only at the mean soft assignment across a batch.
Our analysis here of hard vs soft cluster assignments does not contradict (Caron et al., 2020).
They report better results using soft assignments as training labels, while we show the superiority of regularizing , i.e., encouraging equal numbers of, hard assignments.
The two are different contexts of hard and soft labels.
We also explored training SK (the method used by Caron et al.
using soft assignments as targets, but the results were slightly worse than using hard targets.
C ONCLUSION
This paper proposed a data-augmentation-free method to prevent collapse in online deep clustering.
We frame probabilistically the problem of deciding which clusters to assign a batch of data points to, given the cluster centroids and features of the data points, and we use this framing to derive a concise optimization objective for making hard cluster assignments.
We then describe an algorithm to approximately solve this optimization problem and demonstrate empirically on four datasets that this method outperforms existing methods, both in better preventing collapse and in leading to better clustering performance.
Finally, we analyze how the cluster assignment distribution is affected by our partition support method and previous comparable methods.
The analysis suggests that regularizing the soft assignments, as is done by existing works, is not sufficient to prevent collapse, and that a better approach is to regularize the hard assignments, as is done by our method.
R