PaperlayerView exact revision on arXiv

Source-linked AI summary

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

Chen Liang, Simiao Zuo, Qingru Zhang, Pengcheng He, Weizhu Chen, Tuo Zhao

arXiv:2210.01351v3cs.CLcs.AIcs.LG

TL;DR

Layer-wise distillation can be difficult when a smaller student must mimic a larger teacher’s representations, including information redundant to the target task. TED uses task-aware filters to align task-relevant knowledge, and the paper reports improvements in continual pre-training and fine-tuning evaluations. Its task-specific setup requires a new student for each new task and may not suit tasks with few or no training samples.

  • Problem

    Layer-wise distillation can under-fit smaller students facing a teacher capacity gap and can require them to mimic hidden information redundant to the target task.

    Source: §1 · p. 1

  • Method

    TED trains task-aware filters to extract target-task knowledge from teacher and student hidden representations, then aligns the filtered representations layer by layer.

    Source: §1 · pp. 1–2, §3.1 · p. 3

  • Results

    TED reports significant improvements over existing methods across continual pre-training and fine-tuning evaluations, including GLUE and SQuAD benchmarks.

    Source: §1 · p. 2

  • Takeaways & Limitations

    Task-aware filters select target-task knowledge and reduce the knowledge gap, making layer-wise distillation easier for the student.

    Source: §1 · p. 2, §8 · p. 9

    Task-specific distillation requires a new student for each new task and may not suit tasks with zero or few training samples.

    Source: §7 · p. 9

Abstract

from arXiv · show

Layer-wise distillation is a powerful tool to compress large models (i.e. teacher models) into small ones (i.e., student models). The student distills knowledge from the teacher by mimicking the hidden representations of the teacher at every intermediate layer. However, layer-wise distillation is difficult. Since the student has a smaller model capacity than the teacher, it is often under-fitted. Furthermore, the hidden representations of the teacher contain redundant information that the student does not necessarily need for the target task's learning. To address these challenges, we propose a novel Task-aware layEr-wise Distillation (TED). TED designs task-aware filters to align the hidden representations of the student and the teacher at each layer. The filters select the knowledge that is useful for the target task from the hidden representations. As such, TED reduces the knowledge gap between the two models and helps the student to fit better on the target task. We evaluate TED in two scenarios: continual pre-training and fine-tuning. TED demonstrates significant and consistent improvements over existing distillation methods in both scenarios. Code is available at https://github.com/cliang1453/task-aware-distillation.

1. Introduction

Layer-wise distillation can leave a smaller student under-fitted and burden it with teacher information that is redundant for the target task. TED uses task-aware filters to focus layer-wise alignment on task-specific knowledge and reports improvements across continual pre-training and fine-tuning evaluations.

  • Layer-wise distillation can under-fit the student when capacity gaps make representation matching dominate task loss, while redundant teacher information can compete with useful knowledge.Source: §1 · p. 1
  • TED introduces task-aware filters at corresponding teacher and student layers to extract task-specific knowledge and reduce the knowledge gap.Source: §1 · pp. 1–2
  • TED trains filters and task-specific heads with model parameters frozen, then jointly trains the student and its filters to align filtered representations.In Stage II, the teacher and its filters remain fixed.Source: Figure 1 · p. 2, §1 · p. 2
  • TED is evaluated in continual pre-training with a 6-layer GPT-2 student and 12-layer teacher, and in fine-tuning with DeBERTaV3-xsmall and DeBERTaV3-base.The reported continual pre-training evaluations include zero-shot and transfer learning settings.Source: §1 · p. 2
  • TED achieves significant improvements on the GLUE benchmark and the SQuAD v1.1/2.0 question-answering datasets.Source: §1 · p. 2

2. Background

The background contrasts prediction-based distillation with layer-wise distillation, which also aligns intermediate representations. It describes their objectives and how student layers can be mapped and projected to teacher layers.

  • Knowledge distillation trains a smaller student to match a teacher’s output predictions, with prediction distance often measured by KL-divergence.Offline distillation keeps the trained teacher fixed while optimizing the student.Source: §2 · pp. 2–3
  • Layer-wise distillation aligns student and teacher hidden representations at mapped intermediate layers to capture knowledge not reflected in their output predictions.The teacher layer mapping M(k) determines which teacher layer corresponds to student layer k.Source: §2 · p. 3
  • A learnable linear projection maps student hidden representations into the teacher’s representation space for layer-wise matching.The layer-wise distance uses mean-squared error.Source: §2 · p. 3
  • The student objective can combine task loss, prediction-distillation loss, and layer-wise loss, with nonnegative hyperparameters weighting the distillation terms.The formulation can be extended with losses on attention scores or embedding layers.Source: §2 · p. 3

3. Method

TED uses task-aware filters to extract target-task knowledge from corresponding teacher and student layers. It trains the filters first, then fine-tunes the student and its filters to match filtered representations.

3. Method

  • TED is a two-stage framework that learns task-aware filters and then minimizes discrepancies between teacher and student filtered representations.Source: §3 · p. 3

3.1. Stage I: Training Task-aware Filters

  • Stage I equips corresponding teacher and student layers with filters that turn hidden representations into target-task predictions, trained using task loss while model parameters stay fixed.Filter architectures may be linear projections, multilayer perceptrons, or Transformer stacks; the filters aim to remove redundant information while retaining task-useful information.Source: §3.1 · pp. 3–4

3.2. Stage II: task-aware Layer-wise Distillation

  • Stage II freezes the teacher and its filters, removes filter heads, and fine-tunes the student and its filters to minimize discrepancies between filtered representations.This relaxes conventional layer-wise matching, which requires matching unfiltered representations regardless of their relevance to the target task.Source: §3.2 · p. 4

4. Language Modeling

The language-modeling experiments evaluate TED in continual pre-training and downstream transfer, using a smaller GPT-2 student. TED improves performance, especially on Open WebText, while layer-wise distillation does not always outperform KD.

4.1. Data

  • Open WebText3 supplies continual pre-training data, with zero-shot and transfer evaluation on LAMBADA and WikiText-103.Source: §4.1 · p. 4

4.2. Models

  • The setup uses a 12-layer, 125M-parameter GPT-2 teacher and a 6-layer, 82M-parameter student initialized from selected teacher layers.Source: §4.2 · p. 4

4.3. Training

  • Stage I trains teacher filters with model weights frozen and transfers mapped filters; Stage II trains the student and filters, with KD and LWD as baselines.Source: §4.3 · pp. 4–5

4.4. Main Results

  • TED improves GPT-2 student performance, especially on Open WebText, while LWD does not always outperform KD.Table 1 covers zero-shot and transfer performance.Source: Table 1 · p. 5, §4.4 · p. 5

5. Natural Language Understanding

On GLUE and SQuAD, TED improves over the best distillation baseline across most evaluated tasks. A separate BERT experiment reports comparable performance with benefits on three NLU tasks.

5.1. Data

  • The evaluation covers nine GLUE NLU tasks and the SQuAD v1.1 and v2.0 question-answering datasets.Source: §5.1 · p. 5

5.2. Models

  • The teacher is a fine-tuned DeBERTaV3-base model (183M, 12 layers), and the student is fine-tuned DeBERTaV3-xsmall (70M, 12 layers).Source: §5.2 · p. 5

5.3. Training

  • Stage I trains filters with teacher and student weights frozen; Stage II trains the student and its filters using Eq. 7.Source: §5.3 · p. 5

5.4. Main Results

5.5. BERT Experiments

  • A separate comparison uses a 66M six-layer BERT-base student distilled from a fine-tuned 109M 12-layer teacher, with comparable performance and benefits on three NLU tasks.Source: §5.5 · p. 6

6. Analysis

The analysis tests whether TED’s filters capture task-specific knowledge, ease distillation, and help across capacity gaps and design choices. Results support task-aware filtering, while showing that greater filter complexity offers little benefit.

6.1. Filters Capture Task-Specific Knowledge

  • Target-task-trained filters give consistent gains over LWD, whereas different-task filters yield smaller, more variable gains.The authors suggest filters trained on more similar tasks may perform better; on RTE, MNLI- and QNLI-trained filters outperform SST-2-trained filters.Source: §6.1 · p. 7, Table 5 · p. 7, §6.2 · p. 7

6.2. TED Alleviates Under-fitting and Eases Distillation

  • TED converges faster with lower target-task training loss and has lower-magnitude, lower-variance distillation loss than LWD.The reported loss plots cover GPT-26 on Open WebText and DeBERTaV3-xsmall on MNLI.Source: §6.2 · p. 7, Figure 2 · p. 7

6.3. Contribution of the Filters

  • Trained teacher filters improve distillation when student filters can learn to match their outputs; trained student filters further improve performance over random initialization.The study also evaluates removing filters and compares initialization on GLUE and Open WebText.Source: §6.3 · pp. 7–8, Table 6 · p. 7

6.4. TED Alleviates the Capacity Gap Issue

  • TED maintains comparable performance with a DeBERTaV3-large teacher, while LWD drops to 79.4 on RTE and TED gains 0.4 points on SST-2.The comparison uses a 435M DeBERTaV3-large teacher and a 70M DeBERTaV3-xsmall student.Source: §6.4 · p. 8, Table 7 · p. 8

6.5. Filter Architectures

  • Adding filter nonlinearity slightly improves zero-shot performance, transfer performance is insensitive to architecture, and greater filter complexity offers little benefit.The compared designs include linear projections, two-layer GeLU perceptrons, and filters initialized from subsequent model layers.Source: §6.5 · p. 8, Table 8 · p. 8

6.6. Design of Distillation Loss

  • Using the Stage I task-specific heads yields prediction distributions for KL-divergence distillation, which improves over baselines and can preserve task-specific knowledge.This alternative replaces the hidden-representation MSE with KL-divergence between prediction distributions.Source: §6.6 · p. 8, Table 9 · p. 8

6.7. Hyper-parameter Study

  • TED shows consistent gains across a wide range of α2 values on MNLI-m and SQuAD v1.0.Figure 3 reports DeBERTaV3-xsmall performance under different α2 settings.Source: §6.7 · p. 9, Figure 3 · p. 9

7. Discussion

TED adds a moderate training cost but leaves inference speed unchanged. The discussion identifies task-specific distillation’s scalability and data constraints, and explains why direct layer-wise distillation is challenging for large language models.

7. Discussion

  • TED’s Stage I filter training adds about 10% computational cost over LWD, while deployment inference speed remains the same because filters are discarded.Filter parameters are around 2%–4% of model parameters, and filter training requires no back-propagation on model parameters.Source: §7 · p. 9
  • Task-specific distillation requires a new student for each task and is unsuitable for tasks with zero or few training samples.The authors propose exploring multi-task distillation as a possible direction, using task-specialized filters and routing inputs to relevant filters.Source: §7 · p. 9
  • Direct TED application to LLMs is challenging because closed models lack accessible layer-wise representations and LWD cost scales with depth and hidden dimension.The authors suggest generating controllable task-relevant input-output samples from an LLM teacher as a possible alternative for student distillation.Source: §7 · p. 9

8. Conclusion

TED addresses the difficulty of layer-wise distillation by learning task-aware filters that extract task-specific knowledge before matching filtered representations. The experiments indicate that these filters capture task-specific knowledge and ease distillation.

8. Conclusion

  • TED learns task-aware filters to extract task-specific knowledge from teacher and student representations before minimizing the discrepancy between their filtered outputs.The approach targets the difficulty of having a student mimic a much larger teacher.Source: §8 · p. 9
  • The filtered representations contain more task-relevant signals for the student to learn.TED first learns the filters, then minimizes the discrepancy between their outputs.Source: §8 · p. 9
  • Experiments verify that the filters capture task-specific knowledge and ease layer-wise distillation.Source: §8 · p. 9

A.1.1. DATA

The appendix describes the datasets, model configurations, training procedures, and initialization recommendations used across the paper’s experiments.

A.1.1. DATA

  • Open WebText is used for language-modeling data, with near-duplicate filtering, byte-level BPE tokenization, and a maximum input sequence length of 1024.The data subsection describes Open WebText as an open-source effort to reproduce OpenAI’s WebText dataset.Source: §A.1.1 · p. 13

A.1.2. TRAINING

  • GPT-2 base has 12 layers, 12 attention heads, and 768-dimensional hidden states; training uses mixed precision on eight 80G Nvidia A100 GPUs.Hyperparameters are listed in Table 10.Source: §A.1.2 · p. 13, Table 10 · p. 13

A.2.1. DATA

  • The NLU evaluation uses GLUE’s nine tasks and SQuAD v1.1/2.0, whose v2.0 set includes unanswerable questions.SQuAD examples come from approximately 500 Wikipedia articles.Source: §A.2.1 · p. 13

A.2.2. MODEL

  • For downstream tasks, the teacher is fine-tuned DeBERTaV3-base and the student is pre-trained DeBERTaV3-xsmall.The teacher receives a target-task classification head on its last layer.Source: §A.2.2 · p. 13

A.2.3. TRAINING

  • Fine-tuning the student before Stage I empirically improved distillation, while the reported comparison found initialization with fine-tuned versus pre-trained weights did not largely influence final performance.The authors hypothesize that target-task initialization helps student filters capture more task-relevant knowledge; Table 11 compares the two initializations.Source: §A.2.3 · p. 14, Table 11 · p. 14

A.2.4. BERT EXPERIMENTS

  • In the BERT experiment, a fine-tuned BERT-base12 teacher supplies selected layers to initialize a six-layer BERT-base6 student, and its trained filters are copied to the student.The teacher is fixed while its filters are trained in Stage I; Stage II distills the student and its filters.Source: §A.2.4 · p. 14

A.3.1. EXPERIMENTS IN SECTION 6.2

  • The setup pairs a fine-tuned 24-layer DeBERTaV3-large teacher (435M parameters) with a 12-layer DeBERTav3-xsmall student (70M parameters); in Stage I, model parameters are fixed while their filters are trained.In Stage II, the student model and its filters are distilled.Source: §A.3.1 · pp. 14–15

A.4. Discussion on the Model Initialization

  • Initialization guidance varies by setting: reuse a matching student when available; otherwise choose scratch pretraining, teacher-layer transfer, or direct Stage II according to budget and architecture.For pre-training and fine-tuning, the appendix gives separate cases; it cautions that poor model initialization can leave task-aware filters without sufficient task-relevant knowledge.Source: §A.4 · pp. 15–16
Loading 2210.01351v3…