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

TL;DR

This paper introduces a novel knowledge integration framework for pre-trained language models that merges multiple teacher models into a versatile student without human annotations, improving performance and generalization.

Contribution

It proposes the MUKI framework that uses model uncertainty and instance-wise re-weighting to effectively merge diverse teacher models into a single student model.

Findings

MUKI outperforms baseline methods on benchmark datasets.

It effectively merges heterogeneous and cross-lingual teacher models.

Demonstrates strong generalization capabilities.

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.