Press "Enter" to skip to content

知识蒸馏在推荐精排中的应用与实践

 

作者 | 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的差异公式为:

 

其中,注意力权重的分配公式为:

 

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

 

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

Be First to Comment

发表回复

您的电子邮箱地址不会被公开。 必填项已用*标注