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. 发现了一系列更优的配置和重要结论，例如：较大的学生模型通常能从预训练蒸馏中获益更多，而更大的教师模型并不一定能保证更好的蒸馏结果。

背景知识与设计原则

目标函数形式化
给定文本 $\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）

• 逻辑值截断（Logits Truncation）：因为存储整个词表的逻辑值需要极大的磁盘空间，作者设计了一种两阶段的 top-$p$-$k$ 截断方法：先使用 top-$p$ 截断，接着使用 top-$k$ 截断。当逻辑值分布尖锐时，top-$p$ 截断就足够了；当分布较均匀且存在长尾非平凡值时，top-$k$ 截断作为二次截断发挥作用。与基础的 top-$p$ 和 top-$k$ 相比，top-$p$-$k$ 方法显著减少了存储空间。实验固定 $k=100$ 研究不同 $p$ 的影响，发现不同的 $p$ 值带来的性能提升相似。因为学生LLM主要捕捉逻辑值的质量部分，这表明可以使用较小的 $p$ 来进一步减少存储空间。接着固定 $p=0.95$ 研究不同 $k$ 的影响，发现所有 $k$ 值都能带来提升，其中 $k=50$ 效果最好。值得注意的是，当 $k=1$ 时（等同于AFM预训练中采用的方法【21，Gunter 等人+2024+Apple intelligence foundation language models】），也带来了性能提升，这可能是因为教师LLM在预训练语料库中进行了隐式的噪声过滤。总体而言，采用较小的 $p$ 和 $k$ 可以在保证性能提升的同时节省存储空间。
  
  

• 温度 $\tau$：较低的温度会使逻辑值分布变尖锐，较高的温度会使分布更均匀。实验首先测试了静态 $\tau$，发现较低温度（$\tau \leq 2.0$）带来相似的提升，而较高温度（$\tau \geq 5.0$）提升有限，说明从过于均匀的分布中学习对学生模型效率不高。随后探索了自适应温度（根据每个样本动态调整），测试了NormKD【8，Chi 等人+2023+Normkd】和WTTM【56，Zheng and YANG+2024+Knowledge distillation based on transformed teacher matching】。作者还实现了一个紧凑版的自适应温度方法AdaKD，对较尖锐的教师逻辑值施加高温度，对不尖锐的施加低温度。分别使用标准差（AdaKD$_{SD}$）和熵（AdaKD$_H$）来衡量尖锐度。AdaKD$_H$ 的温度计算公式为：$\tau_H = \tau_{\mathrm{max}} - (\tau_{\mathrm{max}} - \tau_{\mathrm{min}}) \times \frac{H}{H_{\mathrm{max}}}$。结果表明，AdaKD$_H$ 在自适应方法中表现最好，但与静态温度 $\tau=0.5$ 相比，自适应温度并未显示出显著的额外收益。

损失函数选择（Loss Selection）

• 蒸馏损失函数 $L$：实验比较了负对数似然（NLL）、Kullback-Leibler散度（KLD）和均方误差（MSE）损失（通过设置 $\alpha=1$ 排除LM损失干扰）。结果显示，使用NLL和KLD训练的LLM均优于仅使用LM损失的基线。虽然KLD整体优于NLL，但NLL在MMLU和C-Eval等更具挑战性的数据集上表现更好。相反，使用MSE损失训练的学生LLM性能显著下降。这一现象与之前的LLM研究一致【39，Muralidharan 等人+2024+Compact language models via pruning and knowledge distillation】，但与图像分类领域的发现相悖【30，Kim 等人+2021+Comparing kullbackleibler divergence and mean squared error loss in knowledge distillation】，表明LLM的预训练蒸馏涉及新的训练动态。
• $\mathcal{L}_{\mathrm{lm}}$ 和 $\mathcal{L}_{\mathrm{kd}}$ 的组合：在固定 $\mathcal{L}_{\mathrm{kd}}$ 为NLL损失的前提下，首先测试了静态 $\alpha$（0.0至1.0）。随着 $\alpha$ 增加，性能先升后降，在 $\alpha=0.9$ 时达到最佳。这表明虽然高比例的蒸馏损失能提升性能，但保留适当比例（约10%）的LM损失能进一步增强效果。接着探索了 $\alpha$ 的动态调度策略：（1）线性增加（Linear Inc，从0到1）或线性减少（Linear Dec，从1到0）。（2）周期性变化（Period，每四个批次设为0.9，其余为0【28，Kiefel and Shah+2024+Lokilm】）。（3）Warmup-Stable-Decay（WSD）非线性调度【24，Hu 等人+2024+Minicpm】，记为WSD-$\alpha$。具体为：预热阶段从0升至1.0，保持1.0，最后余弦衰减至0（预热比例10%，衰减比例1%）。同时引入了WSD学习率调度器（WSD-LR）。结果表明：线性减少优于线性增加，说明早期引入更多KD损失更有利；WSD-LR通常能带来收益；WSD-$\alpha$ 结合 WSD-LR 取得了最佳性能。WSD-$\beta$（即对LM损失比例进行WSD调度）效果有限，这证明了在保持最大学习率时使用KD损失能有效提升模型性能。

缩放定律（Scaling Law）

• 模型大小：采用9B和32B的教师模型来蒸馏330M、670M、1.9B、3.8B和6.8B的学生模型。结果表明：（1）越大的学生模型通常从预训练蒸馏中获益越多。（2）从更大的教师模型蒸馏并不一定带来更好的性能。这可能是由于师生模型之间存在容量差距【35，Mirzadeh 等人+2020+Improved knowledge distillation via teacher assistant】；【18，Gou 等人+2021+Knowledge distillation: A survey】。从压缩角度看，大模型压缩信息更有效【11，Deletang 等人+2024+Language modeling is compression】，导致小模型更难学习。实验表明，当学生模型大小达到教师模型的约10%或以上时，预训练蒸馏是有效的，并且随着比例增加，收益随之增长。
  

• 语料库大小：使用GLM-4-9B作为教师，在5000亿（500B）token上蒸馏1.9B和3.8B学生模型，每10000步（约83B token）保存一次检查点。结果显示：与仅用LM损失训练的模型相比，预训练蒸馏在整个预训练过程中始终产生提升；蒸馏带来的增益在初期增加，随后收敛并略有下降，但在预训练结束时仍然显著。这表明PD不仅提高了训练效率，还提升了学生模型的性能上限。
  

离线或在线（Offline or Online）

• 探讨了逻辑值的获取方式。离线指从预训练好的教师模型获取；在线指在教师模型预训练的同时保存生成的逻辑值。在线方式的优势在于没有额外的推理成本，且类似于课程学习【46，Soviany 等人+2022+Curriculum learning: A survey】。作者从头预训练了一个9B教师模型（400B token）并实时保存逻辑值。使用前100B和后100B token分别训练了LLM-Online-100B-L和LLM-Online-100B，结果表现很差，原因是教师模型远未收敛，逻辑值包含大量噪声。随后调整策略，设置 $\alpha=0.1$ 并使用 top-0.95-50 截断训练了 LLM-Online-100B*，其性能略好于仅用LM损失的基线，但仍低于使用离线逻辑值的模型。这表明未收敛教师的逻辑值也能起作用，但最佳实践是使用教师预训练后期的逻辑值。

实验环境

• 数据集：涵盖英语语言理解与常识推理（HellaSwag, WinoGrande, PIQA, MMLU）、中文语言理解与常识推理（KBQA, C3, C-Eval）以及数学数据集（GSM8k）。
• 模型架构参数：学生模型规模从330M到6.8B。以1.9B模型为例：隐藏层大小2048，FFN隐藏层大小6912，24层，16个注意力头，2个查询组（采用GQA），不共享词嵌入与输出权重。教师模型采用GLM-4-9B和GLM-4-32B。
• 硬件配置：所有实验均在Nvidia H800 GPU上进行。存储100B token的截断逻辑值需约15TB磁盘空间；在线保存400B token逻辑值消耗约180TB空间。
• 软件配置：采用Adam优化器，批大小为2048，最大序列长度为4096。余弦学习率调度器（最大学习率 $6 \times 10^{-4}$，最小 $6 \times 10^{-5}$，1%预热）。模型以BFLOAT16格式训练。由于预训练模型在某些数据集上接近随机猜测，评估前在10B高质量指令微调数据上进行了监督微调（SFT）。

实验结果

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中获益更多，而更大的教师模型并不一定能保证更好的结果。这些探索有望为未来大语言模型的预训练蒸馏实践提供指导。

附录与补充细节

• SFT微调细节：在评估前的SFT阶段，混合使用了10B高质量指令微调数据和额外的10B预训练文本语料。指令数据仅对回复部分计算LM损失。采用256批大小，余弦学习率（最大 $4 \times 10^{-5}$，最小 $4 \times 10^{-6}$，1%预热）。
• 评估细节：HellaSwag、WinoGrande、PIQA和KBQA采用zero-shot评估；C3和C-Eval采用5-shot；MMLU采用6-shot；GSM8k采用8-shot。采样温度设为0。
• 自适应温度参数：NormKD设置超参数 $T_{\mathrm{norm}}=1.0$ 且 $\alpha=0.5$。WTTM设置 $\gamma=0.1$ 且 $\beta=1.0$。AdaKD$_H$ 中，最大熵 $H_{\mathrm{max}}$ 在1000万个token上估计为4.8，设置 $\tau_{\mathrm{max}}=2.0$，$\tau_{\mathrm{min}}=0.1$。
• 在线教师预训练：从头训练的9B教师模型采用1728批大小，4096最大序列长度，余弦学习率（最大 $6 \times 10^{-4}$，最小 $6 \times 10^{-5}$，1%预热）。由于成本高昂，仅使用了400B token。
• 局限性：受限于计算成本，本文未探索预训练蒸馏中不同因素之间的交互作用（即因素组合）。目前的控制变量实验已经产生了高昂的计算成本。
• 伦理考量：严格遵守所有使用模型和数据集的版权许可；数据已适当匿名化；使用ChatGPT进行部分句子的改写和语法检查。