Wasserstein Distance Rivals Kullback-Leibler Divergence for Knowledge Distillation

Authors: jiaming  lv
haoyuan  yang
peihua  li
Venue: NeurIPS_cc_2024_Conference
Published: 2024-05-05 03:28:20

Abstract: Since pioneering work of Hinton et al., knowledge distillation based on Kullback-Leibler Divergence (KL-Div) has been  predominant, and recently its variants have achieved compelling performance.  However, KL-Div  only compares probabilities of the corresponding category between the teacher and student while lacking a mechanism for cross-category comparison.  Besides, KL-Div is problematic when applied to intermediate layers, as it cannot handle non-overlapping distributions and is unaware of geometry of the underlying manifold. To address these downsides, we propose a methodology of Wasserstein Distance (WD) based knowledge distillation. Specifically, we propose  a logit distillation method called WKD-L based on discrete WD, which performs cross-category comparison of  probabilities and thus can explicitly leverage  rich interrelations among categories.  Moreover, we introduce a feature distillation method called WKD-F, which  uses a parametric method for modeling feature distributions and adopts continuous WD  for transferring knowledge from  intermediate layers. Comprehensive evaluations  on image classification and object detection  have shown  (1) for logit distillation WKD-L outperforms very strong KL-Div variants; (2) for feature distillation WKD-F is superior to the KL-Div counterparts and state-of-the-art competitors.

Main paper: 
 introduction: 
 Knowledge distillation (KD) aims to transfer knowledge from a high-performance teacher model with large capacity to a lightweight student model. In the past years, it has attracted ever increasing interest and made great advance in deep learning, enjoying widespread applications in visual recognition and object detection, among others [1] . In their pioneering work, Hinton et al. [2] introduce Kullback-Leibler divergence (KL-Div) for knowledge distillation, where the prediction of category probabilities of the student is constrained to be similar to that of the teacher. Since then, KL-Div has been predominant in logit distillation and recently its variants [3; 4; 5] have achieved compelling performance. In addition, such logit distillation methods are complementary to many state-of-the-art methods that transfer knowledge from intermediate layers [6; 7; 8] .Despite the great success, KL-Div has two downsides that hinder fully transferring of the teacher's knowledge. First, KL-Div only compares the probabilities of the corresponding category between the teacher and student, lacking a mechanism to perform cross-category comparison. However, real-world categories exhibit varying degrees of visual resemblance, e.g., mammal species like dog and wolf look more similar to each other while visually very distinct from artifact such as car and bicycle. Deep neural networks (DNNs) can distinguish thousands of categories [9] and thus are (a) Real-world categories exhibit rich interrelations (IRs) in feature space, e.g., dog is near other mammal while far from artifact like car. We quantify pairwise IRs as feature similarities among categories. Best viewed by zooming in.(b) For logit distillation, discrete WD performs cross-category comparison by exploiting pairwise IRs, in contrast to KL-Div that is a category-to-category measure and lacks a mechanism to use such IRs (cf. Figure 2 ). For feature distillation, we use Gaussians for distribution modeling and continuous WD for knowledge transfer. well-informed of such complex relations among categories, as shown in Figure 1a . Unfortunately, due to its category-to-category nature, the classical KD [2] and its variants [3; 4; 5] are unable to explicitly utilize this rich cross-category knowledge.Secondly, KL-Div is problematic for distilling knowledge from intermediate layers. Deep features of an image are generally high-dimensional and of small size, so being populated very sparsely in the feature space [10, Chap. 2] . This not only makes non-parametric density estimation (e.g., histogram) that KL-Div requires infeasible due to curse of dimensionality, but also leads to non-overlapping discrete distributions that KL-Div fails to deal with [11] . One may turn to parametric, continuous methods (e.g., Gaussian) for modeling feature distributions. However, KL-Div and its variants have limited ability for measuring dis-similarity between continuous distributions, as it is not a metric [12] and is unaware of geometric structure of the underlying manifold [13] .The Wasserstein distance (WD) [14] , also called Earth Mover's Distance (EMD) or optimal transport, has the potential to address the limitations of KL-Div. The WD between two probability distributions is generally defined as the minimal cost to transform one distribution to the other. Several works have made exploration by using WD for knowledge transfer from intermediate layers [15; 16] . Specifically, they measure the dis-similarity of a mini-batch of images between the teacher and student based on discrete WD, which concerns comparison across instances in a soft manner, failing to make use of relations across categories. Moreover, they mainly quest for non-parametric method for modeling distributions, behind in performance state-of-the-art KL-Div based methods.To address these problems, we propose a methodology of Wasserstein distance based knowledge distillation, which we call WKD. This methodology is applicable to logits (WKD-L) as well as to intermediate layers (WKD-F) as shown in Figure 1b . In WKD-L, we minimize the discrepancy between the predicted probabilities of the teacher and student using discrete WD for knowledge transfer. In this way, we perform cross-category comparison that effectively leverages interrelations (IRs) among categories, in stark contrast to category-to-category comparison in the classical KL-Div. We propose to use Centered Kernel Alignment (CKA) [17; 18] to quantify category IRs, which measures the similarity of features between any pair of categories.For WKD-F, we introduce WD into intermediate layers to condense knowledge from features. Unlike the logits, there is no class probability involved in the intermediate layers. Therefore, we let the student directly match the feature distributions of the teacher. As the dimensions of DNN features are high, the non-parametric methods (e.g., histogram) are infeasible due to curse of dimensionality [10, Chap. 2], we choose parametric methods for modeling distributions. Specifically, we utilize one of the most widely used continuous distributions (i.e., Gaussian), which is of maximal entropy given 1stand 2nd-moments estimated from features [19, Chap. 1] . WD between Gaussians can be computed in closed form and is a Riemannian metric on the underlying manifold [20] .We summarize our contributions in the following.• We present a discrete WD based logit distillation method (WKD-L). It can leverage rich interrelations among classes via cross-category comparisons between predicted probabilities of the teacher and student, overcoming the downside of category-to-category KL divergence. • We introduce continuous WD into intermediate layers for feature distillation (WKD-F). It can effectively leverage geometric structure of the Riemannian space of Gaussians, better than geometryunaware KL-divergence. • On both image classification and object detection tasks, WKD-L perform better than very strong KL-Div based logit distillation methods, while WKD-F is supervisor to the KL-Div counterparts and competitors of feature distillation. Their combination further improves the performance. 
wd for knowledge transfer: 
 Given a pre-trained, high-performance teacher model T, our task is to train a lightweight student model S that can distill knowledge from the teacher. As such, supervisions of the student are from both the ground truth label with the cross entropy loss and from the teacher with distillation losses to be described in the next two sections.Interrelations (IRs) among categories. As shown in Figures 1a and 4 , real-world categories exhibit complex topological relations in the feature space. For instance, mammal species are nearer each other while being far away from artifact or food. Moreover, features of the same category cluster and form a distribution while neighboring categories have overlapping features and cannot be fully separated. As such, we propose to quantify category IRs based on CKA [18] , which is a normalized Hilbert-Schmidt Independence Criterion (HSIC) that models statistical relations of two sets of features by mapping them into a Reproducing Kernel Hilbert Space (RKHS) [21] .Given a set of b training examples of category C i , we compute a matrix X i ∈ R u×b where the k th column indicates the feature of example k that is output from the DNN's penultimate layer. Then we compute a kernel matrix K i ∈ R b×b with some positive definite kernel, e.g., a linear kernel for which K i = X T i X i where T indicates matrix transpose. Besides the linear kernel, we can choose other kernels such as polynomial kernel and RBF kernel (cf. Section A.1 for details). The IR between C i and C j is defined as:EQUATIONHere H = I-1 b 11 T is the centering matrix where I indicates the identity matrix and 1 indicates an all-one vector; tr indicates matrix trace. IR(C i , C j ) ∈ [0, 1] is invariant to isotropic scaling and orthogonal transformation. Note that the cost to compute the IRs can be neglected since we only need to compute them once beforehand. As the teacher is more knowledgeable, we compute category interrelations using the teacher model, which is indicated by IR T (C i , C j ).Besides CKA, cosine similarity between the prototypes of two categories can also be used to quantify IRs. In practice, the prototype of one category can be computed as the average of the features of the category's examples. Alternatively, the weight vectors associated with the softmax classifier of a DNN model can be regarded as prototypes of individual categories [22] .Loss function. Given an input image (instance), we let z = z i ∈ R n be the corresponding logits of a DNN model where i ∈ S n △ = {1, • • • , n} indicates the index of i th category. The predicted category probability p = p i is computed via the softmax function σ with a temperature τ , i.e., p i = σ z τ i △ = exp(z i /τ )/ j∈Sn exp(z j /τ ). We denote by p T and p S the predicted category probabilities of the teacher and student models, respectively. The classical KD [2] is an instancewise method, which measures the discrepancy between p T and p S given the same input image: EQUATIONKL-Div (2) only compares predicted probabilities corresponding to the same category between the teacher and student, essentially short of a mechanism to perform cross-category comparison, as shown in Figure 2 . Though during gradient backpropagation probability of one category affects probabilities of other categories due to the softmax function, this implicit effect is insignificant and, above all, cannot explicitly exploit rich knowledge of pairwise interrelations as described in (1) .In contrast to KL-Div, WD performs cross-category comparison and thus naturally makes use of category interrelations, as shown in Figure 1b (left). We formulate discrete WD as an entropy regularized linear programming [23] :D WD (p T , p S ) = min qij i,j c ij q ij + ηq ij log q ij (3) s.t. q ij ≥ 0, j q ij = p T i , i q ij = p S j , i, j ∈ S n ,where c ij and q ij respectively indicate transport cost per mass and the transport amount while moving probability mass from p T i to p S j ; η is a regularization parameter. We define the cost c ij by converting the similarity measure IRs to a distance measure according to the commonly used Gaussian kernel [10, Chap. 6] , i.e., c ij = 1 -exp(-κ(1 -IR T (C i , C j ))), where κ is a parameter that can control the degree of sharpening of IR. The smaller IR T (C i , C j ) in the feature space, the less cost is needed for transport between the two categories. As such, the loss function of WKD-L isEQUATIONRecent work [3] discloses target probability (i.e., the probability of the target category) and the nontarget ones play different roles: the former concerns the difficulty of training examples while the latter containing prominent "dark knowledge". It has been shown that this separation is helpful to balance their roles and improves greatly over the classical KD [3; 4] . Inspired by them, we also consider a similar separation strategy. Let t be index of target category and z T \t = z T i ∈ R n-1 , i ∈ S n \{t} be the teacher's logits of non-target categories. We normalize z T \t as previously, obtaining the teacher's non-target probabilities p T \t = p T i . In this case, our loss functions consist of two terms:EQUATIONwhere λ is the weight.As the features output from intermediate layers of DNN are of high dimension and small size, the non-parametric methods, e.g., histogram and kernel density estimation, are infeasible. Therefore, we use one of the widely used parametric methods (i.e., Gaussian) for distribution modeling.Feature distribution modeling. Given an input image, let us consider feature maps output by some intermediate layer of a DNN model, whose spatial height, width and channel number are h, w and l, respectively. We reshape the feature maps to a matrix F ∈ R l×m where m = h × w and the i th column f i ∈ R l indicates a spatial feature. For these features, we estimate the 1st-momentµ = 1 m i f i and the 2nd-moment Σ = 1 m i (f i -µ)(f i -µ) T .We model feature distribution of the input image by a Gaussian with mean vector µ and covariance matrix Σ as its parameters:EQUATIONwhere| • | indicates matrix determinant.We estimate Gaussian distribution of the teacher directly from its backbone network. For the student, as in previous arts [24; 25; 26] , a projector is used to transform the features so that they are compatible in size with the features of the teacher. Then the transformed features produced by the projector are used to compute the student's distribution. We select Gaussian for distribution modeling as it is of maximal entropy for given the 1st-and 2nd-moments [19, Chap. 1] and has a closed form WD that is a Riemannian metric [20] .Loss Function. Let Gaussian N T △ = N (µ T , Σ T ) be feature distribution of the teacher. Similarly, we denote by N S the student's distribution. The continuous WD between the two Gaussians is defined asEQUATIONwhere f T and f S are Gaussian variables and ∥ • ∥ indicates Euclidean distance; the joint distribution q is constrained to have marginals N T and N S . Minimization of Eq. ( 7) leads to the following closed form distance [14] :EQUATIONHereEQUATION2 ) where superscript 1 2 indicates matrix square root. As the covariance matrices estimated from highdimensional features are often ill-conditioned [27] , we add a small positive number (1e 5) to the diagonals. We also consider diagonal covariance matrices, for which we haveD cov (Σ T , Σ S ) = δ T -δ S 2, where δ T is a vector of standard variances formed by the square roots of the diagonals of Σ T . We later compare Gaussian (Full) and Gaussian (Diag) which have full and diagonal covariance matrices, respectively. To balance role of the mean and covariance, we introduce a mean-cov ratio γ and define the loss asEQUATIONWe can use the strategy of spatial pyramid pooling [28; 29; 6] to enhance representation ability. Specifically, we partition the feature maps into a k × k spatial grid, compute a Gaussian for each cell of the grid and then match per cell the Gaussians of the teacher and student.KL-Div [30] and symmetric KL-Div (i.e., Jeffreys divergence) [31] , both having closed form expressions for Gaussians [32] , can be used for knowledge transfer. However, they are not metrics [12] , unaware of the geometry of the space of Gaussians [13] that is a Riemannian space. Conversely, D WD is a Riemannian metric that measures the intrinsic distance [20] . Note that G 2 DeNet [33] proposes a metric between Gaussians that leverages the geometry based on Lie group, which can be used to define distillation loss. Besides Gaussian, one can also use Laplace and exponential distributions for modeling feature distributions. Finally, though histogram or kernel density estimation are infeasible, one can still model feature distribution with probability mass function (PMF) and accordingly use discrete WD to define the distillation loss. Details on these methods can be found in Section A.2. 
related works: 
 We summarize KD methods related to ours and show their connections and differences in Table 1 .KL-Div based knowledge distillation. Zhao et al. [3] disclose the classical KD loss [2] is a coupled formulation that limits its performance, and thereby propose a decoupled formulation (DKD) that consists of a binary logit loss for the target category and a multi-class logit loss for all non-target categories. Yang et al. [4] propose a normalized KD (NKD) method, which decomposes the classical KD loss into a combination of the target loss (like the widely used cross-entropy loss) and the loss of normalized non-target predictions. WTTM [5] introduces Rényi entropy regularizer without temperature scaling for student. In spite of competitive performance, they cannot explicitly exploit relations among categories. By contrast, our Wasserstein distance (WD) based method can perform cross-category comparison and thus exploit rich category interrelations.WD based knowledge distillation. The existing KD methods founded on WD [15; 16] mainly concern cross-instance matching for feature distillation, as shown in Figure 3 a global and a local contrastive loss. The former loss minimizes mutual information (via dual form of WD) between the distributions of the teacher and student; the latter loss minimizes Wasserstein distance for the penultimate layer only, where the set of features of the mini-batch images are matched between the teacher and student. Lohit et al. [16] independently propose a similar cross-instance matching method called EMD+IPOT, in which knowledge is transferred from all intermediate layers and discrete WD is computed via an inexact proximal optimal transport algorithm [34] . The differences of our work from them are twofold: (1) They fail to leverage category interrelations which our WKD-L can make full use of; (2) They are concerned with cross-instance matching based on discrete WD, while our WKD-F involves instance-wise matching with continuous WD.Other arts based on statistical modeling. NST [35] is among the first to formalize feature distillation as a distribution matching problem, in which the student mimics the distributions of the teacher based on Maximum Mean Discrepancy (MMD). They show that the 2nd-order polynomial kernel performs best among the candidate kernels of MMD, and that the activation-based attention transfer (AT) [36] is a special case of NST. Yang et al. [37] propose a novel loss function, which transfers the statistics learned by the student back to the teacher based on adaptive instance normalization. Liu et al. [6] propose inter-channel correlations (ICKD-C) to model feature diversity and homology for better knowledge transfer. Both NST and ICKD-C can be regarded as Frobenius norm based distribution modeling via 2nd-moment along spatial-and channel-dimension, respectively, as shown in Figure 3 (right). However, they fail to utilize the geometric structure of the 2nd-moment matrices, which are symmetric positive definite (SPD) and form a Riemannian space [38; 39] . Ahn et al. [40] introduce variational information distillation (VID) based on mutual information. VID assumes that feature distributions are Gaussians, and its loss boils down to mean square loss (i.e., FitNet [24] ) if Gaussians are further assumed to have unit variance. 
experiments: 
 We evaluate WKD for image classification on ImageNet [41] and CIFAR-100 [42] . Also, we evaluate the effectiveness of WKD on self-knowledge distillation (Self-KD). Further, we extend WKD to object detection and conduct experiment on MS-COCO [43] . We train and test models with PyTorch framework [44] , using a PC with an Intel Core i9-13900K CPU and GeForce RTX 4090 GPUs.Image classification. ImageNet [41] contains 1,000 categories with 1.28M images for training, 50K images for validation and 100K for testing. In accordance with [25] , we train the models for 100 epochs using SGD optimizer with a batch size of 256, a momentum of 0.9 and a weight decay of 1e 4. The initial learning rate is 0.1, divided by 10 at the 30th, 60th and 90th epochs, respectively. We use random resized crop and random horizontal clip for data augmentation. For WKD-L, we use POT library [45] for solving discrete WD with η = 0.05 and 9 iterations. For WKD-F, the projector has a bottleneck structure, i.e., a 1×1 Convolution (Conv) and a 3×3 Conv both with 256 filters followed by a 1×1 Conv with BN and ReLU to match the size of teacher's feature maps.Target←|→Non-target Top- Object detection. MS-COCO [43] is a commonly used benchmark for object detection that contains 80 categories. Following the common practice, we use standard split of COCO 2017 with 118K images for training and 5K images for validation. As in ReviewKD [29] , we employ the framework of Faster-RCNN [47] with Feature Pyramid Network (FPN) [48] on Detectron2 platform [49] . As in previous arts [50; 29; 51] , we use the detection models officially trained and released as the teachers, while training the student models whose backbones are initialized with the weights pre-trained on ImageNet. The student networks are trained in 180K iterations with a batch size of 8; the initial learning rate is 0.01, decayed by a factor of 0.1 at 120K and 160K iterations.We analyze key components of WKD-L and WKD-F on ImageNet. We adopt ResNet34 as a teacher and ResNet18 as a student (i.e., setting (a)), whose Top-1 Accuracies are 73.31% and 69.75%, respectively. See Section C.1 for analysis on hyper-parameters, e.g., temperature and weight.How WD performs against KL-Div? We compare WD to KL-Div with (w/) and without (w/o) separation of target probability in Table 2a . For the case of without separation, WD (w/o) improves over KL-Div (w/o) by 1.0%; for the case of with separation, WD (w/) outperforms non-trivially KL-Div (w/) based DKD and NKD. The comparison above clearly shows that (1) WD performs better than KL-Div in both cases and (2) the separation also matters for WD. As such, WD with separation of target probability is used across the paper.How to model category interrelations (IRs)? Table 2b compares two methods for IR modeling, i.e., CKA and cosine. For the former, we assess different kernels while for the latter we evaluate prototypes with classifier weights or class centroids. We note all WD-based methods perform much better than the baseline of KL-Div. Overall, IR with CKA performs better than IR with cosine, indicating it has better capability to represent similarity among categories. For IR with CKA, RBF kernel is better than polynomial kernel, while linear kernel is the best so is used throughout.Full covariance matrix or diagonal one? As shown in Table 3a (3rd and 4th rows), for Gaussian (Full), WD performs better than G 2 DeNet [33] , which suggests that the former metric is more appropriate for feature distillation. When using WD, Gaussian (Diag) (5th row) produces higher accuracy than Gaussian (Full). We conjecture the reason is that high dimensionality of features makes estimation of full covariance matrices not robust [52] ; in contrast, for Gaussians (Diag), we only need to estimate 1D variances for univariate data of single dimension. Besides, Gaussian (Diag) is much more efficient than Gaussian (Full). So we use Gaussians (Diag) throughout the paper. (c) Distillation position and grid scheme.Table 3 : Ablation analysis of WKD-F for image classification (Acc, %) on ImageNet.How to model distributions? In Table 3a , we compare different parametric methods for knowledge distillation, including Gaussian, Laplace, exponential distribution, as well as separate 1st moment and 2nd moment. Note that NST [35] adopts spatial 1st-and 2nd-moment while ICKD-C [6] uses channel 2nd moment. Besides, we compare to non-parametric method based on PMF.For Gaussians (Diag), KL-Div and symmetric (Sym) KL-Div produce similar accuracies which are both lower than WD. The reason may be that KL-related divergences are not intrinsic distances, failing to exploit geometric structure of the manifold of Gaussians. For statistical moments, we see that channel-wise moments perform better than spatial-wise ones. For channel-wise representations, 1st-moment outperforms 2nd-moment, suggesting that the mean plays a more important role. At last, the non-parametric method of PMF underperforms the parametric method of Gaussians.Besides Gaussian (Diag), We can also use univariate Laplace or exponential functions to model distributions of each component of features. For them KL-Div can be computed in closed-form [53] but WD is an unsolved problem. When using KL-Div, Gaussian (Diag) achieves better performance than both Laplace and exponential distributions, highlighting Gaussian as a more suitable option among these parametric alternatives. Further, the Gaussian (Diag) combined with WD yields superior performance compared to KL-Div, suggesting advantage of the Riemannian metric.Instance-wise or cross-instance matching? Our WKD-F is an instance-wise matching method based on continuous WD, while WCoRD and EMD+IPOT concern cross-instance matching for a mini-batch of images based on discrete WD. As seen in Table 3b , WCoRD produces an accuracy much higher than EMD+IPOT, which may be attributed to its extra global contrast loss based on mutual information; WKD-F outperforms the two counterparts by a large margin of 1.0%, which suggests the advantage of our strategy. Note that our WKD-F runs remarkably faster than the two counterparts that rely on optimization algorithm to solve discrete WD.Distillation position and grid scheme. We evaluate in Table 3c Table 4 compares to existing works in two settings. Setting (a) involves homogeneous architecture, where the teacher and student networks are ResNet34 and ResNet18 [9] , respectively; setting (b) concerns heterogeneous architecture, in which we set the teacher as ResNet50 and the student as MobileNetV1 [57] . Refer to Section C.2 for hyper-parameters in Settings (a) and (b).For logit distillation, we compare our WKD-L with KD [2] , DKD [3] , NKD [4] , CTKD [54] and WTTM [5] . Our WKD-L performs better than the classical KD and all its variants in both settings. Particularly, our WKD-L outperforms WTTM, a very strong variant of KD, which additionally introduces a sample-adaptive weighting method. This suggests Wasserstein distance that performs cross-category comparison is superior to category-to-category KL-Div. For feature distillation, we compare to FitNet [24] , CRD [25] , ReviewKD [29] and CAT [55] . Our WKD-F improves ReviewKD, previous top-performer, by ∼0.9% in the setting (a) and ∼0.6% in the setting (b) in terms of top-1 accuracy; this comparison indicates that, for knowledge transfer, matching of Gaussian distributions is better than matching of features. Finally, combination of WKD-L and WKD-F further improves and outperforms strong competitors, including CRD+KD [25] , DPK [7] , FCFD [8] and KD-Zero [56] .More results of combination about WKD-L or WKD-F can be found in Table 9 . We evaluate WKD in the settings where the teacher is a CNN and the student is a Transformer or vice versa. We use CNN models including ResNet (RN) [9] , MobileNetV2 (MNV2) [58] and ConvNeXt [59] , as well as vision transformers that involve ViT [60] , DeiT [61] , and Swin Transformer [62] . The setting of hyper-parameters can be found in Section C.5.For logit distillation, we compare WKD-L to KD [2] , DKD [3] , DIST [63] and OFA [46] . As shown in Table 6 , WKD-L consistently outperforms state-of-the-art OFA for transferring knowledge from Transformers to CNNs or vice versa. For feature distillation, we compare to FitNet [24] , CC [64] , RKD [65] and CRD [25] . WKD-F ranks first across the board; notably, it significantly outperforms the previous best competitors by 2.1%-3.4% in four out of five settings. We attribute superiority of WKD-F to our distribution modeling and matching strategies, i.e., Gaussians and Wasserstein distance. We posit that, for knowledge transfer across CNNs and Transformers that yield very distinct features [46] , WKD-F is more suitable than raw feature comparisons as in FitNet and CRD. We implement our WKD in the framework of Born-Again Network (BAN) [66] for self-knowledge distillation (Self-KD). Specifically, we first train an initial model S 0 using ground truth labels. Then we distill, using WKD-L, the knowledge of S 0 into a student model S 1 with the same architecture as S 0 . For the sake of simplicity, we do not perform multi-generation distillation, such as training a student model S 2 with S 1 as the teacher, etc. We compare with existing methods in two settings, as shown in Table 8 . In RN101→RN18 setting, the teacher is ResNet101 and the student is ResNet18; in RN50→MNV2, the teacher and student are ResNet50 and MobileNetV2 [58] , respectively. For logit distillation, our WKD-L significantly outperforms the classical KD [2] and is slightly better than DKD [3] . For feature distillation, we compare with FitNet, FGFI [50] , ICD [51] and ReviewKD [29] ; our WKD-F improves ReviewKD, the previous top feature distillation performer, by a non-trivial margin in both settings. Finally, by combining WKD-L and WKD-F, we achieve performance better than DKD+ReviewKD [3] . When additional bounding-box regression is used for knowledge transfer, our WKD-L+WKD-F improves further, outperforming previous state-of-the-art FCFD [8] . 
conclusion: 
 The Wasserstein distance (WD) has shown evident advantages over KL-Div in several fields such as generative models [11] . However, in knowledge distillation, KL-Div is still dominant and it is unclear whether WD will outperform. We argue that earlier attempts on knowledge distillation based on WD fail to unleash the potential of this metric. Hence, we propose a novel methodology of WD-based knowledge distillation, which can transfer knowledge from both logits and features. Extensive experiments have demonstrated that discrete WD is a very promising alternative of predominant KL-Div in logit distillation, and that continuous WD can achieve compelling performance for transferring knowledge from intermediate layers. Nevertheless, our methods have limitations. Specifically, WKD-L is more expensive than KL-Div based logit distillation methods, while WKD-F assumes features follow Gaussian distribution. We refer to Section F in Appendix for detailed discussion on limitations and future research. Finally, we hope our work can shed light on the promise of WD and inspire further interest in this metric in knowledge distillation. 
experiment statistical significance: 
 Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments? Answer: [Yes] Justification: We provide the mean accuracy and standard deviation over three runs on CIFAR-100 in Table 6 and Table 11 .Guidelines:• The answer NA means that the paper does not include experiments.• The authors should answer "Yes" if the results are accompanied by error bars, confidence intervals, or statistical significance tests, at least for the experiments that support the main claims of the paper. • The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions). • The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.) • The assumptions made should be given (e.g., Normally distributed errors).• It should be clear whether the error bar is the standard deviation or the standard error of the mean. • It is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have a 96% CI, if the hypothesis of Normality of errors is not verified. • For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are out of range (e.g. negative error rates). • If error bars are reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables in the text. 
experiments compute resources: 
 Question: For each experiment, does the paper provide sufficient information on the computer resources (type of compute workers, memory, time of execution) needed to reproduce the experiments? Answer: [Yes] Justification: We report the computational resources used for our experiments in Section 4 and the training latency of our method in Table 5 .Guidelines:• The answer NA means that the paper does not include experiments.• The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage. • The paper should provide the amount of compute required for each of the individual experimental runs as well as estimate the total compute. • The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't make it into the paper). 
code of ethics: 
 Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics https://neurips.cc/public/EthicsGuidelines? Answer: [Yes] Justification: All our research strictly adheres to the NeurIPS Code of Ethics, and we have discussed the potential societal impact in Section G. Guidelines:• The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.• If the authors answer No, they should explain the special circumstances that require a deviation from the Code of Ethics. • The authors should make sure to preserve anonymity (e.g., if there is a special consideration due to laws or regulations in their jurisdiction). 10. Broader Impacts Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed? Answer: [Yes] Justification: We have discussed this in Section G. Guidelines:• The answer NA means that there is no societal impact of the work performed.• If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does not address societal impact. • Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.• The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is a direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster. • The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology. • If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML). 
safeguards: 
 Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?Answer: [NA] Justification: All our experiments are conducted on open-source datasets. Therefore, our paper does not poses such risks.Guidelines:• The answer NA means that the paper poses no such risks.• Released models that have a high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters. • Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images. • We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make a best faith effort.12. Licenses for existing assets Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?Answer: [Yes]Justification: We have cited the data, code and models used in the paper, and all the assets we utilized are open-source.Guidelines:• The answer NA means that the paper does not use existing assets.• The authors should cite the original paper that produced the code package or dataset.• The authors should state which version of the asset is used and, if possible, include a URL. • The name of the license (e.g., CC-BY 4.0) should be included for each asset.• For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided. • If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswithcode.com/datasets has curated licenses for some datasets. Their licensing guide can help determine the license of a dataset. 
 
 Appendix: A.1 Interrelations (IRs) among Category for WKD-L Visualization of category interrelations. We select in random fifty images from each of one hundred categories randomly chosen on the training set of ImageNet. Then we feed them to a pre-trained ResNet50 model and extract from the penultimate layer features that are projected to 2D space using t-SNE. The 2D embedding points are shown in Figure 4a where different categories are indicated by different colors. For intuitive understanding, inspired by karpathy1 , we display features by the corresponding images at their nearest 2D embedding locations in Figure 4b . It can be seen that the categories exhibit complex topological relations (distances) in the feature space, e.g., mammal species are nearer each other while far from artifact or food. The relations encode abundant information and are beneficial for knowledge distillation. In addition, the features of the same category cluster and form a (unknown) distribution that often overlaps with those of neighboring categories.The observation suggests that it is more desirable to model the interrelations with statistical method. Figure 4 : Visualization of interrelations among 100 categories in feature space. The categories exhibit complex topological relations, where features of the same category cluster and form a distribution that often overlaps with those of neighboring categories.Quantization of IRs with CKA. We propose to use CKA for modeling category IRs as it can effectively characterize similarity of deep representations [17] . CKA is normalized HSIC [18] that measures statistical dependence of random variables (features) by mapping them into a RKHS with some positive definite kernels. Recall that, for category C i , we have a matrix X i ∈ R u×b of u-dimensional features of b training images. As such, for the commonly used kernels, we have linear kernelK lin i = X T i X i , the polynomial kernel K poly i = (X T i X i + 1) k where k ∈ {2, 3, 4}, and RBF kernel K rbf i = exp(- Di 2α 2 Med(Di) )where the bandwidth α ∈ {0.2, 0.4, 0.6} andD i = 2 diag(X T i X i )1 T sym -2X T i X i .For a matrix A, Med(A) denotes the median of all of its entries, diag(A) denotes a vector formed by its diagonals, and (A) sym = 1 2 (A+A T ). Quantization of IRs with cosine similarity. Besides CKA, cosine similarity between the prototypes of two categories is used to quantify category interrelations. The prototype of a category can be naturally computed as feature centroid of this category's training examples, i.e., x i = 1 b X i 1. Alternatively, the weight vectors associated with the softmax classifier can be used as prototypes [22] . Specifically, if the weight matrix of the last FC layer is W ∈ R u×n where n is number of total classes, then its i th column w i can be regarded as the prototype of category C i , i.e., x i = w i .Recall that, for an input image, we have 3D feature maps output by some layer of a DNN, whose spatial height, width and channel number are h, w and l, respectively. We reshape the feature maps to a matrix F ∈ R l×m where m = h × w; we denote the i th column as a feature f i ∈ R l , while the j th row (after transpose) as a feature f j ∈ R m . In WKD-F, we estimate channel 1st-moment µ ∈ R l and 2nd-moment Σ ∈ R l×l :EQUATIONwhich are used to construct a parametric Gaussian N (µ, Σ). For measuring difference between Gaussians, we use Wasserstein distance (WD) that is a Riemannian metric. We prefer Gaussians (Diag) that have diagonal covariances to Gaussians (Full) that have full covariances, as they are much more efficient and have better performance as shown in Table 3a .G 2 DeNet. Different from WD, Wang et al. [33] propose a method called G 2 DeNet, which considers Lie group of the space of Gaussians and embeds a Gaussian into the space of SPD matrices:EQUATION)Thus the difference between two Gaussians is defined as the Euclidean distance of the Gaussian embeddings of the teacher and student models.ICKD-C and NST. ICKD-C [6] uses raw channel 2nd-moment for exploring inter-channel correlations of features, i.e., M = 1 m m i=1 f i f T i . Instead of channel-wise statistics, NST [35] estimates raw spatial 1st-moment µ ∈ R m and raw spatial 2nd-moment M ∈ R m×m for describing distributions:EQUATIONKL-Div between Gaussians (Diag).Let µ = µ 1 , • • • , µ l T be mean and δ = δ 1 , • • • , δ l T bevariance of Gaussians (Diag). In this case, both KL-Div and symmetric KL-Div have closed forms:EQUATIONEQUATIONδ T i ) 2 +( 1 δ S i ) 2 )+( δ S i δ T i ) 2 +( δ T i δ S i ) 2 -2.Here T and S denote the teacher and student, respectively.We assume components of featuref = [f 1 , • • • , f l ] T ∈ R l are statistically independent. Then the Laplace distribution of f is L(µ, ν) = i 1 2νi exp(-|fi-µi| νi), where µ is the mean and ν = ν 1 , • • • , ν l T is the scale parameter. The KL-Div takes the following form [53] :EQUATIONKL-Div between exponential distributions. Under independence assumption, the exponential distribution of featuref is E(β) = i β i exp(-β i f i ), where β = β 1 , • • • , β l T is the rate parameter.For exponential distributions, the KL-Div can be written as [53] :EQUATIONNon-parametric PMD. Though non-parametric methods, e.g., histogram or kernel density estimation, are infeasible due to curse of dimensionality, we can still use probability mass function (PMF) for distribution modeling. Specifically, given a set of features f i , i = 1, . . . , m, the PMF of the features is of the form:EQUATIONwhere ψ(f i ) denotes Kronecker function that is equal to one if f = f i and zero otherwise. Let p T f /p S f be the PMFs of the teacher/student. We can use discrete WD to measure their discrepancy, i.e., D WD (p T f , p S f ); the transport cost c ij between two features f T i and f S j is computed as 1-cos(f T i , f S j ). Note that KL-Div is inapplicable to PMF as it cannot handle non-overlapping distributions [11] .It is worth noting that (1) the space of Gaussians is a Riemannian space, on which WD is an intrinsic distance [20] while KL-Div or its symmetric version are not [12] and thus are unaware of the geometry [13] ; (2) the space of either channel or spatial 2nd-moments is a manifold of symmetric, positive definite matrices rather than a Euclidean space [38; 39] , so the Frobenius norm is not an intrinsic distance [35; 6] and fails to exploit the geometric structure of the manifold.The logit-based WKD-L is formulated as an entropy regularized linear programming, which can be solved fastly by Sinkhorn algorithm [23] . Let n be the dimension of the predicted logits, the complexity of WKD-L can be written as O(Dn 2 log n) [70] . Here D = ∥C∥ 3 ∞ ϵ is a constant, where ∥C∥ ∞ indicates the infinity norm of the transportation cost matrix C = [c ij ], and ϵ > 0 indicates a prescribed error. In contrast, the computational complexity of KL-Div is O(n). Despite its high complexity, WKD-L can be computed efficiently as Sinkhorn algorithm is highly suitable for parallel computation on GPU [23] . For feature-based WKD-F, the dominant cost lies in computation of means and variances. Given a set of m features f i of l-dimension, the means can be computed by global average pooling that takes O(ml) time; the complexity of variances is also O(ml), as it can be obtained by element-wise square operations followed by global average pooling.We adopt the setting (a) for ablation on ImageNet, in which the teacher is ResNet34 and the student is ResNet18.Hyper-parameters of WKD-L. The total loss function of WKD-L consists of the cross-entropy loss L CE , the target loss L t and WD-based logit loss L WKD L . Following [3; 4] , we set the weights of the two former losses to 1 across the paper. As such, our hyper-parameters includes the temperature τ , the weight λ of L WKD L , the sharpening parameter κ that controls smoothness of IRs for transport cost in WD, and the regularization parameter η (set to 0.05) in discrete WD. As simultaneous optimization of them is computationally infeasible, we first analyze the effect of temperature by fixing the sharpening parameter and weight to 1 and 30, respectively. As seen in Figure 5 (top left), across all temperatures WKD-L performs better than the baseline of KD [2] by large margins while achieving the highest accuracy when the temperature is 2. Next, we fix the temperature and weight, analyzing the effect of sharpening parameter. From Figure 5 (top right) we see that performance variation is not large against the sharpening parameter and the best performance is obtained when it is equal to 1. Subsequently, by fixing the best temperature and sharpening parameters, we evaluate the weight of WKD-L. As Figure 5 (bottom left) shows, its accuracy varies smoothly against the weight and achieves the best accuracy when the weight is 30. Finally, Figure 5 (bottom right) illustrates that the performance varies smoothly as a function of the regularization parameter η over a reasonably wide range; for consistency, we set η to 0.05 in all experiments throughout this paper.Hyper-parameters of WKD-F. The loss function of WKD-F method contains the cross-entropy loss L CE and WD-based feature loss L WKD F . As in previous work, we set the weight of the cross-entropy loss to 1. Therefore, there are two hyper-parameters, i.e., the mean-cov ratio γ and the weight of L WKD F . We first analyze the effect of the former by fixing the latter to 2e-2. Table 6 (left) shows the performance as a function of the mean-cov ratio. We see that in the whole range of the mean-cov ratio WKD-F is clearly better than the baseline of FitNet [24] ; the best accuracy is obtained when its value is 2, which suggests that the means play a more important role than the covariances. Next, we fix the mean-cov ratio and evaluate the effect of the weight. As seen in Table 6 (right), WKD-F is much better than the baseline of FitNet across all weights and achieves the best performance when the weight is 0.02. The setting (a) involves homogeneous architecture, where the teacher and student networks are ResNet34 and ResNet18, respectively. For WKD-L, the weights of L CE , L t and L WKD L are 1, 1 and 30, respectively; the temperature τ = 2 and sharping parameter κ = 1. For WKD-F, we set the weights of L CE and L WKD F to 1 and 2e 2, respectively. We adopt features from Conv5_x and Gaussian (Diag) with mean-cov ratio γ = 2. The setting (b) concerns heterogeneous architecture, in which the teacher is ResNet50 and the student is MobileNetV1. In this setting, the weight of L WKD L is set to 25 while that of L WKD F is 1e 3, and the remaining hyper-parameters are identical to those in the setting (a).How WKD complements other KD methods? We combine WKD-F with a state-of-the-art logitbased knowledge distillation (i.e., NKD), and combine WKD-L with a state-of-the-art feature-based method (i.e., ReviewKD). The results are shown in Table 9a and Table 9b , respectively. We can see that NKD+WKD-F improves over individual NKD and WKD-F, which suggests our WKD-F is complementary to NKD. Notably, NKD+WKD-F is slightly inferior to WKD-L+WKD-F. In addition, WKD-L+ReviewKD improves over ReviewKD but underperforms WKD-L. We conjecture that, due to the large gap (∼ 0.9%) between ReviewKD and WKD-L, the combination hurts WKD-L.Will separation of target probability in WKD-L still help when combined with WKD-F? To answer this question, we integrate WKD-L without separation into WKD-F, which is called WKD-L (w/o)+WKD-F. From Table 9c , we can see that it is slightly inferior to WKD-L (w/)+WKD-F in which the separation scheme is used in WKD-L. This suggests that the benefit due to the separation scheme decreases but still persists in the logit+feature approach. (c) Effect of separation scheme on WKD-L+WKD-F.Table 9 : Analysis on different combinations of logit and feature distillation methods.In Table 4 of the main paper, we conduct comparison in the ordinary setting, which is formalized in CRD [25] and adopted by most methods. However, there exist some works which adopt Settings (a) and (b) but with non-trivially different setups, e.g., DIST [63] , NKD [4] , MGD [71] , and DiffKD [72] . Specifically, their students and/or teachers have performance higher than those in the ordinary settings, which make them somewhat advantageous when comparing with the metric of Top-1 accuracy. For a fair comparison, we propose to use gains of the distilled student over the vanilla student in Top-1 accuracy, which is denoted by ▲ . They can more faithfully indicate how much knowledge the student has learned from the teacher. The comparison results are shown in Table 10 .For logit distillation, WKD-L outperforms NKD, a strong KL-Div based variant, by 0.68% and 0.93% in setting (a) and setting (b), respectively. Instead of KL-Div, DIST matches predicted probabilities based on Pearson correlation coefficients, which exploits both instance-wise and cross-instance knowledge, in contrast to WKD-L that uses instance-wise knowledge only. Nevertheless, WKD-L performs better than it by 0.43% in setting (a) and a large margin of 1.19% in setting (b). For feature distillation, we compare to MGD that randomly masks the student's features and then forces the student to generate the teacher's features. The accuracies of WKD-F are over 1.0% higher than those of MGD in both settings. Regarding logit+feature distillation, we contrast with MGD+WSLD and DiffKD. DiffKD denoises the student features through a diffusion model trained by teacher features whose computation cost is high. WKD-L+WKD-F surpasses DiffKD by 0.55% and 1.33% in setting (a) and setting (b), respectively. [25] . ▲ represents the gains of the distilled student over the vanilla student.Additionally, MLKD [73] makes use of a stronger image augmentation, i.e., RandAugment [74] (cf. their implementation), for improving performance. Vanilla KD [75] studies the great potential of vanilla KD in a very different setting, in which optimizer with much longer epochs, diverse and very strong augmentations along with more regularization methods are used. In contrast, we and many of the state-of-the-art methods follow the ordinary setting formalized in CRD [25] .For a fair comparison, we follow OFA [46] and separately tune the hyper-parameters for different settings. For WKD-L, we set the temperature to 2 and the sharpening parameter to 1, while searching the optimal weight of L WKD L in [50, 200] with a step of 50. For WKD-F, the projector is simply a linear layer. We set mean-cov ratio to 2, and perform grid search for the weight of L WKD F in {1, 2, 4, 8}×1e 2. The spatial grid for computing Gaussians (Diag) is set to 1×1. We report average accuracy and standard deviation after three runs in Table 6 ; the results of all competing methods are duplicated from OFA.C.6 Knowledge Distillation within CNN Architectures on CIFAR-100 Experimental Setup. We follow the setting of CRD [25] , where the networks cover ResNet [9] , Wide-ResNet (WRN) [76] , VGG [77] , MobileNetV2 [58] , and ShuffleNetV1 (SNV1) [78] . All models are trained for 240 epochs via the SGD optimizer with a batch size of 64, a momentum of 0.9 and a weight decay of 0.0005. The initial learning rate is 0.01 for MobileNetV2 and ShuffleNetV1 and 0.05 for the remaining networks, divided by 10 at the 150th, 180th, and 210th epochs, respectively. As in CRD [25] , the projector is a 1×1 Conv or 4×4 transposed Conv both with BN and ReLU.For a fair comparison, we follow the practice of DKD [3] , CAT [55] , ReviewKD [29] , FCFD [8] and WTTM [5] , i.e., tuning hyper-parameters separately for different architectures. For WKD-L, we perform grid search for the weight of L WKD L in [50, 800] with a step of 50, the temperature τ in {4, 8} and the sharpening parameter κ in {0.5, 1}. For WKD-F, we perform grid search for the weight of L WKD F in {1, 2, . . . , 50}×1e 2 and mean-cov ratio γ in {2, 3, . . . , 8}. The spatial grid for computing Gaussians (Diag) is searched in {1×1, 2×2, 4×4}. We report the average accuracy and standard deviation of our method after three runs.Results. The comparison results are shown in Table 11 . For logit distillation, WKD-L outperforms the classical KD by large margins in all settings, and surpasses the competitors in 5 out of 6 settings across homogeneous and heterogeneous architectures. For feature distillation, WKD-F invariably achieves better results than EMD-based counterparts (i.e., WCoRD and EMD+IPOT) and 2ndmoment based counterpart (i.e., NST); meanwhile, WKD-F is also very competitive, compared to other state-of-the-art methods. For logit+feature distillation, WKD-L+WKD-F ranks first in 4 out of 6 settings across the board. Overall, our methods have comparable or lower standard deviation, as opposed to the competing methods, which suggests that our method is statistically robust. Following [25; 3] , we visualize the difference of correlation matrices of student and teacher logits. Figure 7a shows visualization results in two settings on CIFAR-100 [42] . In setting of ResNet32x4→ResNet8x4, the teacher and student are ResNet32x4 and ResNet8x4, respectively; the setting VGG13→VGG8 consists of a teacher of VGG13 and a student of VGG8. On the whole, in both settings the colors of WKD-L are much lighter than those of KD [2] , which indicates WKD-L produces correlations matrices more similar to the teacher than KD. As the differences of correlation matrices capture inter-class correlations [25] , smaller differences of WKD-L suggest better-informed cross-category relations than KL divergence can be learned.In addition, we visualize whether WKD-F can learn distributions that are more similar to the teacher.To that end, we use high-level features output from the last Conv layer of a network model for computing distributions, as they encode the most discriminative information for classification. For each validation image of the i th category, we compute WD between feature distributions of the teacher and student. Hence, for n categories each with k validation images, we obtain a n×k matrix of distribution matching. Figure 7b shows visualization of two settings on CIFAR-100. We can see that overall WKD-F demonstrates smaller discrepancies with the teacher than [24] , suggesting it can better mimic the teacher's distributions. We use Grad-CAM [79] to compute class activation maps using features output from the last Conv layer. Figure 8 shows, for three example images, CAMs of different models, including the teacher, vanilla student, distilled models by KD, WKD-L, FitNet and WKD-F. It can be seen that WKD-L and WKD-F have more similar CAMs with the teacher than KD and FitNet, and localize more accurately the important regions of objects. The comparison suggests that WKD-L and WKD-F can learn features with better representation capability. The framework of Faster-RCNN [47] with Feature Pyramid Network (FPN) [48] is much more complicated than that of image classification. It contains the backbone network, Region Proposal Network (RPN), Feature Pyramid Network (FPN), and detection head consisting of a classification branch and a localization branch. For feature distillation, besides RoIAlign layer, FPN can also be used for knowledge transfer [29] .For WKD-L, we use discrete WD to match the probabilities predicted by the classification branches of the teacher and student. We set the temperature τ = 1 and sharpening parameter κ = 2, and set all the weights of L CE , L t and L WKD L to 1 for RN101→RN18. All hyper-parameters of RN50→MV2 are the same as RN101→R18 but κ that is set to 1. For WKD-F, we transfer knowledge from features straightly fed to the classification branch, i.e., features output by the RoIAlign layer. We let RoIAlign generate a high resolution of 18×18 feature maps to exploit more features, and choose a 4×4 spatial grid for computing Gaussians. For both settings, we set the weight of L CE to 1, and that of L WKD F to 5e 3 with mean-cov ratio γ = 2; as in FCFD [8] , we use a projector of 3×3 Conv with BN.We analyze key component of WKD-F specific to object detection on MS-COCO [43] with ResNet101 as the teacher and ResNet18 as the student.Where to distill features: RoIAlign or FPN? Among five stages (P2-P6) of FPN, we select a single P3 with a spatial grid of 16×16 for extracting Gaussians; this option produces the best result among the candidates, and combination of multiple stages brings us no benefit. We compare to the baseline of FitNet [24] and the results are shown in Table 12a . When distilling features of FPN, WKD-F significantly outperforms FitNet (over 2% mAP). By moving the distillation position to RoIAlign layer, the performance of both methods improve, while WKD-F still performs better than FitNet by a non-trivial margin.How spatial size of RoIAlign features affect performance? In Faster-RCNN, the RoIAlign layer outputs standard 7×7 feature maps. To exploit more spatial information, we let RoIAlign layer output feature maps of higher resolution. Table 13 : Our WKD benefits state-of-the-art KD methods for object detection.We assess whether our methodology is complementary to two state-of-the-art methods, i.e., ReviewKD and FCFD. To this end, for ReviewKD, we add the losses of WKD-L and WKD-F into the original loss functions; the weights of our two losses and hyper-parameters are same as those in Section E.1.For FCFD, we replace respectively the original KD loss and feature distillation loss (i.e., FitNet loss) by the losses of WKD-L and WKD-F whose weights are identical to those specified in the main paper. The results are shown in Table 13 . We can see that, by integrating our methods, both FCFD and ReviewKD improve by non-trivial margins, which indicates that our methodology is parallel to them and thus can enhance their performance.Besides common logit distillation and feature distillation, FCFD [8] additionally uses BB regression for knowledge transfer and achieves state-of-the-art performance. Here we also introduce it into our methodology. Specifically, for each proposal from RPN of the student [47] , both student and teacher perform BB regression, predicting separately a localization offset vector (LOV) o, from which we obtain predicted bounding box B of the target class. Then the distillation loss is written asEQUATIONwhere L 2 indicates the square of Euclidean distance between the student's LOV vector o S and the teacher's one o T . We use the Distance-IoU loss L DIoU defined by Zheng et al. [80] , which measures the Intersection over Union (IoU) between two bounding-boxes B S and B T with a penalty term. The constant ξ is to balance the two losses and is set to 20 throughout.The cost of our WKD-L is higher than KL-Div based methods due to the regularized linear programming [23] . However, it is affordable as shown in Table 5 and will benefit from advance of faster algorithms for solving WD [45] . Potentially, WKD-L can generalize to be a label smoothing regularization method. Specifically, besides the CE loss, one introduces an additional loss that measures WD rather than KL-Div between the logits and a uniform distribution [81] , while computing IRs using embeddings of textual category names as 'prototypes' via ready-made LLMs [82] . Besides BAN based Self-KD in Section 4.5, it is promising to further generalize WKD-L for teacher-free knowledge distillation, e.g., by using customized soft labels [4] and IRs based on weights of the softmax classifier.Our WKD-F models distribution of features with Gaussians. As deep features of DNNs are generally of high-dimension and small-size, accurate estimation of covariance matrices is difficult [52] . Therefore, exploration of robust and efficient methods for estimating Gaussian may further improve the performance of WKD-F. Besides, the Gaussians may be sub-optimal for modeling feature distributions. We address limitations of Kullback-Leibler Divergence for knowledge distillation (KD), and can obtain stronger, lightweight models suitable for resource-limited devices. Our methodology is very promising, readily applicable to a variety of visual tasks, e.g., image classification and object detection. We hope our work sheds light on importance of WD and inspires future exploration of it in the field of knowledge distillation.With breakthroughs of large-scale pre-training, multimodal large language models (LLMs) like GPT-NeurIPS Paper Checklist • The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper. • The authors are encouraged to create a separate "Limitations" section in their paper.• The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be. • The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on a few datasets or with a few runs. In general, empirical results often depend on implicit assumptions, which should be articulated. • The authors should reflect on the factors that influence the performance of the approach.For example, a facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or a speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon. • The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size. • If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness. • While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an important role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations. • The experimental setting should be presented in the core of the paper to a level of detail that is necessary to appreciate the results and make sense of them. • The full details can be provided either with the code, in appendix, or as supplemental material.https://cs.stanford.edu/people/karpathy/cnnembed[83] have excelled in many visual tasks. Our methodology can be potentially applied to transfer knowledge from LLMs to smaller ones for specific visual or language tasks, allowing for improved performance while preserving fast inference cost. Consequently, researchers as well as practitioners can more easily benefit from advanced technologies of LLMs, facilitating their widespread use and broader accessibility.Currently, the theoretical quest for KD is limited. Due to the black-box nature of the distillation process, the student distilled using our methods may inevitably inherit harmful biases from the teacher[84]. Moreover, the reduction in deployment cost may lead to more potential harms of model abuse. This highlights the necessity for more widespread efforts to regulate the use of artificial intelligence techniques including knowledge distillation. 


Decision: The paper provides an interesting perspective of WD in distillation and a new algorithm for knowledge distillation. After initial rating, the authors provide a nice rebuttal which addressed all the issues mentioned by the reviewers. Now, all the reviewers though there is no further issues with this paper and the AC also learns toward an acceptance for this submission.

Reviews: review: paper_summary: paper_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.
strength_weakness: strength_weakness: 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: questions: None
limitations: limitations: The cost expensiveness has been mentioned.
paper_summary: paper_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.
strength_weakness: strength_weakness: * 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: 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?
limitations: limitations: * please check the questions and weaknesses.
paper_summary: 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.
strength_weakness: strength_weakness: 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: 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.
limitations: limitations: The authors discussed the limitations.
paper_summary: paper_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.
strength_weakness: strength_weakness: (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: 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?
limitations: limitations: Limitations are discussed in the paper

Comments: title: Response by Reviewer Zit9
comment: Dear Reviewer F2zQ,

Thank you for your positive feedback. We are delighted to hear that our rebuttal has effectively addressed your concerns. The responses to your constructive comments will be incorporated into the revised paper. We would greatly appreciate it if you could consider reflecting this in your updated score.
Thanks for the responses from the authors.

My concerns have been addressed.
Dear Reviewer Zit9,

Thank you for your positive feedback on our rebuttal and for raising your score. We are pleased to hear that your concerns have been satisfactorily addressed. The responses to your constructive comments will be incorporated into the revised paper.
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!
Dear Reviewer 5GY4,

We are pleased to learn that we have effectively addressed your concerns. We appreciate your very positive comments on the soundness, presentation and contribution of our paper.
Thanks for the rebuttal. It clarifies my questions and I will maintain the score.
Dear reviewers, could you share your feedback to the authors after reading their response? Does the authors' response address your concerns well?