Pre-training Distillation for Large Language Models: A Design Space Exploration

📄 arXiv: 2410.16215v1 📥 PDF

作者: Hao Peng, Xin Lv, Yushi Bai, Zijun Yao, Jiajie Zhang, Lei Hou, Juanzi Li

分类: cs.CL, cs.AI

发布日期: 2024-10-21

💡 一句话要点

探索预训练蒸馏设计空间，提升大语言模型知识迁移效率

🎯 匹配领域: 支柱二：RL算法与架构 (RL & Architecture) 支柱九：具身大模型 (Embodied Foundation Models)

关键词: 大语言模型 知识蒸馏 预训练 模型压缩 迁移学习

📋 核心要点

  1. 现有知识蒸馏方法主要集中在LLM的后训练阶段，忽略了预训练阶段的知识迁移潜力。
  2. 本文提出预训练蒸馏(PD)方法，将知识蒸馏应用于LLM的预训练阶段，提升学生模型的学习效率。
  3. 实验结果表明，PD方法有效，并深入探索了logits处理、损失函数等关键因素对蒸馏效果的影响。

📝 摘要（中文）

知识蒸馏(KD)旨在将知识从大型教师模型迁移到较小的学生模型。以往在大型语言模型(LLM)领域应用KD的工作通常集中在后训练阶段，即学生LLM直接从教师模型生成的指令和相应响应中学习。本文将KD扩展到LLM的预训练阶段，称为预训练蒸馏(PD)。我们首先使用GLM-4-9B作为教师LLM来蒸馏一个1.9B参数的学生LLM，验证了PD的有效性。考虑到蒸馏的关键影响因素，我们系统地探索了预训练蒸馏在四个方面的设计空间：logits处理、损失选择、缩放定律以及离线或在线logits。我们进行了广泛的实验来探索预训练蒸馏的设计空间，并发现了更好的配置和有趣的结论，例如，较大的学生LLM通常从预训练蒸馏中获益更多，而较大的教师LLM并不一定保证更好的结果。我们希望我们对设计空间的探索能够为未来的预训练蒸馏实践提供信息。

🔬 方法详解

问题定义：现有的大语言模型知识蒸馏方法主要集中在后训练阶段，即在指令微调阶段利用大模型生成的数据来训练小模型。这种方法忽略了预训练阶段的知识迁移，可能导致学生模型无法充分学习到教师模型的通用语言能力。因此，如何有效地在预训练阶段进行知识蒸馏，提升学生模型的初始能力，是一个重要的研究问题。

核心思路：本文的核心思路是将知识蒸馏技术应用于大语言模型的预训练阶段。通过让学生模型在预训练阶段学习教师模型的logits输出，从而使学生模型能够更好地学习到教师模型的知识表示和语言模式。这种方法旨在提升学生模型在下游任务上的泛化能力和学习效率。

技术框架：该研究的技术框架主要包括以下几个步骤：1) 使用大型教师模型（如GLM-4-9B）在预训练数据集上进行推理，生成logits输出。2) 使用较小的学生模型（如1.9B参数模型）在相同的预训练数据集上进行训练，并同时计算学生模型的logits输出。3) 使用知识蒸馏损失函数（如KL散度）来衡量学生模型和教师模型logits输出之间的差异，并将其作为训练目标之一。4) 通过优化学生模型的参数，使其logits输出尽可能接近教师模型的logits输出，从而实现知识迁移。

关键创新：该研究的关键创新在于将知识蒸馏技术扩展到大语言模型的预训练阶段。与传统的后训练蒸馏方法相比，预训练蒸馏能够更早地将教师模型的知识融入到学生模型中，从而提升学生模型的初始能力和泛化性能。此外，该研究还系统地探索了预训练蒸馏的设计空间，包括logits处理、损失函数选择、缩放定律等关键因素，为未来的研究提供了有价值的指导。

关键设计：在预训练蒸馏过程中，几个关键的设计选择包括：1) Logits处理：研究了不同的logits处理方法，例如温度缩放，以调整教师模型logits的平滑程度。2) 损失函数：比较了不同的知识蒸馏损失函数，例如KL散度、交叉熵等，以选择最适合预训练阶段的损失函数。3) 缩放定律：研究了教师模型和学生模型大小对蒸馏效果的影响，发现更大的学生模型通常受益更多。4) 离线/在线Logits：探索了使用离线预计算的教师模型logits和在线计算的logits对蒸馏效果的影响。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

该研究通过实验验证了预训练蒸馏的有效性，并系统地探索了预训练蒸馏的设计空间。实验结果表明，更大的学生模型通常从预训练蒸馏中获益更多，而更大的教师模型并不一定保证更好的结果。此外，研究还发现，合适的logits处理方法和损失函数选择对蒸馏效果有显著影响。这些发现为未来的预训练蒸馏实践提供了重要的指导。

🎯 应用场景

该研究成果可应用于各种需要部署小型化大语言模型的场景，例如移动设备、边缘计算设备等。通过预训练蒸馏，可以在保证模型性能的前提下，显著减小模型体积，降低计算成本，从而实现大语言模型在资源受限环境下的高效应用。此外，该研究也为未来大语言模型知识迁移和模型压缩提供了新的思路。

📄 摘要（原文）

Knowledge distillation (KD) aims to transfer knowledge from a large teacher model to a smaller student model. Previous work applying KD in the field of large language models (LLMs) typically focused on the post-training phase, where the student LLM learns directly from instructions and corresponding responses generated by the teacher model. In this paper, we extend KD to the pre-training phase of LLMs, named pre-training distillation (PD). We first conduct a preliminary experiment using GLM-4-9B as the teacher LLM to distill a 1.9B parameter student LLM, validating the effectiveness of PD. Considering the key impact factors of distillation, we systematically explore the design space of pre-training distillation across four aspects: logits processing, loss selection, scaling law, and offline or online logits. We conduct extensive experiments to explore the design space of pre-training distillation and find better configurations and interesting conclusions, such as larger student LLMs generally benefiting more from pre-training distillation, while a larger teacher LLM does not necessarily guarantee better results. We hope our exploration of the design space will inform future practices in pre-training distillation.