SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Lujun Li†* Peijie Dong† Anggeng Li† Zimian Wei Ya Yang
HKUST HKUST(GZ) Huawei NUST CityU
{lilujunai, dongpeijie98} @ gmail.com,{anggeng.li,yya9} @outlook.com, weizimian16@nudt.edu.cn
Knowledge distillation (KD) has emerged as an effective technique for compressing models that can enhance the lightweight model. Conventional KD methods propose various designs to allow student model to imitate the teacher better. However, these handcrafted KD designs heavily rely on expert knowledge and may be sub-optimal for various teacher-student pairs. In this paper, we present a novel framework, KD-Zero, which utilizes evolutionary search to automatically discover promising distiller from scratch for any teacher-student architectures. Specifically, we first decompose the generalized distiller into knowledge transformations, distance functions, and loss weights. Then, we construct our distiller search space by selecting advanced operations for these three components. With sharpness and represent gap as fitting objectives, we evolve candidate populations and generate better distillers by crossover and mutation. To ensure efficient searching, we employ the loss-rejection protocol, search space shrinkage, and proxy settings during the search process. In this manner, the discovered distiller can address the capacity gap and cross-architecture challenges for any teacher-student pairs in the final distillation stage. Comprehensive experiments reveal that KD-Zero consistently outperforms other state-of-the-art methods across diverse architectures on classification, detection, and segmentation tasks. Noticeably, we provide some practical insights in designing the distiller by analyzing the distiller discovered. Codes are available in supplementary materials.
Deep Neural Networks (DNNs) have achieved great success in tackling a variety of visual recognition tasks [31, 17, 66]. Despite the appealing performance, the prevailing DNN models usually have large numbers of parameters, leading to heavy costs of memory and computation. Conventional techniques such as Neural Architecture Search [16, 6, 26, 15] and quantizing networks to use low-bit parameters [11, 53, 49] have proven to be effective for mitigating this computational burden. Recently, Knowledge Distillation (KD) has been widely used for training compact and efficient neural networks by transferring the knowledge lied in the logits [23] or features [57] from a large, pre-trained teacher to a smaller student.
Recently, although KD has made significant progress in the handcrafted designs, there are still some limitations to its practical
Figure 1: Illustration of search space in KD-Zero.
*Corresponding author, † Co-first authors, equal contribution.
Figure 2: Top-1 gain  $(\%)$ . of WRN-16-2 (left) & ResNet-20 (right) distilled via various capacities teacher models on CIFAR-100.
Figure 3: Top-1 gain  $(\%)$ . of various student with KD-Zero (left). Various KD methods combined with KD-Zero (right) for ResNet-20 on C-100.
application for different scenarios from three aspects: (1) Capac-
ity gap problem: Traditional KD methods do not always distill better from stronger teachers because of the large teacher-student gap [27, 46]. Larger and more accurate teacher models tend to be overconfident and fail to improve students [79]. To alleviate this, architecture-level methods use the assistant model [46] or architecture search [45], introducing additional training budgets. Other KD techniques use optimized designs or manually modified distillers to reduce the disparity. However, these techniques rely on expert knowledge and require careful tuning, which suffers from poor generality and efficiency. (2) Cross-architecture issue. Teacher-student networks with different architectural styles suffer from mismatches in position dependence, receptive fields, and feature dimensions. As shown in Figure 2, while handcrafted methods (e.g., KD [23], DIST [27, 46] and WSLD [79]) provide slight gains in some cases, they still are weak in most cross-architecture pairs<sup>2</sup>. (3) Manual tuning difficulties. Different choices of knowledge sources, distillation functions, and hyperparameters can largely affect the performance of KD. However, designing and training a proper distillation setting requires trial and error, substantial effort, and experiments. Thus, two questions are raised: (1) How to efficiently discover the optimal distillation strategies without expert knowledge? (2) How to reduce the teacher-student gap with different capabilities and architectures?
"We may hope that machines will eventually compete with men in all purely intellectual fields."
—Alan Turing
As this well-said quote goes, recently, machine learning approaches have successfully replaced human experts in algorithm design, architecture search, and drug discovery. For the first problem, inspired by Auto-Zero [55], we decide to automate the process of designing the distillation functions for the first time in the field of knowledge distillation. For the second question, we systematically review previous studies of KD design for distillation gaps. Our observations reveal that: (1) Most of the distillers can be divided into three components with various key elements: knowledge transformation, distance function, and loss weights. (2) Some essential operations, like normalize ops [3] and mask ops [71], play important roles in reducing the teacher-student gap. (3) Optimal distillation functions can be effectively integrated with additional strategies (e.g., feature aggregation [50], projector ensemble [9], N-to-One match [67], and cross-layer mapping [28, 13] in Figure 3).
Based on the above analysis, we present KD-zero, an automated search framework that utilizes evolutionary algorithms from scratch to efficiently discover the best distiller without manual design. Specifically, our framework is organized into three parts: search space, search algorithm, and acceleration strategy. Firstly, we establish the search space with basic transform operations, distance functions, and loss weights (see Figure 1). For example, we select operators normalized in different dimensions (e.g., batchnorm, normHW,C,N), various types of activation functions (e.g., exp, relu, sigmoid, pow2), multi-scale process and spatial-wise/channel-wise mask transforms and other advanced operations in the knowledge transformation. Our distance function options include smooth  $\ell_1$ ,  $\ell_1$ ,  $\ell_2$ ,  $\ell_{KL}$ ,  $\ell_{hard}$ ,  $\ell_{Cosine}$ ,  $\ell_{Pearson}$  and  $\ell_{Correlation}$  distance. Options in the loss weight part include various values for loss factors, temperature factors, and weight calibration strategies. In this way, our search space includes over 10,000 candidates covering the existing SOTA KD methods and designs. Then, we construct a calculation graph for these candidates and use selected features, representations, and logits from the teacher-student network as input. Based graph structure, we
2In our paper, W40-2, R32×4, R8×4, MV2, SV1, and SV2 stand for WRN-40-2, ResNet32×4, ResNet8×4, MobileNetV2, ShuffleNetV1, and ShuffleNetV2. DeiT-Ti [61] T2T-ViT-7 [72], PVT-Ti [64].
initialize the candidate total groups from the search space, crossover, and mutate them according to the evaluation results. We employ loss-rejection protocol and search space shrinkage for search efficiency to filter out weak candidates. With early-stop proxy settings, we achieve at least  $40\times$  acceleration during the distiller search. For distillation gap reduction, we take the representation gap and sharpness gap between teacher-student as the fitting objectives besides the accuracy metric of the validation set. Finally, we distill student architectures with discovered distiller, and our KD-Zero surpasses existing KD approaches by a large margin without prior knowledge (see Figure 3).
In principle, our KD-Zero differs from previous hand-designed KD methods, opening new doors to automated distillation designs. Its merits can be highlighted in three aspects: (1) Effective. KD-Zero effectively reduces the teacher-student gap by presenting a general distiller search space and adaptive evolutionary search for different teacher-student pairs. KD-Zero extends the KD formulation and allows for additional gains with extra design besides distillers. In addition, it reduces human bias and ensures that the resulting distillers are optimized for the target problem or dataset. (2) Efficient. KD-Zero increases efficiency in practice by a series of flexible, systematic, and efficient search procedures without additional laborious tuning. By contrast, other manual methods with fixed KD forms involve complex parameter tuning with additional training time and resources. (3) Insightful. KD-Zero undertakes an in-depth analysis of the existing advanced distillation designs, with the aim of exploring their potential combination to produce numerous novel distillers. KD-Zero provides guidelines for practical applications and develops a new research direction. We hope our efforts on the automated design of distillers could facilitate future research for automated KD works to some extent. In summary, our contributions are:
- To alleviate architecture & capability gaps of teacher-student, we present KD-Zero, the first auto-search framework for evolving best distillers from scratch to our best knowledge.
- We present a comprehensive distiller search space, including advanced operations on transformations, distance functions, and loss weights. Then, we evolve the distiller search with performance and sharpness & represent the gap as fitting objectives. In addition, We achieve significant search acceleration via loss-rejection protocol & space shrinkage, and proxy settings.
- We conduct extensive experiments on classification, detection, and segmentation. KD-Zero performs state-of-the-art in multiple datasets and architectures (e.g., CNN and vision transformer). Specifically, ResNet-18 and MobileNet with KD-Zero achieve  $72.17\%$  and  $73.02\%$  Top-1 accuracy on ImageNet, outperforming KD by  $1.51\%$ ,  $2.34\%$ , respectively.
Knowledge Distillation. The idea of teacher-student learning is first proposed in pioneering explorations [1, 2], and the formal definition is defined by the original KD [23]. Subsequent efforts explore on different knowledge (e.g., intermediate feature representations [38, 43, 37, 33, 36], sample relationships [48, 60]) and applications [18, 14, 35]. Compared to KD for distillation gap. Previous methods propose assistant teachers, architecture search, KD designs on transformations [27], distance functions [59], and weight-tuning [42, 34] for this problem. However, such KD designs rely on expert knowledge and tuning, and their performance can fluctuate significantly across different situations. In contrast to these methods, KD-Zero develops automated searches for distillers to address these difficulties that do not require additional architecture modification and manual KD design. Compared to Meta-KDs. These works [13, 42] only focus on hyperparameter tuning and involve complex optimization challenges. In contrast, our approach searches for complementary distiller design besides hyperparameters establishing a new paradigm for KD research and application.
Automated Machine Learning. AutoML [81, 80] aims to automate Network Architecture Search (NAS) and HyperParameter Optimization (HPO), making them more accessible to non-experts. NAS chooses architecture rather than KD designs. Compared to HPOs [54, 69], they generally focus on the hyperparameters on training configurations. Recent methods search for loss formulation [39, 32]. In contrast to these methods, we present a new complex search space for distiller design in transforms, distances, loss weights, and new search objectives and accelerations according to the KD task.
Figure 4: Overview of KD-Zero. In the search phase, we first randomly sample candidate distillers to initialize the population and evaluate their validation performance and sharpness & represent the gap between teacher-student. Then, we perform loss-rejection protocol & space shrinkage to remove weak individuals and crossover & mutation to generate new populations within promising ones. Finally, we pick up the best-performing distiller for distillation.
Table 1: Specific operations in KD-Zero. More details of their formulas are available in the Appendix.
Transform	norm-based: batchnorm, min - max, normHW,C,N, softmaxHW,C,N, logsoftmaxHW,C,N activation-based: exp, mish,leaky,relu,tanh,sigmoid,pow2,pow4,log,sqrt scale-based: scale, multi - scale, scaler1,r2,localr1,r2,r4, batch, channel attention&mask-based: drop,satt,natt,catt,mask, other: no,bmm, mm
Distance	no-norm loss: smooth ℓ1, ℓ1, ℓ2, ℓKL, ℓhard; norm loss: ℓCosine, ℓPearson, ℓCorrelation
Weight	calibration: entropy, focal, sim, conf, no; weight values: 0.01,..., 100; τ values:1,4,8
In this section, we first illustrate the design of our distiller search space, the search process, the acceleration, and the fitting objectives. Then we analyze the search results and give some guidelines. Finally, we analyze the student distilled via KD-Zero and expansion for different distillation scenarios. The pipeline of our approach is shown in Figure 4.
3.1 Search Space for Distillers Discovery
Search space structure. In KD, the student student  $S$  is distilled with the fixed teacher  $T$  by minimizing:
$$
\mathcal {L} _ {K D} = \tau^ {2} \times \mathcal {W} _ {f} \times \mathcal {W} _ {C a l} \times \mathcal {D} \left(\mathcal {T} \left(f _ {S} / \tau\right), \mathcal {T} \left(f _ {T} / \tau\right)\right), \tag {1}
$$
where  $\mathcal{W}_f$  and  $\mathcal{W}_{Cal}$  is the loss weights factor and calibration [79],  $\tau$  is the temperature factor,  $\mathcal{T}$  is transformations,  $\mathcal{D}(\cdot ,\cdot)$  is distance function measuring the knowledge difference.  $f_{T}$  and  $f_{S}$  are outputs (e.g., features, embeddings, and logits) of the teacher-student. Following this general KD formulation, our search space consists of different types of operations (see Table 1) in transformations, distance functions, and loss weights parts. Then, we use a computation graph to represent each candidate, in which the input nodes are different types of knowledge and the intermediate nodes are primitive operations. In addition, we assign three transform options as transform-1  $\rightarrow$  transform-2  $\rightarrow$  transform-3 for the transformation part to ensure effective processing of the input knowledge.
Insight of space design. Our space design enjoys multiple merits. (1) Comprehensive & flexible: KD-Zero contains key elements of most existing KD methods, including the normalize, mask,
attention, and focal-calibration KDs [3, 79, 27] for distillation gaps. In addition, we introduce basic operations in loss design to improve the flexibility and diversity of the search. (2) Extended & Innovative: We extend advanced KD design for omni-dimensional and various knowledge. For example, we expand the normalized operation for batch, channel, and spatial dimensions, which can be used for logits, embeddings, and feature knowledge. Thus, such a unified and flexible search space would provide some inspiration for KD design.
3.2 Evolution Procedure, Acceleration and Objective
Evolution process. Our EA starts with an initial population of candidate distillers evaluated for their fitness based on our distillation gap-related multi-objectives. The algorithm then iteratively evolves the population over generations using genetic operators, such as selection, crossover, and mutation, to generate better solutions. Specifically, distillers are first randomly generated to form the initial population. Each candidate solution in the population is evaluated using multiple fitness functions that measure teacher-student gaps. Based on these evaluations, we select the best-performing individuals from the population to create a new population for the next generation. Then, we apply crossover to the selected individuals to create new offspring and use mutation to the new offspring to introduce diversity into the population, which helps to explore new distillers.
Search acceleration As the search space is sparse with many unpromising distillers, we employ several strategies to accelerate: (1) Loss-rejection protocol. We filter out candidates with excessive loss values or collapsed optimization during the search. (2) Search space shrinkage. We reduce sampling probabilities for the operations frequently in loss rejection and tail candidates with search iterations. (3) Proxy settings. With diverse and informative knowledge learned from a teacher, student models offer advantages in terms of faster training speeds. Based on these properties, we employ early stop the training process once the student model performs well enough to determine the quality of the candidate distillation. Nevertheless, proxy settings can also introduce evaluation uncertainties, and we alleviate this issue by introducing multiple distillation gap metrics in the following sections.
Fitting objectives. To accurately evaluate each distiller and reduce the distillation gap, we include cross-entropy loss, sharpness-gap [20] on prediction, and CKA-gap [52] on representation between teacher-student as the multi-objectives. Specifically, we conduct a training-free evolutionary search algorithm to efficiently discover the optimal distiller  $\alpha^{*}$  from search space  $\mathcal{A}$ , as:
$$
\alpha^ {*} = \arg \min  _ {\alpha \in \mathcal {A}} \left(\mathcal {L} _ {C E} \left(f _ {S}, Y\right) + \overbrace {\left(\log \left(\exp \left(f _ {S}\right)\right) - \log \left(\exp \left(f _ {T}\right)\right)\right)} ^ {\text {S h a r p n e s s - g a p}} - \overbrace {\frac {\mathrm {H S I C} \left(f _ {S} , f _ {T}\right)}{\sqrt {\mathrm {H S I C} \left(f _ {S} , f _ {S}\right) \mathrm {H S I C} \left(f _ {T} , f _ {T}\right)}}} ^ {\text {C K A - g a p}}\right),
$$
where  $\mathcal{L}_{CE}$  is the regular cross-entropy objective with labels  $Y$ , the sharpness metric is the logarithm of the exponential sum of logits, and Centered Kernel Alignment (CKA) metric is normalized from Hilbert-Schmidt Independence Criterion (HSIC) [19] on high-level feature.
Figure 5: Probability of operations within each search part, which counted from the Top-3 searched distillers for all teacher-student pairs in the CIFAR-100 experiment.
3.3 Results Analysis and Practical Guidance
Figure 5 present searched distillers for different models. Based on these results, some practical guidance for KD designs can be summarized as follows:
- For knowledge input, feature knowledge enjoys superiority while logits and embedded knowledge share identical status. Searched distillers with feature knowledge take more proportion, and ablation studies on vanilla feature knowledge also surpass logits and embedding knowledge.
This observation also aligns with the existing conclusion that feature-based KD outperforms other KD on detection tasks [27, 63].
- For transformations, the normalized-based operations play a key role in most optimal distillers, and its vanilla performance outperforms other types of transformations in ablation studies. These observations illustrate that normalization benefits distillation, consistent with the current KD methods. Scale-based and activation-based operations are also crucial for the KD because they are often present in optimal distillers and enjoy good vanilla performance in ablation studies.
- For distance functions, normalized-based distances such as  $\ell_{\text{Pearson}}$  and  $\ell_{\text{correlation}}$  enjoy better results in ablation studies and are often adopted by teacher-student pairs across architectures. This suggests that these normalized-based distances can reduce distillation gaps in complex scenarios. In addition, some simple distance functions (e.g.,  $\ell_1$ ,  $\ell_2$  and  $\ell_{KL}$ ) also appear in the optimal distiller. These may be attributed to these distillers employing advanced transformation operations for input knowledge.
- For loss weights, adjusting the temperature values is helpful for different scenarios, and the optimal weight values are generally between 1 and 10 based on search results and ablation studies. In addition, the focal weight calibration outperforms the no-weight calibration in ablation studies, and it is often used for some optimal distillers. In summary, we should employ smaller weight values and actively explore different weight calibration and temperature values to reduce distillation gaps under different teacher-student pairs.
3.4 Distilling Student via Discovered KD-Zero functions
After search, the discovered distiller  $\mathcal{L}_{KD}$  is used for distilling student  $f_{S}$  combined with cross-entropy  $\mathcal{L}_{CE}$  via label  $Y$  in single teacher  $f_{T}$  or multiple teacher  $f_{T_i}$  KD as:
$$
\mathcal {L} _ {\text {s i n g l e}} = \mathcal {L} _ {C E} \left(f _ {S}, Y\right) + \mathcal {L} _ {K D} \left(f _ {S}, f _ {T}\right), \quad \mathcal {L} _ {\text {m u l t i p l e}} = \mathcal {L} _ {C E} \left(f _ {S}, Y\right) + \sum_ {i = 1} ^ {N} \mathcal {L} _ {K D} \left(f _ {S}, f _ {T _ {i} ^ {\prime}}\right). \tag {3}
$$
Extended to various distillation designs & scenarios KD-Zero focuses on the distiller's search and employs simple  $1 \times 1$ -Conv for channel alignment. Recently, some KD methods have proposed other designs in feature aggregation [50], projector ensemble [9], N-to-One match [67], and cross-layer mapping [28, 13]. As shown in Table 2, KD-Zero can combine well with them by replacing their default  $\ell_2 / \ell_{KL}$  losses. In addition, the distiller search of KD-Zero also benefits different distillation scenarios. Specifically, Self-KDs, online KDs, and multi-teacher KDs with KD-Zero achieve extra gains than the default settings. Also, KD-Zero can combine architecture-level methods (e.g., Assistant [46]) to reduce distillation gaps further.
Table 2: Top-1 (%) accuracy of KD-Zero combined with different KD designs and scenarios for ResNet-20 and WRN-16-2 on CIFAR-100.
Net	Method	Different KD design	Different KD scenarios
Review [50]	PE [9]	NORM [67]	L2T-ww [28]	ONE [76]	BYOT [75]	DML [77]	AVE-MKD [58]	AE-MKD [58]	Assistant [46]
R20	ℓKL or ℓ2	71.89	71.36	71.55	70.89	70.77	70.37	70.92	71.24	71.36	71.06
KD-Zero	72.27	72.18	72.00	72.03	71.35	70.98	72.02	72.22	72.35	71.89
W-16-2	ℓKL or ℓ2	76.20	76.02	75.65	75.50	74.25	74.12	75.33	75.22	75.68	75.35
KD-Zero	76.62	76.47	76.26	76.36	75.65	75.21	76.45	76.72	76.78	76.05
Why can KD-Zero bridge the teacher-student gap? The answer is intuitive: In KD-Zero, the search space contains many operators for distillation gap reduction, and search objectives are directly designed to reduce the prediction and representation gap between teacher-student. For example, We approximate the sharpness gap using a Taylor second expansion [20]:
$$
G _ {g a p} = \log \left(\exp \left(f _ {T}\right)\right) - \log \left(\exp \left(f _ {S}\right)\right) \approx \log \left(1 + f _ {T} + \frac {1}{2} f _ {T} ^ {2}\right) - \log \left(1 + f _ {S} + \frac {1}{2} f _ {S} ^ {2}\right), \tag {4}
$$
Following Hindon's assumption [23] that the logits of each training sample are approximately zero-meaned, i.e.,  $\bar{f}_T,\bar{f}_S = 0$ . So the gap can be rewritten as  $\log \left(1 + \frac{1}{2}\text{Var}(f_T)\right) - \log \left(1 + \frac{1}{2}\text{Var}(f_S)\right)$  [20]. Our options like normalized transform [3, 27], and weight calibrations [79] can effectively reduce the variance of teacher-students' outputs and minimize the distillation gap. In addition, the student models distilled with KD-zero enjoy the merits: (1) Fast convergence and superior performance. As shown in Figure 6, KD-Zero has surpassed the best accuracies of KD
Figure 6: Comparison of training curves (left) of ResNet-20 on CIFAR-100 and ResNet-18 on ImageNet (right).
Figure 7: Penultimate-layer visualization (left), logits-correlation map of teacher-student (ResNet-110/20) via KD (middle) & KD-Zero (right).
Figure 8: Comparison on the Grad-CAM++ [4] visualization results between the features of the baseline student model, the student model trained with KD and our KD-Zero, and the teacher model. Results are obtained on ImageNet with ResNet18 (left) and MobileNet (right).