P$^2$OT: Progressive Partial Optimal Transport for Deep Imbalanced Clustering

We appreciate the diligent assessment by reviewer qtAk, as well as the recognition of our technical contributions and our proposed methods. We hope the concerns raised could be addressed by the following response.

The technical choice about Eq.(13) is unique.

It is possible to apply the Lagrange multiplier to impose a hard equality constraint. But resulting OT problem can not be easily solved by fast scaling algorithm, making it less practical for large-scale application problems. By contrast, we introduce a virtual cluster and weighted KL constraint to transform Eq.5-6 into a form similar to Eq.1, which can be solved by the fast scaling algorithm.

Regarding the additional hyperparameter, we assign $\lambda_{:K}$ as 1 for all datasets without any tuning. And for $\lambda_{K+1}$, we believe that setting it to a sufficiently large value is adequate. In practice, directly assigning the $\frac{\lambda_{K+1}}{\lambda_{K+1}+\epsilon}$ as 1 in Proposition 2 without any further adjustment proves effective.

Comparing with OT

We thank the reviewer qrAk for pointing out this problem. We have added the results of OT[1] to Table 2, and our method shows considerable improvement over naive OT [1].

As for the comparison of SPICE on balance datasets, please refer to General Response.

[1] Yuki M. Asano, Christian Rupprecht, and Andrea Vedaldi. Self-labelling via simultaneous clustering and representation learning. In International Conference on Learning Representations (ICLR), 2020.

Comparison with representation learning.

Please refer to General Response.

We appreciate the thorough evaluation by reviewer xbkw and express our sincere gratitude for xbkw's interest in the problem and the method we proposed. We hope the following response could address the concerns raised.

Great computation cost

The unit of the horizontal axis is K (kilo), and the analysis is conducted on iNature1000. Thus, when there are 40000 points and 1000 clusters, it takes only 1 second, which is practical for mini-batch training in most cases.

Comparisons on Balanced Dataset

Please refer to the general response.

This paper does not learn representations.

We learn the representation. Specifically, we employ an ImageNet-1K unsupervised pre-trained ViT-B16 model and fine-tune the last block, which strikes a balance between complexity and flexibility. We have updated Appendix F to make it clear.

We thank the reviewer's suggestion. We note that K-means clustering is also affected by data imbalance, leading to inferior results and we provide detailed experimental analysis in the second General Response.

Evaluation Metric

We appreciate the importance of a good evaluation and recognize that it is an open question in imbalanced deep clustering. However, we disagree that a uniform clustering result is beneficial for evaluation. As evident from the results below, the naive OT, which enforces equality uniform distribution constraints, yields inferior results.

In addition, regarding Hungarian matching, it is true that on imbalanced training sets, tail classes may be less important during the matching process. However, when evaluating on balanced test sets (Tables S8-S10), all classes are equally important during Hungarian matching. Therefore, we can mitigate the bias of Hungarian matching by evaluating on balanced test sets, where we still observe significant improvements.

We appreciate reviewer xbKw's suggestion, but we believe that kNN is not a suitable evaluation metric for imbalanced clustering due to the following reasons:

1. kNN evaluation requires label information from the training set, which is not available in deep clustering.
2. kNN is usually an evaluation of representation, while deep clustering learns not only representation but also clustering head.

In conclusion, we suggest evaluating on balanced test set to alleviate the bias of Hungarian matching and evaluate the generalization of different methods.

Huang et al. achieve uniform constraints by maximizing the inter-cluster distance between prototypical representations to avoid degenerate solutions. Unlike [1], we achieve that by utilizing the KL constraint in the pseudo-label generation process. We have updated Appendix D to discuss more on [1] with respect to uniform constraints.

[1] Zhizhong Huang, Jie Chen, Junping Zhang, and Hongming Shan. Learning representation for clustering via prototype scattering and positive sampling. IEEE Transactions on Pattern Analysis and Machine Intelligence, 2022.

Questions

1. We have corrected the typo.
2. The DINO is pre-trained on ImageNet-1K, therefore we have not adopted the ImageNet-1K dataset as a training set. In the training, we only fine-tuned the last transformer block of ViT-B16.
3. Please refer to General Response Part 2.
4. Figure 2 is the visualization of the feature before the cluster head. We have updated the manuscript to make it clear.

REUTERS-10K dataset

As shown in Table 1 of [1], the REUTERS-10k dataset contains only 10,000 points and 4 classes, which is significantly smaller than the ImageNet-R and iNature500 datasets used in our experiments. Additionally, it is a dataset intended for document clustering, whereas our manuscript focuses on image clustering.

[1] Xie, Junyuan, Ross Girshick, and Ali Farhadi. "Unsupervised deep embedding for clustering analysis." International conference on machine learning. PMLR, 2016

Why it can be called an unbalanced OT problem?

When $F_1$ and $F_2$ are equality constraints, Eq. 1 is typically referred to as "classical" optimal transport. To generalize "classical" optimal transport, unbalanced OT [1,2,3] defines the problem on arbitrary positive measures and some divergence functions. In our manuscript, we refer to Eq. 3 as an unbalanced OT in order to distinguish it from "classical" optimal transport.

[1] Lenaic Chizat, Gabriel Peyré, Bernhard Schmitzer, and François-Xavier Vialard. Scaling algorithms for unbalanced optimal transport problems. Mathematics of Computation, 87(314):2563–2609, 2018.

[2] Matthias Liero, Alexander Mielke, and Giuseppe Savar ́e. Optimal entropy transport problems and a new Hellinger–Kantorovich distance between positive measures. Inventiones Mathematicae, pages 1–149, 2015.

[3] Séjourné, Thibault, Gabriel Peyré, and François-Xavier Vialard. "Unbalanced Optimal Transport, from theory to numerics." Handbook of Numerical Analysis 24 (2023): 407-471.

Mini-Batch

Please refer to General Response 3.

NMI explanation

Our method does achieve better embedding but due to the calculation of NMI, we achieve modest improvement.

Baselines tend to merge tail classes into head classes, whereas our approach tends to divide head classes into multiple clusters, resulting in a more even class distribution (Figure S7). The formula for the Normalized Mutual Information (NMI) is $\frac{2I(Y;C)}{H(Y)+H(C)}$, where Y represents the ground truth and C denotes the cluster label. Despite our method achieving superior clustering performance, as indicated by a higher $I(Y;C)$, it also results in a relatively uniform distribution, thereby increasing $H(C)$. This leads to only a modest improvement or decrease in the NMI metric.

Explanation of Figure 4

In Figure 4, we show the weighted precision and recall based on our selection strategy, which generates a weight, $Q1_K$, for samples. As training progresses, we increase $\rho$ using a sigmoid ramp-up strategy and select more samples, which inevitably includes more noisy samples, leading to a decrease in weighted precision and recall after 10 epochs. However, the precision and recall of the entire dataset are steadily increasing, illustrating our method learns meaningful clusters gradually. We have analyzed this phenomenon in the ablation study and mentioned that the current ramp-up function may not be optimal, leaving more advanced designs for future work.

We appreciate the valuable feedback and thoughtful comments from reviewer H8zS. We have diligently analyzed the suggestions and concerns raised, and we hope the following response could address them. However, we respectfully disagree with the assessment that "The technical contributions are overclaimed.".

Some claims ... are confusing ...

We appreciate that reviewer H8zS points out the confusing statements. We clarify that our method with $KL$ constraint can generate imbalanced pseudo-labels instead of KL can generate imbalanced pseudo-labels. It's worth noting that while the KL constraint is defined for a uniform distribution, it serves as a relaxed constraint compared to the typical equality constraint used in [1]. Consequently, the optimal Q in our P$^2$OT should take into account both the inherent imbalanced class distribution in P and a prior uniform distribution. This enables our method to:

1. generate pseudo labels reflecting imbalanced characteristics akin to P
2. avoid degenerate solutions where all samples are assigned to a single cluster.

For clarity, we have updated the confusion statement in the paper as follows:

1. The sentence "the KL divergence-based uniform distribution constraint empowers our method to avoid degenerate solutions and generate imbalanced pseudo-labels" should be changed to: "Notably, the $KL$ divergence-based uniform distribution constraint empowers our method to avoid degenerate solutions. And our approach with $KL$, which represents a relaxed constraint compared to an equality constraint, facilitates the generation of imbalanced pseudo-labels."
2. The sentence "demonstrating that KL constraint enables our P$^2$OT to generate imbalanced pseudo labels." should be changed to: "demonstrating that our P2OT with $KL$ constraint can generate imbalanced pseudo labels."

[1] Yuki M. Asano, Christian Rupprecht, and Andrea Vedaldi. Self-labeling via simultaneous clustering and representation learning. In International Conference on Learning Representations (ICLR), 2020.

Eq.5 - Eq.6 are not necessary.

We respectfully disagree with Reviewer H8zS's assessment of this weakness for the following reasons:

1. Eq.5-6 formulates the confident sample selection and imbalanced pseudo-label generation as a novel partial optimal transport problem, which is the central component of our method.

2. Virtual clusters are commonly employed in addressing partial optimal transport (POT) problems. However, Eq.3-4 does not constitute a POT problem. Therefore, introducing Eq.8 directly after Eq.3-4 through the introduction of a virtual cluster makes it challenging to interpret and understand.

In conclusion, we believe that Eq.5-6 is necessary.

[1] Luis A Caffarelli and Robert J McCann. Free boundaries in optimal transport and monge-ampere obstacle problems. Annals of Mathematics, pp. 673–730, 2010.

[2] Laetitia Chapel, Mokhtar Z Alaya, and Gilles Gasso. Partial optimal transport with applications on positive-unlabeled learning. Advances in Neural Information Processing Systems, 33:2903–2913, 2020.

Thank you so much for taking the time to review our submission. We appreciate your detailed feedback and appraisal of our work, which we have taken into careful consideration in our rebuttal response. As the rebuttal process is coming to an end, we would be grateful if you could acknowledge receipt of our responses and let us know if they address your concerns. We remain eager to engage in any further discussions.

Different representation learning baselines

We appreciate the reviewer's suggestion and conduct additional experiments with the representation learning baselines, including K-means clustering on DINO's representation and [2].

Note that we exclude reporting results from [1] and [3] for the following reasons:

1. [1] prunes the ResNet, making it challenging to transfer to ViT
2. [3] has not yet open-sourced the code for kernel density estimation, which is crucial for its success.

For a fair comparison, we adopt the pretrained ViT-B16 and finetune the last block of the ViT-B16 with the latest BCL [2]. The results of BCL and the original DINO are presented below.

For both BCL and DINO, which focus on representation learning, we apply K-means to cluster the imbalanced training set in the representation space. The results indicate the following:

1. DINO achieves satisfactory results, particularly on CIFAR100, where the dataset is akin to the ImageNet-1K used for pre-training
2. BCL exhibits marginal improvement or even a decrease compared to DINO, as DINO serves as a robust unsupervised pre-trained model
3. Our method significantly outperforms baseline and unsupervised pre-training models across a wide range of metrics and datasets.

These results demonstrate the effectiveness of our approach in handling imbalanced data.

We have appended the results to Appendix J and updated the self-supervised long-tailed learning into related work.

[1] Jiang, Z., Chen, T., Mortazavi, B.J. and Wang, Z., 2021, July. Self-damaging contrastive learning. In International Conference on Machine Learning (pp. 4927-4939). PMLR.

[2] Zhou, Z., Yao, J., Wang, Y.F., Han, B. and Zhang, Y., 2022, June. Contrastive learning with boosted memorization. In International Conference on Machine Learning (pp. 27367-27377). PMLR.

[3] Liu, H., HaoChen, J.Z., Gaidon, A. and Ma, T., 2021. Self-supervised learning is more robust to dataset imbalance. In International Conference on Learning Representation, 2022.

[4] Mathilde Caron, Hugo Touvron, Ishan Misra, Hervé Jégou, Julien Mairal, Piotr Bojanowski, and Armand Joulin. Emerging properties in self-supervised vision transformers. In Proceedings of the IEEE/CVF international conference on computer vision, pp. 9650–9660, 2021.

Implementation of Mini-Batch OT

To address the limitation of the mini-batch, we introduce a memory buffer that stores a substantial number of sample predictions (e.g., 5120) to ensure the existence of minority clusters. In particular, before inputting data into P$^2$OT in each iteration, we concatenate predictions from the memory buffer with the current batch predictions to enhance stability while preserving efficiency. Additional details can be found in the last sentence of Section 4.3 and Appendix F.