Source-linked AI summary
Knowledge Distillation from Internal Representations
Gustavo Aguilar, Yuan Ling, Yu Zhang, Benjamin Yao, Xing Fan, Chenlei Guo
TL;DR
Standard knowledge distillation may fail to transfer a large teacher’s internal knowledge because matching output probabilities does not guarantee matching internal representations. The paper distills BERT’s internal representations into simplified students using matched-layer objectives and multiple training algorithms. Across four GLUE datasets, the method consistently outperforms standard knowledge distillation.
Problem
Output-only distillation does not guarantee that a smaller student acquires the teacher’s internal representations, which may limit transfer of the teacher’s intended generalization capabilities.
Method
The paper distills internal representations from multiple layers of a large BERT model into a simplified version using formalized objectives and several distillation algorithms.
Results
The internal-distillation student consistently outperforms standard knowledge distillation across four GLUE benchmark datasets.
Takeaways & Limitations
Adding internal representations makes knowledge distillation more powerful than using soft-label distillation alone in the reported GLUE experiments.
Takeaways & Limitations
The experimental setup initializes the student from general-purpose BERTbase parameters rather than the fine-tuned teacher’s parameters.
Abstract
from arXiv · showhide
Knowledge distillation is typically conducted by training a small model (the student) to mimic a large and cumbersome model (the teacher). The idea is to compress the knowledge from the teacher by using its output probabilities as soft-labels to optimize the student. However, when the teacher is considerably large, there is no guarantee that the internal knowledge of the teacher will be transferred into the student; even if the student closely matches the soft-labels, its internal representations may be considerably different. This internal mismatch can undermine the generalization capabilities originally intended to be transferred from the teacher to the student. In this paper, we propose to distill the internal representations of a large model such as BERT into a simplified version of it. We formulate two ways to distill such representations and various algorithms to conduct the distillation. We experiment with datasets from the GLUE benchmark and consistently show that adding knowledge distillation from internal representations is a more powerful method than only using soft-label distillation.
Introduction
Large transformer models deliver strong NLP performance but are costly to train, run, and deploy. The paper addresses the limits of output-only distillation by transferring the teacher’s internal representations into a smaller student.
- Motivation: Transformer models achieve state-of-the-art results across many NLP tasks but often contain hundreds of millions of parameters.Their size slows training and inference and increases memory consumption, hindering production use under strict resource constraints.
- Knowledge Distillation: Knowledge distillation trains a simplified student from a cumbersome teacher’s output probabilities.Soft-label distributions also encode similarities among classes beyond one-hot ground-truth labels.
- Research Gap: Large teacher–student gaps make it difficult for students to approximate the teacher, even when output probabilities are available.The paper argues that output probabilities expose only part of the abstraction captured by the teacher, leaving internal knowledge difficult for the student to infer.
- Proposed Direction: The paper distills internal representations at multiple teacher layers so the student can learn and compress hidden-layer abstractions.The proposed approach is evaluated by comparing students trained with internal distillation against homologous models trained with ground-truth labels, soft labels, or both.
Related Work
Prior work extends knowledge distillation to internal representations, transformer compression, curriculum-style training, and complementary compression techniques. This paper positions its formulation as applying distillation directly to internal representations.
- Knowledge Distillation Extensions: Knowledge distillation has been extended from single teachers to multi-task networks and from model ensembles to task-specific networks.These extensions preserve the teacher–student compression principle while changing the source or destination of distilled knowledge.
- Knowledge Distillation Extensions: This paper extends knowledge distillation by applying the same principle to internal representations.The formulation differs from extensions that focus primarily on outputs or task arrangements.
- Internal-Representation Distillation: FITNET compares internal representations at matched teacher and student points, using a convolutional regressor when their dimensionalities differ.The regressor maps representations into a shared space but adds parameters to the student-training process.
- Transformer Compression: DistilBERT distills during pretraining, whereas related work also applies hidden-state and soft- or hard-label losses during task-specific fine-tuning.The cited transformer-based methods differ in training stage and loss composition.
- Curriculum Learning: Curriculum learning teaches increasingly complex tasks in stages but requires preliminary-task annotations and careful task ordering.The paper distinguishes its approach from curriculum learning’s dependence on annotated build-up tasks.
- Complementary Compression: Quantization reduces parameter precision, while pruning removes weights with limited impact on performance.The paper describes both techniques as complementary to its internal-distillation method.
Methodology
The methodology distills BERT’s internal knowledge into shallower BERT students through matched-layer losses and several layer-training schedules. It combines attention-distribution matching with [CLS] representation matching.
- Method Overview: The method formalizes internal knowledge distillation for transformer models and proposes multiple algorithms for conducting it.The section first reviews standard knowledge distillation, then defines internal-representation objectives and training procedures.
- Matching Internal Representations: The student learns from teacher soft labels while matching internal representations so multiple teacher layers can be compressed into one student layer.The illustrated setup compresses two teacher layers into one student layer while preserving linguistic behavior across those layers.
- Internal Distillation Loss: Internal distillation uses KL divergence across self-attention probabilities and cosine similarity between teacher and student [CLS] activation vectors.The method studies BERT and simplified versions with fewer transformer layers.
- Attention Matching: Minimizing divergence between teacher and student self-attention distributions is intended to preserve linguistic behavior in the student.The loss compares the full row-wise probability distributions across corresponding attention matrices.
- Activation Matching: Cosine similarity matches the teacher and student [CLS] hidden vectors to align activation behavior through the network.This term complements attention matching by targeting the weighted hidden vectors passed to upper layers.
- Distillation Algorithms: The method supports simultaneous all-layer distillation, progressive bottom-up distillation, and stacked distillation that retains earlier layer losses.Progressive distillation optimizes one layer at a time, while stacked distillation accumulates losses as training moves upward.
- Distillation Algorithms: Progressive and stacked schedules advance to upper layers after a layer-specific epoch limit or cosine-loss threshold.Both schedules may also be combined with classification loss throughout training.
Experiments and Results
Experiments evaluate BERTbase and six-layer BERT students across four GLUE datasets using progressively richer distillation objectives. Internal-representation distillation consistently outperforms the undistilled and standard soft-label baselines, while achieving similar performance with substantially fewer parameters.
- Experimental setup: Experiments compare BERTbase with BERT6 students across CoLA, QQP, MRPC, and RTE using multiple distillation configurations.The study includes no-distillation, soft-label, KL-divergence, cosine-similarity, progressive, and stacked internal-distillation settings.
- Statistical significance: Internal-KD results differ significantly from standard KD across datasets, with p-value less than 1e-3.The reported significance comparison concerns the internal distillation results versus the standard knowledge-distillation method.
- Development results: Internal distillation consistently outperforms both undistilled BERT6 and standard distillation across all four datasets.Standard distillation does not always beat undistilled BERT6, but the proposed internal method does.
- Test results: Internal distillation also performs consistently better than standard soft-label distillation on the test set.Table 2 selects the best internal-KD technique for BERT6 according to development-set results.
- Parameter-performance trade-off: Performance decays as layers are removed, but internal distillation is more resilient than standard distillation at smaller model sizes.With one transformer layer, compression is too aggressive to provide an additional internal-distillation boost.
Analysis
The analysis evaluates internal distillation across parameter reduction, training-data size, convergence, attention behavior, and prediction errors. Internal KD generally preserves performance with fewer parameters and outperforms soft-label KD, especially with less data, while producing representations closer to the teacher.
- Performance vs. Parameters: BERT6 reduces transformer-layer parameters by about 50% while reaching 91.38 F1 on QQP versus 91.45 F1 for BERTbase.BERTbase has around 109.4M total parameters, while BERT6 has 43.1M transformer-layer parameters.
- Performance vs. Parameters: Internal KD remains more resilient than standard KD as the number of student layers decreases, although both methods degrade with more aggressive compression.With one transformer layer, the compression is too high for internal distillation to provide an additional boost over standard distillation.
- The Impact of Data Size: Internal KD consistently outperforms soft-label KD on QQP, with its advantage increasing as the training-data size decreases.The performance gap is small with large datasets but tends to favor internal KD more strongly with less data.
- Student Convergence: Progressive and stacked internal distillation experiments support teaching and compressing internal knowledge without updating the classification layer.Figure 4 compares convergence across epochs and marks layers that have been completely optimized; after L6, only the classification layer is trained.
- Inspecting the Attention Behavior: For one example, internal KD yields a lower cross-layer attention KL-divergence than standard KD: 0.085 versus 2.229.Figure 5 compares attention for head 8 in layer 5 and reports the KL-divergence across matching teacher-student layers.
- Error Analysis: Internal KD better matches the teacher’s right-prediction behavior, reaching 97.9% agreement when the teacher is correct versus 95.7% for standard KD.When the teacher is wrong, internal KD makes more mistakes than standard KD, but its overall student score is significantly higher.
Conclusions
The paper extends knowledge distillation by teaching smaller models both the teacher’s output probabilities and internal representations. Across four GLUE datasets, this approach preserves similar performance, compresses internal behavior, and consistently outperforms standard KD.
- Conclusions: The paper proposes extending KD by revealing the teacher’s internal representations in addition to its output probabilities.The method aims to compress the teacher’s internal behavior into smaller models while preserving similar performance.
- Conclusions: Across four GLUE datasets, internal representation distillation consistently outperforms standard knowledge distillation.The conclusion reports this benchmark-wide comparison as validation of the method’s effectiveness.