关于teacher-student(知识蒸馏)的一些思考与总结

最近看了几篇有关teacher-student架构的paper，感觉收获挺大的，今天就来总结一下。这个teacher-student模式架构主要的目的就是用来进行深度学习模型的压缩，属于model compression领域中的一种比较流行的做法。因为深度学习下为了能够获得更好的准确率，训练出的网络往往结构比价复杂，而对于一些在线预测任务来说，复杂的模型结构不利于线上预测任务的快速响应需求，因此模型压缩的需求应运而生。故在该模型框架下，teacher结构相当于原始的复杂的深度神经网络结构，student则是一种轻量级的网络结构；因此teacher会有更高的预测准确率，它会指导student到达在简化参数之后最好的模型效果。既然已经介绍了teacher-student网络的工作过程，下面就来介绍一些学习到的一些知识（主要参考了阿里、爱奇艺等公司的在推荐系统精排阶段的一些落地实践，以及大神张俊林对于知识蒸馏方法在推荐系统工业落地的一些思考总结）。

第一个参考来源是Hinton大神2015年发表的paper《Distilling the Knowledge in a Neural Network》。整体来说这篇paper由于发表的比较早，所以他整体模型的思路也比较简单。首先使用训练数据集，通过构建比较复杂的网络结构来学习到一个teacher network，然后使用这个学习到的teacher network重新对训练数据集进行预测，
生成soft的结果概率分布，$q_i=\frac{exp(z_i/T)}{{\sum }_{j}exp(z_j/T)}$。其实论文里提到，在这个地方有两种可选方案：1 只是对training data预测soft概率分布；2 对所有的label数据和unlabel数据预测soft概率分布。通过实验证明了，第一种方式准确度更高，因此选用了第一种方案。

其中这个$T$是一个缩放因子，这个值越高说明预测结果的概率分布变的越soft。在distillinged即student网络进行训练的时候，其损失函数来自两个地方：1 使用了T进行缩放的teacher产生的soft概率分布的cross entropy；2 来自该样本真实类别属性的hard cross entropy，即$L=\alpha \ast L_{soft}+(1-\alpha )\ast L_{hard}$。

其实在该paper中，关于teacher网络对于student网络的指导，仅仅只是在网络输出的结果部分，并且以soft的类别概率分布的方式体现出来。这样student在进行学习的时候，相比于只提供类别的归属信息1或者0来说，它可以知道更多的信息，（虽然两个样本都被划为了1类，只能说明它们被预测为1类的概率大于被预测为0类的概率，但是它们被划分为1类的强弱信息是不知道的）。

第二个参考来源是阿里妈妈精准定向广告推荐组在2018年AAAI上发表的paper《Rocket Launching: A Universal and Efficient Framework for Training Well-performing Light Net》,这篇paper也用到了teacher-student架构。整个模型如下所示：

其中左边的Booster net就对应了teacher network，右边的Light Net就对应了student network。假设light net网络的softmax输出形式为：$p(x)=softmax(l(x))$，booster net网络的softmax输出形式为：$q(x)=softmax(z(x))$。那么该模型的损失函数如下所示：
$L=H(y,p(x))+H(y,q(x))+\lambda ||l(x)-z(x)|{|}^2$。和标准的teacher-student 架构不同的是在该paper提出的架构中Lighter和booster是联合训练的，这么做的原因是：student network可以学习teacher network整个的优化过程，而不仅仅是一个最终优化好的结果。同时在使用误差梯度逆向传播对网络进行更新的时候，损失函数中的$\lambda ||l(x)-z(x)|{|}^2$部分只对student network参数部分进行更新（这样才是teacher 去指导student），过程如下所示：

从中还可以发现，在该paper提出的网络结构中，light net和booster net还共享了部分底层的网络参数。（整个神经网络模型，从功能的角度来看，可以分为表示层+判别层，网络的底层主要进行表示层工作，因此具有share的特性。）

其实，在衡量student和teacher网络输出差异部分还可以有其他的选择，作者也在实验中进行了对比，结果肯定是他采用的这种方式效果是最好的，其他两种方式如下所示，给大家一个参考，毕竟不同场景下最优的策略可能是不一样的：
1 $||p(x)-q(x)|{|}^2$
2 $H(\frac{p(x)}{T},\frac{q(x)}{T})$

第三个参考的来源是爱奇艺公司在精排阶段采用的知识蒸馏模型，整体架构图如下所示：

其实从整理来看和阿里妈妈提出的蒸馏架构是相似的，都在teacher网络和student网络在输入部分共享了Embedding部分参数，都在模型最后一层的输出部分增加了logits形式的损失函数$||l(x)-z(x)|{|}^2$，但是需要注意的是有两点改进的地方：
1 除了网络最后的输出部分，中间的MLP隐藏侧部分也增加了对应的损失函数，也就是说除了要保证teacher和student在输出决策部分是相似的，也要保证中间隐藏层单元的输出也是相似的；
2 既然teacher网络和student网络在输入Embedding共享，而中间的MLP隐藏层参数一致（只有一致的情况下，才能针对隐藏层部分增加损失函数），那么如果体现出teacher网络优于student网络的地方呢？就是图中用红框框起来的Feature Interaction layer，这部分是teacher网络所拥有但是student网络没有的地方，在实际操作中这部分可以设计相对复杂一些。

最后参考的就是张俊林大神对于蒸馏网络在推荐系统特别是在召回阶段的一些思考和展望。其实给我印象、启发最大有以下几点：
1 虽然爱奇艺和阿里妈妈都是在精排阶段采用了蒸馏网络的思想，其实在召回/粗排阶段也可以采用蒸馏网络架构，而将召回/粗排作为student，那么对应后续的精排网络就可以作为teacher进行指导，这样可以使得前置两个环节（召回/粗排）的优化目标和推荐任务的最终优化目标保持一致；

2 就是在训练student的时候，并不是传统的直接使用logits结果信息进行辅助，即所说的Without-Logits 方案，即根据精排模型输出的物品最后呈现的顺序来进行辅助训练，排名靠前的肯定是精排模型认为越重要的，这里构造样本就有了几种方式，即point-wise，pair-wise和list-wise，其中point-wise是作者亲自试验了并证明是有一定效果的方法：
所谓point-wise来进行构造样本，即取一个超参数K，认为精排模型输出排名前k个商品是正样本，排名后面的是负样本（个人认为，负样本集合可以再从全局负采样一些也许不错）。同时作者注意到，即使都是排名前K的样本，这个排序的顺序往往是有意义的，也就是说越靠前的样本越重要，就越“正”，于是在损失函数部分，根据样本位置的不同，给每一个损失函数都进行了加权处理，一种示例的加权方式如下：$LossWeight=(1+\alpha \ast 1/position)$，position取值越小，样本排序越靠前，$LossWeight$的值越大。
至于pairwise和listwise方式其实大同小异，这里就不再赘述了。

其实关于teacher-student架构的paper还有不少，不过大概的核心思想和实现方式基本大同小异，比方说2015年ICLR的paper《fitnets: hints for thin deep nets》，也是利用类似的方式把一个shallow and wide的网络变成deep and thin的网络，它不仅仅是对最终分类的概率输出进行了比对，同时对网络的中间层神经元的输出进行了比对，具体的方式可以去原文中学习。