ICLR'24 | MiniLLM：大模型的知识蒸馏

知识蒸馏与模型压缩

量化和蒸馏是两种典型的模型压缩方法，或者说是减少模型推理阶段计算代价的方法。其中，知识蒸馏将待压缩的模型作为教师模型，将体积更小的模型作为学生模型，让学生模型在教师模型的监督下进行优化，最后使用学生模型代替原本的教师模型。在大模型时代之前，知识蒸馏广泛地应用在BERT这类结构的压缩中，代表工作有DistillBERT、PKD-BERT、TinyBERT、MobileBERT。将知识蒸馏应用在生成模型的工作感觉不多。正好看到这篇工作，特来阅读一下。

黑盒知识蒸馏 / 白盒知识蒸馏

作者首先将大模型上的 知识蒸馏工作划分为两类 ，黑盒知识蒸馏和白盒知识蒸馏。它们的区别是后者除了可以获取教师模型产生的文本，还可以获取教师模型输出的概率分布和中间层的隐状态。这里的黑盒知识蒸馏主要是指拿GPT-4产生的数据去训练自己的模型，其缺点是只能利用模型产生的文本做监督，知识蒸馏的效率比较低。

在这篇文章中，作者重点探究了在大模型的知识蒸馏过程中，如何妥当地利用教师模型输出的概率分布。（中间层的隐状态没用上）

前向KL散度

目前利用概率分布的典型方法是使用KL散度。对于输入文本 ，分别采样教师模型和学生模型的概率分布 。然后计算KL散度：

其中 要么是真实值、要么是从教师模型中采样的。接下来，我们就可以通过最小化KL散度来优化学生模型了。

作者将上面的方法称之为前向KL散度，并指出该方法会在 表达能力不足（学生模型能力差）的时候，倾向于过度估计 的void regions（应该就是指概率较小的那些词）。 具体为什么，文章没有具体讲。不过看了几篇文章都说前向KL不好。大概是学生模型差的时候， 太小，这导致KL散度不好优化。

反向KL散度

作者指出在计算KL散度的时候，应该从学生模型中采样 。这种情况下KL散度的计算如下：

) = 𝐾 𝐿 [ 𝑞 𝜃 ‖ 𝑝 ] = − 𝐸 𝑥 ∼ 𝑝 𝑥 , 𝑦 ∼ 𝑞 𝜃 log ⁡ 𝑝 ( 𝑦 | 𝑥 ) 𝑞 𝜃 ( 𝑦 | 𝑥 ) . 这两个KL散度的不同可以理解为：前向KL散度是从语料中采样一个prompt，然后从教师模型中采样一个response，接下来让学生模型逼近教师模型在response上的条件概率；后向KL散度，则是从学生模型中采样这个response，让教师模型根据自己的偏好来对采样的response进行指导。

进一步改进

在上述损失函数 的基础上，作者进行了三点实现上的改进。

• 单步分解 这是将每步的生成质量从损失的梯度中单独提出来，以减少训练时的方差并加速收敛。

• 教师指导的采样 前面提到反向KL散度中 是从学生模型中采样的。在这个改进中，作者在采样 时混合教师和学生模型的分布，即 ，其中 是超参数，文章中设为0.2。相应地作者在损失函数的计算上也进行了修改，添加了一个重要性权重。

• 长度正则化 作者指出当前的损失容易导致蒸馏后的模型产生较短的序列，因此增加了一个正则化到损失函数中。

经过上面的改进，最后的损失函数变为

实施细节

训练过程中用到了两个数据集，预训练数据集和指令数据集（15k）。首先在指令数据集上对学生模型进行有监督的微调，然后在指令数据集上进行知识蒸馏。此外，作者还在蒸馏过程中，结合了学生模型在预训练数据集上的损失。