PaperlayerView exact revision on arXiv

Source-linked AI summary

AD-KD: Attribution-Driven Knowledge Distillation for Language Model Compression

Siyue Wu, Hongzhan Chen, Xiaojun Quan, Qifan Wang, Rui Wang

arXiv:2305.10010v1cs.CL

TL;DR

Knowledge distillation can compress pre-trained language models, but existing approaches may neglect teacher reasoning and data-specific knowledge. AD-KD transfers token-level attribution knowledge, including attributions for all possible teacher predictions, to the student. In BERT experiments on GLUE, it outperforms several state-of-the-art baselines.

  • Problem

    Existing KD methods may imitate teacher behavior without capturing its reasoning and may overlook data-specific knowledge.

    Source: §1 · p. 1

  • Method

    AD-KD transfers token-level teacher attributions computed with Integrated Gradients, using top-K filtering and multi-view maps for all possible predictions.

    Source: §1 · p. 2, §3.2.1 · p. 4

  • Results

    AD-KD outperforms several state-of-the-art baselines in BERT experiments on GLUE.

    Source: §1 · p. 2

  • Takeaways & Limitations

    The authors conclude that input attribution is a knowledge source for distillation, and report effectiveness for AD-KD and its components.

    Source: §6 · p. 9

    Experiments cover BERT at different scales, not other model structures, and task-specific rather than task-agnostic distillation.

    Source: Limitations · p. 9

Abstract

from arXiv · show

Knowledge distillation has attracted a great deal of interest recently to compress pre-trained language models. However, existing knowledge distillation methods suffer from two limitations. First, the student model simply imitates the teacher's behavior while ignoring the underlying reasoning. Second, these methods usually focus on the transfer of sophisticated model-specific knowledge but overlook data-specific knowledge. In this paper, we present a novel attribution-driven knowledge distillation approach, which explores the token-level rationale behind the teacher model based on Integrated Gradients (IG) and transfers attribution knowledge to the student model. To enhance the knowledge transfer of model reasoning and generalization, we further explore multi-view attribution distillation on all potential decisions of the teacher. Comprehensive experiments are conducted with BERT on the GLUE benchmark. The experimental results demonstrate the superior performance of our approach to several state-of-the-art methods.

1 Introduction

Knowledge distillation aims to compress large pre-trained language models, but existing methods can miss teacher reasoning and data-specific knowledge. AD-KD transfers token-level attribution knowledge, including views of all possible teacher predictions, and outperforms several state-of-the-art baselines in BERT experiments on GLUE.

1 Introduction

  • Pre-training on large unlabeled corpora followed by task-specific fine-tuning yields strong NLP performance, while growing model scale complicates deployment.Knowledge distillation seeks a smaller student that approaches the teacher's performance, including in low-resource scenarios.Source: §1 · p. 1
  • Existing methods focus on what the teacher predicts rather than why, and may overlook data-specific rationale while emphasizing model-specific knowledge.The introduction links these gaps to student reasoning and generalization limitations.Source: §1 · pp. 1–2
  • AD-KD uses Integrated Gradients and top-K filtering to produce token-level attribution scores, then transfers views for all possible teacher predictions.The approach treats attribution as reasoning knowledge complementary to soft-label knowledge.Source: §1 · p. 2
  • AD-KD outperforms several state-of-the-art baselines in BERT experiments on GLUE.The authors report superior performance and validate the approach on the benchmark.Source: §1 · p. 2

2 Related Work

Prior work transfers teacher outputs or intermediate representations through response-, feature-, and relation-based distillation. Attribution methods assign importance to input or intermediate features, and this paper applies attribution to knowledge distillation.

2.1 Knowledge Distillation

  • Response-based KD transfers final outputs, while feature- and relation-based methods align intermediate representations or their relations.The section distinguishes approaches that use hidden representations, attention matrices, and pairwise or structural relations.Source: §2.1 · p. 2

2.2 Attribution

  • Attribution methods assign importance scores to features; this paper investigates their use in KD, where attribution had not been explored.The section describes occlusion-, gradient-, propagation-, and perturbation-based approaches, noting occlusion's repeated-forward-pass cost and gradient methods' single-forward-process requirement.Source: §2.2 · pp. 2–3

3 Methodology

AD-KD transfers token-level attribution knowledge from teacher to student, using Integrated Gradients to represent prediction rationale. It builds multi-view maps, distills their relative token importance, and combines that loss with label and logit losses.

3.1 Preliminary

  • Integrated Gradients approximates feature attribution by summing gradients along a path, with magnitude indicating importance and sign indicating the effect on the target output.A larger number of summation steps usually gives a better approximation.Source: §3.1 · p. 3

3.2 AD-KD

  • AD-KD extracts teacher attribution maps and transfers attribution-based knowledge to the student.The framework uses an input example, extracts attribution maps from the teacher, then transfers that knowledge to the student.Source: §3.2 · p. 3, Figure 2 · p. 4

3.2.1 Attribution Maps

  • Teacher maps score each token’s embedding dimensions with IG; AD-KD filters low-scoring dimensions and keeps the top K, while the student uses all dimensions.The student’s all-dimension extraction is intended to reduce training difficulty and let it learn from the teacher’s noiseless map.Source: §3.2.1 · pp. 3–4, §3.1 · p. 3
  • AD-KD extracts attribution maps for every possible teacher prediction, giving the student a more comprehensive view of the teacher’s soft-label distribution.The teacher attribution scores reflect token importance; their signs are omitted because dimensions within a token embedding can counteract each other when combined.Source: §3.2.1 · pp. 3–4

3.2.2 Attribution Distillation

  • The attribution distillation loss matches normalized teacher and student maps, preserving relative token importance rather than absolute attribution magnitude.Normalization is used to support smoother distillation when teacher and student attribution magnitudes differ early in training.Source: §3.2.2 · pp. 4–5

3.2.3 Overall Objective

  • The overall objective combines ground-truth cross-entropy, response-based logit loss, and attribution-driven distillation loss.The logit loss is defined using KL divergence between the teacher and student output distributions.Source: §3.2.3 · p. 5

4 Experimental Settings

The experiments evaluate AD-KD on eight GLUE tasks against response-, feature-, and relation-based distillation baselines. The setup uses a BERTbase teacher and a six-layer BERT student, with task-specific hyperparameter searches.

4.1 Datasets

  • Evaluation covers eight GLUE tasks, reporting accuracy, F1, Matthews correlation, or Spearman correlation according to the task.Accuracy is reported for MNLI, SST-2, QNLI, QQP, and RTE; MRPC uses F1, CoLA uses Matthews correlation, and STS-B uses Spearman correlation.Source: §4.1 · p. 5

4.2 Baseline Methods

  • Baselines include response-based Vanilla KD and PD, feature-based PKD and TinyBERT, and relation-based CKD and MGSKD.MiniLM and MobileBERT are excluded because their two-stage settings differ; MGSKD is reimplemented with a six-layer student without external span-level knowledge.Source: §4.2 · p. 5

4.3 Implementation Details

  • The experiments fine-tune BERTbase as teacher and use a six-layer, 768-hidden-unit, 12-head BERT student, tuning learning rate, α, β, τ, and K.The main results use one IG step because of computational overhead; results for m from 1 to 8 are also reported.Source: §4.3 · p. 6

5 Results and Analysis

AD-KD improves on several baselines, with attribution distillation contributing most to performance. The analyses also examine multi-view attribution, student size, retained dimensions, IG steps, attribution layers, and loss weights.

5.1 Main Results

  • 1.0 and 1.9 points over CKD and MGSKD, respectively, on development sets, and 0.9 points over MGSKD on test sets.AD-KD outperforms all baselines on most datasets.Source: §5.1 · p. 6

5.2 Ablation Study

  • Removing attribution distillation causes an obvious performance drop; removing cross-entropy or logit distillation also degrades results, but attribution contributes most.The ablation compares loss-term variants on GLUE development sets.Source: §5.2 · pp. 6–7, Table 2 · p. 7
  • On MNLI, each single-view attribution beats vanilla KD, while combining all views yields further improvement.The comparison covers the entailment, contradiction, and neutral labels.Source: Figure 4 · p. 7, §5.2 · p. 7
  • AD-KD consistently outperforms vanilla KD across student scales on MRPC and QNLI.Source: Figure 5 · p. 7, §5.2 · p. 7

5.3 Impact of Top-K

  • On STS-B, performance first improves with K and slightly declines beyond 600; on QNLI, it improves almost monotonically as K increases.The reported dataset sizes are 7k for STS-B and 108k for QNLI.Source: §5.3 · p. 7, Figure 6 · p. 8

5.4 Impact of IG Steps

  • As IG steps m increases, AD-KD performance fluctuates; m=1 offers a better performance–computation tradeoff than costly task-specific tuning.Although some settings surpass the default and even the teacher, finding each task’s optimum is costly.Source: §5.4 · pp. 7–8, Figure 7 · p. 8

5.5 Attribution Distillation Layer

  • Uniform mapping performs best among intermediate-layer methods, but none beats input-layer attribution, and combining input and intermediate layers adds no improvement.Source: Table 3 · p. 8, §5.5 · p. 8

5.6 Impact of α and β

  • Across MRPC and QNLI, performance is best when α is around 0.9; increasing β first improves performance to a peak, then degrades it, with its optimum varying by task.The authors select α close to 1; β is more task-sensitive.Source: Figure 8 · p. 8, §5.6 · p. 8

6 Conclusion

AD-KD distills input-attribution knowledge from teacher to student, using top-K attribution maps and multi-view attribution to transfer model knowledge. Experiments and ablations report the effectiveness of AD-KD and its components.

6 Conclusion

  • AD-KD is a language-model compression framework that transfers input-attribution knowledge from teacher to student.The paper describes input attribution as vital yet easy to transfer between the models.Source: §6 · pp. 8–9
  • Top-K selects input-token attribution maps to reduce noise, while multi-view attribution supports more comprehensive distillation.Source: §6 · p. 9
  • Experiments, including ablations, show the effectiveness of AD-KD and its components.The paper also suggests occlusion-based attribution or chain-of-thoughts as possible extensions to black-box models with unavailable parameters.Source: §6 · p. 9

Limitations

The paper identifies open questions about attribution methods, model structures, and distillation settings: experiments use Integrated Gradients and BERT, and cover task-specific distillation.

Limitations

  • Whether attribution methods other than Integrated Gradients perform better when combined with knowledge distillation remains unclear.Source: Limitations · p. 9
  • AD-KD was evaluated on BERT at different scales, but its effectiveness on other model structures has not been validated.Source: Limitations · p. 9
  • Experiments cover task-specific distillation; applying AD-KD to task-agnostic distillation remains an open direction.Source: Limitations · p. 9

A.1 Details of Datasets

The appendix describes the GLUE datasets used for evaluation and reports experimental hardware and relative training costs.

A.1 Details of Datasets

  • Evaluation covers eight GLUE tasks: CoLA, SST-2, MRPC, QQP, MNLI, QNLI, RTE, and STS-B.These comprise single-sentence, paraphrase, inference, and similarity tasks.Source: §A.1 · p. 13

A.2 Hyperparameter Settings

  • Experiments used GeForce RTX 2080 Ti GPUs; AD-KD runs 1.2 to 3 times slower than vanilla KD, while student inference speeds are equal.The stated training slowdown is due to extra back-propagation.Source: §A.2 · p. 13

B More Discussion

The discussion distinguishes attribution maps from attention matrices as distillation targets, arguing that attention has limitations as a global input-level signal.

B More Discussion

  • Attention and attribution both reflect input-token contributions to predictions, but attention correlates with attribution locally in specific layers and heads rather than globally.Source: B More Discussion · p. 13
  • Adversarial training can yield alternative attention weights with nearly unchanged predictions, so students need not share the teacher’s attention matrices.Source: B More Discussion · p. 13
  • AD-KD uses gradient-based attribution maps, which the paper reports as better than attention matrices used by baselines.Source: B More Discussion · p. 13

C.1 Results on MultiRC

On long-text MultiRC, AD-KD improves on distillation baselines and performs on par with the teacher; additional analyses report that attribution knowledge generalizes to unseen data and show more teacher-consistent token focus in two examples.

C.1 Results on MultiRC

  • 0.97% over vanilla KD and 0.38% over MGSKD, with AD-KD performance on par with the teacher on the MultiRC development set.The evaluation uses SuperGLUE MultiRC, whose maximum sequence length is 512, compared with 128 for GLUE.Source: Table 6 · p. 13, §C.1 · p. 13

C.2 Overfitting Study

  • Development-set attribution gaps track training-set trends as β changes, suggesting learned attribution knowledge generalizes to unseen data and overfitting tends not to happen.The analysis covers MRPC and QNLI.Source: §C.2 · p. 14, Figure 9 · p. 14

C.3 Case Study

  • In two QNLI examples, AD-KD's attribution maps align more closely with the teacher's, focusing on “Louisiana” and not “billion,” while vanilla KD makes mistakes.The examples illustrate learning what to focus on and what to ignore.Source: §C.3 · p. 14, Figure 10 · p. 15
Loading 2305.10010v1…