From Mimicking to Integrating: Knowledge Integration for Pre-Trained Language Models (2022)

Abstract: Investigating better ways to reuse the released pre-trained language models (PLMs) can significantly reduce the computational cost and the potential environmental side-effects. This paper explores a novel PLM reuse paradigm, Knowledge Integration (KI). Without human annotations available, KI aims to merge the knowledge from different teacher-PLMs, each of which specializes in a different classification problem, into a versatile student model. To achieve this, we first derive the correlation between virtual golden supervision and teacher predictions. We then design a Model Uncertainty–aware Knowledge Integration (MUKI) framework to recover the golden supervision for the student. Specifically, MUKI adopts Monte-Carlo Dropout to estimate model uncertainty for the supervision integration. An instance-wise re-weighting mechanism based on the margin of uncertainty scores is further incorporated, to deal with the potential conflicting supervision from teachers.Experimental results demonstrate that MUKI achieves substantial improvements over baselines on benchmark datasets. Further analysis shows that MUKI can generalize well for merging teacher models with heterogeneous architectures, and even teachers major in cross-lingual datasets.

## introduction large-scale pre-trained language models (plms), such as bert @xcite , roberta @xcite and t5 @xcite have recently achieved promising results after fine-tuning on various natural language processing (nlp) tasks. many fine-tuned plms are generously released for facilitating researches and deployments. reusing these plms can greatly reduce the computational cost of retraining the plm from scratch and alleviate the potential environmental side-effects like 1 our code is available at https://github.com/ lancopku/muki . part of the work was done while yankai lin and peng li were working at tencent. and knowledge integration (ki). kd assumes that the student performs predictions on the identical label set with the teacher, while ki trains a student model that is capable of performing classification over the union label set of teacher models. ## knowledge integration for plms in this section, we first give the task formulation for knowledge integration, followed by the elaboration on the proposed muki framework. ## error analysis the muki is built on the assumption that the model uncertainty estimation can faithfully reflect the ability of teacher model. we perform an error analysis to investigate when this assumption will fail, i.e, the estimated model uncertainty misguides the teacher selection. specifically, we probe the label distribution of instances on which muki assigns a higher uncertainty score to the correct teacher on thuc-news. as shown in the left part of figure 5 , the uncertainty-based teacher selection only fails on a small portion of training examples, i.e., 2.8% of total training instances. interestingly, the labels of these instances are not uniformly distributed, e.g., estate and tech. have higher error rates than other classes. we further plot the label confusion matrix of an oracle model that is fine-tuned with labeled data with all categories, in the right part of figure 5 . we find that there are classes that can even confuse the oracle model, e.g., instances of estate are tending to be classified into politics and finance, indicating that the mis-identification is partially due to the inherent class similarity. these findings suggest that the performance muki can be limited when the integrating teacher models whose classes are highly correlated. as a remedy, the proposed instance re-weighting mechanism is an effective ad-hoc strategy, as the average teacher uncertainty margin v(x) is 0.04, which can greatly reduce the negative impact of these instances. developing better model uncertainty estimation techniques like incorporating more class information into the estimation process is also promising. besides, the uncertainty estimation can also be influenced by the model capacity. as deeper models tend to produce more confident predictions, the model selection based on uncertainty scores will favor stronger teacher models when there exist particularly weak teacher models. in such a case, applying model calibration techniques like temperature scaling according to the teacher model size @xcite before estimating the uncertainty can be beneficial. ## conclusion in this paper, we explore knowledge integration for plms to promote better model reuse. we present muki, which integrates teacher predictions according to the model uncertainty estimated via monte-carlo dropout, and dynamically adjusts the instance contribution according to the uncertainty margin. extensive results on benchmark datasets demonstrate that muki can substantially outperform strong baselines, and perform well in challenging settings such as merging heterogeneous teachers. further investigation shows that muki can be extended to sequence labeling. in the future, we are interested in developing better integration frameworks for more complex tasks.