Research map / Label-efficient and robust learning / Knowledge distillation
Wasserstein Distance Rivals Kullback-Leibler Divergence for Knowledge Distillation
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Proposes Wasserstein distance based knowledge distillation: WKD-L (discrete WD over teacher/student class probabilities with category interrelations from CKA) for logits and WKD-F (closed-form continuous WD between channel-wise Gaussians) for intermediate features, evaluated on ImageNet, CIFAR-100 and MS-COCO.
Problem
KL divergence in knowledge distillation compares only same-category probabilities (no cross-category comparison) and is ill-suited to intermediate-layer features, since high-dimensional sparse features yield non-overlapping distributions and KL is not a metric or geometry-aware.
Approach
Replace KL divergence with Wasserstein distance: for logits, use discrete WD with a ground cost derived from CKA-based category interrelations so probability mass can move across categories; for features, model each feature map as a Gaussian (channel-wise mean and covariance) and use the closed-form Gaussian WD, a Riemannian metric.
Key results
- On ImageNet ResNet34->ResNet18, WKD-L reaches 72.49% Top-1 vs 71.03% for KD, and WKD-L+WKD-F reaches 72.76%.
- WKD-F reaches 72.50% Top-1 on ImageNet (ResNet34->ResNet18), improving ReviewKD (71.61%) by about 0.9%.
- On MS-COCO RN101->RN18, WKD-F gets 37.21 mAP vs 36.75 for ReviewKD; WKD-L+WKD-F reaches 37.49 mAP.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Use WKD-L instead of NKD for logit distillation | ImageNet distillation, ResNet34 → ResNet18 and ResNet50 → MobileNetV1 · ImageNet | Top-1 accuracy: 71.96 -> 72.49; 72.58 -> 73.17 | Helped |
| Use WKD-F instead of ReviewKD for feature distillation | ImageNet distillation, ResNet34 → ResNet18 and ResNet50 → MobileNetV1 · ImageNet | Top-1 accuracy: 71.61 -> 72.50; 72.56 -> 73.12 | Helped |
| Combine WKD-L and WKD-F | ImageNet distillation, ResNet34 → ResNet18 and ResNet50 → MobileNetV1 · ImageNet | Top-1 accuracy: 72.49 -> 72.76; 73.17 -> 73.69 | Helped |
| Use WKD-L instead of OFA for logit distillation | CNN–Transformer cross-architecture transfer · CIFAR-100 | Top-1 accuracy: 80.54 -> 81.42±0.22; 80.15 -> 80.81±0.21; 78.45 -> 79.04±0.05; 75.76 -> 76.11±0.18; 78.32 -> 78.94±0.17 | Helped |
| Use WKD-F instead of the strongest listed feature baseline | CNN–Transformer cross-architecture transfer · CIFAR-100 | Top-1 accuracy: 78.87 -> 81.57±0.12; 77.71 -> 81.12±0.24; 78.14 -> 79.11±0.07; 69.79 -> 73.27±0.22; 72.63 -> 74.80±0.13 | Mixed |
| Replace KL-Div with WD for logit distillation without target/non-target separation (ablation) | WKD-L · ImageNet, ResNet34 → ResNet18 | Top-1 accuracy: 71.03 -> 72.04 | Helped |
| Add target/non-target probability separation (ablation) | WKD-L · ImageNet, ResNet34 → ResNet18 | Top-1 accuracy: 72.04 -> 72.49 | Helped |
| Use linear CKA rather than cosine class centroids to model interrelations (ablation) | WKD-L · ImageNet, ResNet34 → ResNet18 | Top-1 accuracy: 72.02 -> 72.49 | Helped |
| Use diagonal rather than full Gaussian covariance with WD (ablation) | WKD-F · ImageNet, ResNet34 → ResNet18 | Top-1 accuracy: 72.37 -> 72.50 | Helped |
| Use Gaussian (Diag) with WD rather than Gaussian (Diag) with KL-Div (ablation) | WKD-F · ImageNet, ResNet34 → ResNet18 | Top-1 accuracy: 71.75 -> 72.50 | Helped |
| Use a 2×2 grid rather than a 1×1 grid for feature distribution matching (ablation) | WKD-F · ImageNet, ResNet34 → ResNet18 | Top-1 accuracy: 72.50 -> 72.40 | Hurt |
| Match features from Conv4x and Conv5x rather than Conv5x alone (ablation) | WKD-F · ImageNet, ResNet34 → ResNet18 | Top-1 accuracy: 72.50 -> 72.44 | Hurt |
Surprising findings
- Earlier discrete-WD distillation methods (EMD+IPOT) underperform FitNet on ImageNet (70.46 vs 70.53), suggesting cross-instance discrete WD alone is not sufficient.
- Gaussian (Diag) with WD outperforms Gaussian (Full) in accuracy, attributed to ill-conditioned full covariance estimates in high dimensions.
Limitations
- WKD-L is more computationally expensive than KL-Div logit distillation (about 1.3x latency).
- WKD-F assumes features follow a Gaussian distribution.
- Experiments use a single PC with RTX 4090 GPUs; the paper does not report results beyond ResNet/ViT/MobileNet-scale models.
Experimental setup
Datasets: ImageNet CIFAR-100 MS-COCO
Benchmarks: ImageNet CIFAR-100 MS-COCO
Metrics: Top-1 accuracy Top-5 accuracy mAP AP50 AP75 training latency
Baselines: KD DKD NKD CTKD WTTM FitNet CRD ReviewKD CAT CRD+KD DPK FCFD
Builds on: Hinton knowledge distillation DKD NKD WTTM CKA WCoRD EMD+IPOT
Closest papers by meaning
- CrossKD: Cross-Head Knowledge Distillation for Object Detection (CVPR 2024 · similarity 0.87)
- DCSF-KD: Dynamic Channel-wise Spatial Feature Knowledge Distillation for Object Detection (AAAI 2025 · similarity 0.87)
- Decoupled Kullback-Leibler Divergence Loss (NeurIPS 2024 · similarity 0.86)
- DFD: Distilling the Feature Disparity Differently for Detectors (ICML 2024 · similarity 0.86)
- Scaled Decoupled Distillation (CVPR 2024 · similarity 0.86)
- VkD: Improving Knowledge Distillation using Orthogonal Projections (CVPR 2024 · similarity 0.86)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.