居委会大爷|别再蒸馏3层BERT了!变矮又能变瘦的DynaBERT了解一下( 二 )


在MHA中 , 我们以为不同的head抽取到了不同的特征 , 因此每个head的作用和权重肯定也是不同的 , intermediate中的神经元连接也是 。 假如直接按照粗暴裁剪的话 , 大概率会丢失重要的信息 , 因此作者想到了一种方法 , 对head和神经元进行排序 , 每次剪枝掉不重要的部门 , 并称这种方法为Netword Rewiring 。
对于重要程度的计算参考了论文[5] , 核心思想是计算去掉head之前和之后的loss变化 , 变化越大则越重要 。
利用Rewiring机制 , 便可以对注意力头和神经元进行排序 , 得到第一步的teacher模型 , 如图:
居委会大爷|别再蒸馏3层BERT了!变矮又能变瘦的DynaBERT了解一下
本文插图

要留意的是 , 固然跟着参数更新 , 注意力头和神经元的权重会变化 , 但teacher模型只初始化一次(在后文有验证增加频率并没带来太大晋升) 。 之后 , 每个batch会练习四种student模型 , 如图:
居委会大爷|别再蒸馏3层BERT了!变矮又能变瘦的DynaBERT了解一下
本文插图

蒸馏的终极loss来源于三方面:logits、embedding和每层的hidden state 。
2、深度自适应 Adaptive Depth
训好了width-adaptive的模型之后 , 就可以训自适应深度的了 。 浅层BERT模型的优化实在比较成熟了 , 主要的技巧就是蒸馏 。 作者直接使用训好的作为teacher , 蒸馏裁剪深度后的小版本BERT 。
对于深度 , 系数 , 设层的深度为[1,12] , 作者根据去掉深度为d的层 。 之所以取是因为研究表明最后一层比较重要[6] 。
最后 , 为了避免灾难性遗忘 , 作者继承对宽度进行剪枝练习 , 第二阶段的练习方式如图:
居委会大爷|别再蒸馏3层BERT了!变矮又能变瘦的DynaBERT了解一下
本文插图

2
实验
根据练习时宽度和深度的裁剪系数 , 作者终极可得到12个大小不同的BERT模型 , 在GLUE上的效果如下:
居委会大爷|别再蒸馏3层BERT了!变矮又能变瘦的DynaBERT了解一下
本文插图

可以看到 , 剪枝的BERT效果并没有太多下降 , 并且在9个任务中都超越了BERT-base 。 同时 , 这种灵活的练习方式也给BERT本身的效果带来了晋升 , 在与BERT和RoBERTa的对比中都更胜一筹:
居委会大爷|别再蒸馏3层BERT了!变矮又能变瘦的DynaBERT了解一下
本文插图

另外 , 作者还和DistillBERT、TinyBERT、LayerDrop进行了实验对比 , DynaBERT均获得了更好的效果 。
在消融实验中 , 作者发现在加了rewiring机制后准确率均匀晋升了2个点之多:
居委会大爷|别再蒸馏3层BERT了!变矮又能变瘦的DynaBERT了解一下
本文插图

3
结论
本篇论文的立异点主要在于Adaptive width的练习方式 , 考虑到后续的裁剪 , 作者对head和neuron进行了排序 , 并利用蒸馏让子网络学习大网络的知识 。
总体来说仍是有些点可以挖的 , 好比作者为什么选择先对宽度进行自适应 , 再宽度+深度自适应?这样的好处可能是在第二阶段的蒸馏中学习到宽度自适应过的子网络知识 。 但直接进行同时练习不可以吗?仍是但愿作者再验证一下不同顺序的差距 。
为了简化 , 作者在宽度上所做的压缩比较简单 , 之后可以继承尝试压缩hidden dim 。 另外 , ALBERT比拟原始BERT实在更适合浅层Transformer , 也可以作为之后的尝试方向 。
参考文献:
[1]LayerDrop: https://arxiv.org/abs/1909.11556
[2]DistillBERT: https://arxiv.org/abs/1910.01108
[3]TinyBERT: https://arxiv.org/abs/1909.10351
[4]DynaBERT: https://www.researchgate.net/publication/340523407_DynaBERT_Dynamic_BERT_with_Adaptive_Width_and_Depth


推荐阅读