Less is More: Task-aware Layer-wise Distillation for Language Model Compression

Large pre-trained language models have achieved state-of-the-art performances in many natural language processing tasks (wang2018glue; rajpurkar2016squad). However, their deployment in resource-limited scenarios is hindered by their huge number of parameters (raffel2019exploring; radford2019language; brown2020language; he2020deberta; he2023debertav). Knowledge Distillation (KD) (hinton2015distilling) is a powerful tool to compress large models (i.e., teacher models) into small ones (i.e., student models) with a minimal loss of performance. This approach trains the student to match the output predictions of the teacher.

However, such a last-layer-only distillation approach does not exploit the intermediate layers of the teacher, which contain rich semantic and syntactic knowledge. To leverage such knowledge, researchers have proposed a layer-wise distillation approach, which trains the student to match the hidden representation of the teacher at each layer (sun2019patient; jiao2019tinybert; sun2020mobilebert; hou2020dynabert; zuo2022moebert). Such an approach often improves the generalization performance of the student model.

Nevertheless, layer-wise distillation faces two major challenges. First, the student may struggle to mimic the hidden representations of the teacher due to their large capacity gap. This often leads to large discrepancies between their hidden representations. Consequently, model training/optimization often favors reducing such large discrepancies over the training loss of the student (i.e., the target task’s loss such as cross-entropy), resulting in an under-fitted student model. Second, mimicking the hidden representations may not be beneficial for the target task’s learning. This is because the hidden representations of the teacher often contain redundant information (dalvi2020analyzing; durrani2020analyzing). Given the limited capacity of the student, such redundant information may compete with the useful information for distillation, hindering the useful knowledge from being distilled. Our empirical observations show that for some tasks, layer-wise distillation only marginally outperforms standard KD (Table LABEL:tb:deberta_glue).

To address these challenges, we propose a novel layer-wise distillation method, TED (Task-aware layEr-wise Distillation), which distills task-specific knowledge from the teacher to the student. We design a pair of task-aware filters for each layer of the teacher and student¹¹1For simplicity, we assume that the student and teacher are of the same depth (number of layers) but different widths. The case of different depths will be elaborated in Section 2.. Each filter is a neural network with a task-specific head (e.g., a linear soft-max layer for classification), and is trained to extract the predictive knowledge from the hidden representation of the corresponding model. Figure 1 illustrates the training procedure of TED, which consists of two stages:

The task-aware filters serve as a selection mechanism that reduces the knowledge gap between the teacher and the student and encourages the distillation of task-specific knowledge. This makes distillation easier for the student.

We evaluate TED on two settings: continual pre-training and task-specific fine-tuning. In the continual pre-training setting, we distill a $6$-layer GPT-2 student model ($82$M) from a $12$-layer GPT-2 teacher model ($125$M) (radford2019language). We show that TED outperforms existing methods in both zero-shot and transfer learning settings on various downstream tasks (paperno2016lambada; merity2016pointer). In the task-specific fine-tuning setting, we distill a DeBERTaV3-xsmall student model ($70$M) from a DeBERTaV3-base teacher model ($183$M) (he2023debertav). We demonstrate that TED achieves significant improvement on the GLUE benchmark (wang2018glue) and the SQuAD v1.1/2.0 question answering datasets (rajpurkar2016squad; rajpurkar2018know).

The rest of the paper is organized as follows: Section 2 briefly reviews the background; Section 3 presents our proposed method; Section LABEL:sec:lm presents experiments on language modeling; Section LABEL:sec:nlu presents experiments on natural language understanding; Section LABEL:sec:analysis presents analysis of models; and Section LABEL:sec:conclusion discusses and concludes the paper.

Transformer-based Language Models. The Transformer architecture is a powerful neural network design for modeling sequential data, such as natural language (vaswani2017attention; devlin2018bert; radford2019language; he2023debertav). It consists of multiple layers that are stacked on top of each other. Each layer performs two operations: a multi-head self-attention mechanism and a two-layer feed-forward neural network. We use $f(\cdot;\Theta)$ to denote a Transformer-based model $f$ that has a set of parameters $\Theta$, where $f$ takes an input sequence $x$ from the input sample space $\mathcal{X}$ and produces an output prediction. We define the loss function $\mathcal{L}(\Theta)=\mathbb{E}_{x\sim\mathcal{X}}[\ell(f(x;\Theta))]$, where $\ell$ is the target task loss. For example, $\ell$ is the causal language modeling loss for generative models (i.e., $\sum_{t=1}^{|x|}\log p(x_{t}|x_{<t};\Theta)$).

Knowledge Distillation is a powerful approach to compress large models (i.e., teacher models) into smaller models (i.e., student models) by transferring knowledge from the former to the latter (hinton2015distilling). The student is trained to mimic the output predictions of the teacher. Specifically, we denote the teacher as $f_{t}(\Theta_{t})$ and the student as $f_{s}(\Theta_{s})$ and consider the following optimization problem:

where $\mathcal{D}_{\rm pred}(\Theta_{t},\Theta_{s})$ is the distillation loss, a distance metric between the output predictions of the teacher and the student. For example, $\mathcal{D}_{\rm pred}$ can be the KL-divergence: $\texttt{KL}(f_{t}(\Theta_{t})/T,f_{s}(\Theta_{s})/T)$, where $T>0$ is the temperature that controls the softness of the prediction probability distributions (hinton2015distilling). A commonly adopted distillation scheme is the offline distillation, where the teacher is fully-trained and fixed, and the student is optimized based on Eq. 1.

Layer-wise Distillation. In large Transformer-based models, the output predictions of the models may not capture all the semantic and syntactic knowledge encoded in the intermediate layers. Therefore, researchers propose a layer-wise distillation approach, which aligns the hidden representations of the student and the teacher at each layer (romero2014fitnets; sun2019patient; sun2020mobilebert; jiao2019tinybert; hou2020dynabert; zuo2022moebert; liang2023homodistil). Specifically, we denote the hidden representation at the $k$-th layer of a $K$-layer student as $H_{s}^{k}\in\mathbb{R}^{|x|\times d_{s}}$, and at the $M(k)$-th layer of the teacher as $H_{t}^{M(k)}\in\mathbb{R}^{|x|\times d_{t}}$. Here $|x|$ is the sequence length; $d_{s}$ and $d_{t}$ are the hidden dimensions of the student and the teacher, respectively. $M(\cdot)$ is a layer mapping function that determines from which layer in the teacher that a student layer should distill. For example, if we set $M(k)=2k$, the student would distill from every other layer in the teacher. The layer-wise distillation loss is defined as:

Here $\texttt{MSE}(\cdot,\cdot)$ is the mean-squared error, $W_{s}^{k}\in\mathbb{R}^{d_{s}\times d_{t}}$ is a randomly initialized and learnable linear projection that projects $H_{s}^{k}$ into the same space as $H_{t}^{M(k)}$, and $\mathcal{W}_{s}=\{W_{s}^{k}\}_{k=1}^{K}$. In practice, the student is often optimized using multiple distillation losses, e.g.,

where $\alpha_{1},\alpha_{2}\geq 0$ are hyper-parameters. Besides the intermediate layers, distilling knowledge from the attention scores and the embedding layers can also improve the distillation performance (sun2020mobilebert; jiao2019tinybert; wang2020minilm; wang2020minilmv2). Eq. 3 can be further extended by adding such losses.

We introduce TED, a two-stage training framework that uses task-aware filters to distill knowledge from a teacher to a student. The task-aware filters are neural networks that learn to extract task-specific knowledge from the hidden representations of the teacher and the student. In the first stage, we add a task-aware filter to each layer of the teacher and the student. We train these filters using the task-specific loss while keeping the model parameters frozen. In the second stage, we fine-tune the student and its filters by minimizing the discrepancy between the filtered representations of the teacher and the student.