Source-linked AI summary
Wasserstein Contrastive Representation Distillation
Liqun Chen, Dong Wang, Zhe Gan, Jingjing Liu, Ricardo Henao, Lawrence Carin
TL;DR
Knowledge distillation transfers teacher information to a more compact student, but KL-based methods may miss structural knowledge and feature generalization across different tasks can be limited. WCoRD combines dual-form global contrastive transfer with primal-form local feature matching. Experiments report consistent student-performance improvements across a variety of tasks.
Problem
KL-based distillation may miss structural information in teacher features, while feature generalization can be limited when teacher and student address different tasks.
Source: §1 · p. 1, §3.1 · p. 3Method
WCoRD uses the dual Wasserstein form for global contrastive transfer and the primal form for local feature-distribution matching.
Source: §1 · p. 1, §3.1 · p. 4, §3.2 · p. 4Results
Across a variety of tasks, WCoRD consistently improves student-model performance.
Source: §6 · p. 8Takeaways & Limitations
WCoRD provides a unified framework that combines global contrastive learning with local feature-distribution matching for knowledge distillation.
Source: §1 · p. 2, §5.1 · p. 6
Abstract
from arXiv · showhide
The primary goal of knowledge distillation (KD) is to encapsulate the information of a model learned from a teacher network into a student network, with the latter being more compact than the former. Existing work, e.g., using Kullback-Leibler divergence for distillation, may fail to capture important structural knowledge in the teacher network and often lacks the ability for feature generalization, particularly in situations when teacher and student are built to address different classification tasks. We propose Wasserstein Contrastive Representation Distillation (WCoRD), which leverages both primal and dual forms of Wasserstein distance for KD. The dual form is used for global knowledge transfer, yielding a contrastive learning objective that maximizes the lower bound of mutual information between the teacher and the student networks. The primal form is used for local contrastive knowledge transfer within a mini-batch, effectively matching the distributions of features between the teacher and the student networks. Experiments demonstrate that the proposed WCoRD method outperforms state-of-the-art approaches on privileged information distillation, model compression and cross-modal transfer.
1. Introduction
Knowledge distillation can transfer teacher knowledge, but prediction-based KL objectives may miss structural information and intermediate-feature methods may not generalize across teachers and students. WCoRD uses both Wasserstein forms for global contrastive transfer and local feature-distribution matching.
1. Introduction
- Real-world deployment and training face constraints from growing model sizes, restricted data access, scarce specialized labels, and limited cross-modal data use.The cited settings include computationally restricted phones, sensors, and edge devices, alongside private, classified, or institutional datasets.Source: §1 · p. 1
- KL-based distillation can transfer limited information from low-dimensional predictions, and those predictions may be irrelevant when teacher and student address different tasks.Intermediate representations offer richer signals, but existing methods often compare poorly with basic KD and rely on copying teacher behavior that may not generalize to the student.Source: §1 · p. 1
- WCoRD combines dual-form global contrastive transfer with primal-form local feature-distribution matching to improve knowledge transfer and generalization.The dual form maximizes mutual information between teacher and student representations; the primal form matches their feature distributions while allowing coupling across multiple mini-batch examples.Source: §1 · pp. 1–2
- The framework is evaluated on model compression, cross-modal transfer, and privileged information distillation.The authors describe comprehensive benchmark experiments for the first two settings and a real-world dataset for privileged information distillation.Source: §1 · p. 2
2. Background
Standard knowledge distillation trains a student with teacher feedback, including a term that aligns their output distributions. WCoRD instead motivates Wasserstein-based transfer using feature representations, with primal and dual forms providing complementary distribution-matching formulations.
2.1. Knowledge Distillation
- Standard KD supplements label supervision by encouraging the student’s output distribution to resemble the teacher’s, but prediction distributions collapse structured hidden features.The passage also notes that KL divergence can be numerically unstable and may overemphasize small differences, although KL-based objectives remain effective and popular.Source: §2.1 · p. 2
- WCoRD uses penultimate-layer teacher and student features and introduces global contrastive learning and local distribution matching as Wasserstein-based transfer mechanisms.The stated aim is to provide a general, principled distillation framework while avoiding some drawbacks of KL-based approaches.Source: §2.1 · p. 2
2.2. Wasserstein Distance
- Contrastive learning moves similar samples closer and different samples farther apart in feature space, and the paper extends this idea using Wasserstein distance.The passage identifies Wasserstein distance as also known as Earth Mover’s Distance or Optimal Transport Distance.Source: §2.2 · p. 2
- The Wasserstein distance is presented in primal and dual forms, with the primal formulation using couplings between distributions and the dual formulation using a 1-Lipschitz function.The primal form uses a point-wise cost and a joint distribution over paired samples; in the dual form, the function may be parameterized as a neural network.Source: §2.2 · pp. 2–3
3. Method
WCoRD combines dual-form global contrastive transfer with primal-form local feature-distribution matching. The two objectives are described as complementary: each can provide a feature space that supports the other.
3.1. Global Contrastive Knowledge Transfer
- The dual-form GCKT objective uses congruent and incongruent teacher–student feature pairs to estimate a lower bound on mutual information.Congruent pairs use the same input for both networks; incongruent pairs use different inputs. A memory buffer supplies negatives from prior mini-batches, and the critic is constrained to be 1-Lipschitz using spectral normalization.Source: §3.1 · pp. 3–4
3.2. Local Contrastive Knowledge Transfer
- The primal-form LCKT matches feature distributions within a mini-batch, allowing a student feature to couple with multiple teacher examples rather than only one.The transport problem is solved with entropy-regularized Sinkhorn iterations; gradients through the iterative transport-plan calculation can be ignored.Source: §3.2 · p. 4
3.3. Unifying Global and Local Knowledge Transfer
- WCoRD combines global matching of joint and marginal feature distributions with local optimal-transport matching of teacher and student feature marginals.The global objective maximizes a lower bound on mutual information, while the local objective minimizes a Wasserstein-distance objective.Source: §3.3 · pp. 4–5
- GCKT and LCKT are presented as complementary: each can provide a more constrained or better feature space for the other’s distribution-matching objective.LCKT targets marginal-distribution discrepancy, while GCKT targets the joint distribution against the product of marginals.Source: §3.3 · p. 5
4. Related Work
The related work reviews output- and feature-based knowledge distillation, prior uses of optimal transport in deep learning, and contrastive learning’s connection to mutual information. It positions WCoRD as applying both Wasserstein forms together to knowledge distillation.
Knowledge Distillation
- Prior knowledge-distillation methods use teacher output distributions or intermediate representations, including attention transfer, representation-space methods, and contrastive objectives.The section also describes alternative teacher–student arrangements and an interpretation of distillation as label smoothing.Source: §4 · p. 5
Optimal Transport
- Optimal transport has been used in GANs and other deep-learning tasks, but the authors describe this work as the first to apply it to knowledge distillation using both primal and dual forms.The dual form uses a 1-Lipschitz critic, while the primal form can be solved with Sinkhorn.Source: §4 · p. 5
Contrastive Learning
- Contrastive objectives have been interpreted as maximizing a lower bound on mutual information, although the role of mutual information versus loss form remains unsettled.The related work connects WCoRD’s global transfer to Wasserstein dependency measures while distinguishing its joint use of primal and dual forms for distillation.Source: §4 · p. 5
5. Experiments
The experiments evaluate WCoRD on model compression, cross-modal transfer, and privileged-information distillation. Results report gains over CRD in compression and cross-modal settings, and over student and CRD baselines in privileged distillation.
5.1. Model Compression
- 48.63% with shared architectures and 43.27% with different architectures are the reported average relative improvements of WCoRD over CRD.WCoRD+KD usually improves performance further, though WCoRD alone performs better for ResNet50 → MobileNetV2 and ResNet50 → VGG8.Source: Table 1 · p. 6, §5.1 · p. 6, Table 3 · p. 7
- In WCoRD’s LLCKT-loss weight ablation (λ1 = 0.8), with standard deviations reported over five runs, WCoRD consistently outperforms GCKT and CRD for λ2 in (0, 0.2], while ResNet-8x4 performs best at λ2 = 0.05.Source: §5.1 · p. 6
- For ImageNet with a ResNet-34 teacher and ResNet18 student, WCoRD+KD achieves the best student performance; WCoRD improves over CRD by 44.4% on Top-1 error and 23.08% on Top-5 error.The reported metrics are error rates, for which lower is better.Source: Table 2 · p. 6, §5.1 · p. 7
5.2. Cross-Modal Transfer
- WCoRD outperforms CRD on features from the 2nd–4th residual blocks in cross-modal transfer.The setup transfers from a luminance-trained Tiny-ImageNet teacher to a chrominance-trained STL-10 student and evaluates frozen student features by linear probing.Source: §5.2 · pp. 7–8
5.3. Privileged Information Distillation
- In privileged-information distillation, WCoRD improves over the basic student by 2.34% and over CRD by 0.96% on GA identification.The teacher is pretrained on a different OCT dataset and classification task, while the student is trained on OCT-GA; KL divergence cannot be used in this setting.Source: §5.3 · p. 8
6. Conclusions
WCoRD is a knowledge-distillation framework that combines Wasserstein contrastive learning with feature-distribution matching. Experiments across varied tasks show consistent improvements in student performance.
6. Conclusions
- WCoRD combines Wasserstein-based contrastive learning with an additional feature-distribution matching term.The framework is presented for knowledge distillation.Source: §6 · p. 8
- Experiments on a variety of tasks show that WCoRD consistently improves student-model performance.Source: §6 · p. 8
- Federated learning and adversarial robustness are identified as possible future applications.Source: §6 · p. 8
A. WCoRD Algorithm
The WCoRD algorithm processes a mini-batch through teacher and student feature extraction, global and local contrastive transfer, and a task-specific loss.
A. WCoRD Algorithm
- The algorithm takes a mini-batch of samples and extracts features from the teacher and student networks.Source: A. WCoRD Algorithm · p. 11
- It stores previously computed features in a memory buffer, then maximizes the GCKT loss and minimizes the LCKT loss over the student parameters.Source: A. WCoRD Algorithm · p. 11
- The algorithm also minimizes a task-specific loss over the student parameters.Source: A. WCoRD Algorithm · p. 11
B.1. Baseline Methods
The evaluation compares WCoRD with a broad set of distillation baselines, with baseline hyperparameters following the CRD setup.
B.1. Baseline Methods
- The compared methods include FitNets [44], Knowledge Distillation [25], Attention Transfer [62], NST [27], FSP, and PKT [40].Source: §B.1 · p. 11
- Other baselines are Factor Transfer [29], Similarity-Preserving KD [55], Correlation Congruence [41], VID [1], RKD [39], and activation-boundary transfer [24].Source: §B.1 · p. 11
- Contrastive Representation Distillation (CRD) [53] via NCE is also compared, and baseline hyperparameters follow CRD [53].Source: §B.1 · p. 11
B.2. Model Architectures
The experiments use several network families, including residual, VGG, ShuffleNet, MobileNetV2, and InceptionNet architectures.
B.2. Model Architectures
- The architecture set includes WRN [63], with WRN-d-w denoting depth d and width factor w.Source: §B.2 · p. 11
- CIFAR-style ResNet-d [21] uses three groups of basic blocks with 16, 32, and 64 channels; ResNet-8x4 and ResNet-32x4 use four times the width.Source: §B.2 · p. 11
- ImageNet-style ResNet [21] uses bottleneck blocks and more channels, while the experiments use MobileNetV2 [46] with a width multiplier of 0.5.Source: §B.2 · pp. 11–12
- The VGG networks [47] are adapted from ImageNet versions, ShuffleNets [66] [52] are adapted to 32x32 inputs, and InceptionNet-v3 [51] is used as the privileged-distillation teacher.Source: §B.2 · p. 12
C. Additional Results
Additional results show that WCoRD performs better in further evaluations and that its GCKT module contributes to performance. The model is fairly robust to choices of λ1 when λ2 is fixed at 0.1, but LCKT alone does not achieve very good performance.
C. Additional Results
- WCoRD achieves better performance in additional evaluations combining baseline distillation methods with KD.Table 8 also reports standard deviations for CRD and WCoRD, with or without KD.Source: C. Additional Results · p. 12
- With LCKT fixed at λ2 = 0.1, WCoRD is fairly robust to different choices of λ1.The tested λ1 values range from 0 to 1.0.Source: C. Additional Results · p. 12
- Models using LCKT without the GCKT module do not obtain very good performance.Source: C. Additional Results · p. 12