Paper: OTMatch: Improving Semi-Supervised Learning with Optimal Transport (ICLR-2024)

Reviewer comment (Novelty):
The novelty is limited. Using OT strategy to address semi-supervised learning has been proposed in previous work [1]. The proposed method does not show significant improvement compared to existing work and lacks the necessary experimental comparisons.

Author response:
Thank you for your comments. **While OT was not the first time being used in semi-supervised learning, there are significant differences between our method and previous work [1], as we pointed out in the related work section in the initial manuscript.** We would like to provide a more detailed discussion as follows. First, we leverage OT theory to improve the loss function, differing from [1], which utilizes OT for refining pseudo-label assignments, reflecting distinct objectives. Second, our method has computational efficiency and ease of implementation, in contrast to [1], which introduces computational burdens. Third, [1] lags behind in accuracy compared to modern threshold-based pseudo-labeling strategies, whereas our method achieves SOTA results on benchmarks among threshold-based pseudo-labeling methods. These distinctions underscore the novelty and superiority of our approach, and we have incorporated the above discussion into the revised version of our paper.

Reviewer comment (Evaluation):
The improvement of proposed method on the evaluated dataset is not significant, and there is no comparison regarding its impact on convergence speed and the loss curve.

Author response:
Thank you for your comments. **We would like to emphasize that the performance improvement achieved by combining our OT loss with the pseudo-labeling approach in semi-supervised learning is absolutely noteworthy.** This is particularly evident in scenarios where the number of labeled samples is extremely limited. Compared to the state-of-the-art method FreeMatch, we achieve a 3.18% and 3.46% reduction in error rate on CIFAR-10 with 1 label per class and STL-10 with 4 labels per class. Our improvement is achieved solely through a simple optimal transport loss, with almost no additional training cost introduced.

We also provide training details about OTMatch to enhance understanding. The core of our method lies in learning a cost matrix. Therefore, at the initial stages of training, OTMatch achieves performance similar to FreeMatch since the cost matrix initialized may lack prior information about classes. As training progresses and the cost matrix continuously updates, typically in the later stages around the predefined 1.04 million iterations, OTMatch begins to show the potential to outperform FreeMatch. 
We have provided a detailed showcase of how the optimal performance of FreeMatch and OTMatch evolves on ImageNet after one million iterations, as follows.
|  Iterations | 978000  |1011000|1013000|  1027000  | 1043000  |
|  - | -  |  -  | -  | -  | ----  |
| FreeMatch Best Accuracy  |0.5795|0.58104|0.58364|0.59026|0.59428|

|  Iterations | 976000  |1007000|1013000|1022000| 1030000|1035000 | 1038000  |1040000  |
|  -  | -  |  -  | ----  | ----  | ----  | ----  | ----  | ----  | 
| OTMatch Best Accuracy|0.57792|0.58364|0.58714|0.59262|0.59938|0.60188|0.60484|0.6071|

It can be observed that OTMatch achieves performance comparable to FreeMatch before one million iterations but gradually surpasses FreeMatch later.

Additionally, we conduct a more in-depth investigation, specifically examining the impact of discarding the cost matrix midway through training. In particular, we experiment with training OTMatch on ImageNet for one million iterations, followed by removing the OTLoss directly or resetting the cost matrix to ones. Then we observe how their performance changes after the next 1000 iterations. The results are shown below.

|  | original  |remove|reset|
|  -  | -  |  -  |  -  |
| Accuracy|0.57832|0.56125|0.54927|

It is evident that discarding or resetting a learned cost matrix significantly affects performance. This demonstrates that the cost matrix has indeed learned crucial information about inter-class relationships, highlighting the effectiveness of our approach.

Reviewer comment (Experiments):
Ablation study on \(\epsilon\) is missing? How did you determine the appropriate \(\epsilon\) value for OT?

Author response:
Thank you for your question. We want to clarify that we do not use the entropic regularized method to compute OT. Instead, for semi-supervised learning, the initial (unregularized) OT loss has a closed-form solution which has O(K) running time faster than the Sinkhorn algorithm for entropic regularized OT. **Thus, we employ the closed form in our method and do not need to determine $\epsilon$.**

Reviewer comment (Reproducibility):
I strongly recommend presenting a structured algorithm or providing a comprehensively defined overall loss function for better clarity.

Author response:
**Thank you for your suggestion and we have added an algorithm box to our revised paper in Appendix C.**

Reviewer comment (Writing):
In the text following the second formulation, "q_ui" should be corrected to "Q_ui."

Author response:
Thank you for pointing out the typo. **We have corrected it in our revised paper.**

Reviewer comment (Presentation):
The depiction in Figure 1 is ambiguous. What constitutes the input for the OT loss?

Author response:
Thank you for your feedback. As we depicted in the initial manuscript, the input probabilities are the one hot pseudo-label from a weakly augmented image and the probability prediction of a strongly augmented image. **As OT needs a cost, we use the classification head vectors to calculate the cost.** You can see more clearly from our algorithm.

Reviewer comment (Experiments):
How can we obtain the ground truth matching matrix \( T \) in formulation (1)?

Author response:
Thank you for your feedback. **As we depicted in the initial manuscript, when we tackle labeled samples, we can organize the labeling function as an empirical matching matrix where each row depicts a labeled sample and the only non-zero column in this row is its true label.**

Reviewer comment (Theory):
Why was KL divergence chosen as the loss function?

Author response:
Thank you for your feedback. The reasons are two-fold. One is the empirical matching matrices and the ground truth matrix can be seen as a probability distribution, KL divergence is a typical way to measure differences between probabilities. **The other is that we adopt an (inverse) optimal transport way to understand FreeMatch, where in the definition of inverse optimal transport, it adopts KL divergence.**

Reviewer comment (Writing):
I am unclear about the statement in the paper that suggests the objective is "to push \( f(x) \) in the direction of \( wk \)." Can this be clarified?

Author response:
Thank you for your feedback. Recall that in calculating logits from representation, logits are the inner product of representation with each classification head. **If the sample is from class k, a very confident classification will be logit k is far larger than other logits.** As logit k is the inner product of representation $f(x)$ and $w_k$, making this inner product larger geometrically means that $f(x)$ is closer to the direction of $w_k$.

Reviewer comment (Theory):
Why the class relationship is helpful for semi-supervised learning? Is there any analysis or discussion?

Author response:
Thank you for your question. Semi-supervised learning usually has very few labelled samples compared to the unlabeled ones. When there are only very limited annotated data, the predicted pseudo-label will be less reliable. More importantly, the pseudo-label method will cause a “overconfidence” issue [r1], which means the model will usually fit on the confident but wrong pseudo-labels, resulting in poor performance. In pseudo-label based semi-supervised learning, cross-entropy loss is usually used, but it only considers pushing the prediction to their respective pseudo-labels and does not take the relationship between classes into account. Thus using cross-entropy loss alone will make the model fit on wrong label. **By introducing optimal transport loss that takes the class-relationship into account, we expect it to overcome this over-confidence drawback of cross-entropy loss in semi-supervised learning.** Intuitively, even when the pseudo label is wrong, as the wrong classified classes usually have close semantic relationship with the correct class, introducing an optimal transport loss will use this semantic relationship to update the cost and may help mitigate the over-confidence effect.

[r1]: Semi-Supervised Learning with Multi-Head Co-Training, Mingcai Chen, Yuntao Du, Yi Zhang, Shuwei Qian, Chongjun Wang, AAAI 2022

Reviewer comment (Experiments):
If the class relationship is important for semi-supervised learning, is there any easier method to exploit the relationship?

Author response:
Thank you for your question. **We consider a method "easy" based on its simplicity, involving small complexity (both in terms of time and space), as well as ease of implementation. Regarding time complexity, our optimal transport loss achieves O(K), representing a relatively low computational cost.** This ensures the efficiency of our method theoretically, a feature lacking in some related works such as FreeMatch [r2]. Moreover, we conduct additional experiments in Appendix D.3 on real running time and observe that our method indeed has a very close runtime to previous related works. (i.e., the running time increases by 0.02s compared to FreeMatch during every iteration). Concerning space complexity, our method introduces only a constant additional memory overhead, resulting in O(1) space complexity. In terms of implementation, our method leverages semantic information among classes with just one line of code (see Appendix C), making it straightforward to implement. In conclusion, our method is designed to be sufficiently easy to effectively exploit class relationships in semi-supervised learning.

[r2] FreeMatch: Self-adaptive thresholding for semi-supervised learning.

Reviewer comment (Theory):
The proposed unsupervised loss is difficult to compute. Although the authors claim in some conditions, the computational complexity can be reduced, whether these conditions are satisfied in real-world tasks is hard to know.

Author response:
Thank you for your feedback. **It's important to clarify that, in all our experiments where pseudo-labels are one-hot vectors, the O(K) time complexity of our proposed method is strictly satisfied without any additional conditions.** This is supported by lemma 4.1, the condition is met when one probability vector has all but one of its components smaller than the corresponding components of another probability vector. This theoretical guarantee ensures reduced computational complexity in all real-world tasks considered in this paper. On the other hand, the calculation of the proposed loss only involves a single line code (see Appendix C), making it easy to implement at minimal cost. Moreover, additional experiments highlight the small running time of our method (see above A2). Therefore, the proposed loss is not only easy to implement but also theoretically efficient.

Reviewer comment (Evaluation):
When there are many classes (such as ImageNet or CIFAR-100), the performance improvement is limited.

Author response:
Thank you for your comments. First of all, our method achieves SOTA performance across all benchmarks, ranging from small to large datasets. The consistent improvements observed on all datasets strongly emphasize the superiority of our method over previous works.
Secondly, in terms of ImageNet top-1 accuracy, our method exhibits a 1.28% improvement over the current SOTA method FreeMatch -- equivalent to FreeMatch's improvement over the previous SOTA method FlexMatch. **Notably, such gains on ImageNet are significant, not only due to the dataset's scale but also because our method is both fast and easy to implement.**
Thirdly,  our loss could be seamlessly inserted into pseudo-label based algorithms as a universal add-on to improve performance (for example FlexMatch, see Appendix D.1). This indicates that our method is ready for integrating with future SOTA methods to achieve additional improvement.
In summary, while the performance improvement on datasets with many classes may appear smaller than those with a smaller number of classes, our method's performance improvement on datasets with many classes is not limited compared to other strong baselines.

Reviewer comment (Experiments):
As the authors stated, the O(K) complexity of the proposed method comes from the mild assumptions. Providing actual runtime would be helpful to justify this statement further.

Author response:
Thank you for your valuable feedback. **It's crucial to highlight that, in all our experiments where pseudo-labels are one-hot vectors, the O(K) time complexity of our proposed method strictly holds without any assumptions.** This theoretical guarantee of time complexity distinguishes our approach from others, such as FreeMatch, which lack such assurances. Additionally, in response to your query, we also conduct experiments to measure the actual runtime of our method and results are in Appendix D.3. **We observe that the running time increases by 0.02s compared to FreeMatch during each iteration.** Our findings indicate that the method's runtime closely aligns with that of Freematch, providing practical evidence supporting the efficiency of our method.

Reviewer comment (Experiments):
What's the loss weight of the proposed loss term, How is it affecting training? Ablation study of it is missing.

Author response:
Thank you for your detailed questions. In our experiments, we set the loss weight as 0.5. **To address your inquiry, we conducted additional experiments specifically to assess how it impacts the model's performance.** The results are presented in the table below, where the loss weight is denoted by $\lambda$.
| $\lambda$ | 0.3 | 0.4 | 0.5 | 0.6 | 0.7 |
| --- | --- | --- | --- | --- | --- |
| Error rate | 4.91 | 4.95 | 4.72 | 4.97 | 4.98 |

From the table, we observe that when lambda is set to 0.5, our methods achieves the best performance. As a reference, other related works such as FreeMatch, the value of the loss weight is much more sensitive to the final performance. We have incorporated this ablation study into the revised version of our paper's Appendix D.2.

Reviewer comment (Novelty):
The contrastive learning-based methods such as SimMatch and CoMatch also consider the relation between classes. What's difference of optimal transport to this?

Author response:
Thank you for your question. While both our OTMatch and contrastive learning-based methods consider the relationship between classes, there are some crucial distinctions. **Our OTMatch focuses on aligning the class classification probabilities of two augmented views, following the line of work such as FixMatch, FlexMatch, and FreeMatch.** In contrast, contrastive learning-based methods emphasize the consistency between two batches of augmented views. To be more specific, our OTMatch calculates each optimal transport loss exclusively involving the two augmented views. In contrast, contrastive learning-based methods such as SimMatch and CoMatch align the two batches by utilizing the representation similarity between samples in the batch. As a result, contrastive learning-based methods necessitate an additional branch, apart from the one calculating sample-wise consistency. Therefore, our OTMatch is orthogonal to contrastive learning-based methods and can be combined with them. We have included citations for SimMatch and CoMatch in the revised version along with the above detailed discussion in Related work and Appendix B.1.