NLP中的预训练语言模型(四)—— 小型化bert(DistillBert, ALBERT, TINYBERT)

  bert之类的预训练模型在NLP各项任务上取得的效果是显著的,可是由于bert的模型参数多,推断速度慢等缘由,致使bert在工业界上的应用很难普及,针对预训练模型作模型压缩是促进其在工业界应用的关键,今天介绍三篇小型化bert模型——DistillBert, ALBERT, TINYBERT。git

一,DistillBertgithub

  论文:DistilBERT, a distilled version of BERT: smaller, faster, cheaper and lighter 算法

  GitHub:暂无网络

  DistillBert是在bert的基础上用知识蒸馏技术训练出来的小型化bert。总体上来讲这篇论文仍是很是简单的,只是引入了知识蒸馏技术来训练一个小的bert。具体作法以下:架构

  1)给定原始的bert-base做为teacher网络。框架

  2)在bert-base的基础上将网络层数减半(也就是从原来的12层减小到6层)。函数

  3)利用teacher的软标签和teacher的隐层参数来训练student网络。性能

  训练时的损失函数定义为三种损失函数的线性和,三种损失函数分别为:学习

  1)$L_{ce}$。这是teacher网络softmax层输出的几率分布和student网络softmax层输出的几率分布的交叉熵(注:MLM任务的输出)。优化

  2)$L_{mlm}$。这是student网络softmax层输出的几率分布和真实的one-hot标签的交叉熵

  3)$L_{cos}$。这是student网络隐层输出和teacher网络隐层输出的余弦类似度值,在上面咱们说student的网络层数只有6层,teacher网络的层数有12层,所以我的认为这里在计算该损失的时候是用student的第1层对应teacher的第2层,student的第2层对应teacher的第4层,以此类推。

  做者对student的初始化也作了些工做,做者用teacher的参数来初始化student的网络参数,作法和上面相似,用teacher的第2层初始化student的第1层,teacher的第4层初始化student的第2层。

  做者也解释了为何减少网络的层数,而不减少隐层大小,做者认为在现代线性代数框架中,在张量计算中,下降最后一维(也就是隐层大小)的维度对计算效率提高不大,反却是减少层数,也提高计算效率。

  另外做者在这里移除了句子向量和pooler层,在这里也没有看到NSP任务的损失函数,所以我的认为做者也去除了NSP任务(主要是不少人证实该任务并无什么效果)。

  总体上来讲虽然方法简单,可是效果仍是很不错的,模型大小减少了40%(66M),推断速度提高了60%,但性能只下降了约3%。

 

二,ALBERT 

  论文:ALBERT: A LITE BERT FOR SELF-SUPERVISEDLEARNING OF LANGUAGE REPRESENTATIONS 

  GitHub:https://github.com/brightmart/albert_zh

  ALBERT主要是从模型架构上作的改变,能极大的减少模型大小,可是没有提高推断速度。ALBERT主要作了三点改变:Factorized embedding parameterization ,Cross-layer parameter sharing ,Inter-sentence coherence loss 。咱们先来介绍这三个改变:

  Factorized embedding parameterization

  在bert中采用的是embedding层的嵌入词向量大小E等于隐层大小H,可是做者认为embedding层只是作了词嵌入,在这一层词与词之间是相互独立的,词嵌入后获得的向量包含的信息也仅仅只有当前词的信息,所以向量长度不须要那么大(咱们经常使用的word2vec向量长度通常不超过300),可是隐层由于会和其余词计算self-attention,所以隐层词对应的向量是包含了上下文信息的,此时含有的信息是很是丰富的,用小的向量容易丢失信息,须要将隐层大小设大一点。针对这样的分析,做者认为像bert中embedding层和隐层大小设置为相等是不合理的(bert-base中,E=H=768)。并且在embedding层,有一个大的词嵌入矩阵V x E。在这里V为vocab size(一般比较大,bert中就是20000多),所以当E很大时,这里的参数就很是多,又基于上面的分析,E能够不用设这么大,所以做者在这里作了一个矩阵分解,将矩阵V x H(E)分解为两个小的矩阵V x E,E x H,E << H。在这里再也不将E=H,而是将E设置为一个远小于H的值,而后再通过一个矩阵E x H将词向量维度映射到H。

  Cross-layer parameter sharing

  再Transformer结构中,咱们也能够选择共享一些参数,不如共享self-attention中的参数,或者是共享FFN层的参数,在这里做者共享了encoder中全部的参数(包括self attention和FFN),结合上面的矩阵分解,模型参数获得了极大的减少,具体结果以下图:

    

  另外这种层之间的参数共享机制,也让各层的输入输出的L2距离和余弦距离变化的更加平滑,以下图所示:

    

  Inter-sentence coherence loss

  上面两种改变主要是在减少模型的参数,而这里的改变主要是提升模型在下游任务上的性能,做者认为在bert中的NSP任务其实是有问题的,当初设计这个任务时,是但愿能学到句子一致性的,但事实上因为当初在构造句子对负例时,是从不一样的document中选择的句子,所以网络并不须要学到句子的一致性,而只要学到句子的主题(由于不一样的document,主题极可能不同),就能够判断两个句子是不是负例,所以做者认为NSP任务只学到了句子的主题,而主题分类是一个浅层语义的NLP任务。

  所以在这里做者提出了一种新的任务SOP(句子顺序预测),在这里构造正例和bert中同样,而负例就是将正例中两个句子的顺序颠倒,这里应为负例中的句子都来自于同一文档,因此经过句子的主题是没法区分正例和负例的,须要理解句子的深层语义才能区分,做者也给出了实验证实:

    

   如上面表中所示,学习SOP(86.5)任务,也能在NSP(78.9)任务上取得不错的效果,可是学习NSP(90.5)任务,是没法在SOP(52.0)任务上取得效果的。

  以上就为这篇论文改变的三个点,所以矩阵分解和参数共享,模型参数减少了不少,可是由于层数没有变,在推断时的计算量并无降低,所以推断速度没啥提高,可是训练时的速度有提高。引入SOP任务也能提高模型的效果。此外,ALBERT中的mask方式,采用的是n-gram mask,其实和SpanBert中的span mask是同样的,只是span中将mask的最大长度控制在10,而这里将最大长度控制在3。另外优化算法采用了LAMB优化器。

  最后上一个ALBERT的模型效果:

    

    

   这里的ALBERT是用的ALBERT-xxlarge,其参数以下:

    

   增大了隐层大小到4096,但实际上在这里做者证实了并非全部的模型增大隐层大小,模型的效果都会提高,好比做者在bert-large上证实了提高隐层大小到2048时,模型降低不少:

    

 

三,TINYBERT

  论文:TINYBERT: DISTILLING BERT FOR NATURAL LANGUAGE UNDERSTANDING

  GitHub:暂无

  TINYBERT也是采用了知识蒸馏的方法来压缩模型的,只是在设计上叫distillBert作了更多的工做,做者提出了两个点:针对Transformer结构的知识蒸馏和针对pre-training和fine-tuning两阶段的知识蒸馏。

  做者在这里构造了四类损失函数来对模型中各层的参数进行约束来训练模型,具体模型结构以下:

    

   做者构造了四类损失,分别针对embedding layer,attention 权重矩阵,隐层输出,predict layer。能够将这个统一到一个损失函数中:

    

   上面式子中$\lambda_m$表示每一层对应的系数,$S_m$表示studnet网络的第m层,$T_{g(m)}$表示teacher网络的第n层,其中$n = g(m)$。而且有$g(0) = 0$,$g(M+1) = N+1$,0表示embedding layer,M+1和N+1表示perdict layer。

  针对上面四层具体的损失函数表达式以下:

  attention 权重矩阵

    

   h为multi attention中头数

  隐层输出

    

   由于student网络的隐层大小一般会设置的比teacher的小,所以为了在计算时维度一致,这里用一个矩阵$W_h$将student的隐层向量线性映射到和teacher一样的空间下。

  embedding layer

    

   $W_s$同理上。

  以上三种损失函数都采用了MSE,主要是为了将模型的各项参数对齐。

  predict layer

    

   predict layer也就是softmax层,在这里的损失函数是交叉熵,t是温度参数,在这里设置为1。

  以上四种损失函数是做者针对transformer提出的知识蒸馏方法。除此以外做者认为除了对pre-training蒸馏以外,在fine-tuning时也利用teacher的知识来训练模型能够取得在下游任务更好的效果。所以做者提出了两阶段知识蒸馏,以下图所示:

    

 

   本质上就是在pre-training蒸馏一个general TinyBERT,而后再在general TinyBERT的基础上利用task-bert上再蒸馏出fine-tuned TinyBERT。

  做者给出了TinyBERT的效果:

    

   另外做者也给出了四种损失对最终结果的贡献:

    

   还有就是关于$n = g(m)$这个式子中$g(m)$怎么选择,假设student的层数为4层,这里的$n = g(m) = 3m$,做者将这种称为Uniform-strategy。另外做者还和其余的$g(m)$作了对比:

    

   Top-strategy指用teacher最后4层,Bottom-strategy指用前面4层,其实这里的映射函数,我感受可能还有更优的方案,例如取平均,或者用attention来作,可能效果会更好。