Contrastive Representation Distillation

TL;DR

Knowledge Distillation often treats outputs independently, missing the teacher’s representational structure. Contrastive Representation Distillation (CRD) introduces a contrastive objective that maximizes a lower bound on the mutual information between teacher and student representations, capturing inter-feature dependencies with a memory-bank-based negative sampling strategy. CRD achieves state-of-the-art performance across model compression, cross-modal transfer, and ensemble distillation, outperforming KD and many prior distillers, and can be effectively combined with KD for further gains. The approach links distillation and representation learning, offering scalable improvements and practical benefits for transferability and ensembling scenarios.

Abstract

Often we wish to transfer representational knowledge from one neural network to another. Examples include distilling a large network into a smaller one, transferring knowledge from one sensory modality to a second, or ensembling a collection of models into a single estimator. Knowledge distillation, the standard approach to these problems, minimizes the KL divergence between the probabilistic outputs of a teacher and student network. We demonstrate that this objective ignores important structural knowledge of the teacher network. This motivates an alternative objective by which we train a student to capture significantly more information in the teacher's representation of the data. We formulate this objective as contrastive learning. Experiments demonstrate that our resulting new objective outperforms knowledge distillation and other cutting-edge distillers on a variety of knowledge transfer tasks, including single model compression, ensemble distillation, and cross-modal transfer. Our method sets a new state-of-the-art in many transfer tasks, and sometimes even outperforms the teacher network when combined with knowledge distillation. Code: http://github.com/HobbitLong/RepDistiller.

Figures (6)

• Figure 1: The three distillation settings we consider: (a) compressing a model, (b) transferring knowledge from one modality (e.g., RGB) to another (e.g., depth), (c) distilling an ensemble of nets into a single network. The constrastive objective encourages the teacher and student to map the same input to close representations (in some metric space), and different inputs to distant representations, as indicated in the shaded circle.
• Figure 2: The correlations between class logits of a teacher network are ignored by regular cross-entropy. Distillation frameworks use "soft targets" hinton2015distilling which effectively capture such correlations and transfer them to the student network, leading to the success of distillation. We visualize here the difference of correlation matrices of student and teacher logits, for different student networks on a CIFAR-100 knowledge distillation task: (a) Student trained without distillation, showing that the teacher and student cross-correlations are very different; (b) Student distilled by attention transfer zagoruyko2016paying; showing reduced difference (see axis); (c) Student distilled by KL divergence hinton2015distilling, also showing reduced difference; (d) Student distilled by our contrastive objective, showing significant matching between student's and teacher's correlations. In this visualization, we use WRN-40-2 as teacher and WRN-40-1 as student.
• Figure 3: Top-1 classification accuracy on STL-10 using chrominance image ($ab$ channel in $Lab$ color space). We initialize the chrominance network randomly or by distilling from a luminance network, trained with large-scale labeled images. We evaluate distillation performance by (a) linear probing and (b) fully finetuning.
• Figure 4: Distillation from an ensemble of teachers. We vary the number of ensembled teachers and compare KD with our CRD by using (a) WRN-16-2 and (b) ResNet20. Our CRD consistently achieves lower error rate.
• Figure 5: Effects of varying the number of negatives, shown in (a), or the temperature, shown in (b).