Knowledge Distillation (KD) is a common method for transferring the ``knowledge'' learned by one machine learning model (the teacher) into another model (the student), where typically, the teacher has a greater capacity (e.g., more parameters or higher bit-widths). To our knowledge, existing methods overlook the fact that although the student absorbs extra knowledge from the teacher, both models share the same input data -- and this data is the only medium by which the teacher's knowledge can be demonstrated. Due to the difference in model capacities, the student may not benefit fully from the same data points on which the teacher is trained. On the other hand, a human teacher may demonstrate a piece of knowledge with individualized examples adapted to a particular student, for instance, in terms of her cultural and interests. Inspired by this behavior, we design data augmentation agents with distinct roles to facilitate knowledge distillation. Our data augmentation agents generate distinct training data for the teacher and student, respectively. We focus specifically on KD when the teacher network has greater precision (bit-width) than the student network. We find empirically that specially tailored data points enable the teacher's knowledge to be demonstrated more effectively to the student. We compare our approach with existing KD methods on training popular neural architectures and demonstrate that role-wise data augmentation improves the effectiveness of KD over strong prior approaches. The code for reproducing our will be made publicly available. Background and Motivation. In the educational psychology literature, it is generally considered beneficial if teachers can adapt curricula based upon students' prior experiences (; ; ;). These vary widely depending on students' cultural s, previous educational experiences, interests, and motivations. To help students with different prior experiences to comprehend, memorise, and consolidate a piece of knowledge, teachers may provide extra and customized teaching material during their teaching processes. For instance, when teaching the concept of the color pink, a teacher may choose flamingos, sakura (cherry blossoms), or ice cream cones as the example, depending on a student's . Knowledge distillation (KD) is a common framework for training machine learning models. It works by transferring knowledge from a higher-capacity teacher model to a lower-capacity student model. Most KD methods can be categorized by how they define the knowledge stored in the teacher (i.e., the "soft targets" of training as defined in existing literature). For instance, originally proposed KD for neural networks, and they define the output class probabilities (i.e., soft labels) generated by the teacher as the targets for assisting the training of students. In a follow up work, defined the soft targets via the feature maps in the teacher model's hidden layers. To train a student network with KD effectively, it is important to distill as much knowledge from the teacher as possible. However, previous methods overlook the importance of the medium by which the teacher's knowledge is demonstrated: the training data points. We conjecture that there exist examples, not necessarily seen and ingested by the teacher, that might make it easier for the student to absorb the teacher's knowledge. Blindly adding more training examples may not be beneficial because it may slow down training and introduce unnecessary biases . The analogy with how human teachers adjust their teaching to their students' particular situations (e.g., with the feedback gathered from the students during teaching) suggests that a reasonable yet uninvestigated approach might be to augment the training data for both the teacher and student according to distinct policies. In this paper, we study whether and how adaptive data augmentation and knowledge distillation can be leveraged synergistically to better train student networks. We propose a two-stage, rolewise data augmentation process for KD. This process consists of: training a teacher network till convergence while learning a schedule of policies to augment the training data specifically for the teacher; distilling the knowledge from the teacher into a student network while learning another schedule of policies to augment the training data specifically for the student. It is worth noting that this two-stage framework is orthogonal to existing methods for KD, which focus on how the knowledge to be distilled is defined; thus, our approach can be combined with previous methods straighforwardly. Although our proposed method can in principle be applied to any models trained via KD, we focus specifically on how to use it to transfer the knowledge from a full-precision teacher network into a student network with lower bit-width. Network quantization is crucial when deploying trained models on embedded devices, or in data centers to reduce energy consumption . KD-based quantization jointly trains a full-precision model, which acts as the teacher, alongside a low-precision model, which acts as the student. Previous work has shown that distilling a full-precision teacher's knowledge into a low-precision student, followed by fine-tuning, incurs noticeable performance degradation, especially when the bit-widths are below four . We show that it is advantageous to use adaptive data augmentation to generate more training data for the low-precision network based on its specific weaknesses. For example, low-precision networks may have difficulties learning rotationrelated patterns, 1 and the data augmentation agent should be aware of this and generate more such data points. One positive side-effect for demonstrating the effectiveness of our method is that the improvement brought by our proposed method is more significant compared to the experiments on all full-precision models. Knowledge distillation. KD is initially proposed for model compression, where a powerful wide/deep teacher distills knowledge to a narrow/shallow student to improve its performance . In terms of the definition of knowledge to be distilled from the teacher, existing models typically use teacher's class probabilities and/or intermediate features (; ; ;). Among those KD methods that utilize intermediate feature maps, Relational KD (RKD) considers the intra-relationship in the same feature map, while Multi-Head KD (MHKD) and KD using SVD (KD-SVD) utilize the inter-relationship across feature maps. By contrast, we propose to incorporate both the intra-and inter-relationships within and across feature maps. Automated data augmentation. Manually applying data augmentation rules such as random rotating, flipping, and scaling are common practices for training neural models on image classification tasks (; . Several recent works attempt to automate the data augmentation process. Generative adversarial networks and Bayesian optimization have been used for this process. augment training data in the learned feature space by injecting noise and interpolation. learn how to combine pairs of images for data augmentation. AutoAugment searches for the optimal data augmentation policies (e.g., how to rotate) based on reinforcement learning. However, the search process is computationally expensive. Population-based augmentation (PBA) uses an evolution-based algorithm to automatically augmenting data in an efficient way. In contrast to previous approaches, we study the effect of the training data for KD and propose to use automatic data augmentation to train the student better from her teacher. PBA , as an evolutionary search algorithm, learns a dynamic per-epoch schedule of augmentation policies, denoted as A. Since this schedule is epoch-based, it will re-create the augmented dataset every epoch. More concretely, PBA begins with a population of models that are trained in parallel on a small subset of the original training data. The weights of the worst performing models in the population are replaced by those from better performing models (i.e., exploitation), and the policies are mutated to new ones within the pre-defined policy search space (i.e., exploration). After training, PBA usually keeps the learned augmentation schedule of policies but discards the elementary parameters of the models. A different model (e.g., a larger one) can then use the learned schedule to improve its training on the same task. Following the notations in , a KD method aims to minimize the objective function where λ is a hyper-parameter to balance the impact of the KD loss term. In this paper, for classification tasks,, y truth ), where X refers to training sample space, y truth ∈ Y are the ground-truth labels, F S (·) is the student network, and H(·) denotes the cross-entropy. The KD term can be defined as where F(·) is the function of the network and l(·) is a loss function to compute the difference between the teacher network and the student network. For KD methods that use soft labels , the objective can be defined as where F final (x i) is the feature map of the final layer. There exist some KD methods that utilize the intermediate feature maps in complementary ways. For example, Relational KD considers the intra-relationships. That is, given the feature map of layer j, the KD loss can be formulated as: where Φ(·) refers to the potential function measuring the pairwise relationship inside a feature map from student network or teacher network and F j (x i) is the feature map of layer j (which may include the final logits layer). Therefore, this feature-based KD method includes the benefits of using soft labels. On the other hand, some works consider the interrelationships, where the KD term can be formulated as: Here, ϕ(·) measures the inter-relationship between feature maps of different layers, i.e. k = j. In this work, we use DoReFa 2 to quantize both weights and activations. The quantization function Q(·) is defined as: where r is the full-precision value, r q indicates the quantized value, n bits refers to the number of bits to represent this value. With this quantization function, the quantization on weights w is defined as: The back-propagation is approximated by the straight-through estimator and the partial gradient ∂l ∂r w.r.t. the loss l is computed as: 4 THE PROPOSED METHOD Our proposed method has two stages, which will be described in the following subsections. In the first stage, we train a teacher network, denoted as N T, with the help of PBA-based augmentation. In the second stage, we further distill the knowledge from N T (pre-trained in the first stage) to the student network, denoted as N S, while learning another augmentation schedule to augment the training data for N S. In general, a teacher can provide better training signals for the student if the teacher's performance increases . As shown in Fig. 1, we apply PBA to learn a dynamic per-epoch schedule of augmentation policies, A T, for N T on a small subset of training data. That is, the augmentation agent's training signal is defined as the feedback of N T's accuracy on a subset of the dataset. After this, we use the discovered schedule A T to augment the whole training dataset and re-train N T on it till convergence. KD methods have shown to be effective at improving the performance of lower-capacity networks using the knowledge from higher-capacity networks. In order to take advantage of this functionality, we apply the KD methods together with data augmentation in stage-β, as shown in Fig. 2. More concretely, we first use PBA to learn an epoch-based augmentation schedule A S for N S on a subset of the dataset. Different from the schedule A T learned in stage-α, A S is learned based on the feedback (i.e., accuracy) from N S, which is trained with KD. In other words, N S receives additional training signals from N T that is pre-trained in stage-α. We then use the learned A S to augment the whole training dataset, and re-train N S on it with the distilled knowledge from N T. Note that, because the learned schedule is epoch-based, we do not use the discovered schedule A T from stage-α to augment the training data as initialization. Figure 2: Concept diagram (stage-β) to augment training datapoints for both N T and N S. The N T has been pre-trained using the method shown in Section 4.1, and is fixed during training. The augmentation agent in stage-β is designed to learn schedules of polices that are different from those learned in stage-α, and thus the agent only receives the feedback from N S. When N S is a low-precision network, following , we share the same network architecture 3 between N T and N S. When N S is a full-precision network, it will have fewer layers compared to N T. We evaluate our approach on two benchmark datasets: CIFAR-10 and CIFAR-100. CIFAR-10 consists of 60,000 32×32 color images in 10 classes, with 6000 images per class and CIFAR-100 has 100 classes with 600 images per class. Both have 50,000 training images and 10,000 test images. We search over a "reduced" CIFAR-10/CIFAR-100 with 4,000 training images and 36,000 validation images, which is the same as in . All the data augmentation models are run with 16 total trials to generate augmentation schedules. Following PBA, in stage-α, we run PBA to create schedules over separate models and then transfer the CIFAR-10 policy to CIFAR-100. However, for student network training in stage-β, we empirically use the respective "reduced" dataset. The data augmentation approaches for the baselines include random crop and horizontal flipping operations. Following , our policy search space has a total of 15 operations, each having two magnitude and discrete probability values. We use discrete probability values from 0% to 100%, in increments of 10%. Magnitudes range from 0 to 9. The models we evaluate on include AlexNet and ResNet18. The number of epochs is 200 and the batch size is set to 128. For full-precision network N T, the learning rate starts from 0.1 and is decayed by 0.1 after every 30% of the total epochs. we use SGD with a Nesterov momentum optimizer. The weight decay is set to 5 × 10 −4. For quantization, the learning rate is set to 10 −3 and is divided by 10 every 30% of the total epochs. We use the pre-trained teacher network model as the initial point of student network. We use a smaller weight decay 10 −5 assuming that less regularization is needed by the lower-precision networks. Following DoReFa , the first layer and last layer are not quantized. , during training, we gradually transit the student from learning based on the teacher to training based on the ground-truth labels. This heuristic provides the student with more rich training signals in the early stage but does not force the student to strictly mimic the teacher's behaviors. As for the implementation, we decay the balancing hyper-parameter λ in the KD loss by 0.5 every 60 epochs. As mentioned in Section 3.2, there exist complementary KD methods considering both intra-and inter-relationships within and across feature maps. A natural question is if it would be beneficial to combine them to further boost the performance together with data augmentation. Therefore, we propose a simple extension to these complementary KD methods, dubbed as II-KD, by incorporating intra-relationships inside the feature map and inter-relationships across different feature maps. We incorporate the two relationships into the final objective function as follows: where we only use a single balancing hyper-parameter λ between the original loss and the distillation loss, which does not introduce extra hyper-parameters. More precisely, our KD method incorporates components of three conventional KD methods: RKD , MHGD and KD-SVD . As shown in Eq., we add the three KD terms together with equal coefficients. We use the loss function l(·) following their approaches. For the back-propagation, we clip the gradient for KD loss as in KD-SVD, because this will smoothly post-processes the gradient to limit the impact of KD loss in training. For AlexNet we select the feature maps of ReLU layers after the convolution/max pooling layer. For ResNet18, we select the feature maps of the last ReLU layer of each residual block. We evaluate our proposed KD extension on CIFAR-100 with ResNet18 for different bit-width settings by comparing with various KD methods. For the baseline methods, we use their default settings with a fixed and pre-trained teacher network in the training stage and λ = 1 for the knowledge distillation loss. We set λ = 0.4 for II-KD in Eq., as we have two KD terms. Tab. 1 reports the on various augmented KD methods. We observe that our proposed methods clearly outperforms the other KD methods on all the settings, though the improvements over MHGD and KD-SVD are not huge. The also reveal that only relying soft labels is not as effective as utilizing multiple supervising signals from the teacher. In this subsection, we aim to answer this question: is our two-stage role-wise augmentation with KD effective for network quantization? We conduct experiments on CIFAR-10 and CIFAR-100 datasets under full-precision, 4-bit, and 2-bit settings. From Tab. 2, we can observe that training with learned data augmentation schedules does not improve the performance of low-precision networks too much. Similar to the obtained in , transferring knowledge from the full-precision to the low-precision student usually helps the training of students, which is especially obvious on the CIFAR-100 dataset. Tab. 2 also clearly shows that our proposed pipeline consistently improves the performance of the lowprecision student networks. For example, the 4-bit N S is comparable with full-precision reference without loss of accuracy for CIFAR-10 and with loss of accuracy within 1.0% on CIFAR-100. When decreasing the precision to 2-bit, the are still promising as compared with other baselines, even though there is a performance gap between the 2-bit and the full-precision models. For instance, our approach usually outperforms the strong baseline, only using II-KD, by more than 1.0%. Here we aim to answer this question: how effective is it if we use A T, learned based on the feedback from N T in stage-α, to dynamically augment the training dataset and train N S on it? Tab. 3 reports the accuracy comparison with different KD methods and augmentation schedules. We can clearly see that augmenting the training dataset for N S with A S consistently outperforms those using the transferred schedules A T among different KD methods. This observation is consistent with our assumption that N S has her own optimal augmentation schedule, A S, that is different from A T for N T. In particular, blindly applying the teacher augmentation schedule A T may negatively influence Table 2: Accuracy (%) on CIFAR-10 and CIFAR-100 datasets with different bit-widths. Vanilla Training for 4-bit and 2-bit refers to training a network based on DoReFa from scratch without learned data augmentation. Teacher after Stage-α refers to using learned schedules discovered by PBA to re-train N T as described in Section 4.1. Student with only II-KD refers to training N S using II-KD but without the learned data augmentation. Student after Stage-β refers to training N S using II-KD and the learned data augmentation. For Vanilla Training and Teacher after Stage-α, we report the accuracy of N T, and for the rest we report the accuracy of N S. the training of N S as compared to only using KD. For example, the learned schedule based on the teacher A T degrades the performance of N S by 0.58% for AlexNet on CIFAR-100 as compared to applying KD methods, as shown in Tab. 2. To analyze the difference on the discovered schedules between N T (i.e., full-precision ResNet18) and N S (i.e., 4-bit ResNet18), we report their augmented schedules quantitatively in terms of normalized probability and magnitude on CIFAR-100 in Fig. 3. We normalize the probability of each epoch by dividing the maximal summation of probabilities for all operations across all epochs. It can be seen that the discovered schedules A S for N S is quite different from A T for N T. In particular, for A T, there is an emphasis on Brightness, Posterize, Rotate, Sharpness and TranslateY, while A S cares more about Contrast, ShearX and TranslateY. Furthermore, we observe that the probability and magnitude increase as the epoch evolves. For A S, in the beginning, KD plays a more important role, and there is no augmentation operation before about epoch 50. As the training continues, the augmentation policies become more important. One possible reason is that, for lowprecision networks, KD methods can provide rich training signals such that data augmentation does not help in the early training phases. Furthermore, we observe that, compared to A T, the schedule for student A S evolves more smoothly in the sense that the policy updating frequency is lower. For example, the probability and magnitude values change about every 40 epochs for student, while the policies for teacher update about every 15 epochs. One possible reason is that for the low-precision N S, KD methods make the training process more smooth and it is not necessary to change the augmentation policies too frequently. This is consistent with the observations shown in Tab. 2 that KD can already provide useful training signals. Also, this validates our assumption that N S has her own optimal augmentation schedule A S that is different from A T. (a) Normalized plot of operation probability parameters over time for the teacher network NT. (b) Operation magnitude parameters over time for the teacher network NT. (c) Normalized plot of operation probability parameters over time for the student network NS. (d) Operation magnitude parameters over time for the student network NS. Figure 3: Evolution of magnitude and probability parameters in the learned schedules. Each operation appears in the parameter list twice, and we take the mean values of the parameter. This subsection aims to verify the effectiveness of our proposed methods on more conventional settings where both N T and N S are full-precision networks. We select ResNet18 as the teacher, and ResNet8 as the student to check how our proposed methods affect the student network. Tab. 4 shows that our proposed method outperforms the standard baseline training. The discovered augmentation schedule further boosts the performance of shallow N S based on II-KD, though the improvement is not that significant compared with the obtained when N S is low-precision network. This shows that our proposed method can be used for full-precision training tasks. Table 4: Accuracy (%) on CIFAR-100 with full-precision ResNet8 as N S and full-precision ResNet18 as N T under different settings. Vanilla Training refers to training a full-precision network from scratch. Re-Training with PBA refers to using learned schedules discovered by PBA to re-train N T as described in Section 4.1. Student with only II-KD refers to training N S using II-KD but without the learned data augmentation. Student after Stage-β refers to training N S using II-KD and the learned data augmentation. Previous literature on KD focuses on exploring the knowledge representation and the strategies for distillation. However, both the teacher and student learn from the same training data without adapting the different learning capabilities. To address this issue, we propose customizing distinct agents to automatically augment the training data for the teacher and student, respectively. We have extensively studied the effect of combining data augmentation and knowledge distillation. Furthermore, we propose a simple feature-based KD variant that incorporates both intra-and inter-relationships within and across feature maps. We have empirically observed that the student can learn better from the teacher with the proposed approach, especially in the challenging low-precision scenario, and the learned schedules are different for the teacher and student.
We study whether and how adaptive data augmentation and knowledge distillation can be leveraged simultaneously in a synergistic manner for better training student networks.