SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Zhiwei Hao $^{1,2}$ , Jianyuan Guo $^{3}$ , Kai Han $^{2}$ , Yehui Tang $^{2}$ , Han Hu $^{1*}$ , Yunhe Wang $^{2*}$ , Chang Xu $^{3}$
$^{1}$ School of information and Electronics, Beijing Institute of Technology.  $^{2}$ Huawei Noah's Ark Lab.
$^{3}$ School of Computer Science, Faculty of Engineering, The University of Sydney. {haozhw, hhu}@bit.edu.cn, jguo5172@uni.sydney.edu.au, {kai.han, yehui.tang, yunhe.wang}@huawei.com, c.xu@sydney.edu.au
Knowledge distillation (KD) has proven to be a highly effective approach for enhancing model performance through a teacher-student training scheme. However, most existing distillation methods are designed under the assumption that the teacher and student models belong to the same model family, particularly the hint-based approaches. By using centered kernel alignment (CKA) to compare the learned features between heterogeneous teacher and student models, we observe significant feature divergence. This divergence illustrates the ineffectiveness of previous hint-based methods in cross-architecture distillation. To tackle the challenge in distilling heterogeneous models, we propose a simple yet effective one-for-all KD framework called OFA-KD, which significantly improves the distillation performance between heterogeneous architectures. Specifically, we project intermediate features into an aligned latent space such as the logits space, where architecture-specific information is discarded. Additionally, we introduce an adaptive target enhancement scheme to prevent the student from being disturbed by irrelevant information. Extensive experiments with various architectures, including CNN, Transformer, and MLP, demonstrate the superiority of our OFA-KD framework in enabling distillation between heterogeneous architectures. Specifically, when equipped with our OFA-KD, the student models achieve notable performance improvements, with a maximum gain of  $8.0\%$  on the CIFAR-100 dataset and  $0.7\%$  on the ImageNet-1K dataset. PyTorch code and checkpoints can be found at https://github.com/Hao840/OFAKD.
In recent years, knowledge distillation (KD) has emerged as a promising approach for training lightweight deep neural network models in computer vision tasks [1, 2, 3, 4]. The core idea behind KD is to train a compact student model to mimic the outputs, or soft labels, of a pretrained cumbersome teacher model. This method was initially introduced by Hinton et al.. [5], and since then, researchers have been actively developing more effective KD approaches. One notable improvement in KD methods is the incorporation of intermediate features as hint knowledge [6, 7, 8]. These approaches train the student to learn representations that closely resemble those generated by the teacher, and have achieved significant performance improvements.
Existing hint-based approaches have predominantly focused on distillation between teacher and student models with homogeneous architectures, e.g., ResNet-34 vs. ResNet-18 in [9], and Swin-S vs. Swin-T in [8], leaving cross-architecture distillation relatively unexplored. However, there are practical scenarios where cross-architecture distillation becomes necessary. For example, when the objective is to improve the performance of a widely used architecture like ResNet50 [10], as our experiments will demonstrate, distilling knowledge from ViT-B [11] to ResNet50 can readily surpass the performance achieved by distilling knowledge from ResNet152. Cross-architecture distillation thus presents a viable alternative, providing more feasible options for teacher models. It may not always be possible to find a superior teacher model with a homogeneous architecture matching the student's architecture. Moreover, the emergence of new architectures [12, 13] further compounds this challenge, making it difficult to find readily available pre-trained teacher models of the same architecture online. This limitation necessitates the need for researchers to resort to pre-trained teachers of different architectures for knowledge distillation in order to enhance their own models.
The few attempts that have been made in this direction mainly involve using heterogeneous CNN models for knowledge transfer [14, 15, 16]. In a recent study, Touvron et al. [17] successfully trained a ViT student by using a CNN teacher and obtained promising results. These endeavors have sparked our interest in studying the question: whether it is possible to distill knowledge effectively across any architectures, such as CNN [10, 18, 19], Transformer [20, 21, 11, 17, 22], and MLP [13, 23, 24].
In previous studies where teachers and students share similar architectures, their learned representations naturally exhibit similarity in the feature space. Simple similarity measurement functions such as mean square error (MSE) loss, are sufficient for information distillation [6]. However, in the case that the teacher and student architectures are heterogeneous, there is no guarantee of successful alignment of the learned features, due to the fact that features from heterogeneous models reside in different latent feature spaces, as shown in Figure 1. Therefore, directly matching these irrelevant features is less meaningful and even impedes the performance of the student model, as shown in the third row of Table 1, highlighting the difficulty in exploring cross-architecture distillation at intermediate layers.
To address above challenges, we propose a one-for-all KD (OFA-KD) framework for distillation between heterogeneous architectures, including CNNs, Transformers, and MLPs. Instead of directly distilling in the feature space, we transfer the mismatched representations into the aligned logits space by incorporating additional exit branches into the student model. By matching the outputs of these branches with that from teacher's classifier layer in logits space, which contains less architecture-specific information, cross-architecture distillation at intermediate layers becomes achievable. Furthermore, considering that heterogeneous models may learn distinct predictive distributions due to their different inductive biases, we propose a modified KD loss formulation to mitigate the potential impact of irrelevant information in the logits. This includes the introduction of an additional modulating parameter into the reformulated loss of the vanilla KD approach. The modified OFA loss adaptively enhances the target information based on the predictive confidence of the teacher, effectively reducing the influence of irrelevant information. To evaluate our proposed framework, we conduct experiments using CNN, Transformer, and MLP architectures. We consider all six possible combinations of these architectures. Our OFA-KD framework improves student model performance with gains of up to  $8.0\%$  on CIFAR-100 and  $0.7\%$  on ImageNet-1K datasets, respectively. Additionally, our ablation study demonstrates the effectiveness of the proposed framework.
Recently, significant advancements have been made in the design of model architectures for computer vision tasks. Here we briefly introduce two prominent architectures: Transformer and MLP.
Vision Transformer. Vaswani et al. [25] first proposed the transformer architecture for NLP tasks. Owing to the use of the attention mechanism, this architecture can capture long-term dependencies effectively and achieves remarkable performance. Inspired by its great success, researchers have made great efforts to design transformer-based models for CV tasks [26, 27, 28]. Dosovitskiy et al. [11] split an image into non-overlapped patches and projected the patches into embedding tokens. Then the tokens are processed by the transformer model like those in NLP. Their design achieves state-of-the-art performance and sparks the design of a series of following architectures [21, 29].
MLP. MLP has performed inferior to CNN in the computer vision field for a long time. To explore the potential of the MLP, Tolstikhin et al. [13] proposed MLP-Mixer, which is based exclusively on the MLP structure. MLP-Mixer takes embedding tokens of an image patch as input and mixes channel and spatial information at each layer alternately. This architecture performs comparably to the best CNN and ViT models. Touvron et al. [24] proposed another MLP architecture termed ResMLP.
The most advanced CNN, Transformer, and MLP achieve similar performance. However, these architectures hold distinct inductive biases, resulting in different representation learning preferences. Generally, there are striking differences between features learned by heterogeneous architectures.
Knowledge distillation. KD is an effective approach for compressing cumbersome models [30, 31, 32, 33, 34, 35], where a lightweight student is trained to mimic the output logits of a pretrained teacher. Such an idea was originally proposed by Bucila et al. [36] and improved by Hinton et al. [5]. Following works further improve logits-based KD via structural information [37, 38], model ensemble [39], or contrastive learning [40]. Recently, Touvron et al. [17] proposed a logits distillation method for training ViT students. Huang et al. [9] proposed to relax the KL divergence loss for distillation between teacher and student with a large capacity gap. Besides logits, some KD approaches also use intermediate features as hint knowledge. Romero et al. [6] first proposed the hint-based distillation approach. They adopted a convolutional layer to project student features into the space of the teacher feature size, and then constrained the projected feature to be like the teacher feature via mean square loss. Zagoruyko et al. [7] proposed to force the student to mimic attention maps. Yim et al. [41] adopted the flow of solution procedure matrix generated by the features as hint knowledge. There are also many other feature distillation approaches using various hint design [42, 43, 8]. Recently, applying KD to compress dataset has emerged as a popular topic [44, 45, 46].
Although existing hint-based distillation methods have achieved remarkable performance, they assume that student and teacher architectures are homogeneous. When the architectures are heterogeneous, existing approaches may fail because features of the student and the teacher are distinct. How to conduct hint-based distillation under this circumstance is still an open problem.
3.1 Revisit feature distillation for heterogeneous architectures
Discrepancy among architectures. For a long time, CNN is the dominant architecture for CV tasks. Its inductive biases of locality and spatial invariance are strong priors for learning natural images, and restrict CNNs from learning global representations. Vision transformers [11, 21, 47] inherited from the NLP field are attention-based models. This architecture lacks inductive biases inherent to CNNs and thus can learn global representations easily. Moreover, inspired by the success of ViTs, pure MLP-based vision models have been proposed [13, 24], which are also low-priori models.
Knowledge distillation. KD trains a tiny student model to learn knowledge from a large pre-trained teacher model. Logits and features are the most common used types of knowledge in KD.
Logits depict the predictive distribution of a model. Logits-based distillation trains the student to mimic output logits of the teacher, which can be formulated as:
$$
\mathcal {L} _ {\mathrm {K D}} = \lambda \mathbb {E} _ {(\boldsymbol {x}, y) \sim (\mathcal {X}, \mathcal {Y})} \left[ \mathcal {D} _ {\mathrm {C E}} \left(\boldsymbol {p} ^ {s}, y\right) + (1 - \lambda) \mathcal {D} _ {\mathrm {K L}} \left(\boldsymbol {p} ^ {s}, \boldsymbol {p} ^ {t}\right) \right], \tag {1}
$$
where  $(\mathcal{X},\mathcal{Y})$  is the joint distribution of samples and class labels.  $\pmb{p}^{s}$  and  $\pmb{p}^t$  are predictions of the student model and the teacher model on sample  $\pmb{x}$ , respectively.  $\mathcal{D}_{\mathrm{CE}}$  denotes the cross-entropy function.  $\mathcal{D}_{\mathrm{KL}}$  denotes the Kullback-Leibler divergence function.  $\lambda$  is a hyperparameter controls the trade-off between one-hot label  $y$  and soft label  $\pmb{p}^t$ .
hint-based distillation adopts more fine-grained teacher activations at intermediate layers to train the student, which can be formulated as:
$$
\mathcal {L} _ {\mathrm {F D}} = \mathbb {E} _ {(\boldsymbol {x}, \boldsymbol {y}) \sim (\mathcal {X}, \mathcal {Y})} [ \sum_ {i} \| \mathbf {F} _ {i} ^ {T} - \psi \left(\mathbf {F} _ {i} ^ {S}\right) \| _ {2} ], \tag {2}
$$
where  $\mathbf{F}^T$  and  $\mathbf{F}^S$  are features of the teacher and the student at layer  $i$ , respectively.  $\psi$  is a feature projector that maps student features to match dimension of teacher features [6]. Knowledge from logits and features are not opposites. In most cases, feature distillation approaches also adopt logits as an auxiliary supervision.
Figure 1: Similarity heatmap of intermediate features measured by CKA. We compare features from MobileNetV2 (CNN), ViT-Small (Transformer) and Mixer-B/16 (MLP model). The first three figures illustrate similarity between homogeneous models, and the last three illustrate feature similarity between heterogeneous models. Coordinate axes indicate the corresponding layer index.
Challenges in heterogeneous feature distillation. Existing feature distillation methods are designed under the assumption of teacher and student belonging to the same model family. Simply adopt a convolutional projector is able to align their features for distillation. However, heterogeneous models learn diverse features because of there distinct inductive biases. Under this circumstance, directly forcing features of the student to mimic the teacher may impede the performance. Moreover, ViTs use embedded image patches as input, and some architectures even adopt an additional classification token [11, 17], which further aggregates the feature mismatch problem.
Centered kernel alignment analysis.2 To demonstrate the representation gap among heterogeneous architectures, we adopt centered kernel alignment (CKA) [48, 49] to compare features extracted by CNN, ViT, and MLP models. CKA is a feature similarity measurement allowing cross-architecture comparison achievable, as it can work with inputs having different dimensions.
CKA evaluates feature similarity over mini-batch. Suppose  $\mathbf{X} \in \mathbb{R}^{n \times d_1}$  and  $\mathbf{Y} \in \mathbb{R}^{n \times d_2}$  are features of  $n$  samples extracted by two different models, where  $d_1$  and  $d_2$  are their dimensions, respectively. CKA measures their similarity by:
$$
\operatorname {C K A} (\mathbf {K}, \mathbf {L}) = \frac {\mathcal {D} _ {\mathrm {H S I C}} (\mathbf {K} , \mathbf {L})}{\sqrt {\mathcal {D} _ {\mathrm {H S I C}} (\mathbf {K} , \mathbf {K}) \mathcal {D} _ {\mathrm {H S I C}} (\mathbf {L} , \mathbf {L})}}, \tag {3}
$$
where  $\mathbf{L} = \mathbf{X}\mathbf{X}^T$  and  $\mathbf{K} = \mathbf{Y}\mathbf{Y}^T$  are Gram matrices of the features, and  $\mathcal{D}_{\mathrm{HSIC}}$  is the Hilbert-Schmidt independence criterion [50], a non-parametric independence measure. The empirical estimator of  $\mathcal{D}_{\mathrm{HSIC}}$  can be formulated as:
$$
\mathcal {D} _ {\mathrm {H S I C}} (\mathbf {K}, \mathbf {L}) = \frac {1}{(n - 1) ^ {2}} \operatorname {t r} (\mathbf {K H L H}), \tag {4}
$$
where  $\mathbf{H}$  is the centering matrix  $\mathbf{H}_n = \mathbf{I}_n - \frac{1}{n}\mathbf{1}\mathbf{1}^T$ .
Figure 1 shows feature similarity heatmaps measured by CKA. From figures at the first row, models of homogeneous architectures tend to learn similar features at layers of similar positions. On the country, heatmaps at the second row demonstrate that features of heterogeneous architectures are remarkably different. In particular, MobileNet v2 can only learn features similar to shallower layers of ViT-S or Mixer-B/16. Moreover, features of ViT-S and Mixer-B/16 are quite different as well.
From the CKA analysis, directly using hint-based distillation methods for cross-architecture KD is unreasonable. Although only using logits for distillation is feasible, the absence of supervision at intermediate layers may lead to suboptimal results. New designs of distilling at intermediate layers are needed to improve heterogeneous KD.
3.2 Generic heterogeneous knowledge distillation
To bridge the gap of heterogeneous architectures in KD, we propose a one-for-all cross-architecture KD method, which can perform KD between models of any architectures, such as CNN, ViT, and MLP. Figure 2 provides an overview of the proposed framework.
Learning in the logits space. Following the CKA analysis results, heterogeneous models learn distinct representations, which impedes the feature distillation process. Hence, the key for heterogeneous feature distillation is to align the mismatched representations. However, existing approaches
2Detailed settings of our CKA analysis are provided in our supplement.
(c) Our OFA distillation
Figure 2: Comparison of different KD approaches. (a) Logits-based distillation: student only learns from final predictions of the teacher; (b) hint-based distillation: besides final predictions, student also learns to mimic intermediate features of the teacher; and (c) our OFA-KD framework: student learns from final predictions of the teacher via multiple branches, where the OFA distillation loss adaptively adjusts the amount of target information to be learned based on the confidence of teacher. Note: Only three stages are shown for convenience. In our experiment, all models are split into four stages.
directly align dimension of features via a convolutional projector [6], which is not a generic solution for all heterogeneous architectures. Inspired by early-exit model architecture [51], we align features in the logits space by introducing additional exit branches into the student. In particular, we build each exit branch with a feature projector and a classifier layer. On the teacher side, we direct adopt its final output as the knowledge. In the logits space, redundant architecture-specific information is removed, and thus feature alignment across any architectures is achievable. When training the student, the additional exit branches are optimized together with the student backbone. At test time, these branches are removed and render no extra inference overhead.
Adaptive target information enhancement. Although different models learn the same targets in the logits space, their different inductive bias leads them to variant destinations, i.e., their predictive distributions are not the same. For example, a CNN model may predict a sample with similar probabilities of belonging to cars and trucks, as CNNs tend to catch the shared local information of the two classes, e.g., round wheels. In contrast, ViTs may differentiate these samples better, as its attention mechanism is apt to learn the distinct global features, e.g., arc-shaped and rectangle-shaped bodies. This difference results in an inconsistent understanding of the dark knowledge of the teacher and the student. To bridge this gap, we propose to enhance the information of the target class as it always relates to the real class of a sample and brings no concern about information mismatch.
We take the first step by reformulating the original distillation loss to decouple information about the target class:3
$$
\begin{array}{l} \mathcal {L} _ {\mathrm {K D}} = - \log \boldsymbol {p} _ {\hat {c}} ^ {s} - \mathbb {E} _ {c \sim \mathcal {Y}} [ \boldsymbol {p} _ {c} ^ {t} \log \boldsymbol {p} _ {c} ^ {s} ] \\ = - \left(1 + \boldsymbol {p} _ {\hat {c}} ^ {t}\right) \log \boldsymbol {p} _ {\hat {c}} ^ {s} - \mathbb {E} _ {c \sim \mathcal {Y} / \{\hat {c} \}} \left[ \boldsymbol {p} _ {c} ^ {t} \log \boldsymbol {p} _ {c} ^ {s} \right], \tag {5} \\ \end{array}
$$
where we ignored the denominator term in KL-divergence as it does not contribute to the gradient.  $c$  is a predicted class and  $\hat{c}$  indicates the target class. To enhance information from the target class, we add a modulating parameter  $\gamma \geq 1$  to the  $(1 + p_{\hat{c}}^{t})$  term:
$$
\mathcal {L} _ {\mathrm {O F A}} = - \left(1 + \boldsymbol {p} _ {\hat {c}} ^ {t}\right) ^ {\gamma} \log \boldsymbol {p} _ {\hat {c}} ^ {s} - \mathbb {E} _ {c \sim \mathcal {Y} / \{\hat {c} \}} \left[ \boldsymbol {p} _ {c} ^ {t} \log \boldsymbol {p} _ {c} ^ {s} \right]. \tag {6}
$$
Without loss of generality, we omit the hyperparameter  $\lambda$ .
Discussion. To learn the behavior of the OFA loss with different  $\gamma$  values, we take the case of  $\gamma \in \mathbb{Z}_{+}$  as an example expand the binomial term:
$$
\begin{array}{l} \mathcal {L} _ {\mathrm {O F A}} = - \sum_ {k = 0} ^ {\gamma} \binom {\gamma} {k} \left(\boldsymbol {p} _ {\hat {c}} ^ {t}\right) ^ {k} \log \boldsymbol {p} _ {\hat {c}} ^ {s} - \mathbb {E} _ {c \sim \mathcal {Y} / \{\hat {c} \}} \left[ \boldsymbol {p} _ {c} ^ {t} \log \boldsymbol {p} _ {c} ^ {s} \right] \\ = - \left(1 + \boldsymbol {p} _ {\hat {c}} ^ {t}\right) \log \boldsymbol {p} _ {\hat {c}} ^ {s} - \mathbb {E} _ {c \sim \mathcal {Y} / \{\hat {c} \}} \left[ \boldsymbol {p} _ {i} ^ {t} \log \boldsymbol {p} _ {i} ^ {s} \right] - \left(\sum_ {k = 1} ^ {\gamma} \binom {\gamma} {k} \left(\boldsymbol {p} _ {\hat {c}} ^ {t}\right) ^ {k} - \boldsymbol {p} _ {\hat {c}} ^ {t}\right) \log \boldsymbol {p} _ {\hat {c}} ^ {s} \tag {7} \\ = \mathcal {L} _ {K D} + - \left(\sum_ {k = 1} ^ {\gamma} \binom {\gamma} {k} \left(\boldsymbol {p} _ {\hat {c}} ^ {t}\right) ^ {k} - \boldsymbol {p} _ {\hat {c}} ^ {t}\right) \log \boldsymbol {p} _ {\hat {c}} ^ {s}. \\ \end{array}
$$
When  $\gamma = 1$ , the OFA distillation loss is identical to the logits-based KD loss. To study cases of  $\gamma \geq 1$ , we take  $\gamma = 2$  as an example:
$$
\mathcal {L} _ {\mathrm {O F A}, \gamma = 2} = \mathcal {L} _ {K D} + - \left(\boldsymbol {p} _ {\hat {c}} ^ {t} + \left(\boldsymbol {p} _ {\hat {c}} ^ {t}\right) ^ {2}\right) \log \boldsymbol {p} _ {\hat {c}} ^ {s}. \tag {8}
$$
Besides the KD loss term, the OFA distillation loss has an additional positive term that only relates to the target class. If the teacher is confident at the target class, the high order term  $(\pmb{p}_{\hat{c}}^{t})^{2}$  decays slowly. Otherwise, the high order term decays faster, which prevents the student to learn from a less confident teacher and achieves adaptive enhancement of target information. The top-right corner of Figure 2 provides a brief illustration about effect of the high order term in the OFA loss.
Heterogeneous models inherently possess differences in their learning capabilities and preferences. Our proposed adaptive target information enhancement method bridges this gap by mitigating the impact of soft labels when the teacher provides suboptimal predictions. Moreover, the efficacy of distillation between heterogeneous architectures can be further enhanced by combining the multibranch learning paradigm in the aligned logits space.
We conducted comprehensive experiments to evaluate the proposed OFA-KD framework. In this section, we offer a concise overview of our experimental settings. Additional details can be found in our supplementary materials.
Models. We evaluated a variety of models with heterogeneous architectures. Specifically, we employed ResNet [10], MobileNet v2 [18], and ConvNeXt [12] as CNN models, ViT [11], DeiT [17], and Swin [21] as ViT models, and MLP-Mixer [13] and ResMLP [24] as MLP models.
In our OFA-KD framework, incorporating additional exit branches into the student model requires identifying appropriate branch insertion points. For models with a pyramid structure, we consider the end of each stage as potential insertion points, resulting in a total of four points. In the case of other models like ViT, we evenly divide them into four parts and designate the end of each part as feasible insertion points for the branches. Exit branches in CNNs are constructed using depth-width convolutional layers, while exit branches in ViTs and MLPs employ ViT blocks.