DL-Paper精读:MobileBERT

MobileNERT: a Compact Task-Agnostic BERT for Resource-Limited Devices

https://arxiv.org/abs/2004.02984

Background

BERT在NLP领域的地位是举足轻重的,其预训练模型,在多种下游任务的迁移工作中都能给出非常好的效果。但于此同时,BERT也受困于其庞大的模型参数和较慢的运行速度,尤其是在于一些资源受限的移动端部署任务中,BERT的实际应用是非常受限的。

Related work and the limit

目前,针对BERT的压缩研究工作已有很多。比如distill BERT等工作,通过知识蒸馏,来获得一个高精度的小模型。但是这些工作一般都是基于具体下游任务场景的(“task-specifically”),这些工作的一般流程是首先将预训练的BERT模型在具体任务中进行fine-tune,然后再进行蒸馏来获得小模型。该过程是比较耗时且繁琐的,不如直接来fine-tune一个任务无关(“task-agnostic”)的小的预训练模型。但是如果直接预训练一个较窄或较浅的BERT,又会造成较严重的精度损失,无法达到原始BERT的精度。

Novel points

在预训练的阶段来进行知识蒸馏,获得一个压缩版的小BERT(mobileBERT),能够达到可观的精度,同时是task-agnostic的。同时为了能够有效地蒸馏mobileBERT,本文设计了一个带反瓶颈结构(inverted-bottleneck)的近似BERT_large的网络(IB-BERT)。在实际训练中,首先充分训练IB-BERT作为teacher model,然后再对mobileBERT进行知识蒸馏。

Methodology

该部分主要分为结构设计和训练方法两个方面。

结构设计:

包括mobileBERT和IB-BERT两个网络,具体根据以下图表来呈现,图1分别为BERT/IB-BERT/mobileBERT的模块结构图,表1给出了整体的网络参数设计。可以看出IB-BERT和mobileBERT与BERT_L有着相同的层数,即深度相同,但mobileBERT每一层的宽度(width)则缩减了8倍,这也就造成了它相比于BERT_L有着接近20倍的参数缩减,同时还能保持着可观的精度。此外,具体的结构方面的改进主要包括以下4个方面:

bottleneck/inverted-bottleneck结构的使用。由图标可以看出,IB-BERT和MobileBERT的模块前后都添加了一个linear层。其中IB-BERT采用了一个Inverted-Bottleneck,将输入的512维度放大为1024,从而匹配BERT_L的模型尺寸,来获得相近的精度;而MobileBERT则采用了一个Bottleneck结构,用来将输入的512维度缩小为128,进一步缩减模型参数量和计算量。IB-BERT和MobileBERT的模块输入输出都是512,因此能够在每一层进行feature map的蒸馏。

堆叠的FFN操作。在MobileBERT的模块中,对FFN操作进行了多次堆叠。在传统的Transformer结构中,MHA(multi-head attention)和FFN(feed-forward network)都扮演着重要的角色,其中MHA能够提取获得各部分的注意力,而FFN能够增加模型的非线性能力,同时在传统BERT中,这两者的参数量总是保持1:2。但是在MobileBERT的结构设计中,因为宽度的变化,参数比例发了改变,MHA的参数相对变多了。因此文中提出通过堆叠FFN的方法来解决该问题,MobileBERT中,MHA之后有连续4个FFN。

文中通过分析发现LN和Gelu操作对于latency的影响较大,因此提出移除LN层,替换为文中所提出的一种element-wise的线性转换操作NoNorm,同时将Gelu换为Relu。(虽然文中这样提出了,但考察github上的几个实现代码,发现一般还是在使用LN和Gelu)

由表格可以看出,在embedding阶段,为了减少原BERT中embedding的庞大参数量,文中采用了一个 emb_128 + 512_conv的操作来代替。

网络训练

该部分主要针对IB-BERT -> MobileBERT的蒸馏训练设计。主要分为两部分。

Loss设计

由于两个网络中bottleneck/inverted-bottleneck结构的设计,所以文中提出可以在每一层都可以加入feature map transfer(FMT)和Attention Transfer(AT)。其中FMT是通过计算教师网络IB-BERT和学生网络MobileBERT每一层的输出的MSE来获得,AD是通过计算两者MHA中attention的分布的KL散度而获得,具体计算公式参考论文。两种transfer可以如图1中的虚线所示。因此,该预训练阶段的distillation训练的Loss函数,由三部分组成:传统的masked language modeling loss(MLM), next sentence prediction(NSP) loss和新的Knowledge Distillation loss(KD)所组成,公式如下:

训练策略

许多常见的蒸馏训练的技巧,包括 Auxiliary Knowledge Transer,Joint Knowledge Transfer, Progressive Knowledge Transfer等,具体见文中介绍。

Evaluation(benchmark, experiments design)

对NLP任务不够熟悉,因此不对实验部分进行过多的评估,GLUE上的结果如图所示。总的来说,相比于BERT_B,MobileBERT能够在保持相近精度的情况下,实现4.3×的压缩和5.5×的加速。

Thoughts:

1. Describe what the authors of the paper aim to accomplish, or perhaps did achieve.

对BERT进行压缩,给出一个精度较高且任务无关的压缩模型

2. If a new approach/ technique/ method was introduced in a paper, what are the key elements of the newly proposed approach?

通过蒸馏来获得一个高精度的预训练MobileBERT。为了能够获得较好的蒸馏效果,首先设计了一个接近的BERT_L的teacher模型,同时通过巧妙地bottleneck/inverted-bottleneck结构的使用,使得大小模型之间能够实现直接的信息transfer。

此外,详尽复杂的蒸馏策略设计等也对结果有着重要的影响。

©著作权归作者所有,转载或内容合作请联系作者
  • 序言:七十年代末,一起剥皮案震惊了整个滨河市,随后出现的几起案子,更是在滨河造成了极大的恐慌,老刑警刘岩,带你破解...
    沈念sama阅读 216,496评论 6 501
  • 序言:滨河连续发生了三起死亡事件,死亡现场离奇诡异,居然都是意外死亡,警方通过查阅死者的电脑和手机,发现死者居然都...
    沈念sama阅读 92,407评论 3 392
  • 文/潘晓璐 我一进店门,熙熙楼的掌柜王于贵愁眉苦脸地迎上来,“玉大人,你说我怎么就摊上这事。” “怎么了?”我有些...
    开封第一讲书人阅读 162,632评论 0 353
  • 文/不坏的土叔 我叫张陵,是天一观的道长。 经常有香客问我,道长,这世上最难降的妖魔是什么? 我笑而不...
    开封第一讲书人阅读 58,180评论 1 292
  • 正文 为了忘掉前任,我火速办了婚礼,结果婚礼上,老公的妹妹穿的比我还像新娘。我一直安慰自己,他们只是感情好,可当我...
    茶点故事阅读 67,198评论 6 388
  • 文/花漫 我一把揭开白布。 她就那样静静地躺着,像睡着了一般。 火红的嫁衣衬着肌肤如雪。 梳的纹丝不乱的头发上,一...
    开封第一讲书人阅读 51,165评论 1 299
  • 那天,我揣着相机与录音,去河边找鬼。 笑死,一个胖子当着我的面吹牛,可吹牛的内容都是我干的。 我是一名探鬼主播,决...
    沈念sama阅读 40,052评论 3 418
  • 文/苍兰香墨 我猛地睁开眼,长吁一口气:“原来是场噩梦啊……” “哼!你这毒妇竟也来了?” 一声冷哼从身侧响起,我...
    开封第一讲书人阅读 38,910评论 0 274
  • 序言:老挝万荣一对情侣失踪,失踪者是张志新(化名)和其女友刘颖,没想到半个月后,有当地人在树林里发现了一具尸体,经...
    沈念sama阅读 45,324评论 1 310
  • 正文 独居荒郊野岭守林人离奇死亡,尸身上长有42处带血的脓包…… 初始之章·张勋 以下内容为张勋视角 年9月15日...
    茶点故事阅读 37,542评论 2 332
  • 正文 我和宋清朗相恋三年,在试婚纱的时候发现自己被绿了。 大学时的朋友给我发了我未婚夫和他白月光在一起吃饭的照片。...
    茶点故事阅读 39,711评论 1 348
  • 序言:一个原本活蹦乱跳的男人离奇死亡,死状恐怖,灵堂内的尸体忽然破棺而出,到底是诈尸还是另有隐情,我是刑警宁泽,带...
    沈念sama阅读 35,424评论 5 343
  • 正文 年R本政府宣布,位于F岛的核电站,受9级特大地震影响,放射性物质发生泄漏。R本人自食恶果不足惜,却给世界环境...
    茶点故事阅读 41,017评论 3 326
  • 文/蒙蒙 一、第九天 我趴在偏房一处隐蔽的房顶上张望。 院中可真热闹,春花似锦、人声如沸。这庄子的主人今日做“春日...
    开封第一讲书人阅读 31,668评论 0 22
  • 文/苍兰香墨 我抬头看了看天上的太阳。三九已至,却和暖如春,着一层夹袄步出监牢的瞬间,已是汗流浃背。 一阵脚步声响...
    开封第一讲书人阅读 32,823评论 1 269
  • 我被黑心中介骗来泰国打工, 没想到刚下飞机就差点儿被人妖公主榨干…… 1. 我叫王不留,地道东北人。 一个月前我还...
    沈念sama阅读 47,722评论 2 368
  • 正文 我出身青楼,却偏偏与公主长得像,于是被迫代替她去往敌国和亲。 传闻我的和亲对象是个残疾皇子,可洞房花烛夜当晚...
    茶点故事阅读 44,611评论 2 353

推荐阅读更多精彩内容