MiniPLM: Knowledge Distillation for Pre-Training Language Models
Authors: name: Yuxian Gu; name: Hao Zhou; name: Fandong Meng; name: Jie Zhou; name: Minlie Huang
Published: 2024-10-22 17:40:32

Abstract: Knowledge distillation (KD) is widely used to train small, high-performing
student language models (LMs) using large teacher LMs. While effective in
fine-tuning, KD during pre-training faces challenges in efficiency,
flexibility, and effectiveness. Existing methods either incur high
computational costs due to online teacher inference, require tokenization
matching between teacher and student LMs, or risk losing the difficulty and
diversity of the teacher-generated training data. To address these issues, we
propose MiniPLM, a KD framework for pre-training LMs by refining the training
data distribution with the teacher's knowledge. For efficiency, MiniPLM
performs offline teacher LM inference, allowing KD for multiple student LMs
without adding training-time costs. For flexibility, MiniPLM operates solely on
the training corpus, enabling KD across model families. For effectiveness,
MiniPLM leverages the differences between large and small LMs to enhance the
difficulty and diversity of the training data, helping student LMs acquire
versatile and sophisticated knowledge. Extensive experiments demonstrate that
MiniPLM boosts the student LMs' performance on 9 widely used downstream tasks,
improves the language modeling capabilities, and reduces pre-training
computation. The benefit of MiniPLM extends to large pre-training scales,
evidenced by the extrapolation of the scaling curves. Further analysis reveals
that MiniPLM supports KD across model families and enhances the utilization of
pre-training data. Our model, code, and data are available at
https://github.com/thu-coai/MiniPLM.