A Contrastive Knowledge Transfer Framework for Model Compression and Transfer Learning

Let $X=\{x_{i}\}_{i=1}^{B}$ and $Y=\{y_{i}\}_{i=1}^{B}$ denote a set of inputs with a batch size of $B$ and its ground truth labels $Y$, respectively. We define a module as a group of convolution layers. The output representations of the modules from the teacher and student can be described as $\{T_{m}\}_{m=1}^{M}$ and $\{S_{m}\}_{m=1}^{M}$, respectively, where $M$ denotes the number of modules. Similarly, let $T_{h}$ and $S_{h}$ denote the output vectors of the penultimate layer from the teacher and student, respectively.

Figure 1 illustrates the workflow of the proposed Contrastive Knowledge Transfer Framework (CKTF). The optimization objective in CKTF consists of three components: 1) the cross entropy loss with the ground truth labels; 2) the proposed contrastive loss to transfer knowledge from the intermediate representations $\{T_{m}\}_{m=1}^{M}$ and the penultimate layer $T_{h}$ of the teacher to $\{S_{m}\}_{m=1}^{M}$ and $S_{h}$ of the student, respectively; and 3) the distillation loss from other KT methods. The loss function of CKTF is as follows:

$$\small\begin{split}L=\gamma L_{CE}(Y,S_{h})+&L_{CKT}(\{T_{m}\}_{m=1}^{M},\{S_{m}\}_{m=1}^{M},T_{h},S_{h})\\ +&\theta L_{Distill}(T_{h},S_{h})\end{split}$$

(1)

where $\gamma$ equals either 1 or 0 depending on the availability of labels, and $\theta$ is a hyper-parameter that controls the weight of the loss term.

The first loss term in Eq. 1 enforces the supervised learning from labels, which is typically implemented as a cross-entropy loss for classification tasks: $L_{CE}(Y,S_{h})=\sum_{i=1}^{c}[Y_{i}log(S_{h,i})+(1-Y_{i})log(1-S_{h,i})]$, where $c$ denotes the number of classes of the dataset. The second loss term is the proposed contrastive loss that transfers high-dimension structural knowledge from both the intermediate presentations and the penultimate layer via contrastive learning. It works because, as opposed to just transferring information about conditionally independent output class probabilities, the multiple contrastive objectives constructed in $L_{CKT}$ better transfer all the information in the teacher’s representational space (see Section 2.2 for details). The third loss term is used to incorporate existing KT methods into CKTF. For example, for the conventional KD [2], it is defined as the KL-divergence-based loss that minimizes the difference between the teacher’s and student’s soft logits: $L_{Distill}(T_{h},S_{h})=KL(softmax(T_{h}/\rho)||softmax(S_{h}/\rho))$, where $\rho$ is the temperature. In this way, CKTF can help improve the performance of existing KT methods (see Section 3.1 for evaluation results).

Note that, in transfer learning, $\gamma$ (Eq. 1) is set to zero since supervision from labels is not available.

CKTF constructs the contrastive loss across intermediate representations from multiple modules of the teacher and student. Directly using intermediate representations $\{T_{m}\}_{m=1}^{M}$ and $\{S_{m}\}_{m=1}^{M}$ to perform contrastive learning is infeasible, since 1) the dimension between $T_{m}$ and $S_{m}$ might be different, and 2) the huge feature dimension of $T_{m}$ and $S_{m}$ may cause memory issues or significantly increase the training time. In detail, the dimension of $S_{m}$ is calculated as $|S_{m}|=B\times o_{m}^{s}\times(k_{m}^{s})^{2}$, where $B$, $o_{m}^{s}$ and $k_{m}^{s}$ denote the batch size, output dimension, and kernel size of the module $m$ of the student. For example, for ResNet-50 on Tiny-ImageNet with a batch size of 32, the feature dimension of one intermediate module can be: $32\times 1024\times 16^{2}\approx 8.39$ millions, and its teacher may also have a similar level of the feature dimension.

To solve the above problem, CKTF first applies an average pooling over $T_{m}\in R^{B\times o_{m}^{t}\times k_{m}^{t}\times k_{m}^{t}}$ and $S_{m}\in R^{B\times o_{m}^{s}\times k_{m}^{s}\times k_{m}^{s}}$, respectively, and it produces the output $\bar{T}_{m}\in R^{B\times o_{m}^{t}\times 1\times 1}$ and $\bar{S}_{m}\in R^{B\times o_{m}^{s}\times 1\times 1}$, respectively. Then it uses a reshaping function $h(\cdot)$ that changes the 4-D $\bar{T}_{m}$ and $\bar{S}_{m}$ to a 2-D space, yielding $H^{T}_{m}\in R^{B*o^{t}_{m}}$ and $H^{S}_{m}\in R^{B*o^{s}_{m}}$, respectively:

$$\vspace{-3pt}\begin{split}&\bar{S}_{m}=AvgPool(S_{m}),\bar{T}_{m}=AvgPool(T_{m})\\ &H_{m}^{S}=h(\bar{S}_{m}),H_{m}^{T}=h(\bar{T}_{m})\end{split}\vspace{-3pt}$$

(2)

Next, a projection network $g(\cdot)$ takes the presentations $\{H_{m}^{T}\}_{m=1}^{M}$ and $\{H_{m}^{S}\}_{m=1}^{M}$ as the input, and for the module $m$, it produces: $G^{T}_{m}=g(H_{m}^{T})\in R^{B\times d}$ and $G^{S}_{m}=g(H_{m}^{T})\in R^{B\times d}$, respectively, where $d$ denotes the output dimension of the projection network. $g(\cdot)$ used in CKTF is a single linear layer of size $d=128$ followed by the $\ell_{2}$ normalization. Note that $g(\cdot)$ is discarded after training, so we do not change the model architecture. We will show that the linear projection is better than Multi-Layer Perceptron (MLP) projection used in representation learning [15, 16, 17, 18] and discuss the effect of $d$ in Section 3.3.

CKTF constructs the contrastive loss using $\{G^{T}_{m}\}_{m=1}^{M}$ and $\{G^{S}_{m}\}_{m=1}^{M}$. Given a batch of random samples $X=\{x_{i}\}_{i=1}^{B}$, we define positive representation pairs as $(G_{m,i}^{S},G_{m,i}^{T})$, which are the outputs of the student’s and teacher’s module $m$ from the same input sample $x_{i}$, and negative representation pairs as $(G_{m,i}^{S},G_{m,j}^{T})$ from their modules’ outputs given two different data samples $x_{i}$ and $x_{j}$, respectively.

CKTF aims to push closer each positive pair $G_{m,i}^{S}$ and $G_{m,i}^{T}$ for every input $x_{i}$, while pushing $G_{m,i}^{S}$ apart from $\{G_{m,j}^{T}\}_{j=1,j\neq i}^{N}$. $N$ is the number of negative representation pairs. CKTF defines the contrastive loss based on the intermediate representations as follows:

$$\vspace{-1pt}\small L_{MCKT}(G_{m}^{S},G_{m}^{T})=-E\left[log\frac{f(G_{m,i}^{S},G_{m,i}^{T})}{\sum_{j=1}^{N}f(G_{m,i}^{S},G_{m,j}^{T})}\right]\vspace{-0.5pt}$$

(3)

where the function $f(\cdot)$ is similar with that used in [15, 16, 17, 18], specifically, $f(G_{m,i}^{S},G_{m,i}^{T})=\frac{exp(G_{m,i}^{S}G_{m,i}^{T}/\tau)}{exp(G_{m,i}^{S}G_{m,i}^{T}/\tau)+N/N_{d}}$. $N_{d}$ is the number of training samples of the dataset, and $\tau$ is a temperature that controls the concentration level. The previous works [15, 16, 17, 18] use the function for different domains or objectives, such as self-supervised representation learning [17] and density estimation [16], whereas we are the first to construct multiple contrastive objectives on the intermediate representations of image classification models for knowledge transfer. The minimization of the contrastive loss $L_{MCKT}$ is maximizing the lower bound of the mutual information [15, 16, 17, 18] between $\{G_{m}^{T}\}_{m=1}^{M}$ and $\{G_{m}^{S}\}_{m=1}^{M}$.

Similar to $L_{MCKT}$, CKTF constructs the contrastive objective on the outputs of the penultimate layer as:

$$\vspace{-1pt}\small L_{PCKT}(S_{h},T_{h})=-E\left[log\frac{f(S_{h,i},T_{h,i})}{\sum_{j=1}^{N}f(S_{h,i},T_{h,j})}\right]\vspace{-0.5pt}$$

(4)

Finally, the proposed contrastive loss (the second loss term in Eq. 1) is defined as the weighted sum of $L_{MCKT}$ and $L_{PCKT}$:

$$\small L_{CKT}=\alpha_{1}\sum_{m=1}^{M}L_{MCKT}(G_{m}^{S},G_{m}^{T})+\alpha_{2}L_{PCKT}(T_{h},S_{h}).$$

(5)

We compare CKTF for model compression tasks with three baselines: 1) a large model that is uncompressed and directly trained (teacher), 2) a small model directly trained without KT (student w/o KT), and 3) the same small model trained with various KT methods. For the implementation, we use the public CRD code-base [14] to conduct a fair comparison.

Table 1 presents the Top-1 test accuracy of various teacher and student combinations on CIFAR-100 and Tiny-ImageNet. CFKT significantly outperforms all the existing KT methods in all cases. Specifically, CFKT outperforms: 1) the conventional KD [2] by 0.5% to 2.41% (none of the existing methods consistently outperforms KD), 2) the other KT methods by 0.04% to 11.59%, and 3) the related contrastive learning method CRD [14] (the second best in the results) by 0.04% to 0.97%.

Compared to the student trained directly on the data without KT, CFKT is 0.95% to 4.41% better. We also observe that, compared to the student w/o KT, CFKT performs better on Tiny-ImageNet (better than the student w/o KT by 4.41% to 1.41%) than on CIFAR-100 (better than the student w/o KT by 3.95% to 0.95%). This could be because Tiny-ImageNet is more complicated than CIFAR-100 (with more classes and data), resulting in more complicated intermediate representation, whereas CFKT is good at capturing this complicated high-dimension structural knowledge. Further, CFKT enables the small student to achieve comparable performance to the large teacher with only 16% of its original size. This confirms that CFKT is beneficial to on-device image classification applications that require small, high-performance models.

We first train the teacher using the source domain data (Tiny-ImageNet) with ground truth labels. Then, we transfer knowledge from the teacher to the student using the unlabeled data from the target domain (STL-10) with KT methods. Finally, we fine-tune the student (only train its linear classifier) using the training set of STL-10 and evaluate its accuracy on the test set of SLT-10. This is a common practice [22, 23, 14] for evaluating the quality of transfer learning.

We compare CKTF with KD [2] and CRD [14]. The teacher and student can be either the same, e.g., VGG-19/VGG-19, or different, e.g., VGG-19/VGG-8. Figure 2 shows how the Top-1 test accuracy of the student evolves during fine-tuning on STL-10. CKTF converges faster than all the baselines and outperforms them in final accuracy by 0.4% to 4.75%. This result validates that CKTF is advantageous for cross-domain transfer learning, even without labeled data in the target domain.