Knowledge Distillation of Large Language Models

关于大模型系列模型的文章，是针对《Knowledge Distillation of Large Language Models》的翻译。

大模型的知识蒸馏

• 摘要

• 引言部分

• 方法论部分

• 2.1 基于逆向Kullback-Leibler散度的知识蒸馏机制

• 2.2 基于策略梯度的优化方法

• 2.3 基于策略梯度的训练机制

  • 3 实验

  • • 3.1 实验设置
    • 3.2 结果
    • 3.3 分析
    • 3.4 消融实验

  • 4 相关工作

  • 5 结论

摘要

知识蒸馏（KD）作为一种新兴技术，展现出巨大的发展潜力，能够有效降低大型语言模型（LLM）的计算需求。然而，以往的KD方法主要集中在白盒分类模型或小型模型的训练，旨在模仿诸如ChatGPT这样的黑盒模型API。在从白盒生成LLM中提取知识方面，这一问题尚待深入探索，随着LLM的迅速发展，这一问题的重要性日益凸显。在本研究中，我们提出了一种名为MINILLM的新方法，该方法能够从生成的大型语言模型中提取出更小的语言模型。具体而言，我们首先对标准KD方法中的前向Kullback-Leibler散度（KLD）目标进行了调整，将其替换为更适合生成语言模型上的KD的反向KLD目标，以避免学生模型在估计教师分布的低概率区域时出现高估现象。接着，我们开发了一种高效的优化算法来学习这一目标函数。通过在指令执行任务中的大规模实验，我们发现MINILLM模型能够生成更精确的响应，展现出更高的整体质量、较低的曝光偏差、更好的校准效果以及更出色的长文本生成性能。此外，该方法适用于参数范围在120 million到13 billion之间的不同模型族。代码和模型预训练版本已将在https://aka.ms/MiniLLM正式发布。

1 引言

随着大型语言模型（LLM；HZD^+，BHA^+，BMR^+，Ope，CND^+）的快速发展，减少高计算资源需求的一种常见技术是知识蒸馏（KD；HVD），在这种技术中，我们在大型教师模型的监督下训练小学生模型。通常应用两类KD：黑盒KD和白盒KD，其中只有教师预测是可访问的，其中教师参数可供使用[JBMD]。最近，黑盒KD在对LLM API[TGZ^+，CLL^+，WWZ^+，PLH^+]生成的提示响应对的小模型进行微调方面显示出了有希望的结果。随着更多开源LLM[ZRG^+，TLI^+]的出现，白盒KD对研究界和行业部门都变得更有价值，因为学生模型从白盒教师模型中接收到更好的信号，从而可能带来更高的性能。然而，白盒KD方法主要用于小（<1B参数）语言理解模型[SDCW，WWD^+]，而用于生成LLM的白盒KD尚待探索。
在这项工作中，我们研究了LLM的白盒KD。我们认为，标准KD目标[KR16，SST^+]对于以生成方式执行任务的LLM来说是次优的。给定由\theta参数化的教师分布p(y|x)和学生分布q_\theta(y|x)，标准KD目标（包括序列级模型的几个变体）本质上最小化了教师和学生分布之间的近似前向Kullback-Leibler散度（KLD），称为KL[p||q_\theta]，这迫使p覆盖q_\theta的所有模式。对于文本分类任务，KL[p||q_\theta]工作得很好，因为输出空间通常由有限个类组成，使得p(y|x)和q_\theta(y|x)都有几个模式。然而，对于开放文本生成任务，其中输出空间要复杂得多，并且由于模型容量有限，p(y|x)可以包含比q_\theta(y|x)所能表达的模式多得多的模式。最小化正向KLD可能导致q_\theta将不合理的高概率分配给p[MG]的空位区域，并在自由运行生成期间在p下产生非常不可能的样本[Hus]。

为了解决这一问题，我们推荐最小化反向KLD，即KL[q_\theta||p]，这一指标在计算机视觉[KW]和强化学习[CPO^+]领域中被广泛应用。与KL[p||q_\theta]相比，最小化KL[q_\theta||p]会导致q_\theta集中于p的主要模式，并将低概率分配给p的空穴区域[M^+]，如第2.1节所示。在LLM的语言生成中，这意味着学生模型避免过度学习教师分布的长尾变体[HBD^+]，并更加关注生成响应的准确性，这在需要真实性和可靠性的场景中尤为重要[JLF^+]。如第2.2节所示，为了优化\min_\theta KL[q_\theta||p]，我们采用策略梯度[SSMSM]来推导目标的梯度。尽管最近的研究表明，通过策略优化在微调PLM方面取得了成功[OWJ^+，RAB^+]，但我们发现，模型训练仍面临高方差、奖励黑客和生成长度偏差等问题。因此，我们引入了（1）单步正则化以降低方差，（2）教师混合采样以缓解奖励黑客行为，以及（3）长度归一化以消除长度偏差。最后，我们在第2.3节中详细介绍了MINLLM算法。

2 方法

针对条件语言生成任务，模型通过从分布p_x采样的输入x生成响应序列y=\{y_t\}^{T}_{t=1}。我们将KD转化为一个优化问题，以最小化固定教师模型分布p(y|x)与由\theta参数化的学生模型分布q_\theta(y|x)之间的差异。生成模型的标准KD方法特别适合最小化前向KLDKL[p||q_\theta]=\mathbb{E}_{x\sim p_x,y\sim p'}\log\frac{p(y|x)}{q_\theta(y|x)}，其中p'可以是真实数据分布（即单词级KD）或教师分布p（即序列级KD）。然而，尽管KL[p||q_\theta]被广泛采用，但当q_\theta无法覆盖p'[JKH^+]的所有模式时，它在语言生成任务中高估了p的空位区域。LLM的KD方法特别适合这种情况，因为LLM以生成的方式执行各种任务，使得低能力的学生模型不能完美地模仿教师模型或人类的复杂语言生成分布。

2.1 MiniLLM：利用逆向KLD进行知识蒸馏

在本研究中，我们以反向Kullback-Leibler散度作为MINILLM的学习目标，旨在最小化学生和教师分布之间的差异。

2.2 策略梯度优化

梯度推导

单步正则化

2.3 训练算法

MINLLM的训练算法如算法2.3所示。基于具有最低验证损失的训练数据集进行微调的初始点设置用于学生模型，并对(\nabla\mathcal{J})_{Main}引入PPO裁剪策略以增强训练稳定性。值得注意的是，我们在PPO训练过程中未采用值网络和KL正则化以提升训练效率。与[OWJ^+]模型相似，我们在预训练语料库上增加了语言建模损失\mathcal{L}_{PT}这一关键组件。

3 实验

3.1 实验设置

我们通过首先在教学响应数据集\mathcal{D}上微调一个大模型作为教师p来进行实验。然后，我们比较了不同的KD方法，通过评估蒸馏模型的指令跟随性能，在教师的指导下蒸馏\mathcal{D}上的一个较小的学生模型。
基础模型
我们提取了三种不同大小的模型：GPT-2（120M、340M、760M）、OPT（1.3B、2.7B、6.7B）和LLaMA（7B），分别使用GPT-2-1.5B、OPT13B和LLaMA-13B作为每种模型类型的教师。我们还在附录C.1中介绍了使用GPT-J作为教师的结果。
训练
我们从由15K人写指令-响应对组成的databricks-doolly-15k构建训练数据。我们随机分割14K个样本作为训练集\mathcal{D}，并分别留下500个样本进行验证和测试。对于\mathcal{D}_{PT}，我们对GPT-2家族使用OpenWebText，对其他模型使用RoBERTa训练语料库。在整个实验中，我们将教师的混合强度设定为\alpha=0.2。我们在验证集上使用Rouge-L分数来选择超参数，因为它比验证损失更符合人类偏好。更多训练细节见附录B.1。
评估
我们在以下五个指令数据集上评估训练后的模型：

• DollyEval：基于databricks-dolly-15k数据集，我们进行了500个样本的测试集抽样分离。
• SelfInst：我们开发了一个用户友好的指令集合，其中包含了252个样本。
• VicunaEval：在Vicuna评估框架中，我们采用了80个具有挑战性的评估问题。
• S-NI：在[PLH^+]之后，我们将测试集划分为3个子集，其基本事实响应长度分别为[0，5]、[6，10]和[11，+\infty]区间。我们在第3.2节中使用了[11，+\infty]子集，并在第3.3节中对所有子集进行了详细分析。
• UnNI：作为UNNATURALINSTRUCTIONS的核心集合，我们收集了60K个样本。与S-NI类似，我们首先对[11，+\infty]子集进行了评估，然后对附录C.2中的所有子集进行了性能分析。

我们采用两个指标来评估模型生成的响应：

• R-L：Rouge-L分数，用于评估生成内容的准确性。研究表明，[WMA^+]表明Rouge-L在大规模教学跟踪评估中表现出色。
• GPT4：GPT-4的反馈要求对生成的回答进行评估，并分别给予1-10分的评分（具体评分依据见附录B.2）。我们计算了模型回答与基本事实回答得分的总分比值。此评估指标仅适用于DollyEval、SelfInst和VicunaEval。

在所有测试集上，我们对温度参数设置为1的响应进行采样操作，并对不同随机种子下的每个提示进行5代的评估，计算其平均得分。

基于无KD的SFT方法，直接对数据集D上的学生模型进行训练，并采用黄金响应进行监督。在KD框架下，教师分布被用作每个表征过程的监督信号，对\mathcal{D}上的学生模型进行微调训练，也被简称为单词级KD。此外，SeqKD方法基于教师生成的示例数据对学生模型进行微调训练。

3.2 结果

在表1中，我们展示了评估结果，其中共有四个观察结果。首先，通过对比SFT与其他KD方法如KD和SeqKD，我们发现这些标准在大多数情况下成功地从教师模型中提取了知识，获得了更好的Rouge-L和GPT-4反馈分数，这与我们之前的研究结论一致。其次，通过将MINILLM的GPT-4反馈分数与基线进行比较，我们发现，当采用不同基础模型进行训练并在多套评估集上进行测试时，MINILLM在几乎所有情况下都优于基线。这表明MINILLM是一种通用的小型模型提取方法。我们还发现，与基线相比，MINILLM在DollyEval以外的数据集上通常表现更好，这表明该方法具有良好的分布外泛化能力。此外，Rouge-L分数显示，MINILLM生成的响应与基本事实高度重叠且具有极高的精确度。我们注意到，在某些特定评估集如VicunaEval、S-NI和UnNI上，MINILLM的Rouge-L分数甚至超过了教师模型的分数，这与[FLT^+]中的观察结果一致。我们推测，这种现象的原因在于标准教师强制微调引入了推断差异，即暴露偏差。相反，MINILLM采用策略优化方法，从而降低了这种偏差。在第3.3节中，我们对这一偏差进行了深入分析。最后，通过比较不同模型尺寸和模型族的表现，我们发现MINILLM的改进在120M至13B尺寸范围内始终保持一致。这一趋势在图1中得到验证，展示了MINILLM在LLM时代的一致改进和广泛适用性。

3.3 分析

暴露偏差

不同响应长度下的性能
我们研究了当黄金响应长度属于不同范围时模型的性能。在图5中，我们展示了不同KD模型相对于SFT模型在三个S-NI子集上的Rouge-L得分，这些子集由基本事实响应的长度划分。我们可以看到，所有方法在预期短响应（≤5个标记）的提示上都获得了低分数，这可能是因为我们训练集中的大多数响应都是长句，这引入了训练和测试之间的分布变化。此外，这些提示的输出空间相对较小，允许学生模型覆盖教师的大多数模式，因此反向KLD和正向KLD具有相似的性能。对于响应较长（≥6个标记）的提示，由于输出空间复杂，教师分布包含的模式比学生多，这表明了MINILLM相对于标准KD方法的优势。UnNI的类似结果如附录C.2所示。
生成多样性
[CF^+]发现，通过最小化反向KLD优化的模型可能会丢失模式，从而影响生成多样性。我们遵循[PH]从三个方面讨论了生成多样性：（i）在给定提示的情况下生成多个不同的响应。（ii）产生语言上复杂的反应。（iii）生成具有真实数据分布的高覆盖率的内容的能力。对于（i），我们认为，对于许多NLP应用程序，生成一个正确的响应就足够了，尤其是对于那些要求高真实性和可靠性的场景。对于（ii）和（iii），我们使用LLaMA家族的基本模型，在表3中报告了测试集上回答的不同4-gram比例和语言建模损失。我们可以看到，MINILLM在生成的响应中保留了明显的4-gram比例，并且不会导致测试集上的语言建模损失增加太多。

3.4 消融实验

优化策略的影响
我们通过从GPT-2-1.5B模型中提取GPT-2-125M模型，对第2.2节中提出的稳定和加速优化的三种策略进行了消融研究。在表4中，我们报告了每次运行的验证集上的最佳Rouge-L分数以及相应检查点的评估结果。在图6中，我们还绘制了训练期间学生和老师之间的反向KLD，其中线被平滑了32步。我们可以看到，教师混合抽样和长度归一化对稳定训练至关重要。尽管在没有这些策略的情况下，反向KLD也会减少，但我们发现，模型很快学会生成在教师分布中具有高概率的重复、短或无意义的字符串（见附录D中的示例），这被称为奖励黑客。这也导致了表4中的低生成性能。从图6中，我们还观察到，单步正则化有效地减少了训练过程的方差，这也导致了验证和测试集的更高性能。
教师组合强度 \alpha的影响
在图7中，我们使用GPT-2-1.5B、OPT-13B和LLaMA-13B作为教师，绘制了GPT-2-125M、OPT-1.3B和LLaMA-7B验证集上的最佳Rouge-L分数，其中MINILLM中教师组合的强度\alpha不同。\alpha=0.0意味着我们只从学生分布中采样，当\alpha=1.0时，我们完全从教师分布中采样。我们发现，\alpha=0.2通常适用于不同的模型族和大小，较大的模型对\alpha的选择更具鲁棒性。
添加预训练损失的影响
在表5中，我们通过将MINLLM与其变体进行比较来研究在算法2.3中添加预训练丢失的影响，其中去除了预训练语料库上的语言建模损失（w/o PT损失）。我们有一个与[OWJ^+]类似的观察结果，即添加预训练损失有助于保持规范NLP任务的能力，同时保持指令跟随任务的性能几乎不变。

4 相关工作

大型语言模型

知识蒸馏

文本生成中的分布差异度量

5 结论

在本研究中，我们探讨了如何将知识从大型LLM向小型LLM蒸馏的问题。在语言生成任务中，我们发现，仅通过最小化正向KLD进行蒸馏的方法并非最优选择。具体而言，教师模型的输出分布包含更丰富的模式，而正向KLD迫使学生模型在估计教师分布的低概率区域时过于保守。基于此，我们提出了一种新的蒸馏框架，MINILLM，该框架旨在最小化教师和学生分布之间的反向KLD，并为此开发了相应的优化算法。通过在指令跟踪任务中的大量实验，我们发现，MINILLM模型在生成响应的精确度和整体质量方面均优于传统蒸馏方法。此外，MINILLM在处理具有多样性的长文本生成时，展现出更低的曝光偏差，同时保持了更好的校准效果和更高的性能。