【Paper Note】模型蒸馏Model Distilling

这篇论文是Hinton在15年提出的，为了提升模型的有效性，模型的复杂度的不断增加，上线实时提供服务成了难题，而知识蒸馏的思路正好可以解决这个问题，同时模型的效果相比复杂模型也不会下降太多。
论文中以生物中蝴蝶变态发育作类比介绍知识蒸馏：通过不同的形态，完成同样的使命（任务）。
Hinton提出可以通过一个简单模型直接学习复杂模型的概率分布结果，如果one-hot的目标是一种hard-targets，那么这种就是一种soft-targets。

一种方法是直接比较logits来避免这个问题。具体地，对于每一条数据，记原模型产生的某个logits是 $v_i$ ，新模型产生的logits是 $z_i$ ，我们需要最小化
$$\frac{1}{2}(z_i-v_i{)}^2\text{\hspace{1em}(1)}$$

Hinton提出了升温蒸馏的概念，温度就是其中的关键点，升温蒸馏，降温预测，完美。

其中温度T就是用来做平滑的，T越大，平滑力度越大，使得轻量模型学习时可以关注到那些概率很小的类别；T越小，则相反，T=1时，就是平常所见的概率分布。

考虑一个广义的softmax函数：
$$q_i=\frac{exp(z_i/T)}{{\sum }_{j}exp(z_j/T)}\text{\hspace{1em}(2)}$$
可以证明，上面的logit值作为训练目标是这种方法的一种特例，总是可以通过调整T来达到。其中 $T$ 是温度，这是从统计力学中的玻尔兹曼分布中借用的概念。容易证明，当温度 $T$ 趋向于0时，softmax输出将收敛为一个one-hot向量；温度 $T$ 趋向于无穷时，softmax的输出则更「软」。因此，在训练新模型的时候，可以使用较高的 $T$ 使得softmax产生的分布足够软，这时让新模型（同样温度下）的softmax输出近似原模型；在训练结束以后再使用正常的温度 $T=1$ 来预测。具体地，在训练时我们需要最小化两个分布的交叉熵(Cross-entropy)，记新模型利用公式 (2) 产生的分布是 $q$ ，原模型产生的分布是 $p$ ，则我们需要最小化

$$C=-p^T\log q\text{\hspace{1em}(3)}$$

下面计算交叉熵损失对softmax输入的梯度，由链式法则，有：
$$\frac{\partial C}{\partial z}=\frac{\partial C}{\partial q}\frac{\partial q}{\partial z}\text{\hspace{1em}(4)}$$

由于式（3）中的 $p$ 是原模型产生的softmax输出，与 $z$ 无关。
结合式（3）可得：
$$\frac{\partial C}{\partial q_i}=-\frac{p_i}{q_i}\text{\hspace{1em}(5)}$$
所以，
$$\frac{\partial C}{\partial q}=\left[\begin{array}{c}-\frac{p_1}{q_1}\\ -\frac{p_2}{q_2}\\ ⋮\\ -\frac{p_n}{q_n}\end{array}\right]\text{\hspace{1em}(6)}$$

式（4）中，$\frac{\partial q}{\partial z}$ 是一个 $n\times n$ 的方阵，分类讨论可以得到。
记 $Z={\sum }_{k}exp(z_k/T)$，由除法的求到公式，输出 $q_i$ 对输入 $z_j$ 的偏导为：
$$\begin{array}{rl}\frac{\partial q_i}{\partial z_j}& =\frac{1}{Z^2}(Z\frac{\partial exp(z_i/T)}{\partial z_j}-exp(z_i/T)\frac{\partial Z}{\partial z_j})\\ & =\frac{1}{Z^2}(Z\frac{\partial exp(z_i/T)}{\partial z_j}-exp(z_i/T)\cdot \frac{1}{T}exp(z_j/T))\\ & =\frac{1}{Z}\frac{\partial exp(z_i/T)}{\partial z_j}-\frac{1}{T{Z}^2}exp(z_i/T)exp(z_j/T)\\ & =\frac{1}{Z}\frac{\partial exp(z_i/T)}{\partial z_j}-\frac{1}{T}\frac{exp(z_i/T)}{Z}\frac{exp(z_j/T)}{Z}\\ & =\frac{1}{Z}\frac{\partial exp(z_i/T)}{\partial z_j}-\frac{1}{T}{q}_i{q}_j\end{array}\text{\hspace{1em}(7)}$$
对 $\frac{\partial exp(z_i/T)}{\partial z_j}$ 分类讨论得到：
$$\frac{\partial exp(z_i/T)}{\partial z_j}=\{\begin{array}{rcl}\frac{1}{T}exp(z_i/T)& & i=j\\ 0& & i\ne j\end{array}\text{\hspace{1em}(8)}$$

将式（8）带入式（7），得到：
$$\begin{array}{rl}\frac{\partial q_i}{\partial z_j}& =\{\begin{array}{rcl}\frac{1}{T}(\frac{exp(z_i/T)}{Z}-q_i{q}_j)& & i=j\\ -\frac{1}{T}{q}_i{q}_j& & i\ne j\end{array}\\ & =\{\begin{array}{rcl}\frac{1}{T}(q_i-q_i{q}_j)& & i=j\\ -\frac{1}{T}{q}_i{q}_j& & i\ne j\end{array}\end{array}\text{\hspace{1em}(9)}$$

所以，$\frac{\partial q}{\partial z}$ 的形式如下：
$$\frac{\partial q}{\partial z}=\frac{1}{T}\left[\begin{array}{cccc}{q}_1-q_1^2& -q_1{q}_2& \cdots & -q_1{q}_n\\ -q_2{q}_1& q_2-q_2^2& \cdots & -q_2{q}_n\\ ⋮& ⋮& \ddots & ⋮\\ -q_n{q}_1& -q_n{q}_2& \cdots & q_n-q_n^2\end{array}\right]\text{\hspace{1em}(10)}$$

将式（10）带入到式（4）中，得到：
$$\begin{array}{rl}\frac{\partial C}{\partial z}& =\frac{1}{T}\left[\begin{array}{cccc}{q}_1-q_1^2& -q_1{q}_2& \cdots & -q_1{q}_n\\ -q_2{q}_1& q_2-q_2^2& \cdots & -q_2{q}_n\\ ⋮& ⋮& \ddots & ⋮\\ -q_n{q}_1& -q_n{q}_2& \cdots & q_n-q_n^2\end{array}\right]\left[\begin{array}{c}-\frac{p_1}{q_1}\\ -\frac{p_2}{q_2}\\ ⋮\\ -\frac{p_n}{q_n}\end{array}\right]\\ & =\frac{1}{T}\left[\begin{array}{c}-p_1+{\sum }_k{p}_k{q}_1\\ -p_2+{\sum }_k{p}_k{q}_2\\ ⋮\\ -p_n+{\sum }_k{p}_k{q}_n\end{array}\right]\\ & =\frac{1}{T}\left[\begin{array}{c}-p_1+q_1\\ -p_2+q_2\\ ⋮\\ -p_n+q_n\end{array}\right]\\ & =\frac{1}{T}(q-p)\end{array}\text{\hspace{1em}(11)}$$

所以，有：
$$\frac{\partial C}{\partial z_i}=\frac{1}{T}(q_i-p_i)\text{\hspace{1em}(12)}$$

结合（2）式，得到：
$$\begin{array}{rl}\frac{\partial C}{\partial z_i}& =\frac{1}{T}(q_i-p_i)\\ & =\frac{1}{T}(\frac{exp(z_i/T)}{{\sum }_{j}exp(z_j/T)}-\frac{exp(v_i/T)}{{\sum }_{j}exp(v_j/T)})\end{array}\text{\hspace{1em}(13)}$$

使用等价无穷小 $e^x-1\sim x$ 作替换：
$$\begin{array}{rl}\frac{\partial C}{\partial z_i}& \approx \frac{1}{T}(\frac{1+z_i/T}{{\sum }_j(1+z_j/T)}-\frac{1+v_i/T}{{\sum }_j(1+v_j/T)})\\ & =(\frac{1+z_i/T}{N+{\sum }_j{z}_j/T}-\frac{1+v_i/T}{N+{\sum }_j{v}_j/T})\end{array}\text{\hspace{1em}(14)}$$

假设所有logits对每个样本都是零均值化的，
$$\sum _j{z}_j=\sum _j{v}_j=0\text{\hspace{1em}(15)}$$

则有，
$$\begin{array}{rl}\frac{\partial C}{\partial z_i}& \approx \frac{1}{T}(\frac{1+z_i/T}{N}-\frac{1+v_i/T}{N})\\ & =\frac{1}{N{T}^2}(z_i-v_i)\end{array}\text{\hspace{1em}(16)}$$

所以，如果：1. $T$ 非常大，2. logits对所有样本都是零均值化的，则知识蒸馏和最小化logits的平方差(公式（1）)是等价的(因为梯度大致是同一个形式)。实验表明，温度 $T$ 不能取太大，而应该使用某个适中的值，这表明忽略极负的logits对新模型的表现很有帮助(较低的温度产生的分布比较「硬」，倾向于忽略logits中极小的负值)。

同一个样本，用在大规模神经网络上产生的软目标来训练一个小的网络时，因为并不是直接标注的一个硬目标，学习起来会更快收敛。

更巧妙的是，这个样本我们甚至可以使用无标注的数据来训练小网络，因为大的神经网络将数据结构信息学习保存起来，小网络就可以直接从得到的soft target中来获得知识。

这个做法类似学习了样本空间嵌入（embedding）信息，从而利用空间嵌入信息学习新的网络。

随着温度上升，软目标分布更均匀

T参数是一个温度超参数，按照softmax的分布来看，随着T参数的增大，这个软目标的分布更加均匀。

所以：
1.首先用较大的T值来训练模型，这时候复杂的神经网络能够产生更均匀分布的软目标；
2.之后小规模的神经网络用相同的T值来学习由大规模神经产生的软目标，接近这个软目标从而学习到数据的结构分布特征；
3.最后在实际应用中，将T值恢复到1，让类别概率偏向正确类别

Reference：
https://arxiv.org/pdf/1503.02531.pdf
https://zhuanlan.zhihu.com/p/71986772
https://zhuanlan.zhihu.com/p/97522736
https://zhuanlan.zhihu.com/p/39945855
https://zhuanlan.zhihu.com/p/93287223
https://zhuanlan.zhihu.com/p/90049906