SelKD: Selective Knowledge Distillation via Optimal Transport Perspective

Thank you for your thorough review. Your comments were particularly helpful to improve our submission. Hopefully, our response below adequately addresses your concerns.

Q1: The paper lacks concrete details regarding the derivation and experimental implementation.

A1: Sorry for your confusion. For the detailed derivation for open-set SelKD, please refer to Algorithm 1, where $C_s$ and $C_t$ can be regarded as cosine distance of image and label features. And the original algorithm can be found in [1].

For the experimental settings, we train 240 epoch for each model, with learning rate 0.05 for CIFAR-10 and CIFAR-100. In epoch 150, 180, 210, learning rate decays by 90%. We adopt SGD optimizer, with momentum 0.9. The temperature is 4.0. The core code can be found in this anonymous link.

[1] Iterative Bregman Projections for Regularized Transportation Problems.

Q2: The analysis of experiments is somewhat superficial, particularly concerning parameter impacts. Given that optimal transport (OT) can be heavily influenced by parameters, providing visual analyses of how these parameters affect the results would enhance understanding.

A2: Paper [2], which we have cited in the references, reformulates softmax and cross-entropy into an bi-level optimization with optimal partial transport. The paper points out that the temperature is equivalent to the coefficient of entropic regularization.

Although we introduced OT, it is merely to view the KD problem from a higher-level perspective. Apart from entropic regularization coefficient $\epsilon$, we did not introduce any additional hyperparameters. Moreover, $\epsilon$ is equivalent to the softmax temperature, so we can set this hyperparameter based on other KD-related papers.

[2] Understanding and Generalizing Contrastive Learning from the Inverse Optimal Transport Perspective. ICML, 2023

Q3: The paper does not include any experiments related to the efficiency of the proposed method. It would be valuable to assess and report the computational efficiency of SelKD in comparison to traditional methods.

A3: We provide the time required to train 5 open-set subtask-specific small models using SelKD and Vanilla KD based on Resnet-50 for 240 epochs each in the table below.

And as the number of subtasks increases and the performance requirements for the teacher model become higher, the time and resource-saving advantages of SelKD will become more significant.

Q4: A broader range of experiments would strengthen the paper’s claims regarding SelKD's applicability.

A4: Thank you for your advice, and we further add more experiments with different backbones. The results are shown in the following table. In this setting, we adopt WideResNet-40-2 as the teacher model and WideResNet-16-2 as the student model.

For other distillation methods, SelKD is a framework designed to handle new tasks that require transferring subset knowledge. It can be used in conjunction with new distillation methods, as we demonstrate in our experiments, where we provide results combining it with various feature-based distillation methods (FitNet, AT, FT).

Dear Reviewer,

Thank you very much for your valuable response. As the discussion deadline is approaching, we would be grateful if you could kindly provide your feedback on our response to the concerns you previously raised. We would also be happy to address any further questions or suggestions you may have.

We sincerely appreciate your time and attention.

Best regards,

The Authors

Dear Reviewer,

Thank you for your insightful feedback and the kindness in raising the score. We will make sure that our discussion and the new findings are reflected in the final version of our paper.

Best regards,

The Authors

We appreciate your in-depth review and the time you spent providing feedback. We address your concerns as follows:

Q1: Could you clarify if KL divergence or cross-entropy is minimized in the outer optimization of Eq. 8?

The paper briefly mentions feature-level transport but doesn’t explore it. Expanding this aspect could reveal further benefits and justify using OT for representation learning.

A1: Thank you for pointing out the shortcomings in our paper. Here, we briefly introduce the main idea we followed, which is based on bi-level optimization to design the loss function.

We mainly follow [1] that understanding or designing the loss via bi-level optimization:

where ${C}^\theta$ represents the cosine distance for image feature and text/label feature, with parameters $\theta$, and ${Y}$ is the known supervision for learning. As proven in [1], $H({P}) = -\langle{P}, \log {P} - {1}\rangle$ is the entropic regularization with coefficient $\epsilon$. The inner optimization is exactly equivalent to the softmax activation, while the outer optimization corresponds to cross-entropy. Thus we can find the above bi-level optimization equals to InfoNCE loss:

Thus, bi-level optimization is fundamentally a method for designing activation layers or loss functions. In [1,2], modifications to the inner optimization improve the loss. Our work follows this learning framework, but we modify the inner optimization with new constraints to adapt to open-set scenarios, solving it with iterative algorithm [3] to obtain the predicted probability matching matrix. The outer optimization is adjusted to use the original KL Divergence in KD as the loss, resulting in an application in KD problems.

[1] Understanding and generalizing contrastive learning from the inverse optimal transport perspective. ICML, 2023

[2] OT-CLIP: Understanding and generalizing clip via optimal transport. ICML, 2024

[3] Iterative Bregman Projections for Regularized Transportation Problems. SIAM, 2015

Q2: In Figure 2, the green line is shown but not illustrated and switching Tables 2 and 3 order might improve flow.

A2: Sorry for the mistakes, and they should be blue lines. We will re-organize the tables for better flow and fix the mistakes in the updated version.

Q3: The notation, especially in the optimal transport section, could be clearer. Additionally, Algorithm 1 may be challenging for readers unfamiliar with classic OT methods. Providing some background or simplifying this section would improve accessibility.

A3: Thank you for your advice. We will further clarify the algorithm in the updated version. Please refer to [3] Proposition 5 for detailed derivation.

[3] Iterative Bregman Projections for Regularized Transportation Problems. SIAM, 2015

Q4: Information on computation complexity would be helpful to fully understand the method’s practicality and limits.

A4: We provide the time required to train 5 subtask-specific small models using SelKD and Vanilla KD based on Resnet-50 for 240 epochs each with single GPU in the table below.

For the computation cost for each student model, SelKD doesn't have advantages. But as the number of subtasks increases and the performance requirements for the teacher model become higher, the time and resource-saving advantages of SelKD will become more significant.

Q5: The impact of random seed on results is not shown.

A5: We adjusted the random seed and conducted further experiments on the closed-set SelKD of CIFAR-100. The results are shown in the table below.

It is relatively robust to random seeds.

Dear Reviewer,

We are delighted that our responses have successfully addressed most of your concerns. Thank you again for your support and the insightful comments you have provided. We are truly grateful for your decision to maintain a positive score. We will make sure that the outcomes of our discussions are integrated into the final version of our paper.

Best regards,

The Authors

Thank you for your detailed review and constructive comments. We hope to thoroughly address all the concerns you highlighted.

Q1: The experimental validation is limited.

A1: Thank you for your advice, and we further add more experiments with different backbones. The results are shown in the following table. In this setting, we adopt WideResNet-40-2 as the teacher model and WideResNet-16-2 as the student model.

Q2: Since SelKD relies on label information, it is challenging to extend it to semi-supervised and unsupervised settings.

A2: Firstly, If $g(\cdot)$ in Eq. (11) is one-hot function, it degenerates into a label encoder. Therefore, our model can be extended to CLIP-based models.

Secondly, traditional neural networks output logits of size $B\times N$, where $B$ is the batch size and $N$ is the number of classes. In contrast, a CLIP-based model outputs image features of size $B\times k$ and text features of size $N\times k$. By performing the multiplication in Eq. (11), we can obtain logits equivalent to those produced by end-to-end neural networks. Consequently, this ensures that our model remains compatible with traditional semi-supervised or self-supervised methods typically applied to neural networks.

Thirdly, compared to end-to-end models, CLIP has greater potential for self-supervised learning. Specifically, we can adopt the framework mentioned in the SLIP paper [1], which involves training the teacher model using both contrastive loss for image-text alignment and contrastive loss for image self-supervision.

[1] SLIP: Self-supervision meets Language-Image Pre-training. ECCV, 2022

Q3: The open-set SelKD formulation, while promising, lacks guidance on how it could be adapted to real-world scenarios with dynamically changing categories.

A3: Take a real-world scenario as an example. Software services are provided to different users and the number of features offered varies depending on their subscription tier. If we regard each type of service as a model, using traditional KD methods would require training a separate teacher model for each user to distill the required functionality. In contrast, with our SelKD approach, a single large model encompassing all functionalities can be used to distill small models tailored to any specific subset of features.

Q4: Lack of theoretical guarantee.

A4: Thanks for your kind reminder. Indeed, our training process lacks a strong theoretical guarantee, and this is an area we need to explore further in our future work. Our work primarily follows the frameworks proposed in [2,3] for designing the training loss or framework. Regarding the OT algorithm, our theoretical foundation is based on [4], which ensures the convergence of our predictions.

[2] Understanding and Generalizing Contrastive Learning from the Inverse Optimal Transport Perspective. ICML, 2023

[3] OT-CLIP: Understanding and Generalizing CLIP via Optimal Transport. ICML, 2024

[4] Iterative Bregman Projections for Regularized Transportation Problems. SIAM, 2015

Thank you for your time and your critical assessment. Your comments have been very insightful for how our submission can be improved. We hope our responses to your questions below can help to clarify our approach and demonstrate the impact of our work.

Q1: What are the benefits of training multiple students, each tailored to a specific task, when a general student model showcases consistently strong performance across all tasks in vanilla KD?

A1: Knowledge distillation is typically applied to small devices with limited resources and less capable hardware. When the number of classification categories is too large, the performance of small models is significantly worse compared to large models. Small models are more suitable for tasks with fewer categories.

Additionally, in a real-world scenarios where software services are provided to different users, the number of features offered varies depending on their subscription tier. Using traditional KD methods would require training a separate teacher model for each user to distill the required functionality. In contrast, with our SelKD approach, a single large model encompassing all functionalities can be used to distill small models tailored to any specific subset of features.

Q2: SelKD could involve higher computational costs due to the need for training multiple task-specific student models.

A2: We simply divided the tasks to test accuracy, and in our settings, Vanilla KD also requires training student models for each subtask. We provide the time required to train 5 subtask-specific small models using SelKD and Vanilla KD based on Resnet-50 for 240 epochs each in the table below.

And as the number of subtasks increases and the performance requirements for the teacher model become higher, the time and resource-saving advantages of SelKD will become more significant.

Q3: To strengthen the justification for the proposed SelKD in resource-constrained environments, the paper should include additional motivation experiments or theoretical analysis demonstrating that a student model cannot fully receive the knowledge from the teacher model.

A3: Large models are typically used on large devices, while small models are often deployed on smaller devices, such as smartphones. The functionality required on smartphones doesn’t need to fully match that of larger devices. Instead, it only needs to meet specific requirements. In our setting, we leverage knowledge distillation to deploy only certain functionalities to smaller devices, enabling the creation of customized models with varying effects and capabilities. Training a dedicated teacher model for each specific functionality would be highly inefficient.

For example, a \$199 phone requires fewer features (i.e. categories) in its model compared to a \\$599 phone, which demands more functionalities. Instead of training two separate teacher models for the models needed by these two types of devices, training a single large model to perform selective knowledge distillation for both saves resources and time.

Q4: How are the teacher and student models practically deployed/trained? It is challenging to simultaneously train a cumbersome teacher and student on resource-constrained edge device.

A4: The teacher and student models are both trained on devices with sufficient resources, and the student model is then deployed on resource-constrained edge device.

Q5: In Eq. (11), why is the cost matrix calculated through the multiplication of $f(\cdot)$ and $g(\cdot)$? Is there a specific reason or advantage for using this multiplicative form?

A5: For traditional classifiers, $f(\cdot)$ represents the image encoder and $g(\cdot)$ represents the label encoder (e.g., one-hot encoding). For multimodal classifiers based on CLIP, $f(\cdot)$ represents the image encoder while $g(\cdot)$ represents the text encoder. Our formulation generalizes this approach, making it more adaptable to a wider range of methods.

For this multiplicative form, many papers [1,2] adopt it to depict similarity or cost. However, if there is a specific method for cost calculation, it can also be used as a substitute.

[1] Self-supervised Learning of Visual Graph Matching. ECCV, 2022.

[2] Deep Graphical Feature Learning for the Feature Matching Problem. ICCV, 2019.

Q6: How are subtasks determined? Can the proposed SelKD consistently maintain superior performance when the categories within the same subtask significantly differ from each other?

A6: To better understand the real-world scenario, we simply randomly split the class labels. In fact, for datasets like CIFAR-10 and Tiny ImageNet that are split by class indices, the categories in each subtask are inherently quite distinct from one another.

Q7: Some multi-task KD studies should be considered as baselines in the experiments.

A7: The goal of multi-task KD [3,4,5] is to enhance the performance of a student model across multiple tasks by transferring knowledge from a teacher model. In contrast, the goal of our SelKD is to transfer a specific portion of the teacher model's knowledge to the student model, thereby optimizing a specific subset of functionalities. The two approaches have distinct goals and are not directly comparable.

[3] Distilling task-specific knowledge from BERT into simple neural networks. arXiv, 2019

[4] Online knowledge distillation for multi-task learning. WACV, 2023

[5] Multi-task learning with knowledge distillation for dense prediction. ICCV, 2023

Dear Reviewer,

Thank you very much for your valuable response. As the discussion deadline is approaching, we would be grateful if you could kindly provide your feedback on our response to the concerns you previously raised. We would also be happy to address any further questions or suggestions you may have.

We sincerely appreciate your time and attention.

Best regards,

The Authors

Thank you to the authors for their efforts in addressing my comments. Your responses to both my comments and that of the other reviewers have addressed most of my concerns, and I am pleased to raise my rating score.