速度×6,模型缩小20倍:用华为&华科的TinyBERT做搜索,性能达BERT 90%

百家 作者:量子位 2020-01-23 09:37:21
原作:Jack Pertschuk
鱼羊 编译整理
量子位 报道 | 公众号 QbitAI

在各大NLP榜单大杀四方的BERT,自打正式加入搜索引擎中,用户搜索体验确实更丝滑了。

不过,强大往往也意味着复杂。

BERT对算力、内存要求很高,为了使其顺利应对成百万次的搜索,谷歌不得不搭建起专门的Cloud TPU来满足其计算量。

那么,问题来了。不像谷歌那样财大气粗,就用不上BERT搜索了呗?

不知道你是否听说过,华为诺亚方舟实验室和华中科技大学联手打造的TinyBERT

现在,就有人用这个Transformer蒸馏方法,让BERT搜索变快6倍,模型大小则只有原来的1/20

TinyBERT搜索

首先,用微软的MS Marco训练三元组,在PyTorch中训练一个基于BERT的教师网络。

注:MS Marco:微软机器阅读理解数据集,由10万个问答组成

然后,让教师网络训练一个较小的学生BERT网络,这个学生网络只有4个隐藏层,而非标准的12层。

每一层的大小仅为312,以使模型更为轻巧。

在整个模型的最后,采用前馈二进制分类层来生成得分,进行搜索排名。

模型的具体参数如下,加粗部分是与标准BERT不同的部分。


“ attention_probs_dropout_prob”:0.1
“单元格”:{},
 “ emb_size”:312
 “ hidden_act”:“ gelu”,
“ hidden_dropout_prob”:0.1
“ hidden_size”:312
“ initializer_range”:0.02
“ intermediate_size”: 1200
 “ max_position_embeddings”:512
“ num_attention_heads”:12
“ num_hidden_layers”:4
 “ pre_trained”:“”,
“ structure”:[],
“ type_vocab_size”:2
“ vocab_size”:30522 
}

那么,用TinyBERT来做搜索,效果如何呢?

作者用MS Marco数据集对其进行检验。这一数据集采集自Bing的实际使用情况,是最大的公共搜索引擎使用情况数据集。

虽然在搜索效果的改进上,TinyBERT的性能只达到了BERT的90%,但速度提升了5倍,并且模型大小仅为BERT的1/20。

谷歌:用BERT做搜索

2019年10月底,谷歌将BERT加入到他们的搜索引擎当中。

谷歌搜索副总裁Pandu Nayak表示:

对于比较长、会话性比较强的查询,或者在“for”、“to”这样的介词比较重要的语句中,使用了BERT的搜索引擎能够更好地理解查询中单词的上下文。

比如,当输入“2019 brazil traveler to usa need a visa”(2019年前往美国的巴西旅行者需要签证)这样一条查询语句时,“to”这个词起到了关键的联系作用(巴西人→美国)

BERT未加入前,搜索引擎的算法无法理解这种关联,因此返回了有关前往巴西的美国公民(美国人→巴西)的结果。

而有了BERT,搜索引擎就能把握住其中细微而关键的差别了。

TinyBERT

而针对BERT这样的庞大模型难以部署到算力、内存有限的设备中的问题,华为诺亚方舟实验室和华中科技大学合作开源了TinyBERT

应用知识蒸馏(knowledge distillation,KD)的方法,大型教师模型BERT中编码的大量知识,能被很好地迁移到小型学生模型TinyBERT中。

研究人员还提出了一种专门用于TinyBERT的两段式学习框架,分别在预训练阶段和特定任务学习阶段进行transformer蒸馏。

这样,就能确保TinyBERT可以获取教师模型在通用和特定任务知识上的经验。

在GLUE基准上,TinyBERT的性能与BERT相当,仅下降了3个百分点。

并且,模型大小不到不到BERT的1/7,推理速度是BERT的9.4倍。

传送门

博客:
https://towardsdatascience.com/tinybert-for-search-10x-faster-and-20x-smaller-than-bert-74cd1b6b5aec

TinyBERT论文:
https://arxiv.org/abs/1909.10351

TinyBERT开源地址:
https://github.com/huawei-noah/Pretrained-Language-Model/tree/master/TinyBERT

作者系网易新闻·网易号“各有态度”签约作者


AI内参|把握AI发展新机遇

拓展优质人脉,获取最新AI资讯&论文教程,欢迎加入AI内参社群一起学习~

跟大咖交流 | 进入AI社群


量子位 QbitAI · 头条号签约作者



վ'ᴗ' ի 追踪AI技术和产品新动态


喜欢就点「在看」吧 !



关注公众号:拾黑(shiheibook)了解更多

[广告]赞助链接:

四季很好,只要有你,文娱排行榜:https://www.yaopaiming.com/
让资讯触达的更精准有趣:https://www.0xu.cn/

公众号 关注网络尖刀微信公众号
随时掌握互联网精彩
赞助链接