用于模型压缩的知识蒸馏

Knowledge Distillation for Model Compression

仅仅依靠模型准确率（或等效的性能指标）很少能决定最终部署哪个模型。这是因为我们还需要考虑若干运营指标，例如：

• 推理（inference）延迟：模型返回一个预测所花费的时间。
• 模型大小：模型占用的内存。
• 易于扩展性等。
• 在本章中，我将分享一种（带有演示的）技术，叫做知识蒸馏（knowledge distillation），它通常用于压缩机器学习模型，并对上述运营指标有所贡献。

简而言之，其思路是训练一个更小、更简单的模型（称为“学生模型”），让它模仿一个更大、更复杂模型（称为“教师模型”）的行为。

这涉及两个步骤：

• 像通常那样训练教师模型（teacher model）。
• 训练一个与学生模型（student model）输出相匹配的模型。
• 如果把它比作学术上的师生场景，学生可能不如老师表现出色。但通过持续的训练，较小的模型可以（几乎）达到和大模型一样好的效果。以这种方式开发的一个经典模型是 DistilBERT。它是 BERT 的学生模型。
• DistilBERT 比 BERT 大约小 40%，这在体积上是一个巨大的差异。
• 尽管如此，它仍然保留了 BERT 约 97% 的能力。

接下来，我们来看一个演示。

为了节省时间，假设我们已经在 MNIST 数据集上训练好了如下这个 CNN 模型（我在文末提供了完整的 Jupyter notebook，别担心）：

逐 epoch 的训练损失和验证准确率如下所示：

接下来，我们定义一个没有任何卷积层的更简单模型：

作为一个分类模型，其输出将是 <N> 个类别上的一个概率分布：

因此，我们可以训练学生模型，使其概率分布与教师模型的概率分布相匹配。

实现这一点的一种方法是使用 KL 散度作为损失函数（loss function）。

它衡量的是当我们用分布 Q 来近似分布 P 时，损失了多少信息。因此，在我们的情况中：

• P 将是来自教师模型的概率分布。
• Q 将是来自学生模型的概率分布。

损失函数的实现如下：

最后，我们按如下方式训练学生模型：

完成！下图比较了两个模型的训练损失和验证准确率：

当然，如上面高亮的几行所示，学生模型的性能不如教师模型，这是意料之中的。不过，考虑到它仅由简单的前馈层组成，结果仍然相当不错。此外，如下图所示，学生模型比教师模型快大约 35%，对于仅带来约个百分点的测试准确率下降而言，模型推理运行时间的提升是相当显著的。

话虽如此，知识蒸馏最大的缺点之一是，仍然必须先训练一个更大的教师模型才能训练学生模型。但在资源受限的环境中，训练一个大的教师模型可能并不可行。因此，这项技术假定我们至少在开发环境中不受资源限制。在下一章中，让我们再讨论一种压缩机器学习模型并减少其内存占用的技术。