内存用量1/20,速度加快80倍,腾讯QQ提出全新BERT蒸馏框架,未来将开源

腾讯 QQ 团队研究员对 BERT 进行了模型压缩,在效果损失很小的基础上,LTD-BERT 模型大小 22M,相比于 BERT 模型内存、存储开销可降低近 20 倍,运算速度方面 4 核 CPU 单机可以预测速度加速 80 余倍。相关代码和更多结果将在近期开源。

BERT 已经被验证是解决口语化短文本语义量化的极为有效的工具,对于口语化人机交互系统如 FAQ、语音助手等产品意义重大。但受限于模型的高复杂度和高计算量,其产业线上应用进展不如预期,尤其很难满足像 QQ、QQ 空间等亿级用户量产品对性能的需求。

为了解决这个问题,QQ 团队研究人员提出了 Learning to Distill BERT (LTD-BERT) 的模型对 BERT 进行了模型压缩,在效果损失很小的基础上,LTD-BERT 模型大小 22M,相比于 BERT 模型内存、存储开销可降低近 20 倍,运算速度方面 4 核 CPU 单机可以预测速度加速 80 余倍。另外,LTD-BERT 也被验证在下游任务可以保持与 BERT 近似的效果,包括相似度计算、短文本分类、短文本聚类等,其应用场景包括但不限于语义匹配、意图识别、文本聚类分析等。

据机器之心了解,从 2019 年 8 月份在腾讯内部开源至今,LTD-BERT 的效果已经在 QQ、腾讯新闻、腾讯游戏、腾讯看点、腾讯健康等服务海量用户的产品的实际任务中得到验证,确保算法拥有足够的泛化能力和实用性。

此外,该团队还表示 LTD-BERT 相关代码和更多结果将在近期开源。

 简介

2018 年底 Google 提出并开源了 BERT(Bidirectional Encoder Representation from Transformers),对 NLP 的影响已经绵延至今,非常多的 NLP 公开任务都被基于预训练 BERT 的工作刷榜,在工业界大家也逐渐的直接或通过 finetuning 将 BERT 运用到实际业务当中。

但是在这个过程中让大家非常头疼的事情就是因为超大的参数量,BERT 的运算速度和资源开销是很难权衡的问题。GPU 上线速度较快,但是成本很高;CPU 上线的话运算速度较慢,需要做大量的底层优化工作。

为了解决这个问题,QQ 研究团队提出了一种基于知识蒸馏的方法,在尽可能保证预训练 BERT 具备的文本理解能力不丢失的前提下,可以大大降低模型规模,提高预测速度。QQ 研究团队主要针对的是基于从 BERT 得到的 sentence embedding 去完成更上层任务的需求,这也能满足当前对于 BERT 的大部分的需求,囊括了文本分类、文本聚类、相似度计算等等。当然,word-level 的压缩也可以以近似的方法去实现。

图 1. BERT 的输入和嵌入(取自 [1]) 

Bert 中带权重的句嵌入 

既然要基于 BERT 的 sentence embedding 去做拟合,那第一步要确定如何获得 BERT 的 sentence embedding。如上图所示,BERT 对一个完整的句子会加入 [CLS] 用于针对句子层面的上层任务和 [SEP] 作为句子分隔符,通常的做法是用 [CLS] 所在位置的 vector 作为 sentence embedding。因为 [CLS] 参与到了句子层面「句对二元关系预测-是否是下一句」的预训练任务中,所以 [CLS] 位置是能一定程度反映句子的语义的,至少可以反应语句整体的 topic。另一种常用的方法是 averaging 每个字(中文)在 context 下的 embedding 来代表句子,这是在 word embedding 层面经常使用的方法。

但是其实通常来讲,尤其是针对句子层面的任务而言,一个句子的信息更多是依靠 keywords 来提供的,因此 QQ 研究人员提出了利用 IDF 值给句子中字的权重进行加权的方式来获取 BERT 的句向量,在试验中也取得了更好的效果。

具体做法是:

首先,在大量文本上,用某个分词工具分词,计算得到词语粒度的 IDF,即:

然后,对任意一个句子分词后,可得到其每个词语的 weight(即这个词语的 IDF)。不过,鉴于 BERT 得到的是字向量,在此用该词语的权重除以词语长度得到字的权重。最后,每个字的权重加上平滑项,再用每个字的权重除以整个句子所有权重和,得到每个字最终的权重。 

在计算得到句子中每个字的权重后,然后对字的 BERT 向量加权求和,得到加权句向量,即

为什么不直接对字进行权重计算主要是考虑字的语义通常不够明确,且在依赖上下文的同时还依赖组词,那么对于词的 IDF 值更具有实际意义。另外,加权的方式另一种考量是对于具体的任务拥有足够的适配能力,可以通过权重的调整来关注到具体任务更关心的词语。

 LTD-BERT

知识蒸馏是一种模型压缩的常见方法,是 14 年 Hinton 提出的 [2]。其在 Teacher-Student 框架中,将一些复杂、学习能力强的网络学到的特征表示「知识」蒸馏出来,传递给参数量小、学习能力弱的网络。常见的蒸馏一般把 Teacher 和 Student 作用在同样的问题上,可以提供 Student 在 one-shot label 上学不到的 soft label 信息,这些里面包含了类别间信息,以及 Student 小网络学不到而 Teacher 网络可以学到的特征表示『知识』,所以一般可以提高 Student 网络的精度 [3]。

 之前已经提出的一些蒸馏方法 [5,6,7,8],多数是针对文本分类、语义匹配、序列标注等任务上,把 BERT 作为 Teacher,直接将标准的蒸馏过程实施起来,或在此基础上改进 Student 的学习目标。

如下图 2 中所示的 BERT 训练流程,这部分工作更关注在右侧 Supervised 部分。而 BERT 的突破很大程度上依赖与图中左边弱监督过程中从大量文本数据学习到的语义建模能力,因此这里 LTD-BERT 的初衷是希望 Student 模型能够将这部分能力蒸馏过来,所以不对 Student 网络建立目标任务,Student 的目标仅为拟合 Teacher 网络的 sentence vector,不再去关注 Teacher 网络的预训练任务。通过合适的 loss function 设计,实验验证简单的 Student 网络可以较好的拟合 BERT 产生的句向量,并运用于各类 Task。

图 2. BERT 适用在 NLP 任务上的流程.(取自 [4]) 

Distilling 流程及架构

 整个 BERT 压缩的过程分成几个部分,首先对输入的句子做一定的清洗(BERT 本身也有一定的清洗过程),然后如前文提到字层面的 weight 计算需要分词、计算/获取 weight、均摊到 character(如下图中 preprocessing 所示)。经历完预处理后,得到了单字列表和每个字对应的权重,也就是 Learning to Distilling 过程的输入。

训练过程一方面用预训练的 BERT(Google 开源的中文预训练模型)给句子建模,并通过加权平均得到 BERT vector(如图中 training 阶段中的右侧部分所示);另一方面,使用简易的序列模型(本文介绍结果是以 LSTM 为例)同样的方式给句子建模,得到每个字在 context 下的语义,并使用同样的 weight 得到 sentence embedding,进而使用基于向量距离的 loss function 进行训练。通过百万级数据上的向量拟合训练,得到了一个规模只有 BERT 1/20 大小的 LTD-BERT 模型来为句子构建句向量。该句向量可以直接用于语句相似度计算、基于语义的文本聚类,另外 LTD-BERT 也可以像 BERT 一样在实际任务上 finetune 来获得更好的目标数据上的适应性。

这里 QQ 团队提出的方法只优化了最后句子的向量,因为使用的是序列模型来作为 LTD-BERT 的基准模型,原理上是可以直接用这个流程拟合 BERT 给每个汉字的建模结果。但是没有这样做一方面是直接拟合已经可以获得效果不错的 sentence embedding;另一方面是拟合的过程是使用的预生成的 BERT based sentence embedding,如果存下每个词的向量存储开销是巨大的,如果不进行预生成那训练速度会被 BERT 的 inference 速度 block。

图 3. LTD-BERT 训练过程

 LTD Student 模型结构

图 4. LTD-BERT 的 Student 模型

这里详细介绍一下当前使用的 Student 模型(即 LTD-BERT)的模型结构,这里基本是一个 standard 的双向 LSTM,在字 embedding 之上构建正向和逆向的 LSTM 获取上文语义和下文语义下一个汉字的语义,然后使用与 BERT 同样的 weights(图中 w)做 weighted sum,得到一个与 BERT dimension 一样的 vector。这里需要使用双向 LSTM 是因为 BERT 预训练任务对于每个词构建上下文感知的表示,保持建模过程的一致性才能保证更好的拟合效果,因此这里 LTD-BERT 基础模型的唯一限制是需要在为每个汉字建模的时候提供上下文语义的建模,QQ 研究团队也试验了不满足这个假设的基础模型,效果会大打折扣。

效果

从 2019 年 8 月份在腾讯内部开源至今,LTD-BERT 的效果已经在如下业务:QQ、腾讯新闻、腾讯游戏、腾讯看点、腾讯健康等海量用户产品线的上得到验证,包括文本分类、语义匹配、文本聚类等任务,确保算法拥有足够的泛化能力和实用性。具体的任务上,本文选取了一些有代表性的对比数据说明模型的效果,因为涉及具体,这里主要用「类别个数+数据量+数据类型+评测标准」来区分数据的特点。

文本分类

语义匹配

总结来说,在任意句子层面的任务上,该方法的方法可以提供一个规模很小的 LSTM 模型来替代 BERT,无差别的进行使用。并且实验发现当数据量相对较少,尤其是业务数据常见的几 k~几十 k 规模上,LTD-BERT 要明显优于 LSTM 以及前文提到的蒸馏算法,并且效果也更加接近 BERT,通常配合上一些其他的逻辑或者简单方法可以与 BERT 持平,甚至更优于单独 BERT 的效果。尤其对于资源有限或者业务量较大的业务来讲可以有效解决已经尝到 BERT 很香却只能线下看看的问题。

另外,QQ 团队也做了一些公开任务上的测试,以更好对比,后续会将代码开源,并在开源项目中陆续放出更多实验结果。另外,现在 NLP 飞速发展的时代,也会有很多相关的工作思路在一个阶段发展起来。在内部使用 LTD-BERT 的同时,也看到了 DistillBERT[9] 和 Tiny-BERT[10] 陆续推出,这两个工作跟 LTD-BERT 的出发点有一致之处,都是关注在图 2 的左边部分,大层面的不同之处主要有两点:1. LTD-BERT 关注在预训练的模型上,DistillBERT 和 Tiny-BERT 关注在预训练过程中引入蒸馏;2,LTD-BERT 关注不依赖 Transformer 或者 BERT 结构的蒸馏过程,探索更小算量的可能性。

 

[1] BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding.

[2] Distilling the Knowledge in a Neural Network.

[3] 简评 | 知识蒸馏(Knowledge Distillation)最新进展(一).

[4] The Illustrated BERT, ELMo, and co. (How NLP Cracked Transfer Learning).

[5] Distilling Task-Specific Knowledge from BERT into Simple Neural Networks.

[6] Scalable attentive sentence-pair modeling via distilled sentence embedding.

[7] Transformer to cnn: Label-scarce distillation for efficient text classification.

[8] Patient knowledge distillation for bert model compression.

[9] DistilBERT, a distilled version of BERT: smaller, faster, cheaper and lighter.

[10] TinyBERT: Distilling BERT for Natural Language Understanding.

理论知识蒸馏BERT腾讯QQ
暂无评论
暂无评论~