PaperlayerView exact revision on arXiv

Source-linked AI summary

ALP-KD: Attention-Based Layer Projection for Knowledge Distillation

Peyman Passban, Yimeng Wu, Mehdi Rezagholizadeh, Qun Liu

arXiv:2012.14022v1cs.CL

TL;DR

Intermediate-layer distillation is limited when student-teacher depth differences force existing approaches to skip teacher layers and choose alignments. The paper combines teacher layers through attention, and reports that this combination is effective, with attention-based combination needed for competitive results.

  • Problem

    Existing intermediate-layer distillation uses selected teacher layers, risking skipped information and requiring a choice of which layers to align.

    Source: Introduction · p. 1, Problem Definition · p. 2

  • Method

    The proposed attention-based approach combines teacher layers for each student layer and uses all teacher-layer information without searching for a subset.

    Source: Problem Definition · p. 2, Methodology · p. 3

  • Results

    The authors report that combination is effective and that attention-based combination is required for competitive results, although simple concatenation often beats skipping.

    Source: Conclusion and Future Work · p. 7

  • Takeaways & Limitations

    The authors conclude that combination becomes more crucial when teacher and student depths differ more, particularly for distillation from deep teachers.

    Source: Conclusion and Future Work · p. 7

    Evaluation uses validation sets because GLUE test-set labels are not publicly available.

    Source: Experimental Study · p. 4

Abstract

from arXiv · show

Knowledge distillation is considered as a training and compression strategy in which two neural networks, namely a teacher and a student, are coupled together during training. The teacher network is supposed to be a trustworthy predictor and the student tries to mimic its predictions. Usually, a student with a lighter architecture is selected so we can achieve compression and yet deliver high-quality results. In such a setting, distillation only happens for final predictions whereas the student could also benefit from teacher's supervision for internal components. Motivated by this, we studied the problem of distillation for intermediate layers. Since there might not be a one-to-one alignment between student and teacher layers, existing techniques skip some teacher layers and only distill from a subset of them. This shortcoming directly impacts quality, so we instead propose a combinatorial technique which relies on attention. Our model fuses teacher-side information and takes each layer's significance into consideration, then performs distillation between combined teacher layers and those of the student. Using our technique, we distilled a 12-layer BERT (Devlin et al. 2019) into 6-, 4-, and 2-layer counterparts and evaluated them on GLUE tasks (Wang et al. 2018). Experimental results show that our combinatorial approach is able to outperform other existing techniques.

Introduction

Knowledge distillation transfers a teacher’s predictive supervision to a student, and intermediate-layer supervision can extend this beyond final predictions. Existing layer matching may discard teacher information and requires choosing which layers to align, motivating attention-based combination of teacher layers.

Introduction

  • Knowledge distillation couples a student with a teacher during training; the student learns from hard labels and the teacher’s soft predictions.The teacher is expected to provide reliable predictions, while the student typically has a simpler architecture.Source: Introduction · p. 1
  • Intermediate-layer distillation extends teacher supervision beyond final predictions, but existing approaches compare only selected teacher layers with student layers.This can leave information unused when the teacher is deeper than the student.Source: Introduction · p. 1, Problem Definition · p. 2

Problem Definition

  • When the teacher has many more layers, matching a student layer to one teacher layer requires skipping others, creating the skip problem.The skipped layers contain information learned during training.Source: Problem Definition · p. 2
  • Choosing which teacher layers to retain is arbitrary without a selection strategy, and finding the best subset requires exhaustive search or expert guidance.The paper calls this the search problem.Source: Problem Definition · p. 2
  • The proposed alternative combines teacher layers through attention, uses all teacher-layer information, and removes the need to search for a subset.The authors suggest the combination approach may adapt to other tasks.Source: Problem Definition · p. 2

Related Work

Prior work uses internal teacher representations for distillation, but selecting a limited set of layers can discard information and makes alignment choices consequential. Combinatorial distillation combines selected teacher layers, while other work explores progressive training, multiple teachers, or self-teaching.

Related Work

  • Knowledge distillation has been applied to NLP and NLU, including BERT compression through Patient KD (Sun et al. 2019).Earlier work also adapted distillation for sequence-level machine translation (Kim and Rush 2016) and multi-task translation (Freitag, Al-Onaizan, and Sankaran 2017).Source: Related Work · p. 2
  • Patient KD (PKD) adds an intermediate-layer loss to the standard task and distillation losses, aligning student layers with selected teacher layers.Only the teacher layers returned by the alignment function contribute, which can mean substantial information loss for deep teachers and small students.Source: Related Work · p. 2
  • CKD combines a subset of teacher layers to address skipping, but still searches for the best subset and therefore retains the search problem.Its combination uses concatenation, with a projection layer to address dimension mismatch.Source: Related Work · p. 3
  • Aguilar et al. (2020) proposed progressive training for PKD, training lower layers first before shifting training toward upper layers.They argue that how internal layers are trained can matter.Source: Related Work · p. 3
  • Other approaches distill from multiple teachers or use a student’s own checkpoints as teachers, rather than relying solely on a fixed external teacher.Tan et al. (2019) combine translation engines; Wei et al. (2019) select between current and best stored checkpoints.Source: Related Work · p. 3

Methodology

The methodology replaces one-to-one teacher–student layer matching with combinatorial projection. It combines teacher layers using attention, first within buckets and then across all teacher layers for each student layer.

Methodology

  • The bucketed version combines every teacher layer in a student's bucket, avoiding skipped layers but still requiring a choice of bucket boundaries.The authors note that bucket size and membership remain uncertain design choices.Source: Methodology · p. 3
  • Unlike CKD's concatenation, ALP-KD uses attention to combine teacher layers.The paper reports that this combination is more useful.Source: Methodology · p. 3
  • Learnable attention weights represent teacher-layer contributions; experiments use dot products to estimate student–teacher layer similarity.The authors state that other similarity functions could also be used.Source: Methodology · p. 3
  • Each student layer receives an attention-weighted average over all teacher layers, addressing both skipped layers and bucket-boundary search.The all-layer configuration assigns a distinct set of attention weights to each student layer.Source: Methodology · p. 4
  • Training combines task cross-entropy, prediction distillation, and a dedicated ALP-KD layer loss, weighted by β, η, and λ.The dedicated layer loss uses mean-square error between student-layer representations and teacher projections.Source: Methodology · p. 4

Experimental Study

Across eight GLUE tasks, the experiments compare BERT students using intermediate-layer distillation variants; ALP-KD is reported better than RKD, PKD, and CKD variations on average. The qualitative analyses examine learned teacher-layer attention and student representations.

Experimental Study

  • The evaluation uses eight GLUE tasks and a 12-layer BERT teacher, with student models assessed on validation sets.The paper notes that GLUE test labels are not publicly available and test-set evaluation requires leaderboard participation.Source: Experimental Study · p. 4
  • Concatenation-based CKD helps in some cases, but simple skipping performs better in others, and changing bucket overlap does not alter results substantially.This comparison motivates attention-based combination.Source: Experimental Study · p. 5
  • ALP-KD is consistently better than RKD, PKD, and CKD variations and delivers the best average result without layer-alignment adjustments.The authors report that replacing CKD concatenation with attention improves results.Source: Experimental Study · p. 5

Training Deeper/Shallower Models Than 4-Layer Students

  • Among 6-layer students, SALP-NO has the best average score, while SALP is second best and needs no layer alignment before training.Attention-based combination remains superior to simple concatenation in this setting.Source: Training Deeper/Shallower Models Than… · p. 6
  • SALP outperforms the other students for 2-layer models, and the PKD–ALP-KD gap is more visible than for larger students.The authors interpret this as evidence that combining intermediate layers becomes crucial when teacher and student architectures differ substantially.Source: Training Deeper/Shallower Models Than… · p. 6

Qualitative Analysis

  • Attention maps show stronger activity in the first half of the teacher for 2- and 4-layer students, but later layers also signal strongly; 6-layer attention concentrates around middle layers.The authors conclude that nearly all teacher layers participate, making fixed skipping or concatenation a poor fit to the observed attention distributions.Source: Figure 3 · p. 6, Qualitative Analysis · p. 6
  • For one example, the reported distance between a PKD student representation and its teacher counterpart is around 0.20, versus 0.05 for ALP-KD.The authors use cosine similarity to measure distance and report that ALP-KD representations are closer to the teacher.Source: Qualitative Analysis · p. 7

Conclusion and Future Work

The conclusion finds that combining teacher layers with attention is useful for intermediate-layer distillation, particularly as teacher and student depths diverge. It recommends attention over simple concatenation for competitive results and identifies deeper teachers and other tasks as future directions.

Conclusion and Future Work

  • The conclusion reports that attention-based combination is effective and avoids skipping teacher layers.The authors also propose combining other internal components, such as attention heads.Source: Conclusion and Future Work · p. 7
  • The authors consider combining multiple internal components essential when distilling from deep teachers.They identify applying ALP-KD to extremely deep teachers and compact students as future work.Source: Conclusion and Future Work · p. 7
  • Intermediate-layer combination becomes more crucial as the teacher and student differ more in layer count.The authors propose applying ALP-KD to other tasks and designing better attention modules, including ones suited to sparse structures.Source: Conclusion and Future Work · p. 7
  • Attention-based layer combination is required for competitive results, although simple concatenation often outperforms skipping layers.The authors suggest a simple attention mask over all teacher layers when the best bucket strategy is unknown.Source: Conclusion and Future Work · p. 7

GLUE Datasets

The experiments use eight GLUE tasks spanning grammaticality, sentiment, inference, paraphrase, question equivalence, and semantic similarity. The passages describe each task’s dataset and prediction target.

GLUE Datasets

  • CoLA evaluates grammatical acceptability, while SST-2 uses sentence-level positive or negative sentiment labels.CoLA has 8,551 training and 1,043 validation instances; SST-2 has 67,349 training and 872 validation sentences.Source: GLUE Datasets · p. 10
  • MNLI, QNLI, and RTE cover textual inference, including entailment, contradiction, or neutral judgments for MNLI.The experiments use MNLI’s matched version.Source: GLUE Datasets · p. 10
  • MRPC and QQP assess whether sentence pairs or question pairs are semantically equivalent.MRPC uses human annotations; QQP examples are collected from Quora.Source: GLUE Datasets · p. 10
  • STS-B measures semantic similarity between sentence pairs using scores from 1 to 5.The dataset contains 5,749 training and 1,500 validation examples.Source: GLUE Datasets · p. 10

Different Attention Models

The authors compare single- and multi-head attention extensions with ALP-KD’s default dot-product module. Neither tested extension surpasses the dot-product technique.

Different Attention Models

  • The default ALP-KD module uses a simple dot product, and the authors test whether another attention technique can improve performance over PKD and CKD.The comparison is reported in Table 4.Source: Different Attention Models · p. 10
  • In the extension, a student layer is the query, teacher layers are keys, and each key has a dedicated trainable value vector.The authors implement single- and multi-head attention using the key, query, and value setup described by Vaswani et al. (2017).Source: Different Attention Models · p. 10
  • Neither single-head nor 4-head attention outperforms the simple dot-product technique, despite performance improving as head count increases.The authors describe this outcome as unexpected.Source: Different Attention Models · p. 10

Hyper-parameters

The experiments specify common batch and sequence settings, with fine-tuning epochs varying by dataset. Hyperparameter searches vary learning rate, temperature, and, for intermediate-layer methods, loss weights.

Hyper-parameters

  • All models use batch size 32; maximum sequence length is 64 for single-sentence tasks and 128 for sentence-pair tasks.Fine-tuning uses 10 epochs for STS-B, SST-2, and QNLI; 5 for datasets with over 300K training instances; 20 for MRPC and RTE; and 50 for CoLA.Source: Hyper-parameters · p. 10
  • For TBERT and SNKD, the learning-rate and temperature grids are searched; SRKD additionally searches η, with β set to 1 − η and λ fixed at 0.The learning-rate candidates are {1e−5, 2e−5, 5e−5}, and temperatures are {1, 5, 10, 20}.Source: Hyper-parameters · p. 10
  • SPKD, SCKD, and SALP search λ alongside the other hyperparameters, with β set to 1 − η − λ for SPKD.SPKD tests λ values {0.2, 0.5, 0.7}; Tables 5 to 10 list the experiment-specific values.Source: Hyper-parameters · p. 10

Hardware

Models were fine-tuned on a single NVIDIA 32GB V100 GPU, with runtime varying by dataset size.

Hardware

  • Fine-tuning used a single NVIDIA 32GB V100 GPU and took from a few hours to one day, depending on dataset size.Source: Hardware · p. 10

Related papers

Loading 2012.14022v1…