SelKD: Selective Knowledge Distillation via Optimal Transport Perspective

Authors: Liangliang Shi, Zhengyan Shi, Junchi Yan
Primary area: unsupervised, self-supervised, semi-supervised, and supervised representation learning
Keywords: Knowledge Distillation, Inverse Optimal Transport

Abstract:
Knowledge Distillation (KD) has been a popular paradigm for training a (smaller) student model from its teacher model. However, little research has been done on the practical scenario where only a subset of the teacher's knowledge needs to be distilled, which we term selective KD (SelKD). This demand is especially pronounced in the era of foundation models, where the teacher model can be significantly larger than the student model. To address this issue, we propose to rethink the knowledge distillation problem from the perspective of Inverse Optimal Transport (IOT). Previous Bayesian frameworks mapped each sample to the probabilities of corresponding labels in an end-to-end manner, which fixed the number of classification categories and hindered effective partial knowledge transfer. In contrast, IOT calculates from the standpoint of transportation or matching, allowing for the flexible selection of samples and their quantities for matching. Traditional logit-based KD can be viewed as a special case within the IOT framework. Building on this IOT foundation, we formalize this setting in the context of classification, where only selected categories from the teacher's category space are required to be recognized by the student in the context of closed-set recognition, which we call closed-set SelKD, enhancing the student's performance on specific subtasks. Furthermore, we extend the closed-set SelKD, introducing an open-set version of SelKD, where the student model is required to provide a "not selected" response for categories outside its assigned task. Experimental results on standard benchmarks demonstrate the superiority of our approach. The source code is available at: \href{https://github.com/machoshi/SelKD}

Decision: accept (poster)
Reviewer scores: 6, 6, 6, 6

Reviews and discussion:
<thread>
creation date:
1729586121584

thread_type:
official_review

<thread_object>
creator: reviewer_bhk7

<content>
summary:
This paper proposes a novel method called Selective Knowledge Distillation (SelKD), which enhances student model performance by transferring only partial knowledge. The integration of Optimal Transport (OT) methods is applied to both closed-set and open-set classification tasks, demonstrating that SelKD outperforms traditional knowledge distillation methods. Overall, the paper offers new insights and practical solutions for the knowledge distillation field.

strengths:
1. The introduction of SelKD allows for effective partial knowledge transfer, reducing redundant learning seen in traditional KD.
2. The method shows high potential for real-world applications, particularly in resource-constrained environments.

weaknesses:
1. The paper lacks concrete details regarding the derivation and experimental implementation. It would be beneficial to provide specific derivations and related code for the experiments to facilitate reproducibility.
2. 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.
3. 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.
4. The experiments primarily focus on ResNet architectures, which restricts the generalizability of the findings. Incorporating more recent distillation methods from 2023 or 2024, as well as testing on different model architectures and datasets, would better demonstrate the effectiveness and versatility of the proposed method. A broader range of experiments would strengthen the paper’s claims regarding SelKD's applicability.

questions:
See above.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730567356115

thread_type:
official_review

<thread_object>
creator: reviewer_up4c

<content>
summary:
This paper introduces Selective Knowledge Distillation (SelKD), a novel approach that allows student models to learn only selected subsets of knowledge from teacher models. The authors reformulate knowledge distillation through the lens of Inverse Optimal Transport (IOT) and propose both closed-set and open-set variants of SelKD.

strengths:
**Originality**: The paper offers a fresh perspective on selective knowledge distillation by using optimal transport (OT) to map logit relationships between image and text features. This new approach allows the student model to learn only the relevant subset of knowledge, making it distinct from traditional, full-transfer methods. This new perspective opens up important new research directions.

**Clarity**: The paper is well-structured, with a clear motivation and related work sections that are easy to follow. The optimal transport concepts are complex, but most parts are presented clearly.

**Significance**: Great motivation: enabling selective transfer for specific tasks. The open-set extension further adds value by allowing the model to handle unknown categories, making it versatile for real-world applications.

weaknesses:
**Clarity and Notation**: The notation, especially in the optimal transport section, could be clearer. Key terms and symbols, like 
$z$ in Eq. 2 and $\tilde{P}$ in Eq. 7, should be introduced. Additionally, Algorithm 1 may be challenging for readers unfamiliar with classic OT methods. Providing some background or simplifying this section would improve accessibility.

**Computational Complexity**: The paper doesn’t explain the computational cost of the OT-based approach, which may be high given the framework. Information on computation complexity would be helpful to fully understand the method’s practicality and limits.

**Performance**: The impact of random seed on results is not shown. Further testing here could demonstrate robustness and validate improvements.

**Feature-level Transport Potential**: 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.

questions:
1. If each subtask requires a separate teacher model, it naturally raises concerns about parameter count and memory costs. Tables presenting parameter counts should clearly explain any dependencies on task splits to avoid misleading impressions.

2. Could you clarify if KL divergence or cross-entropy is minimized in the outer optimization of Eq. 8? The description seems inconsistent.

3. In Figure 2, the green line is shown but not illustrated. Could you clarify this?

4. Switching Tables 2 and 3 order might improve flow.
</content>
</thread_object>

</thread>
<thread>
creation date:
1731042499373

thread_type:
official_review

<thread_object>
creator: reviewer_o62s

<content>
summary:
The paper proposes Selective Knowledge Distillation (SelKD) for transferring only selected parts of a teacher model’s knowledge to a student model, framed within the Inverse Optimal Transport (IOT) perspective. This method is relevant for scenarios with resource constraints, where distilling only relevant subsets of knowledge is more practical than full-model distillation. The authors present closed-set and open-set SelKD, where the student learns a subset of categories or can reject classes outside its training scope. Experiments on CIFAR and Tiny ImageNet demonstrate SelKD’s advantages over traditional KD techniques, with improved efficiency and comparable performance on targeted tasks.

strengths:
The paper addresses a practical challenge in knowledge distillation, particularly relevant to deploying models on edge devices with limited resources. By reinterpreting KD within an IOT framework, the authors introduce an innovative approach to selective knowledge transfer that has potential for efficient model adaptation in constrained environments. The experiments on CIFAR and Tiny ImageNet datasets validate the method’s utility, showing that SelKD can effectively focus on specific subtasks with lower computational costs.

weaknesses:
1. The experimental validation is somewhat limited, with results provided only on relatively small datasets (CIFAR and TinyImagenet). The model used in the study is relatively small (ResNet50), with significantly fewer parameters than current large-scale pretrained models. I recommend the authors add experiments using larger backbone models such as ViT.
2. Since SelKD relies on label information, it is challenging to extend it to semi-supervised and unsupervised settings.
3. The open-set SelKD formulation, while promising, lacks guidance on how it could be adapted to real-world scenarios with dynamically changing categories.
4. Lack of theoretical guarantee.

questions:
Would the authors consider exploring SelKD in semi-supervised or unsupervised settings where labeled data for all categories may not be available?
</content>
</thread_object>

</thread>
<thread>
creation date:
1731059561501

thread_type:
official_review

<thread_object>
creator: reviewer_9lkc

<content>
summary:
This paper discusses the knowledge distillation (KD) problem from the perspective of inverse optimal transport (IOT), introducing a selective KD (SelKD) framework tailored for classification tasks. Specifically, closed-set SelKD allows the student model trained exclusively on the data relevant to specific tasks, while open-set SelKD enables the student model to identify non-selected knowledge to unassigned tasks. Experimental results on image datasets showcase the superior performance of proposed SelKD over the state-of-the-art methods.

strengths:
This paper formulates the KD problem as a bi-level optimization through the perspective of IOT. Building on the IOT-based formulation, this paper further proposes novel closed-set and open-set SelKD methods that allow the student model to learn specific knowledge from assigned tasks, making SelfKD well-suited for real-world, resource-constrained environments.

weaknesses:
However, there are several weaknesses where this paper can improve:
1.  The advantages of the proposed SelKD over vanilla KD are unclear in Figure 1. 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? SelKD could involve higher computational costs due to the need for training multiple task-specific student models. 

2. 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, or in other words, that the teacher fails to transfer essential task-specific knowledge to the student. 

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

4. In Eqn (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?

5. The data storage and computational capacity are often limited in resource-constrained environments. Thus, the proposed SelKD is expected to exhibit strong robustness against insufficient training data, while also ensuring efficient computational cost and resource utilization.

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

7. Some multi-task KD studies [1-3] should be considered as baselines in the experiments. 

[1] Raphael Tang, et al. "Distilling task-specific knowledge from BERT into simple neural networks." arXiv, 2019.

[2] Geethu Miriam Jacob, et al. "Online knowledge distillation for multi-task learning." WACV, 2023.

[3] Yangyang Xu, et al. "Multi-task learning with knowledge distillation for dense prediction." ICCV, 2023.

questions:
Please refer to the weaknesses mentioned above.
</content>
</thread_object>

</thread>
<thread>
creation date:
1732366813992

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
General Response

comment:
Dear Area Chairs and Reviewers,

We extend our gratitude to all reviewers for their efforts, insightful comments and constructive suggestions. Overall, the reviewers deem our work as "novel/innovative/fresh" (9LKC, o62s, uP4C), "well-structured" (uP4C), "great motivated" (uP4C), "effective" (o62s, bHk7), "showing potential/well-suited for real-world/resource-constrained environments" (9LKC, o62s, uP4C, bHk7) and "opening up important new research directions" (uP4C).

We make every effort to address reviewers’ concerns and provide additional experimental results. We have updated our revised PDF based on the reviewers' suggestions, with specific changes highlighted in blue. The updates we made are as follows:

1. We conduct more experiments under larger-scaled models. The experimental results on CIFAR-100 dataset demonstrate the effectiveness of our SelKD method under closed-set tasks.
2. We give a clearer explanation on the bi-level optimization and further clarify our OT-based method. The outer optimization is equivalent to cross-entropy, while the inner is equivalent to the softmax activation.
3. We further explain the feature extraction functions $f(\cdot)$ and $g(\cdot)$. In particular, we clarify that the text encoder $g(\cdot)$ degenerates into a label encoder when implemented as a one-hot function.
4. We address formatting issues and fix typos to improve the overall layout and readability of the paper.

In our individual rebuttal, we provide thorough explanations and responses to each of the reviewers' questions. We aim for these responses to effectively address the concerns raised by the reviewers, and we are open to further discussion to ensure a thorough evaluation of our work.
</content>
</thread_object>

</thread>


Metareview:
metareview:
This paper discusses the knowledge distillation (KD) problem from the perspective of inverse optimal transport (IOT), introducing a selective KD (SelKD) framework tailored for classification tasks. The paper addresses a practical challenge in knowledge distillation, particularly relevant to deploying models on edge devices with limited resources. By reinterpreting KD within an IOT framework, the authors introduce an innovative approach to selective knowledge transfer that has potential for efficient model adaptation in constrained environments. The experiments on CIFAR and Tiny ImageNet datasets validate the method’s utility, showing that SelKD can effectively focus on specific subtasks with lower computational costs. 

The paper initially has some issues (insufficient experiments and unclear settings). After the rebuttal, I believe most issues have been solved and all reviewers agree to accept.