On partial prototype collapse in clustering-based self-supervised learning

We thank the reviewer for the insightful feedback and valuable comments. We appreciate that the reviewer recognizes the importance of the partial prototype collapse problem and for commenting on our overall presentation and experiments. We address the weaknesses and questions raised by the reviewer below.

W1. Is there any theoretical results to guarantee that minimizing L_KP will make the prototypes to be uniformly distributed in the latent space? Why not just minimizing $- \sum_i \sum_j || \mu_i - \mu_j ||^2$ ?

The objective that you suggested was one of our first thoughts as well but we realized that this can be reduced to something trivial. Since $\mu_i$ and $\mu_j$ are unit vectors, $\sum_i \sum_j || \mu_i - \mu_j ||^2$ = $\sum_i \sum_j 2 - 2 \langle \mu_i, \mu_j \rangle$ = $2 - 2 \langle \sum_i \mu_i, \sum_j \mu_j \rangle$ = $2 - 2 || \sum_i \mu_i ||^2$.

The objective reduces to minimizing the $L^2$-norm of the sum of the prototypes. This cannot solve the problem of partial prototype collapse. Consider a simple example where K/2 prototypes are equal to $\mu$ and the other K/2 prototypes are equal to $- \mu$. This will minimize the above objective perfectly but learns only 2 unique prototype vectors.

For detailed theoretical results related to optimizing the KoLeo objective, kindly refer to [1]. The KoLeo objective is an estimate of the differential entropy. A uniform distribution of prototypes in the latent space results in maximum differential entropy.

[1] Sablayrolles, Alexandre, et al. "Spreading vectors for similarity search." ICLR 2019.

W2. Overall, the improvement is small when compared with the baseline iBOT-vMF. For example, the classification accuracy of imagenet in Table 2 are very close for iBOT-vMF and iBOT-vMF (kp).

Our secondary contribution of mitigating partial prototype collapse is aimed at effectively utilizing the prototypes which we manage to achieve. We study the downstream impact of effectively utilizing all the prototypes - we study uniform and long-tailed datasets under different tasks to understand both the advantages and limitations that arise from this. In the experiments reported in Table 2, we use the same number of initialized prototypes (K=8192) as the baseline. We do this to measure the performance improvement that can be obtained only from the effective utilization of the prototypes and not from making use of additional prototypes (which also changes the memory footprint of the model compared to the baseline; see Table 12 in revised paper for details on memory costs). As we note in section 6.3, we observe that the Imagenet kNN top-1 accuracy continues to improve as the number of prototypes are increased for iBOT-vMF with KoLeo-proto regularization, which is not the case with the baseline. The Imagenet classification performance with our proposed regularization can be improved further by scaling the number of prototypes. We do not investigate this for the longer experiments due to the significantly higher computational cost required for such experiments.

W3. The authors should give more quantitative or qualitative results to demonstrate how the proposed KP regularization can solve the partial prototype collapse problem, e.g. visualization.

The goal of adding the KoLeo-prototype regularization is to mitigate the partial prototype collapse issue. We show in Figure 3a that our proposed regularization is able to achieve this goal, by comparing the number of unique prototypes learned. Through several evaluation experiments we further quantify the downstream impact of this regularization and how it varies for uniform and long-tailed pre-training datasets. We have now added section A.6 in the revised paper to provide an intuitive visual explanation of the partial prototype collapse problem and how the KoLeo-proto regularization impacts the prototypes and the clusters learned by the method.

We thank the reviewer for the useful feedback. We appreciate that the reviewer found our paper to be well presented in terms of writing and organization. We consider our identification and definition of the partial prototype collapse to be one of the key contributions of our work and we appreciate that the reviewer considers it as one of the strengths. Our secondary contribution is to propose KoLeo-proto regularization as a way to mitigate partial prototype collapse. We study the downstream impact of effectively utilizing the prototypes in different scenarios and report both the advantages and the limitations.

W1. The title of this work seems to explore the partial prototype collapse for all clustering-based SSL methods. However, without any clarification, all explorations and discussions are focused on the DINO family of SSL methods by default.

We agree that there is a slight disconnect between the paper title and the experiments in the paper. Note, however, that we consider the main contribution of the paper to be the definition and identification of the partial prototype collapse problem, and this is what the title refers to. This contribution applies to self-supervised clustering methods that use a prototypical formulation which includes SwAV, DINO, EsViT, MSN and iBOT. The DINO family of methods using ViT backbones is the predominant recipe among self-supervised clustering methods. This is our motivation for focusing on the DINO family in the paper. We have now furthermore investigated the SwAV and EsViT pre-trained models using various backbones and find that they also suffer from a partial prototype collapse (added to Table 9 in revised paper). For the SwAV Resnet50 backbone model, only 1669 prototypes are unique out of the 3000 initialized prototypes.

For the investigation of the effect of KoLeo prototype regularization (a secondary contribution) it is true that we focus on DINO and iBOT. The DINO family of methods is the prominent approach to self-supervised clustering in recent times. Several methods such as MSN, EsViT, iBOT are based on the DINO objective. The iBOT formulation is a strong recent baseline among the DINO family of methods and also used as the foundation for DINOv2. The vMF version is chosen because it stabilizes the iBOT training [1] (and provides improved performance). This is our motivation for choosing iBOT-vMF as the baseline for evaluating our second contribution of mitigating partial prototype collapse through KoLeo-prototype regularization.

[1] Govindarajan, Hariprasath, et al. "DINO as a von Mises-Fisher mixture model." ICLR 2023.

W2. The last sentence of the Sec. Conlusion claims the superior learning efficiency of the proposed method. However, there is no experiment to prove that.

There appears to be a misunderstanding regarding what we mean in the last two sentences of the conclusion section. We merely point out that the computational cost of training the DINO family of methods depends (in a non-negligible way) on the number of latent classes, K. Computing probability distributions over K latent classes involves a prototype layer of size $K \times 256$. The probability distributions in DINO are of the size $B \times K$ and in iBOT they are of the size $B \times T \times K$, where T is the number of tokens (T = 197 in the default configuration). The choice of K has a notable impact on the training times and memory use.

We find that around 1000 unique prototypes are used when the number of initialized prototypes are 2048, 4096, 8192 or 10240. For a batch size of 100, using 2048 vs 10240 prototypes consumes GPU memory of 21 GB vs 31.1 Gb respectively. We quantify this for different batch sizes and number of prototypes in section A.5.4 of revised paper. Considering the computational cost associated with using more prototypes, it is important to effectively utilize the initialized prototypes. This is made possible by our proposed regularization, which itself adds negligible computational cost in terms of memory and time.

When pre-training on another dataset, if a smaller number of prototypes are sufficient, then initializing the model with fewer prototypes and effectively utilizing them will bring computational gains compared to setting the number of prototypes to a larger number along with the occurrence of partial prototype collapse.

W3. Although several experiments justify the efficacy of the proposed method on long-tailed fine-grained datasets and few-shot settings, there is no convincing explanations.

Our intuitive explanation is that learning fine-grained clusters encourages the model to learn more informative–but also more low-level and dataset specific–features that are able to discriminate between the fine-grained clusters. This is intuitively useful for a fine-grained classification task, as well as for few-shot learning. However, some of these additional features may not generalize well to other datasets, thus hurting transfer performance.

W4. The proposed method is implicitly built on an assumption that improving the uniformity of the representation space while keep its alignment will produce a better sample distribution. However, without superior distribution priors, the trials of the proposed methods report limited performance improvements on standard SSL evaluations.

Sorry, we do not fully understand what the reviewer meant by “The proposed method is implicitly built on an assumption that improving the uniformity of the representation space while keep its alignment will produce a better sample distribution”. We can provide a better explanation if the reviewer can elaborate this.

We discuss two forms of regularization with different purposes - MLCD regularization and prototype regularization. MLCD regularization is used to prevent complete representation collapse. As a consequence of partial prototype collapse (PPC), the method can cheat and satisfy the MLCD constraint regardless of the choice of prior and MLCD regularization technique by allowing the prototypes to partially collapse. This results in fewer distinct latent classes for the self-supervision task. The prototype regularization is aimed at mitigating partial prototype collapse and learning diverse prototypes. The distribution of representations in the DINO family of methods is shown to correspond to that of a mixture of von Mises-Fisher components [1].

We study the downstream impact of effectively utilizing all the prototypes - we study uniform and long-tailed datasets under different tasks to understand both the advantages and limitations that arise from this. In the experiments reported in Table 2, we use the same number of initialized prototypes (K=8192) as the baseline. We do this to measure the performance improvement that can be obtained only from the effective utilization of the prototypes and not from making use of additional prototypes (which also changes the memory footprint of the model compared to the baseline; see Table 12 in revised paper for details on memory costs). As we note in section 6.3, we observe that the Imagenet kNN top-1 accuracy continues to improve as the number of prototypes are increased for iBOT-vMF with KoLeo-proto regularization, which is not the case for the baseline. The Imagenet classification performance with our proposed regularization can be improved further by scaling the number of prototypes. We do not investigate this for the longer experiments due to the significantly higher computational cost required for such experiments.

[1] Govindarajan, Hariprasath, et al. "DINO as a von Mises-Fisher mixture model." ICLR 2023.

We understand that it has only been a few days since our response. We would like to kindly remind the reviewer that the discussion period is going to end soon. We sincerely appreciate if the reviewer takes our responses into consideration.

Q2. How does the proposed regularization strategy perform in terms of robustness under different hyperparameters?

We mainly experiment with adding the KoLeo-proto regularization to the iBOT-vMF baseline method (with different ViT backbones). Additionally, we have now investigated DINO using a Resnet50 backbone. We also explore using different types of pre-training datasets - uniform and long-tailed. We use the default hyperparameter configurations for all the models, which have been identified by the original works after extensive experimentation. However, note that these hyperparameter configurations vary between ViT-S/16, ViT-B/16 and Resnet50 backbones. This includes differences in multicrop settings, epochs, learning rate, temperature schedule and weight decay. We find the proposed regularization to be robust over these hyperparameter choices.

Q3. The paper claims that the proposed method is more effective in long-tail classification tasks. Does this suggest that long-tail distributions are more likely to lead to Partial Prototype Collapse?

In our experiments, we observe partial prototype collapse in both uniform and long-tailed datasets, not suggesting that it is more likely in the latter. As brought up by reviewer hqpg, a minority collapse problem is identified in supervised long-tail classification [2]. However, the partial prototype collapse occurs in the self-supervised context among latent classes that are learned by the model. For a long-tailed dataset, the partial prototype collapse results in learning more coarse-grained latent classes where the number of unique prototypes are less than the number of classes. Note that this can also occur in fine-grained datasets with many classes. When models use fewer unique prototypes than the number of classes, the learned clusters are likely to have merged several of the fine-grained classes (or classes from the long tail). This can lead to less informative representations which cannot discriminate between fine-grained classes. This can disproportionately affect long-tailed or fine-grained classification datasets.

[2] Fang, Cong, et al. "Exploring deep neural networks via layer-peeled model: Minority collapse in imbalanced training." Proceedings of the National Academy of Sciences, 2021

Q5. Is it necessary to determine the initial number of prototypes before the algorithm is executed?

For existing methods in the DINO family, it is necessary to make this hyperparameter choice before training the method. It can be of interest as stated above to explore ways to avoid this. Previous works have relied on ablation experiments to select the number of prototypes but they were affected by partial prototype collapse. Our work sheds light on the importance and impact of this choice.

Q6. Some details of the experimental setup were not clarified. For instance, what optimizer was used?

We use the adamw optimizer for all the model trainings. The complete experimental setup including the optimizer settings are provided in section A.4 of the appendix.

Q7. This paper only utilizes various ViT-based architectures, does the Partial Prototype Collapse also observed in CNN-based architectures?

Yes, this is also observed in the CNN-based architectures, for instance, we observe partial prototype collapse in the DINO and SWaV models which use the Resnet50 backbone. We also investigate the EsViT models pre-trained using other Transformer variants such as ViL and CvT. This additional investigation is presented in section A.5.1 of the revised paper. We also explore the proposed KoLeo-proto regularization with the DINO method and a Resnet50 backbone and observe improved prototype utilization and performance (see section A.5.3 in the revised paper).

We thank the reviewer for the insightful feedback and valuable comments. We appreciate that the reviewer recognizes our identification of the Partial Prototype Collapse in self-supervised learning methods based on the clustering task. We address the weaknesses and questions raised by the reviewer below.

W1. The employed objective function is rather conventional, which may limit the overall innovation of the study.

We agree that the KoLeo objective for spreading vectors in a latent space [1] is not novel. However, we would like to point out that its usage in the context of learning diverse prototypes is new. Though this solution might be derivative, we believe that identifying and defining the partial prototype collapse issue is a key contribution that is of importance to this family of methods.

[1] Sablayrolles, Alexandre, et al. "Spreading vectors for similarity search." ICLR 2019.

W2. The experiments in the paper mainly compare the proposed strategy with the baselines, without a direct comparison with the existing regularization techniques.

The complete representation collapse is a known problem in self-supervised learning where the representations of all data samples converge to the same vector. This is addressed using various architectural tricks and regularization methods. In clustering-based self-supervised methods, the common way to prevent complete collapse is through MLCD regularization. We unify several existing MLCD regularization techniques in section 3 of the paper and follow-up with comparative experiments in section 6.1. However, the partial prototype collapse is a new problem that we have found in this family of self-supervised learning methods. Since this problem was not known previously, there are no existing regularization techniques to address this specific issue. Our proposed KoLeo-proto regularization can be used in conjunction with different MLCD regularization techniques.

W3. The impact of regularization parameter \lambda on the effectiveness of the proposed regularization method may not have been thoroughly explored in the paper.

We apologize for not including this ablation experiment in the earlier version of the paper. We conducted an ablation experiment on the regularization parameter and found that $\lambda=0.1$ is sufficient to utilize all the prototypes and increasing the strength further does not have any qualitative or quantitative improvements. The ablation experiment is added to section A.5.2 in the revised paper.

W4. The paper heavily relies on textual descriptions, with limited use of graphics to explain the issues or algorithms, which might pose difficulties for readers. AND

Q4. Employing visuals to illustrate Partial Prototype Collapse could enhance readers' comprehension.

We added section A.6 in the appendix to provide an intuitive explanations of our results:

• We provide t-SNE plots of all the learned prototypes to illustrate how the prototypes are impacted by partial prototype collapse and how KoLeo-proto regularization is able to spread the prototypes and hence, learn more diverse prototypes.
• We show t-SNE plots of the same subset of data points based on representations from the baseline iBOT-vMF method and after adding KoLeo-proto regularization (from the iNat18 experiment with ViT-S/16 backbone). We demonstrate that the KoLeo-proto regularization leads to more fine-grained cluster structure for the same data when compared to the baseline. We complement this by presenting some example images belonging to these clusters.

We hope that these newly added illustrations complement the textual explanations in the main paper and improve the readability of the paper.

Q1. How is the number of prototypes determined for different datasets?

Typically, the number of prototypes are determined through ablation experiments in prior literature. However, because of the partial prototype collapse, this hyperparameter previously did not perform its intended purpose of controlling the number of clusters learned by the method. Through KoLeo-proto regularization, we enable effective utilization of the initialized prototypes. Our work makes this hyperparameter relevant and sheds light on the impact of this choice. In our experiments, we use the same number of initialized prototypes as the baselines in order to only demonstrate the impact of their effective utilization (exact numbers used in our experiments are listed in the appendix section A.4.1).

Selecting the number of prototypes for a new unfamiliar dataset is a common challenge for this family of methods. One could design methods that can also learn the optimal number of clusters for a dataset. This is a difficult problem since we are learning representations and the clustering model of the representations simultaneously. We consider this as an interesting direction for future research.

We thank the reviewer for the constructive feedback and valuable comments. We address the weaknesses and questions raised by the reviewer below.

W1. The proposed method mainly works on long-tailed datasets, however, its motivation does not rigorously align with its goal. Although the authors mention that previous works encourage MLCD to match a prescribed uniformed prior distribution, why could KoLeo-proto regularization move away from such a conventional prior? More elaborations are needed to shed light on how the KoLeo-proto regularization could help with the long-tailed problem.

We do not in fact move away from encouraging the MLCD to match a prescribed prior distribution. This could be a uniform (considered by most previous work) or non-uniform (considered by, e.g., PMSN). As a consequence of partial prototype collapse (PPC), the method can cheat and satisfy the MLCD constraint regardless of the choice of prior and MLCD regularization technique by allowing the prototypes to partially collapse. KoLeo-proto regularization is a way to prevent this type of collapse, removing the possibility of “cheating” in this way, which makes the MLCD regularization more effective. This is true for any choice of prior and MLCD regularization technique. It is true that we focus on the common setting with a uniform prior, together with probability centering as MLCD regularization, but the PPC issue is not limited to this case. We find in Table 8 that by choosing a better KL penalty strength, we can also achieve better performance on long-tailed data by encouraging MLCD to match a uniform prior rather than a long-tailed prior (comparing MSN and PMSN). This is our motivation for focusing on the uniform prior in our experiments. Also, note in Table 4 that both MSN and PMSN learn ~3000 unique prototypes out of the 8142 initialized prototypes, despite using different priors (uniform in MSN and long-tailed in PMSN).

We provide a simple example to explain how a model can “cheat” through partial prototype collapse to match the MLCD to a prescribed uniform prior. Consider a dataset of size D and 4 initialized prototypes. Then, consider the following sets of probability distributions corresponding to the D samples when the number of unique prototypes learned by the model are 4 and 2:

1. 4 unique prototypes: $[v_1, v_2, v_3, v_4]$
   • D/4 samples: $[1.0, 0.0, 0.0, 0.0]$
   • D/4 samples: $[0.0, 1.0, 0.0, 0.0]$
   • D/4 samples: $[0.0, 0.0, 1.0, 0.0]$
   • D/4 samples: $[0.0, 0.0, 0.0, 1.0]$
2. 2 unique prototypes: $[v_1, v_1, v_2, v_2]$
   • D/2 samples: $[0.5, 0.5, 0.0, 0.0]$
   • D/2 samples: $[0.0, 0.0, 0.5, 0.5]$

Both configurations of probability distributions achieve a uniform MLCD. However, they use different numbers of unique prototypes. Classifying the data samples into more latent classes is a harder task, as is the case in the first configuration. This encourages the model to learn more informative representations to discriminate between the fine-grained latent classes. We achieve this using the KoLeo-proto regularization that ensures effective utilization of the prototypes.

For a long-tailed dataset, the partial prototype collapse results in learning more coarse-grained latent classes where the number of unique prototypes are less than the number of classes. We confirm this to be the case using the iNaturalist-2018 dataset (see Table 4; column M in the paper). When models use fewer unique prototypes than the number of classes, the learned clusters are likely to have merged several of the fine-grained classes (typically, classes from the long tail). This is mitigated when the prototypes are effectively utilized, leading to more diverse clusters and hence, more informative representations.

Q1. This paper mainly takes iBoT-vMF as its baseline. Considering the simple and elegant form of KoLeo-proto regularization, would it also benefit other baselines either in qualitative (i.e. the uniqueness of prototypes) or quantitative (i.e. the performance) ways?

Yes, the KoLeo-proto regularization can be easily added to several prototypical formulations - which is the case for SWaV and the DINO family of methods including EsViT, iBOT and MSN. We focus our experiments on iBOT which is a recent state-of-the-art method in the DINO family. The vMF extension stabilizes the iBOT training and provides a further performance boost [3]. Even the most recent DINOv2 work builds on top of iBOT. This is our motivation for our choice of iBOT-vMF as the baseline. SWaV, EsViT, iBOT and MSN also suffer from a partial prototype collapse (see Table 1 for iBOT, MSN. See below for EsViT and SWaV; added to Table 9 in the revised paper) and this can be mitigated using the KoLeo-proto regularization. We conducted an additional experiment exploring KoLeo-proto regularization on DINO and used the Resnet50 backbone (to explore a CNN backbone, complementing our previous ViT experiments). The results are in agreement with the results observed for ViT backbones (see below; added to section A.5.3 in the appendix).

[3] Govindarajan, Hariprasath, et al. "DINO as a von Mises-Fisher mixture model." ICLR 2023.

Number of unique prototypes in existing models

ImageNet classification with full data (kNN, linear) using DINO - Resnet50 backbone model

W2. The KoLeo-proto regularization mainly aims to improve the diversity of the learned clusters while reducing partial prototype collapse. Another stream of study [R1], which explores the neural collapse phenomenon that also scatters the prototypes into an ETF geometric structure, also proves to be effective for imbalanced learning. Are there any similar spirits between this KoLeo-proto regularization and the neural-collapse-inspired design? [R1] Inducing Neural Collapse in Imbalanced Learning: Do We Really Need a Learnable Classifier at the End of Deep Neural Network, NeurIPS, 2022.

We thank the reviewer for bringing this work [1] to our notice, where the last linear layer of a supervised classifier is replaced with prototypes initialized using an ETF geometric structure which are not trained. If we apply the KoLeo-proto regularization with a large strength, we would learn a similar ETF structure. However, this makes the implicit assumption that all classes are equally similar/dissimilar to one another, which is a strong assumption. Such a hard constraint may not be a suitable choice in the self-supervised setting, where the objective is to learn good representations. It should be possible for semantically similar clusters to be located nearer in the latent space instead of being separated by a fixed distance. Hence, we argue that a softer penalty through the KoLeo-proto regularization can suit the self-supervised learning scenario better than using a fixed ETF.

Our explanation above (in response to the previous question) for why effective utilization of the prototypes is beneficial for long-tailed datasets is partly related to the minority collapse issue explored in [2]. In the supervised case, [2] finds that the minority classes were merged with one another, thus increasing the similarity between representations of different minority classes. In the self-supervised learning case, since we do not use class labels during pre-training, this happens at a latent class level and the minority class samples can be merged with samples from any other class (head/mid/tail). Also note that PPC is not just limited to long-tailed data (also shown to occur in ImageNet).

[1] Yang, Yibo, et al. "Inducing Neural Collapse in Imbalanced Learning: Do We Really Need a Learnable Classifier at the End of Deep Neural Network?." NeurIPS, 2022.

[2] Fang, Cong, et al. "Exploring deep neural networks via layer-peeled model: Minority collapse in imbalanced training." Proceedings of the National Academy of Sciences, 2021.

W3. The KoLeo-proto regularization does not seem to improve the transfer performance (-0.7% on ViT-base and +0.1% on ViT-small ), as pointed out by the authors in Table 3. The authors explain that it is a common trade-off also existing in other works whose performance thrives in few-shot scenarios but falls short in transfer learning. Unfortunately, there are no further explanations for why this trade-off exists. Also, it seems that KoLeo-proto generally gains negligible improvement on finetuning tasks, as also indicated in Table2, 4. It would be better to elaborate on the results.

This trade-off is a more generic issue not specific to our work, and although we found it interesting to observe that the tradeoff is present also in the current work, we consider a more in-depth investigation to be out-of-scope for this paper. We have not investigated this in sufficient depth to provide a clear explanation. A possible reasoning for reduced transfer performance is that the features learned to discriminate between the fine-grained clusters (which is encouraged by KoLeo-proto regularization) become more specific to the pre-training dataset. These features may not generalize well to other datasets, resulting in reduced transfer performance.