Source-linked AI summary
Patient Knowledge Distillation for BERT Model Compression
Siqi Sun, Yu Cheng, Zhe Gan, Jingjing Liu
TL;DR
Large pre-trained language models such as BERT are effective but computationally costly, creating a need for practical compression. The paper proposes Patient Knowledge Distillation, in which a shallow student learns from multiple intermediate teacher layers through PKD-Last and PKD-Skip. Across multiple NLP tasks, the approach improves performance and training efficiency while maintaining comparable accuracy to large models, though initialization constraints limit one teacher-student setting.
Problem
BERT’s large parameter count makes pre-training and downstream use computationally expensive, motivating compression without sacrificing performance.
Method
Patient Knowledge Distillation trains a shallow student to imitate a large teacher using outputs and selected intermediate layers through PKD-Last and PKD-Skip.
Results
Across seven datasets and four NLP tasks, Patient-KD improves performance and generalization over standard knowledge distillation while increasing training efficiency and reducing storage with comparable accuracy to original large models.
Takeaways & Limitations
Multi-layer distillation can compress BERT into a shallow model while retaining comparable task performance and improving practical efficiency.
Takeaways & Limitations
When BERT-Large teaches BERT6[Base], the patient teacher loss cannot be computed because their [CLS] token embedding sizes differ, so only vanilla KD is evaluated.
Abstract
from arXiv · showhide
Pre-trained language models such as BERT have proven to be highly effective for natural language processing (NLP) tasks. However, the high demand for computing resources in training such models hinders their application in practice. In order to alleviate this resource hunger in large-scale model training, we propose a Patient Knowledge Distillation approach to compress an original large model (teacher) into an equally-effective lightweight shallow network (student). Different from previous knowledge distillation methods, which only use the output from the last layer of the teacher network for distillation, our student model patiently learns from multiple intermediate layers of the teacher model for incremental knowledge extraction, following two strategies: ($i$) PKD-Last: learning from the last $k$ layers; and ($ii$) PKD-Skip: learning from every $k$ layers. These two patient distillation schemes enable the exploitation of rich information in the teacher's hidden layers, and encourage the student model to patiently learn from and imitate the teacher through a multi-layer distillation process. Empirically, this translates into improved results on multiple NLP tasks with significant gain in training efficiency, without sacrificing model accuracy.
1 Introduction
BERT is effective across NLP tasks but costly to train and fine-tune because of its size. Patient Knowledge Distillation compresses BERT by having a shallow student learn from multiple teacher layers, improving efficiency while maintaining comparable accuracy.
- Motivation: BERT-Base has 12 layers and 110 million parameters, making training and fine-tuning computationally expensive.Training from scratch typically takes four days on 4 to 16 Cloud TPUs, while one fine-tuning epoch may take several hours.
- Approach: Patient Knowledge Distillation compresses a large teacher model into a lightweight student without sacrificing performance.The student learns to mimic the teacher’s predictions on training samples, with additional unannotated data supported.
- Approach: Patient-KD distills knowledge from intermediate teacher layers rather than relying only on the final layer.This multi-layer process is intended to extract richer information from the teacher’s deep structure.
- Distillation strategies: PKD-Last teaches from the teacher’s last k layers, whereas PKD-Skip teaches from every k layers.The strategies respectively emphasize knowledge near the top of the teacher and information distributed across lower layers.
- Results: Experiments across seven datasets and four NLP tasks show superior performance and generalization over standard knowledge distillation, with better efficiency and storage reduction.Patient-KD maintains comparable accuracy to the original large models.
2 Related Work
Related work establishes pre-training as a route to transferable language representations and identifies model compression as necessary for practical deployment. Patient knowledge distillation extends prior compression and distillation research by transferring information from large language models to compact students.
- Language model pre-training: Pre-training approaches learn language representations from large unlabeled corpora, with feature-based and fine-tuning approaches forming two broad categories.Feature-based methods include word, sentence, and contextualized representations, while fine-tuning approaches adapt pre-trained models to labeled downstream data.
- Language model pre-training: BERT uses multiple Transformer layers and achieves strong performance on natural language understanding tasks after masked language modeling and next sentence prediction.BERT-Base has 12 Transformer layers and 110 million parameters, while BERT-Large has 24 layers and 330 million parameters.
- Model compression: Large pre-trained language models remain difficult to apply when computational resources are limited, motivating compression into shallower models without sacrificing performance.The paper presents BERT compression as an example and notes possible applicability to XLNet and RoBERTa.
- Model compression: Model compression reduces the size or computational demands of deep networks through approaches such as inference acceleration, training-time reduction, and network pruning.Prior work also shows that intermediate representations from large models can provide hints for improving compact-model training and performance.
- Knowledge distillation: Patient knowledge distillation differs from prior efforts by targeting large-scale language-model compression and transferring knowledge from a teacher to a student.Its motivation builds on evidence that intermediate representations can serve as useful hints for compact models.
3 Patient Knowledge Distillation
Patient Knowledge Distillation compresses a large teacher such as BERT into a shallower student by combining output imitation with intermediate-layer representation matching. It uses patient multi-layer distillation to address vanilla KD saturation while retaining comparable performance with lower computational cost.
- Distillation Objective: Knowledge distillation trains a shallower student to match a large teacher’s performance with much lower computational cost.The teacher is represented by f(x; θ), while the student is g(x; θ′) and learns by imitating teacher outputs on training data.
- Distillation Objective: The teacher produces fixed soft labels, and the student minimizes a distance between teacher and student probability outputs, optionally combined with task-specific cross-entropy.Temperature T controls reliance on the teacher’s soft predictions, while α balances cross-entropy and distillation loss.
- Patient Knowledge Distillation: Vanilla KD uses the teacher’s last-layer logits and reaches comparable training-set performance, but quickly saturates on the test set as epochs increase.This motivates distilling knowledge from intermediate teacher layers rather than relying only on the final layer.
- Patient Knowledge Distillation: Patient Knowledge Distillation adds intermediate-layer learning so the student extracts richer information through the teacher’s deep structure.The framework imitates only the [CLS] representations in intermediate layers, based on their role in predicting final labels.
- Patient Knowledge Distillation: PKD-Skip learns from every k teacher layers, whereas PKD-Last learns from the last k layers.For BERT12-to-BERT6 compression, the selected sets are {2, 4, 6, 8, 10} and {7, 8, 9, 10, 11}; the final teacher layer is omitted because its output is already included in the KD loss.
- Patient Knowledge Distillation: The patient teacher adds a mean-square loss between normalized hidden states to the conventional KD objective.The combined objective weights intermediate-layer feature matching with β, alongside the distillation and task losses.
4 Experiments
Experiments across four NLP task families and seven datasets evaluate Patient-KD using compressed BERT3 and BERT6 students. The approach generally preserves or improves performance while reducing model size and inference time, though results depend on training-data size, initialization, and teacher-student configuration.
- GLUE results: Patient-KD models perform best on almost all GLUE tasks except MRPC, compared with direct fine-tuning and vanilla KD.The evaluation uses BERT3 and BERT6 students distilled from a BERT12 teacher.
- GLUE results: 5 of 7 tasks: the Patient-KD BERT6 student achieves similar performance to the BERT-Base teacher, with those tasks each exceeding 60k training samples.The reported teacher-relative differences are SST-2 -2.3%, QQP -0.1%, MNLI-m -2.2%, MNLI-mm -1.8%, and QNLI -1.4%.
- Patient-KD variants: PKD-Skip performs slightly better than PKD-Last, plausibly because cross-layer sampling captures more diverse low- to high-level representations.Both patient strategies improve over vanilla KD.
- Machine reading comprehension: 1.6%: Patient-KD further lifts RACE performance after vanilla KD already exceeds direct fine-tuning by 4.42%.The comparison is reported on the RACE test set.
- Efficiency: 1.94x and 3.73x: Patient-KD achieves near-linear inference speedups for BERT6 and BERT3, while reducing machine memory by 1.64x and 2.4x, respectively.Inference uses 105k QNLI training samples on a single Titan RTX GPU with batch size 128, maximum sequence length 128, and FP16.
5 Conclusion
Patient-KD compresses large BERT models into shallow networks by distilling knowledge from the teacher’s deep structure. The authors identify pre-training, broader applications, improved loss metrics, and more complex settings as future directions.
- Patient-KD compresses large BERT models into shallow ones through multi-layer distillation from the teacher network.
- Future work includes pre-training BERT from scratch to address initialization mismatch and adapting Patient-KD for pre-training.
- The authors also plan to explore more sophisticated loss distance metrics and applications to multi-task and meta-learning settings.