Wasserstein Distance Rivals Kullback-Leibler Divergence for Knowledge Distillation

Summary

This paper proposes a new algorithm for knowledge distillation by replacing KL-divergence loss with Wasserstein distance loss. The proposed algorithm contains two parts: (1) Logits distillation with Wasserstein distance loss, implementing with an entropy regularized linear programming; (2) Assuming the features obey Gaussian distribution, the Wasserstein distance loss can be solved with parametric tricks. Experiments on ImageNet and CIFAR100 show obvious improvements compared to baselines like KD, DKD, and NKD.

Strengths

(1) The paper is written clearly and is easy to follow. (2) Sound ablations for combinations of the proposed method and previous ones.

Weaknesses

(1) Eq. (3) and Eq. (5) both have hyperparameter $\lambda$. Does the two lambda have the same value in implementation? The ablation for $\lambda$ in Eq. (3) is missing. (2) There are too many hyper-parameters to be tuned. lambda in Eq. (3), lambda in Eq. (5), k for calculating IR, gamma for a trade-off between D_mean and D_cov, weight between WKD-L and WKD-F, temperature. Do all experiments in the paper adopt consistent hyper-parameters? (3) What's the layer features are used for WKD-F? Feature dimensions are usually higher than logits. Moreover, multiple-layer features are often used for feature-based methods. Thus, it seems to be wired that WKD-F (207ms) is even faster than the original KD (215ms). (4) The paper claims that Wasserstein distance rivals Kullback-Leibler divergence for knowledge distillation as indicated by their title. However, as shown in Table 2(a) and Table 2(b), the KL-div method achieves 71.96 in the setting of separating target and non-target. The WD-based method just achieves comparable results with the Polynomial kernel, class centroid, or classifier weight for calculating IR. It seems that the IR is more important than the form of WD formulation.

Questions

(1) Ablation for $\lambda$ in Eq. (3). (2) Explanation and analysis of sensitivity of all hyper-parameters. (3) How to calculate $\mu^T$ and $\mu^S$ in your implementation for Eq. (9)? (4) Does the proposed method work well in the setting of self-KD where the student and teacher share the same architecture?

Rating

5

Confidence

4

Soundness

3

Presentation

3

Contribution

3

Limitations

Limitations are discussed in the paper

Summary

This paper proposes a Wasserstein distance based knowledge distillation method for both the logit distillation and feature distillation settings. The logit based version uses discrete WD to model the discrepancy beteween the prediction probabilites of student and teacher networks. It further uses the separation of target probability to improve performance. The feature based version minimizes the continuous WD between the patch features of an image from the student and teacher networks, under the assumption that they form Gaussian distributions. The covariance term is simplifed to its diagnoals to further improve performance. Comapred to recent baslines, the proposed method shows superior performance on ImageNet classification and COCO object detection.

Strengths

1. Strong performance compared to recent work. 2. The presentation is easy to follow and well structured. Related works are introduced to give good contexts. 3. Extensive comparison between baselines, detailed abalation study, and has extra experiments on distillation across CNNs and transformers in appendix.

Weaknesses

One of the motivation for the WKD-L is the cross-category comparison, however, it is not clear to me why the "cross-category" comparison is attrubted as the source of improvement. Firstly, there is implicit cross-category comparision for KL based methods (line 117), so it is not a differentiator. Secondly, without the story of cross-category comparison, the WDK-F also shows improvement.

Questions

1. Could you explain the statement at line 120 "this implicit effect is insignificant"? 2. Line 38. "... features of an image are ... small size". What does the "small size" of an image referes to? 3. Are the student networks randomly initialized or initialized from a trained weights for baseline methods (line 239)? 4. Line 172 "... partition the feature maps into a kxk spatial grid", and line 289 "use ... 1x1 grid for classification". There is no grid setting for detection in the main text and can mislead readers to think grid is not used at all.

Rating

7

Confidence

3

Soundness

3

Presentation

3

Contribution

4

Limitations

The authors discussed the limitations.

Summary

The paper introduces a novel methodology for knowledge distillation using Wasserstein Distance (WD) instead of the traditional Kullback-Leibler Divergence (KL-Div). The proposed methods include a logit distillation approach (WKD-L) that leverages cross-category comparisons and a feature distillation method (WKD-F) that models feature distributions parametrically. The authors demonstrate that their methods outperform strong KL-Div variants on image classification and object detection tasks.

Strengths

* The introduction of WD in knowledge distillation provides a fresh perspective and addresses the limitations of KL-Div, particularly in terms of cross-category comparisons and handling non-overlapping distributions. * The use of parametric methods for feature distribution modeling, specifically Gaussian distributions, is innovative and effectively leverages the geometric structure of the data.

Weaknesses

* The paper suffers from several writing issues, including grammatical errors and unclear explanations, making it difficult to follow the arguments and methodology at times. * Some of the assumptions made for the application of WD, particularly the choice of Gaussian distributions for feature modeling, may not hold universally. Further justification or exploration of alternative parametric methods would strengthen the work. * The computational complexity of implementing WD-based methods, especially in large-scale scenarios, is not adequately addressed. A comparison of computational costs between WD and KL-Div would be beneficial.

Questions

* Why did you choose Gaussian distributions for feature modeling in the context of WD? Are there other parametric methods that you considered, and how would they compare in terms of performance and feasibility? * How does the computational complexity of your proposed WD-based methods compare to traditional KL-Div based methods? Can you provide a detailed analysis or empirical comparison of the computational costs involved?

Rating

6

Confidence

4

Soundness

3

Presentation

2

Contribution

3

Limitations

* please check the questions and weaknesses.

Summary

The paper proposes the utilization of Wasserstein Distance based distillation as opposed to KLD, as is common in practice. This is because the latter does not facilitate cross-category comparisons. Both logit and feature based variants have been proposed. The comparisons have been shown for both classification as well as detection.

Strengths

Though straightforward look at the empirical gains are not massive, the method is theoretically sound, and the paper is easy to understand. The experiments are decently presented in my opinion. The paper provides a fresh perspective of WD in distillation.

Weaknesses

I would have been happier to see if changing the divergence led to more significant boosts, if at all possible (basis the premise of providing cross category comparisons).

Questions

None

Rating

5

Confidence

4

Soundness

3

Presentation

3

Contribution

3

Limitations

The cost expensiveness has been mentioned.

Response by Reviewer Zit9

I have carefully reviewed the feedback from other reviewers, considered the author’s rebuttal, and global responses, followed the ensuing discussion, and read the paper again . I appreciate the authors' thorough responses, particularly their clarification on W2 and new experimental results during the rebuttal period as well as Q2. Itherefore I will raise my score slightly from 5 to 6. Good luck!