Improved Knowledge Distillation via Full Kernel Matrix Transfer

Knowledge distillation has a long history in ensemble learning and becomes popular for training small-sized neural networks [1, 3, 7, 10, 12, 13, 18, 19, 23]. Various algorithms have been developed to transfer different information from the teacher model to the student model. For example, [7] considers the final output of the teacher model as the soft label and regularizes the similarity between the label distribution output from the student model and that of the soft label from the teacher model. [23] transfers the attention maps from intermediate layers, which provides a way to explore more information from the teacher model. [1] proposes an information-theoretic framework for knowledge transfer which formulates knowledge transfer as maximizing the mutual information between the teacher and the student networks.

Recently, a number of algorithms [10, 13, 14, 20] are proposed to transfer the pairwise similarity for knowledge distillation. However, they focus on developing similarity functions while only transferring the similarity information of pairs within a mini-batch, which can be inefficient for distilling. We, on the other hand, aim to transfer the full matrix directly to achieve a better performance. Furthermore, we provide a theoretical analysis to demonstrate the effectiveness of the proposed method.

Besides classification, some methods are proposed for other tasks, e.g., detection [2] and metric learning [3]. We focus on classification in this work while the proposed method can be easily extended to metric learning that aims to optimize the performance of the embedding layer.

Nyström method is an effective algorithm to obtain a low-rank approximation for a Gram matrix [21]. Given a full Gram matrix, it tries to reconstruct the original one with the randomly sampled columns. The data points corresponding to the selected columns are denoted as landmark points. The approximation error can be bounded even with the randomly sampled landmark points. Later, researchers show that a delicate sampling strategy can further improve the performance [5, 9, 24]. [5] proposes to sample landmark points with a data-dependent probability distribution rather than the uniform distribution. [9] and [24] demonstrate that using clustering centers as landmark points provides the best approximation among different strategies.

Note that Nyström method is developed for unsupervised Gram matrix approximation while we can access the label information in knowledge distillation. In this work, we provide an analysis on the selection criterion of landmark points for Gram matrix transfer and develop a supervised strategy accordingly. Besides, Nyström method is useful only for a single Gram matrix approximation, while we aim to do transferring between two Gram matrices.

Nyström method is prevalently applied to approximate the Gram matrix [5, 9, 21, 24]. We briefly review it in this subsection.

Given a $n\times n$ Gram matrix $K$, we can first randomly shuffle columns and rewrite it as

where $W\in\mathbb{R}^{m\times m}$. Then, a good approximation for $K$ can be obtained as $\widetilde{K}=CW^{+}C^{\top}$, where $C=\left[\begin{array}[]{c}W\\ K_{12}\end{array}\right]\in\mathbb{R}^{n\times m}$ and $W^{+}$ denotes the pseudo inverse of $W$ [21].

Let $K_{k}$ denote the best top-$k$ approximation of Gram matrix $K$ and $k\leq m$. The rank-k approximation derived by the Nyström method can be computed as $\widetilde{K}_{k}=CW_{k}^{+}C^{\top}$, where $W_{k}$ denotes the best top-$k$ approximation of $W$ and $W_{k}^{+}$ is the corresponding pseudo inverse. The performance of the approximation can be demonstrated by the following theorem.

The examples corresponding to the selected columns in $C$ are referred as landmark points. Theorem 1 shows that the approximation is applicable even with randomly sampled landmark points.

With the low-rank approximation, the Gram matrix from a certain layer of the student and teacher network can be computed as

Compared with the original Gram matrix, the partial matrix $C$ has significantly less number of terms. Let $D_{\mathcal{S}}\in\mathbb{R}^{d_{\mathcal{S}}\times m}$ and $D_{\mathcal{T}}\in\mathbb{R}^{d_{\mathcal{T}}\times m}$ denote the landmark points for the student and teacher Gram matrices, then we have

With Theorem 1, it is easy to show that transferring the compact matrix is up-bounding the distance between original Gram matrices from the teacher and the student.

In Corollary 1, the partial Gram matrix will be regularized with the pseudo inverse of $W$. The computational cost of obtaining pseudo inverse is expensive and it can introduce additional noise when the feature space of the student is unstable in the early stage of training. Besides, it still optimizes the difference between two $n\times n$ matrices.

By selecting ingenious landmark points, we can bound the original loss in Eqn. 3.1 solely with $C_{\mathcal{S}}$ and $C_{\mathcal{T}}$ as in the following theorem.

• Proof.

  	$\displaystyle\|K_{\mathcal{S}}-K_{\mathcal{T}}\|_{F}=\|K_{\mathcal{S}}-\tilde{K}_{\mathcal{S}}^{k}-(K_{\mathcal{T}}-\tilde{K}_{\mathcal{T}}^{k})+\tilde{K}_{\mathcal{S}}^{k}-\tilde{K}_{\mathcal{T}}^{k}\|_{F}$
  	$\displaystyle\leq\|K_{\mathcal{S}}-\tilde{K}_{\mathcal{S}}^{k}\|_{F}+\|(K_{\mathcal{T}}-\tilde{K}_{\mathcal{T}}^{k})\|_{F}+\|\tilde{K}_{\mathcal{S}}^{k}-\tilde{K}_{\mathcal{T}}^{k}\|_{F}$
  	$\displaystyle\leq 2\varepsilon+\|C_{\mathcal{S}}W_{k\mathcal{S}}^{+}C_{\mathcal{S}}^{\top}-C_{\mathcal{T}}W_{k\mathcal{T}}^{+}C_{\mathcal{T}}^{\top}\|_{F}$
  	$\displaystyle\leq 2\varepsilon+c\|W_{k\mathcal{S}}^{+}C_{\mathcal{S}}^{\top}-W_{k\mathcal{T}}^{+}C_{\mathcal{T}}^{\top}\|_{F}+c(\|W_{k\mathcal{T}}^{+}\|_{F}\|C_{\mathcal{S}}-C_{\mathcal{T}}\|_{F})$
  	$\displaystyle\leq 2\varepsilon+c^{2}\|W_{k\mathcal{S}}^{+}-W_{k\mathcal{T}}^{+}\|_{F}+2c^{2}\|C_{\mathcal{S}}-C_{\mathcal{T}}\|_{F}$

  Then, we want to show that $\|W_{k\mathcal{S}}^{+}-W_{k\mathcal{T}}^{+}\|_{F}\leq\|C_{\mathcal{S}}-C_{\mathcal{T}}\|_{F}$. Note that $\|W_{k\mathcal{S}}-W_{k\mathcal{T}}\|_{F}\leq\|C_{\mathcal{S}}-C_{\mathcal{T}}\|_{F}$ due to the fact that $W_{k}$ is a partial matrix from $C$. We can prove $\|W_{k\mathcal{S}}^{+}-W_{k\mathcal{T}}^{+}\|_{F}\leq\|W_{k\mathcal{S}}-W_{k\mathcal{T}}\|_{F}$ instead. Let

  $$W_{k\mathcal{S}}=\sum_{j=1}^{k}\alpha_{j}\mu_{j}\mu_{j}^{\top};\quad W_{k\mathcal{T}}=\sum_{j=1}^{k}\beta_{j}\nu_{j}\nu_{j}^{\top}$$

  where $\alpha_{1}\geq\cdots\geq\alpha_{k}>1$ and $\beta_{1}\geq\cdots\geq\beta_{k}>1$, and $M\in\mathbb{R}^{k\times k}$, where $M_{i,j}=\langle\mu_{i}\mu_{i}^{\top},\nu_{j}\nu_{j}^{\top}\rangle$. We have

  $$\|W_{k\mathcal{S}}-W_{k\mathcal{T}}\|_{F}^{2}=\sum_{j}\alpha_{j}^{2}+\sum_{j}\beta_{j}^{2}-2\sum_{s,t}\alpha_{s}\beta_{t}M_{s,t}$$

  $$\|W_{k\mathcal{S}}^{+}-W_{k\mathcal{T}}^{+}\|_{F}^{2}=\sum_{j}\frac{1}{\alpha_{j}^{2}}+\sum_{j}\frac{1}{\beta_{j}^{2}}-2\sum_{s,t}\frac{1}{\alpha_{s}\beta_{t}}M_{s,t}$$

  So

  	$\displaystyle\|W_{k\mathcal{S}}^{+}-W_{k\mathcal{T}}^{+}\|_{F}^{2}-\|W_{k\mathcal{S}}-W_{k\mathcal{T}}\|_{F}^{2}$
  	$\displaystyle=\sum_{j}\frac{1}{\alpha_{j}^{2}}+\sum_{j}\frac{1}{\beta_{j}^{2}}-\sum_{j}\alpha_{j}^{2}+\sum_{j}\beta_{j}^{2}$
  	$\displaystyle+2\sum_{s,t}(\alpha_{s}\beta_{t}-\frac{1}{\alpha_{s}\beta_{t}})M_{s,t}$

  According to the definition of $M$, we have

  $$\forall i,\quad\sum_{j}^{k}M_{i,j}=\mu_{i}(\sum_{j}\nu_{j}\nu_{j}^{\top})\mu_{i}^{\top}\leq 1;\quad\forall j,\quad\sum_{i}^{k}M_{i,j}\leq 1$$

  Since $M$ is a doubly stochastic matrix, we can show that

  $$\sum_{s,t}(\alpha_{s}\beta_{t}-\frac{1}{\alpha_{s}\beta_{t}})M_{s,t}\leq\sum_{s}(\alpha_{s}\beta_{s}-\frac{1}{\alpha_{s}\beta_{s}})$$

  It can be proved by contradiction. If the optimal solution $M^{*}$ has a larger result than the R.H.S., we can denote the first column index of the non-zero off-diagonal element as $j>i$ (i.e., $M_{i,j}$), and the corresponding row index as $k>i$ (i.e., $M_{k,i}$). Let $A_{i,j}=\alpha_{i}\beta_{j}-\frac{1}{\alpha_{i}\beta_{j}}$ and we have

  	$\displaystyle A_{i,j}+A_{k,i}-A_{i,i}-A_{k,j}$
  	$\displaystyle=\alpha_{i}\beta_{j}-\frac{1}{\alpha_{i}\beta_{j}}+\alpha_{k}\beta_{i}-\frac{1}{\alpha_{k}\beta_{i}}$
  	$\displaystyle-(\alpha_{i}\beta_{i}-\frac{1}{\alpha_{i}\beta_{i}})-(\alpha_{k}\beta_{j}-\frac{1}{\alpha_{k}\beta_{j}})$
  	$\displaystyle=-(\alpha_{i}-\alpha_{k})(\beta_{i}-\beta_{j})+(\frac{1}{\alpha_{i}}-\frac{1}{\alpha_{k}})(\frac{1}{\beta_{i}}-\frac{1}{\beta_{j}})$
  	$\displaystyle=-(\alpha_{i}-\alpha_{k})(\beta_{i}-\beta_{j})+\frac{(\alpha_{i}-\alpha_{k})(\beta_{i}-\beta_{j})}{\alpha_{i}\alpha_{k}\beta_{i}\beta_{j}}<0$

  It shows that the assignment with the diagonal element can achieve a larger result than the optimal assignment, which contradicts the assumption.

  With the optimal results from the assignment of $M$, we obtain that

  $$\|W_{k\mathcal{S}}^{+}-W_{k\mathcal{T}}^{+}\|_{F}^{2}-\|W_{k\mathcal{S}}-W_{k\mathcal{T}}\|_{F}^{2}\leq\sum_{j}^{k}(\frac{1}{\alpha_{j}}-\frac{1}{\beta_{j}})^{2}-(\alpha_{j}-\beta_{j})^{2}$$

  For each term, it is easy to show that

  $$\forall j,\quad(\frac{1}{\alpha_{j}}-\frac{1}{\beta_{j}})^{2}-(\alpha_{j}-\beta_{j})^{2}\leq 0$$

  Therefore,

  $$\|W_{k\mathcal{S}}^{+}-W_{k\mathcal{T}}^{+}\|_{F}\leq\|W_{k\mathcal{S}}-W_{k\mathcal{T}}\|_{F}\leq\|C_{\mathcal{S}}-C_{\mathcal{T}}\|_{F}$$

       

Theorem 2 illustrates that minimizing the difference between the partial Gram matrices from the teacher and the student using landmark points can transfer the original Gram matrix from the teacher model effectively. Note that the partial matrices have the size of $n\times m$, where $m$ is the number of landmark points. When $m\ll n$, landmark points $D_{\mathcal{S}}$ and $D_{\mathcal{T}}$ can be kept in GPU memory as parameters of the loss function for optimization, which is much more efficient than transferring the original full Gram matrix. In addition, the simplified formulation depends on the property of eigenvalues from landmark points. Therefore, an appropriate selection is crucial for the efficient transfer. We elaborate our strategy in the next subsection.

We consider a strategy that obtains a single landmark point for each class, which means $m=L$ for a $L$-class classification problem. The desired landmark points should capture the distribution of examples well while keeping a sufficient diversity between each other for large eigenvalues. With the setting that each class has one landmark point, we can demonstrate the selection criterion as follows.

• Proof.

  For an arbitrary pair, we have

  	$\displaystyle\|K_{\mathcal{S}}(\mathbf{x}_{i},\mathbf{x}_{j})-K_{\mathcal{T}}(\mathbf{x}_{i},\mathbf{x}_{j})\|=\|\mathbf{x}^{i\top}_{\mathcal{S}}\mathbf{x}^{j}_{\mathcal{S}}-\mathbf{x}^{i\top}_{\mathcal{T}}\mathbf{x}^{j}_{\mathcal{T}}\|$
  	$\displaystyle=\|(\mathbf{x}^{i\top}_{\mathcal{S}}-\mathbf{d}^{i\top}_{\mathcal{S}}+\mathbf{d}^{i\top}_{\mathcal{S}})\mathbf{x}^{j}_{\mathcal{S}}-(\mathbf{x}^{i\top}_{\mathcal{T}}-\mathbf{d}^{i\top}_{\mathcal{T}}+\mathbf{d}^{i\top}_{\mathcal{T}})\mathbf{x}^{j}_{\mathcal{T}}\|$
  	$\displaystyle\leq\|(\mathbf{x}^{i\top}_{\mathcal{S}}-\mathbf{d}^{i\top}_{\mathcal{S}})\mathbf{x}^{j}_{\mathcal{S}}\|+\|(\mathbf{x}^{i\top}_{\mathcal{T}}-\mathbf{d}^{i\top}_{\mathcal{T}})x^{j}_{\mathcal{T}}\|$
  	$\displaystyle+\|\mathbf{d}^{i\top}_{\mathcal{S}}\mathbf{x}^{j}_{\mathcal{S}}-\mathbf{d}^{i\top}_{\mathcal{T}}\mathbf{x}^{j}_{\mathcal{T}}\|$
  	$\displaystyle\leq e\|\mathbf{x}^{i\top}_{\mathcal{S}}-\mathbf{d}^{i\top}_{\mathcal{S}}\|+e\|\mathbf{x}^{i\top}_{\mathcal{T}}-\mathbf{d}^{i\top}_{\mathcal{T}}\|+\|\mathbf{d}^{i\top}_{\mathcal{S}}\mathbf{x}^{j}_{\mathcal{S}}-\mathbf{d}^{i\top}_{\mathcal{T}}\mathbf{x}^{j}_{\mathcal{T}}\|$

  We obtain the result by accumulating over all pairs.         

According to Theorem 3, we can find that the transfer loss comes from two aspects. Term $A$ in Eqn. 3.2 contains the distance from each example to its corresponding landmark point. Since the corresponding landmark point for $\mathbf{x}_{\mathcal{S}}^{i}$ can be obtained by $\arg\min_{1\leq l\leq L}\{\|\mathbf{x}^{i\top}_{\mathcal{S}}-\mathbf{d}^{l\top}_{\mathcal{S}}\|\}$, we can rewrite the problem of minimizing the original term as

Apparently, this objective is a standard clustering problem. Unlike the conventional Nyström method, which is often in an unsupervised learning setting, we can access the label information in knowledge distillation. When we set the number of clusters to be the number of classes, the landmark point becomes the center in each class and can be computed by averaging examples within the class as

where $y_{i}$ is the class label of the $i$-th example and $n_{l}$ denotes the number of examples in the $l$-th class. In addition, class centers keep the diversity from original examples which can have the large eigenvalues as required by Theorem 2. We also verify the assumption in the experiments.

Term $B$ from Eqn. 3.2, in fact, indicates the difference between the student and teacher Gram matrices defined by the landmark points that is consistent as in Theorem 2.

With landmark points $D_{\mathcal{S}}$ and $D_{\mathcal{T}}$ obtained from optimizing the term $A$, we can formulate the Knowledge Distillation problem by transferring Approximated Gram matrix (KDA) as

where $\ell_{\mathrm{KDA}}(X_{\mathcal{S}}^{\top}D_{\mathcal{S}}-X_{\mathcal{T}}^{\top}D_{\mathcal{T}})=\sum_{i=1,l=1}^{i=n,l=L}\ell(\mathbf{d}^{l\top}_{\mathcal{S}}\mathbf{x}^{i}_{\mathcal{S}}-\mathbf{d}^{l\top}_{\mathcal{T}}\mathbf{x}^{i}_{\mathcal{T}})$ and we adopt $\ell(\cdot)$ as the smoothed $L_{1}$ loss for the stable optimization as

With the KDA loss, we propose a novel knowledge distillation algorithm that works in an alternative manner. In each iteration, we first compute the landmark points with features of examples accumulated from the last epoch by Eqn. 3.3. Then, the KDA loss defined by the fixed landmark points in Eqn. 3.4 will be optimized along with a standard cross-entropy loss for the student. The proposed algorithm is summarized in Alg. 1. Since at least one epoch will be spent on collecting features for computing landmark points, we will minimize the KDA loss after $H$ epochs of training, where $H\geq 1$.

In the conventional KD method [7], only the outputs from the last layer in the teacher model are adopted for the student. By setting an appropriate parameter, [7] illustrates that the loss function for KD can be written as

where $\mathbf{x}_{\mathcal{S}}^{i},\mathbf{x}_{\mathcal{T}}^{i}\in\mathbb{R}^{C}$ denote the logits before the SoftMax operator. With the identity matrix $I$, the equivalent formulation is

Compared to the KDA loss in Eqn. 3.4, the conventional KD can be considered as applying one-hot label vectors as landmark points to transfer the Gram matrix of the teacher network. However, it lacks the constraints on the similarity between each example and its corresponding landmark point as illustrated in Theorem 3, which may degenerate the performance of knowledge distillation.

We perform the ablation study on CIFAR-100 [8]. This data set contains $100$ classes, where each class has $500$ images for training and $100$ for test. Each image is a color image with size of $32\times 32$.

In this subsection, we set ResNet-34 as the teacher and ResNet-18 as the student. During the training, each image is first padded to be $40\times 40$, and then we randomly crop a $32\times 32$ image from it. Besides, random horizontal flipping is also adopted for data augmentation.

In this subsection, we compare the proposed KDA method to baseline methods on CIFAR-100. The results of different methods can be found in Table 2. All experiments are repeated $3$ times and the average results with standard deviation are reported.

First, knowledge distillation methods outperform training the student network without a teacher, which shows that knowledge distillation can improve the performance of student models significantly. By transferring the full Gram matrix before the last FC layer, KDA surpasses RKD by $1.3\%$ when ResNet-34 is the teacher and ResNet-18 is the student. The observation is consistent with the comparison in the ablation study, which confirms that $\|K_{\mathcal{S}}-K_{\mathcal{T}}\|_{F}/\|K_{\mathcal{T}}\|_{F}$ is an appropriate metric to evaluate the transfer loss of the Gram matrix. Moreover, KDA shows a significant improvement on different student networks, which further demonstrates the applicability of the proposed method.

Furthermore, when transferring the Gram matrix after the last FC layer, both of KD and KDA can demonstrate good performance compared with the student model. It is due to the fact that these methods transfer the Gram matrix with landmark points, which is efficient for optimization. Besides, KDA can further improve the performance compared to KD. The superior performance of KDA demonstrates the effectiveness of the proposed strategy for generating appropriate landmark points.

Finally, compared to benchmark methods, KDA can distill the information from different layers with the same formulation in Eqn. 3.4. The proposed method provides a systematic perspective to understand a family of knowledge distillation methods that aim to transfer Gram matrix. If transferring the Gram matrices before and after the FC layer simultaneously, the performance of KDA can be slightly improved as illustrated by “Combo” in Table 2.

Then, we compare different methods on Tiny-ImageNet data set¹¹1https://tiny-imagenet.herokuapp.com. There are $200$ classes in this data set and each class provides $500$ images for training and $50$ for validation. We report the performance on the validation set. Since the size of images in Tiny-ImageNet is $64\times 64$ that is larger than CIFAR-100, we replace the random crop augmentation with a more aggressive version as in [6], and keep other settings the same.

Table 3 summarizes the comparison. We can observe the similar results as on CIFAR-100. First, all methods with the information from a teacher model can surpass the student model without a teacher by a large margin. Second, KDA outperforms other methods no matter in which layer the transfer happens. Note that CIFAR-100 and Tiny-ImageNet have very different images, which demonstrates the applicability of the proposed algorithm in different real-world applications.