知识蒸馏技术

1 简介

知识蒸馏（Knowledge Distillation）是深度学习领域很火的概念，是一种模型压缩、知识迁移的技术。它从一个参数多、结构复杂、维护成本高的大模型（称为教师模型）迁移到参数少、结构简单、维护成本低的小模型（称为学生模型）。

迁移的原因很简单，复杂的大模型虽然能够取得很好的效果，但是参数量过大，推理成本高，计算量过大，需要的硬件资源多，很难在资源小的硬件上运行（尤其是移动设备）。知识蒸馏通过将大模型迁移到小模型，降低资源要求的同时尽可能使模型的效果不下降过多。

2 原理

知识蒸馏有不同的层次：

2.1 输出层标签蒸馏

这是较为常见的知识蒸馏方式。知识蒸馏中的“知识”包括硬标签和软标签。以分类问题为例，机器学习模型的输出结果一般是一个概率分布，表示图片属于各个类别的概率，取概率最大的分类作为最终的结果。原始数据中的标签称为“硬标签”，教师模型输出的标签概率分布称为“软标签”。例如，图片分类问题中，对于一张小狗的图片，硬标签是指这张的图片归属的类别是“狗”，软标签是指教师模型输出这张图片的概率分布，比如“狗”的概率是90%，“猫”的概率是5%，“狼”的概率是5%。学生模型在训练时，除了接收硬标签训练，还接收软标签学习，通过最小化学生模型与教师模型的概率分布之间的距离（例如KL散度），从而学习到教师模型所特有的判断。

2.2 中间层特征蒸馏

教师模型的中间特征也包含有用的信息。学生模型通过输出中间层的特征，计算与教师模型中间特征的相似性或距离，从而学习到教师模型特有的知识。例如Transformer中注意力矩阵包含了上下文之间的关联，CNN的中间特征图反映了图片各个层次（轮廓、形状等）的特征，这些特征都是有用的信息。

2.3 参数蒸馏

直接将教师模型的参数作为学生模型训练和模仿的目标，例如神经网络中的参数。如果教师模型和学生模型结构差异较大或者参数量差异较大，可以通过参数变换或映射，将教师模型较多的参数空间映射到学生模型较低的参数空间。

3 优缺点

总结来看，教师模型相当于给学生模型开小灶，将自身有用的信息，例如思考的过程（中间特征、输出标签）、方法论（参数权重）等提供给学生模型学习，从而使得学生模型利用额外的信息训练，快速缩小与教师模型的差距。

优点：

第一，提高了学生模型推理速度，降低资源要求。

第二，扩充了学生模型的训练方式。通过挖掘教师模型参数中隐藏的知识提高学习效果，这些知识无法通过传统的输出层体现出来。

缺点：

第一，教师模型可能和学生模型的结构差异很大。如果网络是不同的结构，则无法进行蒸馏。

第二，即使结构一致，参数量不一样也同样需要处理。例如教师模型的输出比学生模型的输出多，那么需要一个映射关系才可以做蒸馏。

第三，有可能过拟合。过于强调中间层输出与教师模型的一致性，可能导致过拟合，尤其是学生模型参数较少的情况下。

第四，有些教师模型可能并不开源，无法获取中间结果，只能用标签的概率分布（OpenAI部分GPT模型的API可以提供输出的前5大概率分布）。