Paper: P$^2$OT: Progressive Partial Optimal Transport for Deep Imbalanced Clustering (ICLR-2024)

Reviewer comment (Experiments):
The technical choice about Eq.(13) is unique. The other choice like introducing a hard equality constraint about the mass of the virtual cluster and applying the lagrange multiplier can be also possible. The main concern here is that the authors introduce two hyperparameters for weighted KL: one is for the target clusters, i.e., \lambda, and the other is the large value for the virtual cluster. How is the performance of directly constraining the equality about 1-\rou with lagrange multiplier compared with Eq.(13).

Author 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.

Reviewer comment (Experiments):
It is not clear why the authors have not compared with the clustering with self-labeling by Asano from the perspective of OT. For SPICE, the performance reported in the appendix is also not convincing. How is the performance comparison of SPICE and P^2OT on the balanced datasets. It will be provide more comprehensive comparison here on both imbalanced datasets and balanced datasets to show their pros and cons.

Author response:
## Comparing with OT
| Formulation | CIFAR100 | ImgNet-R | iNature500 |
|:-----------:|:--------:|:--------:|:----------:|
|             |    ACC   |    ACC   |     ACC    |
|      OT     |   28.2   |   18.2   |    24.1    |
|   P$^2$OT   |   38.2   |   25.9   |    32.2    |

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.

Reviewer comment (Experiments):
Great computational cost. In Figure 5, it takes 1 second to estimate the pseudo-labels when K=40. It is impossible in practice with a larger number of clusters, eg, imagenet, and mini-batch training.

Author response:
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.

Reviewer comment (Experiments):
This paper does not learn representations, which may be confusing and needs to be clarified. In addition, this paper uses a large pre-trained model and lacks a simple baseline with representation learning models. For example, BYOL or DNIO pre-trained on ImageNet can be used to extract the representations for subsequent K-means clustering. In such settings, we do not need to determine whether the data is imbalanced, as each sample belongs to a single class or a large number of clusters can be pre-defined.

Author response:
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.

Reviewer comment (Evaluation):
A good evaluation metric should be important for imbalanced deep clustering. Under the imbalanced setting, no samples may be assigned to tail classes during Hungarian matching. As we can see in Figure S7, the predictions of the proposed method are more uniform than baselines. There are more samples assigned to tail classes, though this is not true in training data. It confirms that a uniform clustering result is beneficial for evaluation. I suggest that kNN evaluation in representation learning can be adopted for imbalanced clustering, without the need of Hungarian matching. Due to the uniform constraints, more discussion should be paid to Huang et al., 2022.

Author response:
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.

| Formulation | CIFAR100 | ImgNet-R | iNature500 |
|:-----------:|:--------:|:--------:|:----------:|
|             |    ACC   |    ACC   |     ACC    |
|      OT     |   28.2   |   18.2   |    24.1    |
|   P$^2$OT   |   38.2   |   25.9   |    32.2    |


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.

Reviewer comment (Presentation):
What is visualized in Figure 2? We usually visualize the features before the classifier instead of class predictions. If the results are consistent for the features, we can conclude: more distinct clusters.

Author response:
Figure 2 is the visualization of the feature before the cluster head. We have updated the manuscript to make it clear.

Reviewer comment (Theory):
Some claims in this paper are confusing and need improvement. For instance, (i) "the KL divergence-based uniform distribution constraint empowers our method to avoid degenerate solutions and generate imbalanced pseudo-labels", "demonstrating that KL constraint enables our P2OT to generate imbalanced pseudo labels." Why can the KL constraint generate imbalanced pseudo-labels, given that it is defined for a uniform distribution?

Author response:
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.

Reviewer comment (Theory):
I don't think Eq. 5-Eq. 6 are necessary. You can consider introducing Eq. 8 directly after Eq. 3 by introducing a virtual cluster. If I am wrong, please point it out.

Author response:
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.

Reviewer comment (Experiments):
A very popular real imbalanced dataset is REUTERS-10K[1], which is more challenging than the constructed dataset. How does the proposed method perform on REUTERS-10K?

Author response:
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

Reviewer comment (Theory):
In Eq. 3, the KL constraint is defined for the uniform distribution. Why it can be called an unbalanced OT problem?

Author response:
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.

Reviewer comment (Evaluation):
Figure 2 shows that the embedding of P2OT is well-separated compared to that of other baselines. Can the authors explain why the embedding is much better than others, given that it is only slightly higher than SCAN in terms of NMI?

Author response:
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.

Reviewer comment (Evaluation):
Figure 4 shows that the performance degenerates after 10 epochs. Can you explain the reason behind this?

Author response:
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.