Source-linked AI summary
Reinforced Multi-Teacher Selection for Knowledge Distillation
Fei Yuan, Linjun Shou, Jian Pei, Wutao Lin, Ming Gong, Yan Fu, Daxin Jiang
TL;DR
Large pre-trained NLP models are costly to deploy, while multi-teacher distillation typically uses fixed or equal teacher weights. The paper introduces reinforcement-learning-based instance-level teacher weighting and finds improved student performance across several NLP tasks, with a limitation on small-data settings for a weak student.
Problem
Multi-teacher knowledge distillation commonly fixes teacher weights across the distillation process, despite variation in training examples and student-model capability.
Method
The paper learns a reinforcement-learning policy that dynamically assigns teacher weights from training-example characteristics and teacher outputs to optimize student performance.
Results
Experiments on several NLP tasks show that the reinforced multi-teacher selection strategy substantially improves student-model performance.
Takeaways & Limitations
Dynamic instance-level teacher selection adapts teacher choice to student capability and teacher strengths during distillation.
Takeaways & Limitations
For weak student BERT3 on small-data MRPC, the reward may be non-robust and compromise teacher-selection policy learning.
Abstract
from arXiv · showhide
In natural language processing (NLP) tasks, slow inference speed and huge footprints in GPU usage remain the bottleneck of applying pre-trained deep models in production. As a popular method for model compression, knowledge distillation transfers knowledge from one or multiple large (teacher) models to a small (student) model. When multiple teacher models are available in distillation, the state-of-the-art methods assign a fixed weight to a teacher model in the whole distillation. Furthermore, most of the existing methods allocate an equal weight to every teacher model. In this paper, we observe that, due to the complexity of training examples and the differences in student model capability, learning differentially from teacher models can lead to better performance of student models distilled. We systematically develop a reinforced method to dynamically assign weights to teacher models for different training instances and optimize the performance of student model. Our extensive experimental results on several NLP tasks clearly verify the feasibility and effectiveness of our approach.
Introduction
Large pre-trained NLP models are difficult to deploy because of slow inference and substantial GPU requirements, motivating knowledge distillation. With multiple teachers, fixed or equal weighting ignores instance-specific teacher strengths and student capacity, so this paper proposes reinforcement-learning-based dynamic teacher selection and reports improved student performance.
- Knowledge distillation compresses large teacher models into smaller student models to reduce computation costs and speed inference.
- Existing multi-teacher methods commonly assign equal weights throughout distillation, overlooking differences among teacher strengths.Such uniform weighting simplifies coordination but leaves potential gains from exploiting teacher diversity.
- Teacher performance can vary by training instance, reflecting differences in hypothesis spaces, optimization, initialization, and other factors.The paper therefore motivates assigning teacher weights according to performance on individual cases.
- A stronger teacher does not necessarily produce a better student because student capacity may limit the ability to approximate complex teachers.Table 1 reports that students distilled from weaker BERT-Base outperform those distilled from stronger RoBERTa-Base on the cited tasks.
- The paper introduces instance-level teacher weighting and formulates teacher selection as reinforcement learning driven by training-example features, teacher outputs, and student performance.The study presents this as a systematic approach to coordinating teacher and student models.
- Experiments on GLUE tasks show that the reinforced multi-teacher selection strategy substantially improves student-model performance.The evaluated tasks include sentiment analysis, paraphrase similarity matching, and natural language inference.
Related Work
Earlier multi-teacher distillation methods combine teacher outputs using fixed weights or select teachers randomly at the mini-batch level. This paper instead learns example-dependent teacher weights through a policy.
- Chebotar and Waters use a weighted average of teacher models with weights fixed during training.
- Fukuda et al. combine soft labels with fixed weights or randomly select one teacher at the mini-batch level.
- Previous multi-teacher methods fix the same teacher weight across all training examples, whereas this paper dynamically assigns weights by individual example.
Our Approach
RL-KD dynamically selects and weights teacher models for each training instance within a reinforcement-learning framework. The selected teachers’ outputs train the student, while student performance supplies delayed rewards for updating the selection policy.
- Teacher selection: RL-KD treats teacher selection as reinforcement learning, using instance and teacher information as state and student performance as the optimization return.The policy makes decisions from training-example characteristics and teacher outputs, then is optimized toward student performance.
- State representation: The policy state concatenates an input representation, a teacher’s class-probability vector, and its ground-truth loss, along with other teachers’ predictions.The input representation uses a BERT-Base [CLS] embedding, while other teacher predictions are included as additional state features.
- Teacher selection: Each teacher-specific agent chooses whether to select its teacher for the current instance, with policy probabilities serving as teacher weights.The action is binary, and the policy assigns probabilities that become the weights used in distillation.
- Student training: Selected teachers’ outputs are averaged into the distillation loss used to train the student, alongside the ground-truth loss.The student objective balances soft teacher targets and hard labels through the knowledge-distillation training process.
- Policy optimization: Rewards can use student ground-truth loss, combined ground-truth and distillation losses, or development-set accuracy, and are delayed until a batch finishes.Policy-gradient optimization updates the selector across training episodes until student performance converges.
Experiments
Experiments evaluate teacher and student models across three GLUE task groups using several single-teacher, ensemble, and instance-level selection baselines. RL-KD consistently performs best in most cases, while its reward-based selection is less reliable for BERT3 on the small MRPC dataset.
- Experimental setup: The evaluation covers sentiment classification, paraphrase similarity matching, and natural language inference tasks from GLUE.Prediction accuracy is used as the evaluation metric.
- Experimental setup: The study compares BERT12, RoBERTa12, ALBERT12, and XLNet12 as teachers with BERT3 and BERT6 as student models.The teacher models have 12 transformer layers, while the student models have 3 or 6.
- Baselines: The baselines include single-teacher KD, equal-weight and fixed-weight ensembles, random teacher selection, logistic-regression ensembles, and best-single-teacher selection.The comparisons span fixed-weight, learned-weight, and instance-level teacher strategies.
- Teacher-model results: Individual teachers specialize by task: RoBERTa leads MNLI-mm, MNLI-m, and QNLI, ALBERT leads MRPC and RTE, and XLNet leads SST-2.The results suggest that teachers learn different local optima and biases.
- Student-model results: RL-KD consistently outperforms the other methods in most cases and dynamically selects teachers according to student capability.Among the three reward functions, reward3 performs best by incorporating development-set accuracy.
- Limitations: BERT3 is the only reported exception on MRPC, where limited training data makes its reward less robust and can compromise teacher-selection policy.BERT6 generalizes more strongly and can produce reliable rewards even with relatively small datasets.
- Statistical analysis: Across five runs, RL-KD has variance similar to baseline KD and outperforms baselines with statistical significance.The comparison uses mean, standard deviation, and a two-sided t-test with threshold 0.05.
Conclusions
The paper proposes an RL-based approach for dynamically assigning teacher weights at the instance level in knowledge distillation. Experiments on several NLP tasks verify its effectiveness.
- The approach tackles teacher model selection when multiple teachers are available.
- It dynamically assigns teacher weights for individual training instances.
- The method adapts teacher selection to both teacher strengths and student model capability.
- Extensive experiments on several NLP tasks verify the approach's effectiveness.