作者 | Chilia

哥伦比亚大学 nlp搜索推荐

整理 | NewBeeNLP

知识蒸馏主要分为两类：logits蒸馏和中间特征蒸馏。其中，logits蒸馏指的是在softmax时使用较高的温度系数，提升负标签的信息，然后使用Student和Teacher在高温softmax下logits的KL散度作为loss。中间特征蒸馏就是强迫Student去学习Teacher某些中间层的特征。

为什么要在精排阶段使用知识蒸馏呢？

其中一个就是节约机器资源和成本。虽然在精排阶段我们可以使用复杂模型，但是复杂模型的在线服务相应必然变慢（例如xDeepFM在CPU上推理耗时为WDL的2.5~3.5倍，在使用GPU时只有在大batch下，推理性能才符合要求）。所以，我们可以使用复杂的Teacher模型去指导训练简单的Student模型，在线上部署Student模型。

另外一个原因是为了充分利用优势特征。所谓优势特征就是那种信号强、但是只能离线获得的特征（例如预测CVR时候的"用户停留时长"特征）。为了充分利用这些优势特征，我们可以使用优势特征训一个Teacher，然后让不使用优势特征的Student学习Teacher的logits，再把Student部署到线上。关于优势特征蒸馏的介绍可以看：

下面介绍一些有代表性的论文。

1. Rocket Launching: A Universal and Efficient Framework for Training Well-performing Light Net

这是阿里在AAAI 2018上发表的一篇论文，属于logits蒸馏。模型结构如下：

其中，简单的网络称为轻量网络(Light Net), 复杂的网络称为助推器网络(Booster Net)。和传统的logits蒸馏一样，在训练阶段，Light Net学习Booster Net的soft-target来模仿Booster Net的学习过程；在测试阶段，只使用Light Net来上线预测。

那么，这个模型比传统logits蒸馏的改进是什么呢？

（1）共享参数

上图黄色部分的参数 是共享的，这样Booster Net可以直接更新共享参数，于是Light Net可以直接得到Booster Net的“助推”（"get direct thrust from the booster"）。共享的是低层的参数，这是因为低层参数主要学习的是信息表示、反应的是对输入信息基本的刻画（如embedding层），低层参数共享可以帮助Light Net学习 更好的特征表示 。

（2）同时训练

一般的知识蒸馏都是先训练出一个好的Teacher，然后再从Teacher中"蒸馏"出Student模型。但是这篇文章是同时训练Light Net和Booster Net的。在训练阶段，Booster Net和Light Net一起前行；在测试阶段，Booster Net剥离，Light Net独自前行 -- 这就是题目中"Rocket Launching"的由来。

那么，为什么要这么做呢？一方面，Booster Net 全程 提供soft target信息给轻量网络，从而 指导Light Net的整个求解过程 。

The light model can learn from not only the difference between the target and its temporary outputs, but also the possible path towards the final target provided by a complex model with more learning capability.

训练时，网络的loss为：

loss包含三部分：

• Light Net对true label的交叉熵损失；

• Booster Net对true label的交叉熵损失；

• 两个网络在softmax之前的logits的MSE（hint loss）

另一方面，同时训练可以 缩短训练时间 。

（3）梯度屏蔽 (gradient block)

在训练过程中，hint loss只能用于Light Net的梯度更新、不能用于Booster Net的梯度更新。Booster Net只能从true label中学习知识。也就是说，学生可以向老师学，但是老师不能跟学生学，否则就会被学生拖累。

梯度传播图如下：

后记：2020年爱奇艺公布的排序阶段知识蒸馏模型和这个模型十分相似。爱奇艺公布的模型如下：

其中，红色框是复杂模型Teacher，黄色框是简单模型Student。他们共享了底层embedding，但是Teacher使用了较为复杂的特征交叉层。两者同时训练，Student学习Teacher的logits输出和中间层。

2. Ensembled CTR Prediction via Knowledge Distillation

这篇华为在CIKM2020上发表的论文，使用了logits蒸馏+中间层特征蒸馏的方法。

读过上面的知识蒸馏基础便可以知道，想要把Teacher学到的知识迁移到Student上，有两种方法。一种是logits蒸馏，即在softmax的时候加一个 较高的温度系数 ，来提高负label的概率，获得更多的知识；一种是特征蒸馏，即最小化Teacher和Student 中间层 的差异。有了Teacher带来的更丰富的知识，学生就能够学的比只用true label好得多。"even a simple DNN model can learn well given useful guidance from a good teacher" 。如下图：

上面所说的是只有一个Teacher的情况。我们知道，模型ensemble在机器学习中经常能取得更好的效果。但是如果在线上使用多个精排模型同时预测并不太现实，因为这样的话训练和推理所耗费的资源就要翻倍了。那么能不能从多个Teacher中一起蒸馏呢？

一种直接的想法是，直接把多个Teacher的logits层/中间层做个平均，然后让Student学习这个已经融合好的单一模型。但是，并不是所有Teacher模型都应该具有同等的权重。毕竟，如果一个Teacher学的很差，就会对整体产生干扰。所以，文中采用了 自适应分配注意力权重 的方法，是用Gate Net实现的：

计算Student和Teacher的差异公式为：

其中，注意力权重的分配公式为：

402 Payment Required

其中 是第 个Teacher的logits或者中间层向量。这样，我们就得到了要分配给每个Teacher的权重。

在训练时，可以使用Teacher和Student一起训练的方法，这样训练的速度更快，但是需要把所有Teachers和Student都load进GPU，对GPU资源要求高。也可以选择"先训练好Teacher、再从Teacher中蒸馏Student"的方法。

一起交流

想和你一起学习进步！『NewBeeNLP』目前已经建立了多个不同方向交流群（机器学习 / 深度学习 / 自然语言处理 / 搜索推荐 / 图网络 / 面试交流 / 等），名额有限，赶紧添加下方微信加入一起讨论交流吧！（注意一定o要备注信息才能通过）

- END -

工作没满一年，我跳槽了！【附面经】

2022-01-24

Green Deep Learning：NLP在大模型之外的另一种思路！

2022-01-22

深度学习中 Batch Size 对训练过程的影响

2022-01-05

美团 大规模商品知识图谱的构建与应用

2021-12-31