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

发表时间: 2024-10 · arXiv:2410.16215 (Tsinghua University / Zhipu AI)

原文: https://arxiv.org/abs/2410.16215

Hao Peng, Xin Lv, Yushi Bai, Zijun Yao, Jiajie Zhang, Lei Hou, Juanzi Li
Tsinghua University, Zhipu AI

速读

一句话结论 本文将知识蒸馏技术扩展到大语言模型的预训练阶段，系统探索了逻辑值处理、损失函数选择、缩放定律及在线/离线策略的设计空间，并总结出一套能稳定提升学生模型性能的最优预训练蒸馏配置。

要解决什么问题 现有的知识蒸馏技术大多局限于大语言模型的后训练阶段，即让学生模型直接学习教师模型生成的指令回复。然而，将蒸馏引入消耗算力最大的预训练阶段面临着严重的机制卡点。首先是极端的存储与显存开销：预训练涉及千亿甚至万亿级别的词元（tokens），如果直接保存大词表（如 150k 规模）教师模型输出的全量逻辑值（logits），仅 100B 词元就需要约 58.6 PB 的磁盘空间，工程上完全不可行。其次是训练动态的未知性：在从头预训练时，传统的语言模型损失与蒸馏损失该如何分配，以及均方误差等传统蒸馏损失是否依然有效，都缺乏机制层面的验证。最后是模型容量差距带来的外推失效风险：盲目使用参数量过大的教师模型，反而可能因为学生模型拟合能力不足而导致蒸馏效果下降，业界亟需一份关于预训练蒸馏的缩放定律（Scaling Law）指南。

怎么做的 核心思路是在预训练阶段引入基于逻辑值的知识蒸馏，让学生模型在海量预训练语料上，不仅学习真实的独热标签，还学习教师模型输出的软标签。教师的逻辑值包含了更丰富的词汇分布信息，起到了标签平滑的作用，从而加速收敛并提升性能。为了绕开存储卡点并优化训练，方法设计了三个关键部件。第一是逻辑值截断与归一化。为了将存储需求降低数千倍，作者设计了二阶段的 top-p-k 截断法：先保留累积概率达到 $p$ 的逻辑值，再从中最多保留 $k$ 个，最后使用温度参数 $\tau$ 重新归一化。该过程定义为：

$$F(z) = \operatorname{softmax}\left(\frac{\operatorname{Truncate}(z)}{\tau}\right)$$


第二是损失函数的组合调度。整体优化目标是标准语言模型损失 $\mathcal{L}_{\mathrm{lm}}$ 与蒸馏损失 $\mathcal{L}_{\mathrm{kd}}$ 的加权和：

$$\theta_S^* = \arg\min_{\theta_S} [(1 - \alpha) \mathcal{L}_{\mathrm{lm}} + \alpha \mathcal{L}_{\mathrm{kd}}]$$
其中蒸馏损失函数选用 KL 散度或负对数似然（实验证明均方误差 MSE 会导致性能崩塌）。为了最大化收益，作者引入了 WSD（预热-稳定-衰减）调度策略来动态调整蒸馏权重 $\alpha$：在学习率处于最高值的稳定期，赋予蒸馏损失极高的权重（如 0.9），而在训练末期随学习率一起衰减。这能让模型在早期充分吸收教师知识，后期再靠真实标签微调。第三是逻辑值的获取策略，分为离线（使用已收敛教师模型的输出）和在线（在教师模型预训练的同时实时保存输出）两种，前者质量更高，后者则能节省额外的推理算力。

效果如何 实验在 Nvidia H800 GPU 上进行，采用 GLM-4-9B 和 GLM-4-32B 作为教师模型，从头预训练 330M 到 6.8B 规模的学生模型，预训练数据量覆盖 100B 到 500B tokens。对比的基线方法是仅使用标准语言模型损失进行从头预训练的常规路线（即 LLM-LM）。在 8 个涵盖中英文理解、常识推理和数学的任务（如 MMLU、GSM8k、C-Eval 等）中，量化结果证明了方法的有效性。在 100B tokens 的设置下，使用探索出的最优配置（top-0.95-50 截断、温度 2.0、KL 散度结合 WSD 动态调度），1.9B 规模的学生模型平均准确率从基线的 37.7% 提升至 41.2%。在扩大到 500B tokens 的长周期预训练中，蒸馏带来的性能增益贯穿始终，证明其不仅能加速收敛，还能切实提升模型的性能上限。此外，实验揭示了几个关键的代价与失效场景：首先是模型容量差距导致的失效，使用 32B 教师指导 1.9B 学生的效果，反而不如使用 9B 教师，说明教师并非越大越好；其次，如果采用在线策略获取逻辑值，使用教师模型预训练早期的未收敛输出会导致学生模型性能暴跌 20.9%；最后，尽管采用了极端的截断策略，100B tokens 的逻辑值仍需占用约 15 TB 的磁盘空间，这对存储 I/O 依然是一个不可忽视的工程代价。

主要贡献

知识蒸馏（KD）旨在将大型教师模型的知识转移到较小的学生模型中。以往在大型语言模型（LLMs）领域应用知识蒸馏通常集中在后训练（post-training）阶段，即学生模型直接从指令和教师模型生成的相应回复中学习。本文的核心问题在于探讨是否可以在LLMs的预训练阶段应用知识蒸馏。

为此，本文的研究目标是将知识蒸馏扩展到LLMs的预训练阶段，提出预训练蒸馏（Pre-training Distillation, PD），并系统性地探索其设计空间。本文的创新点包括：

  1. 通过使用GLM-4-9B作为教师模型蒸馏1.9B参数的学生模型进行初步实验，验证了预训练蒸馏的有效性。
  2. 系统性地探索了预训练蒸馏设计空间的四个关键影响因素：逻辑值（logits）处理、损失函数选择、缩放定律（scaling law）以及离线或在线逻辑值。
  3. 发现了一系列更优的配置和重要结论，例如：较大的学生模型通常能从预训练蒸馏中获益更多，而更大的教师模型并不一定能保证更好的蒸馏结果。
图1：预训练1.9B、3.8B和6.8B学生LLM的结果，仅使用LM损失、基础PD配置（§3.1）以及经过我们探索后更好的PD配置（PD*）。详情见附录A.6。
图1：预训练1.9B、3.8B和6.8B学生LLM的结果，仅使用LM损失、基础PD配置（§3.1）以及经过我们探索后更好的PD配置（PD*）。详情见附录A.6。

背景知识与设计原则

目标函数形式化
给定文本 $\mathbf{x} = \{x_t\}_{t=1}^T$，参数化为 $\theta_S$ 的学生LLM，以及参数化为 $\theta_T$ 的教师LLM，预训练蒸馏的目标函数形式化如下：

$$ \theta_S^* = \arg \min_{\theta_S} \mathcal{L} = \arg \min_{\theta_S} [(1 - \alpha) \mathcal{L}_{\mathrm{lm}} + \alpha \mathcal{L}_{\mathrm{kd}}] $$

语言建模损失
其中 $\mathcal{L}_{\mathrm{lm}}$ 表示传统的one-hot语言建模（LM）预训练损失，形式化为：

$$ \mathcal{L}_{\mathrm{lm}} = \frac{1}{T} \sum_{t=1}^T -\log P_{\theta_S}(x_t | \mathbf{x}_{<t}) $$ <p>蒸馏损失
$\mathcal{L}_{\mathrm{kd}}$ 表示蒸馏损失，形式化为：

$$ \mathcal{L}_{\mathrm{kd}} = \frac{1}{T} \sum_{t=1}^T L(P_{\theta_S}(x_t | \mathbf{x}_{<t}), F(P_{\theta_T}(x_t | \mathbf{x}_{<t}))) $$ <p>
其中 $L$ 表示蒸馏损失函数（如Kullback-Leibler散度）。$P_{\theta_S}$ 和 $P_{\theta_T}$ 分别代表学生和教师LLM的概率。

逻辑值处理函数
$F$ 表示对教师LLM逻辑值进行的截断和归一化操作，$\tau$ 是归一化的温度参数：

$$ F(z) = \mathrm{softmax}(\frac{\mathrm{Truncate}(z)}{\tau}) $$

基于上述公式，本文在四个维度上探索预训练蒸馏的设计空间：（1）处理教师逻辑值的方法 $F$（截断方法和温度 $\tau$）；（2）损失函数的选择（蒸馏损失 $L$ 和组合系数 $\alpha$）；（3）预训练蒸馏的缩放定律（模型大小和语料库大小）；（4）获取教师概率的策略（离线生成或在线同步生成）。

方法细节

逻辑值处理（Logits Processing）

损失函数选择（Loss Selection）

缩放定律（Scaling Law）

离线或在线（Offline or Online）

实验环境

实验结果

  1. 初步实验验证（表1）：使用GLM-4-9B蒸馏1.9B模型（100B token，NLL损失，top-0.95-100截断）。LLM-KD的平均得分为38.3，略高于仅用LM损失的LLM-LM（37.7），相对提升1.6%，验证了PD的可行性。
  2. 逻辑值处理实验（图2/图3/表2/表3）：top-$p$-$k$ 截断中，不同 $p$ 值差异不大，$k=50$ 表现最佳。静态温度 $\tau \leq 2.0$ 效果较好。自适应温度 AdaKD$_H$ 表现最好（平均38.8），但相比静态 $\tau=0.5$（平均38.7）提升不明显。
  3. 损失函数选择实验（表4/表5）：KLD损失（平均38.7）优于NLL（38.3），MSE导致性能严重下降（34.9）。静态 $\alpha=0.9$ 表现最佳（39.1）。动态调度中，WSD-$\alpha$ 结合 WSD-LR 达到了最佳性能（平均40.7），相比基线相对提升8.0%。
  4. 缩放定律实验（图4/图5）：模型规模实验证实大参数学生模型获益更多，且9B教师在蒸馏小模型时有时优于32B教师。500B token的扩展实验证实，随着训练数据增加，PD带来的性能提升能够持续保持。
  5. 在线与离线实验（表6）：直接使用未收敛的在线逻辑值（LLM-Online-100B-L，平均29.8）性能崩溃。调整超参数后（LLM-Online-100B*，平均37.9），略高于基线（37.7），但不如离线蒸馏。
  6. 最优配置验证（图1/表13）：结合探索得出的更优配置（PD*：top-0.95-50截断，$\tau=2.0$，KLD损失，WSD-$\alpha$最大值为0.9，WSD-LR，离线逻辑值），在1.9B、3.8B和6.8B规模上均显著优于Vanilla PD和LM基线。

结论

本文系统地探索了预训练蒸馏（PD）的设计空间，涵盖逻辑值处理、损失选择、缩放定律以及获取逻辑值的策略（离线或在线）。通过大量实验，确定了更优的配置方案，并得出了重要结论：较大的学生模型通常能从PD中获益更多，而更大的教师模型并不一定能保证更好的结果。这些探索有望为未来大语言模型的预训练蒸馏实践提供指导。

附录与补充细节