知识蒸馏解决的问题
- 提升模型精度
- 降低模型时延,压缩网络参数
- 标签之间的域迁移
基本框架
知识蒸馏采取Teacher-Student模式:将复杂且大的模型作为Teacher,Student模型结构较为简单,用Teacher来辅助Student模型的训练,Student模型去学习Teacher模型的泛化能力。
Teacher模型不上线,真正部署上线进行预测任务的是Student小模型。
知识蒸馏根据迁移的方法不同分为基于目标蒸馏(也称为Soft-target蒸馏或Logits方法蒸馏)和基于特征蒸馏的算法两个大的方向。
目标蒸馏 Logits方法
训练的目标
限定在分类问题下,传统的神经网络训练方法是定义一个损失函数,目标是使预测值尽可能接近于真实值(Hard- target),在知识蒸馏中,是使用大模型softmax层输出的类别概率作为Soft-target的训练过程。
- Hard-target:原始数据集标注的 one-shot 标签,除了正标签为 1,其他负标签都是 0。
- Soft-target:Teacher模型softmax层输出的类别概率,每个类别都分配了概率,正标签的概率最高。
Teacher模型的作用
知识蒸馏用Teacher模型预测的 Soft-target 来辅助 Hard-target 训练Student模型的方式为什么有效呢?
softmax层的输出,除了正例之外,负标签也带有Teacher模型归纳推理的大量信息,比如某些负标签对应的概率远远大于其他负标签,则代表 Teacher模型在推理时认为该样本与该负标签有一定的相似性。而在传统的训练过程(Hard-target)中,所有负标签都被统一对待。也就是说,知识蒸馏的训练方式使得每个样本给Student模型带来的信息量大于传统的训练方式。
具体方法
几个概念:
- logits:经过DNN网络各种非线性变换,在网络最后Softmax层之前,会得到这张图片属于各个类别的大小数值zi。
- softmax函数:exp(zi)/sum(exp(z)),指数会增大logits之间的差异,对输出结果进行归一化表示。
- 带温度系数的softmax函数:exp(zi/T)/sum(exp(z/T)),当T越大,z/T越小,输出的分布越平滑,负标签所携带的信息相当于被放大。
步骤:
- 训练好Teacher模型;
- 使用高温度系数产生soft-target;
- 使用soft-target和hard-target(ground-truth)同时训练Student模型;
- Student线上做推理,设置温度系数T=1。
第3步中,损失函数是两个交叉熵的加权和L=aLsoft+bLhard
,之所以要加上hard-target的部分,是因为Teacher模型有错误的概率,需要ground-truth进行纠正。
Lhard权重取较小值时,效果比较好,由于Lsoft贡献的梯度大约为Lhard的1/T2,因此在同时使用Soft-target和Hard-target的时候,需要在Lsoft的权重上乘T2的系数,这样才能保证Soft-target和Hard-target贡献的梯度量基本一致。
温度T参数调整
温度T特点:
原始的softmax函数是T=0时的特例;T越小时,概率分布比原始更“陡峭”,也就是说,当T接近于0时,Softmax 的输出值会接近于 Hard-target;随着T的增加,Softmax 的输出分布越来越平缓,信息熵会越来越大。温度越高,softmax上各个值的分布就越平均,思考极端情况,当T=无穷大时,softmax的值是平均分布的。
- 当想从负标签中学到一些信息量的时候,温度T应调高一些;
- 当想减少负标签的干扰的时候,温度T应调低一些。
T的选择和Student模型的大小有关,Student模型参数量比较小的时候,相对比较低的温度就可以了。因为参数量小的模型不能学到所有Teacher模型的知识,所以可以适当忽略掉一些负标签的信息。
在整个知识蒸馏过程中,我们先让温度T升高,然后在测试阶段恢复“低温“,从而将原模型中的知识提取出来,因此将其称为是蒸馏。
特征蒸馏
作用:将深且宽的网络蒸馏成更深但窄的网络
步骤:
第一阶段:首先选择待蒸馏的中间层(即Teacher的Hint layer和Student的Guided layer),由于两者的输出尺寸可能不同,在Guided layer后另外接一层卷积层,使得输出尺寸与Teacher的Hint layer匹配。接着通过知识蒸馏的方式训练Student网络的Guided layer,使得Student网络的中间层学习到Teacher的Hint layer的输出。
第二阶段: 在训练好Guided layer之后,将当前的参数作为网络的初始参数,利用知识蒸馏的方式训练Student网络的所有层参数,使Student学习Teacher的输出。
相关论文
目标蒸馏
Distilling the Knowledge in a Neural Network 》,NIPS,2014。
《Deep Mutual Learning》,CVPR,2018。
《Born Again Neural Networks》,CVPR,2018。
《Distilling Task-Specific Knowledge from BERT into Simple Neural Networks》,2019。
特征蒸馏
《FitNets: Hints for Thin Deep Nets》,ICLR,2015。
《Paying More Attention to Attention: Improving the Performance of Convolutional Neural Networks via Attention Transfer》, ICLR,2017。
《A Gift from Knowledge Distillation: Fast Optimization, Network Minimization and Transfer Learning》,CVPR,2017。
《Learning Efficient Object Detection Models》,NIPS,2017。