Teaching What You Should Teach: A Data-Based Distillation Method

As illustrated in Fig. 1, TST can be divided into three stages, namely Stage I, Stage II, and Stage III. To be specific, Stage I introduces the reasonable priori bias into the neural network-based data augmentation module before formal training, which helps ensure the stability of training in Stage II and Stage III, and improves the student’s ultimate generalization ability. Then, Stages II and III obtain an excellent student by alternately training them. In Stage II, we initialize the learnable magnitudes and probabilities, and search augmented samples that the teacher is good at but the student is not, by minimizing the cross-entropy of the teacher and maximizing the cross-entropy of the student. We expect that the student will serve its deficiencies in Stage III and obtain good generalization ability after training. In particular, Stage III can be viewed as a simple knowledge distillation framework that only applies a traditional cross-entropy loss and a Kullback-Leibler divergence to optimize the student. Therefore, the core of TST is to find more desirable training data for Stage III through Stage I and Stage II. The more detailed algorithmic procedure of TST can be found in Appendix C.

Stage III is a simple knowledge distillation framework, which can be interpreted as follows: given a teacher $f_{T}$ with parameter $\theta_{T}$ and a student $f_{S}$ with parameter $\theta_{S}$, when a training set $\mathcal{X}$ including the original samples and augmented samples are utilized for training, we can sample the mini-batch $\{x^{k}\}^{B}_{k=1}$ ($B$ denotes the batch size) from it and get the student output $\{f_{S}(x^{k})\}^{B}_{k=1}$ and the teacher output $\{f_{T}(x^{k})\}^{B}_{k=1}$ by forward propagation. Let us define the normalized exponential function $p$, i.e., softmax, to calculate the probability distribution of the model output. The goal of knowledge distillation is to minimize the cross-entropy loss $\mathcal{L}_{\textrm{CE}}$ between $\{p(f_{S}(x^{k}))\}^{B}_{k=1}$ and the ground truth label $\{y^{k}\}^{B}_{k=1}$, and the Kullback-Leibler divergence $\mathcal{L}_{\textrm{KL}}$ between $\{p(f_{T}(x^{k})/\tau)\}^{B}_{k=1}$ and $\{p(f_{S}(x^{k})/\tau)\}^{B}_{k=1}$, where $\tau$ refers to the temperature weight. Thus, the total loss function $\mathcal{L}_{\textrm{total}}$ is formulated as

where $w_{\textrm{ce}}$ and $w_{\textrm{kl}}$ are balanced weights. Note that Eq. 1 can also be appended with other distillation loss functions, e.g. Mean Squared Error function ($\mathcal{L}_{MSE}$), to improve the performance of distillation, which is applied in TST for the object detection task.

The goal of TST is to assist the student in overcoming its deficiencies in Stage III by discovering more desirable data that the teacher excels at but the student does not in Stage II. We denote the neural network-based data augmentation module $f_{DE}$ contains a data augmentation meta-encoder set $\{f^{i}_{E}\}_{i=1}^{N}$ with a parameter set $\{\theta^{i}_{E},\theta^{i}_{p}\}_{i=1}^{N}\cup\{\theta^{i}_{m}\}_{i=1}^{N-N_{\textrm{nl}}}$, where $N$ and $N_{\textrm{nl}}$ refer to the total number of all sub-policies and the total number of sub-policies without magnitudes, respectively. Note that each $f_{E}$ has the frozen parameter $\theta_{E}$ containing priori bias, and its other parameters, i.e., the magnitude²²2Magnitudes indicate the strength of the sub-policies transition. Some sub-policies do not have this variable, such as Cutout, Equalize, and Invert. These details will be introduced later. $\theta_{m}$ and the probability $\theta_{p}$, are learnable in Stage II. Therefore, as shown in the upper right corner of Fig. 1, we denote the loss function of TST in Stage II as

where $\hat{x}$ denotes the augmented sample obtained from the original sample after $f_{DE}$, and $\alpha$ and $\beta$ denote the balanced weights. Specific details about $f_{DE}$ will be presented in Sec. 3.5. The optimization objective of Stage II $\mathop{\textrm{min}}\limits_{\{\theta^{i}_{m},\theta^{i}_{p}\}_{i=1}^{N}}\!\mathcal{L}_{\textrm{TST}}$ is to let the student misclassify but let the teacher correctly classify so that the teacher can transfer the most helpful knowledge to students. To better find desirable data that the student cannot classify correctly, we use $\mathcal{L}_{\textrm{CE}}\left(1\!-\!p(f_{S}(\hat{x}^{k})),y^{k}\right)$ instead of $-\mathcal{L}_{\textrm{CE}}\left(p(f_{S}(\hat{x}^{k})),y^{k}\right)$ to avoid non-convex optimization, and the proof can be found in Appendix D. Moreover, the parameters $\{\theta^{i}_{E}\}_{i=1}^{N}$ are frozen because they contain priori bias in favor of distillation, which is imported in Stage I. As demonstrated later in Sec. 4.2, distillation can easily fail if $\{\theta^{i}_{E}\}_{i=1}^{N}$ is not frozen due to the instability of the data augmentation module.

During standard distillation training, Stage II and Stage III alternate. Specifically, TST first employs Stage II to find new samples suitable for distillation, and then utilizes Stage III to improve the student’s generalization ability. Inspired by Carlini and Wagner (2017), we claim that it is easy to search augmented samples in Stage II that match what the teacher is good at but the student is not, while it is difficult for the student to absorb knowledge contained in the augmented samples in Stage III. Therefore, the number of iterations $n_{\textrm{encoder}}$ of Stage II will be much smaller than the number of iterations $n_{\textrm{student}}$ of Stage III. Also, Sec. 4.2 illustrates that the student trained by TST will have better performance when $n_{\textrm{encoder}}\ll n_{\textrm{student}}$.

Noise-based generative models require expensive computational costs and suffer from training instability (more detailed can be found in Appendix D.5). Thus, Stage I is applied to introduce the augmented priori bias into $f_{DE}$ and let $\{\theta^{i}_{p}\}_{i=1}^{N}\cup\{\theta^{i}_{m}\}_{i=1}^{N-N_{\textrm{nl}}}$ be differentiable. Thanks to priori bias, TST is able to distill a strong student in a limited number of iterations. For simplicity, we consider each meta-encoder $f_{E}$ in the set $\{f^{i}_{E}\}_{i=1}^{N}$ as a black box in this paragraph. In Stage I, since Equalize, Invert and Cutout do not have the property of magnitude, they do not need to be fitted and can be called directly in Stage II and Stage III. Of course, their probabilities are learnable and will still be optimized in Stage II. Then, we categorize a series of single data augmentations as follows: (a) magnitude-unlearnable transformations, including Equalize, Invert, and Cutout; (b) learnable affine transformations, including Rotate ShearX, ShearY, TranslateX, and TranslateY; (c) learnable color transformations, including Posterize, Solarize, Brightness, Color, Contrast, and Sharpness Cubuk et al. (2019, 2020). For the learnable affine transformations and color transformations, we apply Spatial Transformer Network Jaderberg et al. (2015) and Color Network for fitting, respectively. Since both types (a) and (b) are fitted in the same way, we define a set of the manual data-augmentation mappings $\{f_{A}^{i}\}_{i=1}^{N-N_{\textrm{nl}}}$ that includes all single data augmentations with magnitude. For each matched pair $f_{A}$ and $f_{E}$ in $\{f_{A}^{i}\}_{i=1}^{N-N_{\textrm{nl}}}$ and $\{f_{E}^{i}\}_{i=1}^{N-N_{\textrm{nl}}}$, we will optimize $\mathcal{L}_{MSE}$ by $n_{\textrm{fitting}}$ iterations. The loss function in Stage I can be formulated as

where $m_{r}\sim\mathcal{U}(0,1)$ refers to the magnitude. $\mathcal{L}_{\textrm{encoder}}$ will be optimized to a very small error bound after $n_{\textrm{encoder}}$ iterations. This guarantees that all meta-encoders are well-trained.

We will describe the whole process of how $x$ goes through $f_{DE}$ and finally becomes $\hat{x}$ in this sub-section. We first describe how $f_{DE}$ calls $\{f^{i}_{E}\}_{i=1}^{N}$ to accomplish the combination of various sub-generated samples, and then we present the details of Spatial Transformer Network Jaderberg et al. (2015) and Color Network.

After normalizing $\{\theta_{m}^{i}\}_{i=1}^{N-N_{\textrm{nl}}}$ and $\{\theta_{p}^{i}\}_{i=1}^{N}$ to $\left[0,1\right]$ by applying the sigmoid activation function, we employ the Relaxed Bernoulli Distribution to sample new magnitudes and probabilities for solving the non-differentiable problem Lim et al. (2019); Li et al. (2020). The process can be formulated as

where $\mathop{Logistic}$ and $\tau_{l}$ stand for the logistic distribution and the temperature, respectively, and $\tau_{l}$ is set as 0.05. Then, we randomly select $N_{A}$ non-repeating numbers $\{Z^{i}\}_{i=1}^{N_{A}}$ from $\{i\}_{i=1}^{N}$. $N_{A}$ can be interpreted as the number of randomly sampled primitives, i.e., a larger $N_{A}$ means that the model is more difficult to identify the augmented samples. $N_{A}$ is set as 4 by default in our experiments. We obtain $\{x^{i}_{E}\}_{i=1}^{N_{A}}$ from $x$ through all meta-encoders and get the final $\hat{x}$ by simple sum denoted as

where $\odot$ denotes the element-wise product. For object detection tasks, the summation in Eq. 5 can easily lead to the undesired result that a single target object turns into multiple target objects. Therefore, we use a separate algorithm for this special task. which is explained in Appendix D.

Spatial Transformer Network $f_{\textrm{STN}}$ and Color Network $f_{\textrm{CN}}$ are two different neural network-based meta-encoders. Each has parameters $\theta_{E}$ to store priori bias learned from Stage I. However, $f_{\textrm{STN}}$ and $f_{\textrm{CN}}$ employ different ways to complete the calculation. As shown in the next Eq., when $x$ and magnitude $m$ are fed into $f_{\textrm{STN}}$, $\theta_{E}$ (can be considered as a vector) is first concatenated with $m$ to form a new vector $\hat{\theta}_{E}$, and then passes through a fully connected layer $\mathcal{FC}:=\mathbb{R}^{\mathbf{\textrm{len}}(\theta_{E})+1}\rightarrow\mathbb{R}^{6}$ to obtain the matrix $A\in\mathbb{R}^{2\times 3}$, denoted as

For $f_{\textrm{CN}}$, its forward propagation is composed of convolution and color transformation. So, $\theta_{E}$ not only has a vector $\theta_{E,V}$ but also a convolution weight $\theta_{E,C}$ to record priori bias. When $x$ and $m$ are fed into $f_{\textrm{CN}}$, ${\theta}_{E,V}$ and $m$ concatenate a new vector $\hat{\theta}_{E,V}$, and we let it through a fully connected layer $\mathcal{FC}:=\mathbb{R}^{\mathbf{\textrm{len}}(\theta_{E,V})+1}\rightarrow\mathbb{R}^{2}$ to obtain the scale and shift parameters, i.e., $\theta_{\textrm{scale}}$ and $\theta_{\textrm{shift}}$. After that, the output of $f_{\textrm{CN}}$ is denoted as

where $\mathcal{C}$ refers to the convolutional layer. In particular, in our experiments, we perform RBD on $A$, $\theta_{\textrm{shift}}$ and $\theta_{\textrm{scale}}$ additionally to guarantee the diversity of data augmentation to prevent $\{\theta_{m}^{i}\}_{i=1}^{N-N_{\textrm{nl}}}$ and $\{\theta_{p}^{i}\}_{i=1}^{N}$ from converging to a locally optimal solution.

We conduct ablation studies in three aspects: (a) the effect of Stage II on the student performance; (b) the effect of different iteration numbers in Stage II; (c) the impact of varying data encoders on TST.

As illustrated in Table 5, no matter what kind of teacher-student pairs or whatever value of $N_{A}$, the augmented sample search strategy of TST, i.e., Stage II, is always beneficial for distillation.

In addition, we suppose that TST may easily find augmented samples that the student is not good at, but it is difficult for the student to absorb the corresponding knowledge fully. Based on this conjecture, TST should perform best in the case that $n_{\textrm{encoder}}\ll n_{\textrm{student}}$. The experimental results in Table 6 verify our assumption. Note that $\frac{|\mathcal{X}|}{B}$ actually refers to the iteration number within an epoch.

At last, we show the impact of different types of data encoders on the performance of TST in Table 7 and analyze it. One of the data encoders, Attack, denotes the direct manipulation of samples in a form similar to PGD Madry et al. (2018), which means that TST must let Stage II and Stage III iterate once each in turn and repeat the process. Based on the findings in Table 7, we can indicate that it is an extremely sensible choice to introduce the augmented priori bias into the data encoder and to freeze the parameters associated with the bias during distillation process.