Title: ROLE-WISE DATA AUGMENTATION FOR KNOWLEDGE DISTILLATION

Abstract: 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 background 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 results will be made publicly available.

1. INTRODUCTION:

Background and Motivation. In the educational psychology literature, it is generally considered beneficial if teachers can adapt curricula based upon students' prior experiences (Bandura, 2002; Brumfiel, 2005; Gurlitt et al., 2006; Slotta & Chi, 2006) . These vary widely depending on students' cultural backgrounds, 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 background.

Knowledge distillation (KD) (Bucilua et al., 2006; Hinton et al., 2014 ) 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, Hinton et al. (2014) 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, Romero et al. (2015) 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 (Ho et al., 2019) . 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: (1) training a teacher network till convergence while learning a schedule of policies to augment the training data specifically for the teacher; (2) 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 (Strubell et al., 2019) . KD-based quantization (Zhuang et al., 2018; Polino et al., 2018) 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 (Zhuang et al., 2018; Polino et al., 2018) . 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.

2. RELATED WORK:

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 (Hinton et al., 2014; Romero et al., 2015) . In terms of the definition of knowledge to be distilled from the teacher, existing models typically use teacher's class probabilities (Hinton et al., 2014) and/or intermediate features (Romero et al., 2015; Park et al., 2019; Seunghyun Lee, 2019; Hyun Lee et al., 2018) . Among those KD methods that utilize intermediate feature maps, Relational KD (RKD) considers (Park et al., 2019) the intra-relationship in the same feature map, while Multi-Head KD (MHKD)(Seunghyun Lee, 2019) and KD using SVD (KD-SVD) (Hyun Lee et al., 2018) 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 (Krizhevsky et al., 2012; He et al., 2016) . Several recent works attempt to automate the data augmentation process. Generative adversarial networks (Ratner et al., 2017) and Bayesian optimization (Tran et al., 2017) have been used for this process. DeVries & Taylor (2017) augment training data in the learned feature space by injecting noise and interpolation. Lemley et al. (2017) learn how to combine pairs of images for data augmentation. AutoAugment (Cubuk et al., 2018) 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) (Ho et al., 2019) 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.

3. PRELIMINARIES:



3.1. POPULATION-BASED AUGMENTATION (PBA):

PBA (Ho et al., 2019) , 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.

3.2. KNOWLEDGE DISTILLATION (KD):

Following the notations in (Park et al., 2019) , a KD method aims to minimize the objective function

EQUATION

)

where λ is a hyper-parameter to balance the impact of the KD loss term.

In this paper, for classification tasks,

L task = xi∈X H(softmax(F final S (x i )), 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

EQUATION

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 (Hinton et al., 2014) , the objective can be defined as

EQUATION

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 (Park et al., 2019) considers the intra-relationships. That is, given the feature map of layer j, the KD loss can be formulated as:

EQUATION

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 (Seunghyun Lee, 2019; Hyun Lee et al., 2018) consider the interrelationships, where the KD term can be formulated as:

EQUATION

Here, ϕ(•) measures the inter-relationship between feature maps of different layers, i.e. k = j. 

3.3. QUANTIZATION:

In this work, we use DoReFa2 (Zhou et al., 2016) to quantize both weights and activations. The quantization function Q(•) is defined as:

EQUATION

)

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:

EQUATION

The back-propagation is approximated by the straight-through estimator (Bengio et al., 2013) and the partial gradient ∂l ∂r w.r.t. the loss l is computed as: ∂l ∂r = ∂l ∂r q

• ∂r q ∂r ≈ ∂l ∂r q

.

(8)

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 .

4.1. STAGE-α:

In general, a teacher can provide better training signals for the student if the teacher's performance increases (Mirzadeh et al., 2019) . 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.

4.2. STAGE-β:

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 . When N S is a low-precision network, following (Furlanello et al., 2018) , we share the same network architecture3 between N T and N S . When N S is a full-precision network, it will have fewer layers compared to N T .

5. EXPERIMENTS:



5.1. SETTINGS:

We evaluate our approach on two benchmark datasets: CIFAR-10 ( Krizhevsky et al., 2009) 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 (Ho et al., 2019) . 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 (Ho et al., 2019) , 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 (Krizhevsky et al., 2012) and ResNet18 (He et al., 2016) . 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 (Zhou et al., 2016) , the first layer and last layer are not quantized.

Following Clark et al. (2019) , 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.

5.2. COMPARING DIFFERENT KD METHODS:

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:

EQUATION

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 (Park et al., 2019) , MHGD (Seunghyun Lee, 2019) and KD-SVD (Hyun Lee et al., 2018) . As shown in Eq. ( 9), 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. ( 9), as we have two KD terms. Tab. 1 reports the results 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 results also reveal that only relying soft labels is not as effective as utilizing multiple supervising signals from the teacher.

Table 1 : Accuracy (%) on CIFAR-100 with ResNet18 using different KD methods. We compare ours with the following methods: Soft labels (Hinton et al., 2014) , DML (Zhang et al., 2018) , RKD (Park et al., 2019) 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 results obtained in (Zhuang et al., 2018) , 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 results 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%.

5.4. COMPARING SCHEDULES:

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 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. 

5.5. ANALYZING THE LEARNED SCHEDULES:

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 . 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 results obtained when N S is low-precision network. This shows that our proposed method can be used for full-precision training tasks. 

6. CONCLUSION:

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 will visualize the learned schedules of policies in Section 5.5.

It should be noted that our proposed method is orthogonal to any particular quantization method.

Note this is not a hard constraint, we choose such strategy to reduce the number of factors that might influence the final performance.

Figures and Tables:

Figure 1: Concept diagram (stage-α) of training a data augmentation agent for the N T .

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 .

(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.

, MHGD (Seunghyun Lee, 2019) and KD-SVD(Hyun Lee et al., 2018).

Table content: <table><tr><td>Bit-Width (Weight / Activation)</td><td colspan="5">Soft labels DML RKD MHGD KD-SVD II-KD</td></tr><tr><td>4/4</td><td>70.48</td><td>72.47 71.84</td><td>73.52</td><td>73.92</td><td>74.21</td></tr><tr><td>2/2</td><td>70.09</td><td>69.72 70.71</td><td>71.80</td><td>72.97</td><td>73.35</td></tr><tr><td colspan="6">5.3 IS ROLE-WISE AUGMENTATION WITH KD EFFECTIVE FOR QUANTIZATION?</td></tr></table>

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(Zhou et al., 2016) 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 .

Table content: <table><tr><td>Methods</td><td/><td>AlexNet CIFAR-10</td><td>AlexNet CIFAR-100</td><td>ResNet18 CIFAR-10</td><td>ResNet18 CIFAR-100</td></tr><tr><td/><td>32-bit</td><td>90.58</td><td>65.80</td><td>93.57</td><td>74.85</td></tr><tr><td>Vanilla Training</td><td>4-bit</td><td>89.72</td><td>60.25</td><td>90.97</td><td>69.81</td></tr><tr><td/><td>2-bit</td><td>88.77</td><td>58.96</td><td>90.00</td><td>67.06</td></tr><tr><td/><td>32-bit</td><td>91.62</td><td>66.40</td><td>94.49</td><td>75.19</td></tr><tr><td>Teacher after Stage-α</td><td>4-bit</td><td>90.06</td><td>60.65</td><td>91.47</td><td>70.24</td></tr><tr><td/><td>2-bit</td><td>89.28</td><td>58.59</td><td>89.99</td><td>67.32</td></tr><tr><td>Student with only II-KD</td><td>4-bit 2-bit</td><td>90.55 89.18</td><td>65.55 63.49</td><td>91.42 90.60</td><td>73.85 72.44</td></tr><tr><td>Student after Stage-β</td><td>4-bit 2-bit</td><td>92.00 90.63</td><td>65.69 64.06</td><td>94.44 93.20</td><td>74.21 73.35</td></tr></table>

Accuracy (%) on CIFAR-100 with 4-bit networks using different KD methods.

Table content: <table><tr><td/><td/><td>AlexNet</td><td/><td/><td>ResNet18</td></tr><tr><td>Methods</td><td colspan="2">DML MHGD</td><td>Ours</td><td colspan="2">DML MHGD II-KD</td></tr><tr><td colspan="2">Schedule based on teacher 61.61</td><td>61.62</td><td colspan="2">64.97 71.78</td><td>69.76</td><td>73.46</td></tr><tr><td colspan="2">Schedule based on student 63.73</td><td>63.47</td><td colspan="2">65.69 72.47</td><td>73.52</td><td>74.21</td></tr></table>

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.

Table content: <table><tr><td/><td colspan="4">Vanilla Training Re-Training with PBA Student with only II-KD After Stage-β</td></tr><tr><td>Accuracy</td><td>74.31</td><td>74.52</td><td>75.52</td><td>75.88</td></tr></table>

Peer review (iclr2020):
This paper takes the idea of Population-Based Augmentation (PBA) and extends it to knowledge distillation (KD). The idea is that the ideal augmentation protocol for training-from-scratch (or in this context teacher training) may not be the best for student networks under a KD loss. I am borderline about this paper and had to pick one, so I landed on Weak Reject. On one hand, I think it’s a really neat idea to apply PBA in this context, and package it as a strage-alpha/stage-beta training procedure. However, the experimental results seem very incremental and I’m not convinced there is a genuine signal there. Experiments: Table 1 offers a great comparison of prior work, as well as the combination of prior work (II-KD). Table 2 tells me that PBA seems incremental both for the teacher and the student. Going from vanilla training to student with II-KD gets you most of the way there, and the primary contribution of this paper just gives you a slight benefit above this. It’s great that a more traditional full-precision comparison was also added in Table 4, but this table also confuses me. First, it was unclear if “vanilla training” referred to the teacher or the student. The number 74.31 is not the same as in Table 2 (74.85), so I assume this means it’s the student? If so, the student is already very close to the teacher, and this is not a great starting point for evaluating KD. The student after stage-beta also outperforms the teacher - something that was mentioned in the related work, but I would like more discussion around it specifically for Table 4. Another thing I was wondering was how important PBA is for the teacher’s ability to be a good teacher. It gives a modest boost in Table 2; what if we skip PBA in the teacher but still do it for the student. This would be interesting to add. Overall, there aren’t that many experiments. The dataset is never more challenging than CIFAR-100. There are also no error bars, which are particularly important when the improvements are small. As for Figure 3, I don’t know if there is anything intuitive we can glean from this. I may just be variation between experiments, as far as I can tell. I think this paper can be made stronger by making the experimental evidence broader, as well as the analysis of why this works stronger. Without these improvements, the reader is left wondering if there really is any significant benefit. We have to remember that PBA is not cheap (perhaps much cheaper than AutoAugment, but more expensive than fixed augmentation). For most practitioners, the complication and compute costs of PBA would probably not be worth adding on top of KD, if the benefits are too modest. Minor: In Table 3, it says *Ours* for AlexNet and *II-KD* for ResNet8. Should the both be the same?