MiniPLM: Knowledge Distillation for Pre-Training Language Models

📄 arXiv: 2410.17215v3 📥 PDF

作者: Yuxian Gu, Hao Zhou, Fandong Meng, Jie Zhou, Minlie Huang

分类: cs.CL

发布日期: 2024-10-22 (更新: 2025-03-19)

备注: ICLR 2025

🔗 代码/项目: GITHUB

💡 一句话要点

提出MiniPLM，通过知识蒸馏提升预训练语言模型效率、灵活性和效果。

🎯 匹配领域: 支柱二：RL算法与架构 (RL & Architecture)

关键词: 知识蒸馏 预训练语言模型 模型压缩 数据增强 离线学习

📋 核心要点

  1. 现有预训练知识蒸馏方法存在计算成本高、需要教师和学生模型分词匹配以及可能丢失教师生成数据难度和多样性的问题。
  2. MiniPLM通过离线教师推理提炼训练数据分布，实现高效、灵活且有效的知识蒸馏，无需在线推理或分词匹配。
  3. 实验表明，MiniPLM提升了学生模型在下游任务和语言建模上的性能，并降低了预训练计算成本，且优势可扩展到更大规模。

📝 摘要（中文）

本文提出了一种名为MiniPLM的知识蒸馏框架，用于预训练语言模型。该框架旨在解决现有方法在预训练阶段知识蒸馏时面临的效率、灵活性和有效性问题。MiniPLM通过使用教师语言模型的知识来优化训练数据分布。为提高效率，MiniPLM执行离线教师推理，允许多个学生语言模型进行知识蒸馏，而无需增加训练成本。为提高灵活性，MiniPLM仅在训练语料库上运行，从而实现跨模型族的知识蒸馏。为提高有效性，MiniPLM利用大型和小型语言模型之间的差异来增强训练数据的难度和多样性，帮助学生语言模型获得通用和复杂的知识。大量实验表明，MiniPLM提高了学生语言模型在9个常见下游任务上的性能，提高了语言建模能力，并减少了预训练计算量。MiniPLM的优势扩展到更大的训练规模，这可以从缩放曲线外推中得到证明。进一步的分析表明，MiniPLM支持跨模型族的知识蒸馏，并提高了预训练数据的利用率。

🔬 方法详解

问题定义：现有知识蒸馏方法在预训练语言模型时，面临效率、灵活性和有效性的挑战。在线知识蒸馏计算成本高昂，因为需要在训练过程中实时进行教师模型的推理。一些方法要求教师和学生模型使用相同的分词器，限制了模型选择的灵活性。此外，直接使用教师模型生成的数据进行训练，可能丢失原始训练数据的难度和多样性，影响学生模型的泛化能力。

核心思路：MiniPLM的核心思路是通过教师模型对原始训练数据进行“提炼”，生成更适合学生模型学习的训练数据分布。具体来说，教师模型用于评估原始数据的质量和难度，并据此调整训练数据的采样概率，从而使学生模型能够更有效地学习知识。这种离线提炼的方式避免了在线推理的计算开销，并且不依赖于教师和学生模型的分词器匹配。

技术框架：MiniPLM框架主要包含以下几个阶段：1) 教师模型推理：使用预训练好的大型教师语言模型对原始训练语料库进行推理，得到每个样本的概率分布或损失值等信息。2) 数据分布调整：根据教师模型的推理结果，调整训练数据的采样概率。例如，可以提高教师模型预测困难的样本的采样概率，或者降低教师模型预测过于简单的样本的采样概率。3) 学生模型预训练：使用调整后的训练数据分布，对小型学生语言模型进行预训练。

关键创新：MiniPLM的关键创新在于其离线数据提炼的方式。与传统的在线知识蒸馏方法相比，MiniPLM避免了在线推理的计算开销，从而可以更高效地进行预训练。此外，MiniPLM不依赖于教师和学生模型的分词器匹配，从而可以更灵活地选择模型架构。通过调整训练数据分布，MiniPLM可以更好地利用原始训练数据的难度和多样性，从而提高学生模型的泛化能力。

关键设计：MiniPLM的关键设计包括：1) 数据难度评估：可以使用教师模型的交叉熵损失或困惑度等指标来评估数据的难度。2) 采样概率调整：可以使用不同的函数来调整采样概率，例如指数函数或sigmoid函数。3) 超参数选择：需要根据具体的任务和数据集选择合适的超参数，例如教师模型的权重、采样概率调整的力度等。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

实验结果表明，MiniPLM在9个下游任务上显著提升了学生模型的性能。例如，在某些任务上，MiniPLM可以将学生模型的性能提升超过2个百分点。此外，MiniPLM还提高了学生模型的语言建模能力，并减少了预训练所需的计算资源。缩放曲线外推显示，MiniPLM的优势可以扩展到更大的训练规模。

🎯 应用场景

MiniPLM可应用于各种需要小型化、高性能语言模型的场景，如移动设备上的自然语言处理、资源受限环境下的文本生成、以及需要快速部署的AI应用。通过知识蒸馏，可以降低模型大小和计算复杂度，同时保持较高的性能水平，加速AI技术在边缘设备和特定领域的落地。

📄 摘要（原文）

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 efficiency, flexibility, and effectiveness issues. 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. In this work, we propose MiniPLM, a KD framework for pre-training LMs by refining the training data distribution with the teacher LM's knowledge. For efficiency, MiniPLM performs offline teacher inference, allowing KD for multiple student LMs without adding training 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 training data difficulty and diversity, helping student LMs acquire versatile and sophisticated knowledge. Extensive experiments demonstrate that MiniPLM boosts the student LMs' performance on 9 common downstream tasks, improves language modeling capabilities, and reduces pre-training computation. The benefit of MiniPLM extends to larger training scales, evidenced by the scaling curve extrapolation. Further analysis reveals that MiniPLM supports KD across model families and enhances the pre-training data utilization. Our code, data, and models can be found at https://github.com/thu-coai/MiniPLM.