Learning to Teach with Student Feedback

In our setting, the teacher model $f(\theta)$ is implemented as a deep encoder such as BERT (Devlin et al. 2019), and the student model $g(\phi)$ is implemented as a lightweight encoder such as a Transformer (Vaswani et al. 2017) encoder with fewer layers.

Given a training sample $(x_{i},y_{i})$ drawn from the course data set $D_{course}=\{x_{i},y_{i}\}_{i=1}^{N}$, the teacher model computes a contextualized embedding $h_{i}^{T}=f(x_{i};\theta)$¹¹1We use the [CLS] token embedding of the last layer as the sentence representation., which is then fed into a Softmax layer to obtain the probability for each category, i.e. the soft targets, that is $y_{i}^{T}=\mathrm{Softmax}(W^{T}h_{i}^{T})$, where $W^{T}$ is a learnable parameter matrix and $W^{T}h_{i}^{T}$ is the logits. For simplicity, $W^{T}h_{i}^{T}$ is abbreviated to $z_{i}^{T}$. In practice, the temperature is often introduced as a hyper-parameter to control the smoothness of the soft targets,

$$y_{i}^{T}=\mathrm{Softmax}(z_{i}^{T}/T)=\frac{e^{z_{i,c}^{T}/T}}{\sum_{c^{\prime}=1}^{C}{e^{z_{i,c^{\prime}}/T}}},$$

(1)

where $C$ denotes the number of classes, $c$ is the correct class.

Similarly, we can obtain the output of the student model $y_{i}^{S}=\mathrm{Softmax}(W^{S}h_{i}^{S})$, where $h_{i}^{S}=g(x_{i};\phi)$. In vanilla knowledge distillation, the KL divergence of the teacher’s prediction and the student’s prediction should be minimized,

$$\mathcal{L}_{KD}=\mathrm{KL}(y^{T},y^{S})=-\frac{1}{N}\sum_{i=1}^{N}\sum_{j=1}^{C}y_{i,j}^{T}\log\frac{y_{i,j}^{S}}{y_{i,j}^{T}},$$

(2)

where $\mathrm{KL}$ means Kullback-Leibler divergence.

In practice, the cross entropy between the student’s prediction and the ground truth is also incorporated to be minimized, i.e.,

$$\mathcal{L}_{CE}^{S}=-\frac{1}{N}\sum_{i=1}^{N}\sum_{j=1}^{C}y_{i,j}\log y_{i,j}^{S}.$$

(3)

Thus the overall loss function for the student model can be written as:

$$\mathcal{L}_{stu}=\lambda\mathcal{L}_{KD}+(1-\lambda)\mathcal{L}_{CE}^{S},$$

(4)

where $\lambda$ is a hyper-parameter to balance the knowledge distillation loss and the cross entropy loss.

The student model is then updated by taking one gradient descent step with the loss function above, i.e.

$$\phi_{t+1}=\phi_{t}-\alpha\nabla_{\phi_{t}}\mathcal{L}_{stu},$$

(5)

where $\alpha$ is the learning rate of the student model.

It is expected that the student with the post-update parameters will generalize well on unseen samples. Therefore, the updated student is then evaluated on another batch of samples drawn from the exam data set $D_{exam}=\{x_{i}^{\prime},y_{i}^{\prime}\}_{i=1}^{M}$. The performance of the student is measured by the cross entropy loss on the batch of exam data, which is called meta loss since it is calculated on the post-update parameter $\phi_{t+1}$ thus provides meta gradient w.r.t. the teacher’s parameters.

	$\displaystyle\mathcal{L}_{meta}$	$\displaystyle=\mathrm{CE}(y^{\prime},y_{t+1}^{\prime S})$		(6)
		$\displaystyle=-\frac{1}{M}\sum_{i=1}^{M}\sum_{j=1}^{C}y_{i,j}^{\prime}\log y_{i,j}^{\prime S}(\phi_{t+1}).$		(7)

Our goal is to minimize $\mathcal{L}_{meta}$ by optimizing the teacher model, therefore the gradient of $\mathcal{L}_{meta}$ is w.r.t. the teacher’s parameters $\theta$, so is a second-order gradient (gradient of gradient). To explicitly show this, we can re-write $y_{i,j}^{\prime S}(\phi_{t+1})$ in the above equation as:

$$y_{i,j}^{\prime S}(\phi_{t+1})=y_{i,j}^{\prime S}(\phi_{t}-\alpha\nabla_{\phi_{t}}\mathcal{L}_{stu}(\theta_{t})),$$

(8)

in which we know that the gradient can flow into $\theta_{t}$ through the knowledge distillation in the course step. Figure 3 shows the computational graph of the forward pass and the backward pass.

Although it is computationally available by conducting an additional backward pass under current deep learning libraries that support automatic differentiation, we use the first-order approximation for its efficiency (Finn, Abbeel, and Levine 2017; Nichol, Achiam, and Schulman 2018).

In the context of meta-learning, the teacher can be regarded as the meta-learner, while $D_{course}$ and $D_{exam}$ can be analogous to support set and query set respectively.

Also, the standard cross entropy loss between the teacher’s prediction and the ground truth can be incorporated,

$$\mathcal{L}_{CE}^{T}=-\frac{1}{N}\sum_{i=1}^{N}\sum_{j=1}^{C}y_{i,j}\log y_{i,j}^{T}.$$

(9)

Thus, the overall loss function for the teacher model is

$$\mathcal{L}_{tea}=\gamma\mathcal{L}_{meta}+(1-\gamma)\mathcal{L}_{CE}^{T},$$

(10)

where $\gamma$ is another hyper-parameter to balance the meta loss and the cross entropy loss.

Thus, the teacher model can also be optimized by taking gradient descent:

$$\theta_{t+1}=\theta_{t}-\beta\nabla_{\theta_{t}}\mathcal{L}_{tea},$$

(11)

where $\beta$ is the learning rate of the teacher model.

Putting the student optimization in course step and the teacher optimization in exam step together, the overall loss function of IKD is defined as

$$\mathcal{L}_{IKD}=\mathcal{L}_{stu}+\mathcal{L}_{tea}.$$

(12)

By iterating the course step and the exam step, the student model learns from the soft targets generated by the teacher while the teacher model learns from the feedback generated by the student.

To take a closer look, in this section we further analyze the expanded form of the meta-gradient, which provides insights into the student’s feedback and sheds light on how to take an approximation to achieve efficient training.

According to the chain rule, we can unfold the meta gradient $\nabla_{\theta_{t}}\mathcal{L}_{meta}$ as follows²²2For the simplicity of derivation, we set $N$ and $M$ to 1, and omit the constant scalar $\lambda$.,

	$\displaystyle\nabla_{\theta_{t}}\mathcal{L}_{meta}$	$\displaystyle=\nabla_{\phi_{t+1}}\mathcal{L}_{meta}\cdot\nabla_{\theta}{\phi_{t+1}}$		(13)
		$\displaystyle=\nabla_{\phi_{t+1}}\mathcal{L}_{meta}\cdot(-\alpha\nabla_{\theta}{\nabla_{\phi_{t}}{\mathcal{L}_{stu}}})$
		$\displaystyle=\alpha\nabla_{\phi_{t+1}}\mathcal{L}_{meta}\cdot\nabla_{\theta}{\nabla_{\phi_{t}}\log y_{t}^{S}\cdot y_{t}^{T}}$
		$\displaystyle=\underbrace{\alpha\nabla_{\phi_{t+1}}\mathcal{L}_{meta}\cdot\nabla_{\phi_{t}}\log y_{t}^{S}}_{\text{feedback}}\cdot\nabla_{\theta_{t}}y_{t}^{T}.$

Note that $\nabla_{\phi_{t+1}}\mathcal{L}_{meta}$ is the Jacobi matrix of shape $|\phi|$ and $\nabla_{\phi_{t}}\log y_{t}^{S}$ is of shape $|\phi|\times C$.³³3For matrix multiplication, $\nabla_{\phi_{t}}\log y_{t}^{S}$ should be in the transpose form, which is omitted here for the simplicity. Therefore, $\nabla_{\phi_{t+1}}\mathcal{L}_{meta}\cdot\nabla_{\phi_{t}}\log y_{t}^{S}$ is actually a weighting vector of shape $C$ to adjust the gradient of teacher’s prediction w.r.t. its parameters $\nabla_{\theta_{t}}y_{t}^{T}$. It is easy to find that the weighting vector only depends on the student’s gradient, therefore it is exactly the student’s feedback.

The behavior relationship between the weighting vector and the teacher’s prediction is analyzed. It is worth noticing that the weighting vector should be detached from the computational graph to forbid the gradient flow into the student. By this, the meta loss will only contribute to the update of the teacher’s parameters.

In practice, the difference between $\phi_{t}$ and $\phi_{t+1}$ is usually very small, which motivates us to take an approximation ${\phi_{t+1}}\approx{\phi_{t}}$ such that ${\nabla_{\phi_{t+1}}\mathcal{L}_{meta}}\approx{\nabla_{\phi_{t}}\mathcal{L}_{meta}}$. By this approximation, Eq. (13) can be simplified as

$$\nabla_{\theta_{t}}\mathcal{L}_{meta}\approx{\alpha\nabla_{\phi_{t}}\mathcal{L}_{meta}\cdot\nabla_{\phi_{t}}\log y_{t}^{S}\cdot\nabla_{\theta_{t}}y_{t}^{T}}.$$

(14)

Besides, note that $\nabla_{\phi_{t}}\mathcal{L}_{meta}$, $\nabla_{\phi_{t}}\log y_{t}^{S}$ and $\nabla_{\theta_{t}}y_{t}^{T}$ are mutually independent, therefore we can calculate them independently and achieve joint optimizing in an efficient way. In our experiments we set $D_{course}=D_{exam}=D_{train}$ for efficient optimization. Though, it should be mentioned that in our framework the course data set can be unlabeled, i.e. $D_{course}=\{x_{i}\}_{i=1}^{N}$. In this case, the cross entropy loss between the student’s prediction and the ground truth is removed such that the loss for the student is exactly the knowledge distillation loss. Thus, only the supervision signal of $D_{exam}$ is required. We leave IKD in this semi-supervised setting as future work.

We conduct experiments on the General Language Understanding Evaluation (GLUE) benchmark (Wang et al. 2019) with the backbone of BERT (Devlin et al. 2019), and on ChemProt dataset (Schneider et al. 2020) and SciCite dataset (Cohan et al. 2019) with backbone of SciBERT (Beltagy, Lo, and Cohan 2019).

Our implementation is based on PyTorch (Paszke et al. 2019). The training and evaluation are performed on a single RTX 2080Ti or RTX 3090 GPU. For downstream tasks in GLUE we first fine-tune the teacher model for 3 epochs with learning rate of 2e-5. For domain tasks the teacher is fine-tuned for 4 epochs. The batch size is 32 for all tasks. The student model is initialized with the parameters from the first $k$ layers of BERT base(12 layer), where $k$ is the number of the student’s layer. In our experiments we mainly evaluate our method when $k=4$ and $k=6$. The total number of parameters of BERT₂₄, BERT₁₂, BERT₆ and Bert₄ are 340M, 110M, 67M and 52M respectively. The batch size is set to 4 for all downstream tasks. We conduct 3 experiments for each dataset and the median is reported.

We report our results on GLUE benchmark. The results are presented in Table 1. We denote BERT_k as a BERT with only $k$ layers. BERT₄-FT is to directly fine-tune the first 4 layers of the pre-trained BERT. BERT₄-KD represents student trained with vanilla knowledge distillation loss, i.e., Eq. (4). We also exhibit the reported result of TAKD (Mirzadeh et al. 2019) and Annealing KD (Jafari et al. 2021) for a direct comparision.

Since IKD allows the teacher to dynamically adjust its soft targets by using the feedback from the student, it would be interesting to see how the output of the teacher changes as the training goes on. In Figure 4 we visualize the output soft targets of the teacher output in different training epochs. We can find that the soft targets are relatively smooth at the beginning and become sharper as the training progresses. This is consistent with the hand-designed smoothing schedules (Dogan et al. 2020; Jafari et al. 2021), which introduce a hyper-parameter to control how quickly the soft targets converge to the one-hot distribution. In contrast, our method automatically learns to adjust the smoothness of the soft targets based on the feedback from the student. In Table 4 we further use entropy to measure the smoothness of output targets in different epochs on full training sets of MRPC, RTE and ChemProt. The average entropy goes down over the training process as shown, which confirms our observation as well.

As mentioned in Implementation Details, the gradient of the meta loss can be viewed as a weighted sum of the gradient of teacher’s prediction w.r.t. its parameters, i.e. $\nabla_{\theta}y^{T}$ (See Eq. (14)). The weighting vector $\alpha\nabla_{\phi}\mathcal{L}_{meta}\cdot\nabla_{\phi}\log y^{S}$ can be regarded as the feedback of the student. Assume $fb=\alpha\nabla_{\phi}\mathcal{L}_{meta}\cdot\nabla_{\phi}\log y^{S}$, Eq. (14) can be written as

$$\nabla_{\theta}\mathcal{L}_{meta}\approx fb\cdot\nabla_{\theta}y^{T}=\nabla_{\theta}\sum_{i=1}^{C}fb_{i}\cdot y_{i}^{T},$$

(15)

where $fb\in\mathbb{R}^{C}$, $y^{T}\in\mathbb{R}^{C}$, $C$ is the number of classes.

In this section, we attempt to analyze the relationship between the student’s feedback $fb$ and the teacher’s prediction $y^{T}$. In particular, we plot $(fb_{i},y^{T}_{i})$ for each $i\in\{1,\cdots,C\}$ in Figure 5. The data points are collected from the 7,500 training steps on CoLA dataset. CoLA is a binary classification task so we can plot $C=2$ data points for each training step such that there are 15,000 data points in total.

Note that The greater the student’s feedback $fb_{i}$ for class $i$, the greater the penalty of teacher’s prediction on class $i$. As shown in Figure 5(a), the student’s feedback is always non-positive for teacher’s prediction on target label, which can also be revealed by the expanded form of the feedback:

$$\alpha\nabla_{\phi}\mathcal{L}_{meta}\nabla_{\phi}\log y^{S}_{c}=-\alpha(\nabla_{\phi}\log y_{c}^{S})^{2},$$

(16)

where $c$ is the target label. Besides, we can see that the student’s feedback is usually non-negative for teacher’s prediction on non-target label. Thus, the teacher’s prediction is always encouraged to approach the one-hot distribution where only the target label is 1, which has been verified in analysis. Also, we find that the range of feedback on target prediction is approximately three times larger than the range of feedback on non-target prediction.

In addition, the student’s feedback can be different even when the teacher’s predictions are the same. To take a deeper look, we demonstrate the value of feedback of the student at different training stage in Figure 5(b). We can see that the student feedback converges to zero as the training goes, which implies that IKD works mainly at the immediate training stage.

Knowledge distillation (KD) was originally proposed in domains other than NLP (Bucila, Caruana, and Niculescu-Mizil 2006; Hinton, Vinyals, and Dean 2015). With the emergence of large-scale pre-trained language models (PLMs) such as GPT (Radford et al. 2019), BERT (Devlin et al. 2019), and RoBERTa (Liu et al. 2019), KD has gain much attention in the NLP community due to its simplicity and efficiency in compressing the large PLMs. Much effort has been devoted to compress BERT via KD, including BERT-PKD (Sun et al. 2019), DistilBERT (Sanh et al. 2019), MobileBERT (Sun et al. 2020), and TinyBERT (Jiao et al. 2020). These methods have achieved great success by exploring informative features to be distilled, such as word embeddings, attention maps, hidden states, and output logits. Different from these work, our proposed IKD focuses on the interaction between teacher and student therefore is orthogonal to these work. Gou et al. (2021) provides a survey of online knowledge distillation where both the teacher and the student are updated simulatenously. Our IKD can be regarded as a special case of online KD.

Optimization-based meta-learning intends to learn a group of initial parameters that can fast converge on a new task with only a few gradient steps, in which the learner and meta-learner share the same architecture. As a representative method, Model-Agnostic Meta-Learning (MAML) (Finn, Abbeel, and Levine 2017) uses a nested loop to find an optimal set of parameters on a variety of learning tasks, such that they can adapt well after a small number of training steps. To tackle the high-order derivative, Finn, Abbeel, and Levine (2017) also propose First-Order MAML (FOMAML) to perform the first-order approximation to reduce the computation. Similar to MAML, Reptile (Nichol, Achiam, and Schulman 2018) utilizes the weights in the inner loop instead of gradient to update the meta parameters and achieves competitive performance on several few-shot classification benchmarks.

Meta-learning has been proved efficient when combined with the teacher-student framework. Liu et al. (2020) employ meta-learning to optimize a label generator that generates dynamic soft targets for intermediate layers. Their method is proposed for self-boosting. Therefore, the teacher and the student are required to share the same architecture and parameters. Pan et al. (2020) focus on training a meta-teacher possessing multi-domain knowledge by metric-based meta-learning and then teach single-domain students. Meta Pseudo Labels (Pham et al. 2020) leverage a teacher network to generate pseudo labels on unsupervised images to enlarge the training set of the student. They aim to train a pseudo labels generator to while IKD dedicates to compress a large model into a smaller one. Thus the pseudo label generator is trained from scratch and shares same architecture with the student while in our IKD the teacher must be a larger pretrained model which leads to the impossibility of implementation of gradient substitution in MAML. Moreover, Meta Pseudo Labels generates hard labels and uses Monte Carlo approximation to obtain gradient while soft targets from the teacher in KD are necessary for passing dark knowledge to student and gradients in IKD are obtained by adopting the chain rule manually.