知识蒸馏基本原理

Posted 2022-09-15 zhiyong_will

1. 概述

蒸馏是一个化学上的词汇，百科上对于蒸馏的解释为：“蒸馏是一种热力学的分离工艺，它利用混合液体或液-固体系中各组分沸点不同，使低沸点组分蒸发，再冷凝以分离整个组分的单元操作过程，是蒸发和冷凝两种单元操作的联合”。原理如下图所示：

蒸馏的目的是在不同的温度下提出了特定的成分。说回到知识蒸馏（knowledge distillation），其是模型压缩的一种常用的方法，最早得到推广的版本是由Hinton在2015年[1]提出并应用在分类任务上。与蒸馏的目的一致，在知识蒸馏中，希望通过提取性能更好的大模型的监督信息，构建一个小模型，同时使得小模型具有较好的性能和精度，而此处的大小模型成为Teacher，小模型称为Student。

2. 知识蒸馏的基本原理

2.1. 知识蒸馏原理

随着计算能力的不断提升，现在的模型也越来越大，网络也越来越深，结构变得异常复杂，这带来了模型准确率的提升，同时，计算复杂度也随之提升。知识蒸馏就是一种有效的模型压缩的方法，同时能够使得压缩后的模型的效果并未下降太多。

在知识蒸馏中，首先需要有一个大模型，也称为Teacher模型，该模型的特点是模型复杂，此时需要对该模型压缩，得到一个较小的模型，也称为Student模型，在蒸馏的过程中，将Teacher模型中学习到的“知识”迁移到Student模型中，以使得Student模型具有与Teacher模型一致的效果。具体的过程如下图所示：

2.2. 知识蒸馏学习过程

对于一个完整的知识蒸馏过程，有两个模型，分别为Teacher模型和Student模型，通过学习将已经训练好的Teacher模型中的知识迁移到小的Student模型中。其具体过程如下图所示[2]：

对于Student模型，其目标函数有两个，分别为蒸馏的loss（distillation loss）和自身的loss（student loss），其最终的损失函数为：

$$L=\alpha L_{soft}+\beta L_{hard}$$

其中，$L_{hard}$为student模型自身的损失，对于分类问题来说，可以通过交叉熵计算$L_{hard}$：

$$L_{hard}=-\sum _{i=1}^n{c}_{i}log\left(q_i\right)$$

其中，$c_i$为样本的真实标签，对于分类问题来说即为0或者1，$q_i$为Student模型的输出。通常，$q_i$可以通过softmax计算得到：

$$q_i=\frac{exp\left(z_i\right)}{{\sum }_{j}exp\left(z_j\right)}$$

对于softmax的计算，是在网络的logits结果$z_i$上，在softmax计算后得到的概率分布毁放大logits，会使得类目之间的差异变大。因此在知识蒸馏中，通常在logits的基础上加上一个温度变量$T$，来对logits结果缩放：

$$q_i=\frac{exp\left(z_i/T\right)}{{\sum }_{j}exp\left(z_j/T\right)}$$

当$T=1$时即为正常的输出，上述的$L_{hard}$即在$T=1$的情况下计算得到。对于$L_{soft}$的计算，通常有两种方式，一种是计算softmax输出结果的差异，另一种是直接比较logits结果的差异。

对于softmax结果的差异，由于softmax的结果是概率分布，因此可通过交叉熵计算分布之间的差异：

$$L_{soft}=-\sum _{i=1}^n{p}_{i}log\left(q_i\right)$$

其中，$p_i$为Teacher模型的输出，$q_i$为Student模型的输出。且输出是在$T=t$的情况下计算得到。

对于logits结果的差异，可以直接比较Teacher网络和Student网络输出logits的平方差，即：

$$L_{soft}=\sum _{i=1}^n{\left(v_i-z_i\right)}^2$$

其中，$v_i$为Teacher模型的logits输出，$z_i$为Student模型的logits输出。

3. 总结

知识蒸馏通过对Teacher模型的压缩得到效果接近的Student模型，由于网络模型复杂度的减小，使得压缩后的Student模型的性能得到较大提升。

参考文献

[1] Hinton G , Vinyals O , Dean J . Distilling the Knowledge in a Neural Network[J]. Computer Science, 2015, 14(7):38-39.

[2] Knowledge Distillation

[3] 【经典简读】知识蒸馏(Knowledge Distillation) 经典之作

[4] 一分钟带你认识深度学习中的知识蒸馏