文 | Cheney
编 | 智商掉了一地

既让模型跑得快又好，又要空间占用少，真的有这样效果爆表的操作存在吗？

在当前疫情反复的情况下，大家平时出门用健康宝刷脸的频率变得越来越高。如果每次人脸识别都需要等很久的话，那也太让人焦虑啦，甚至会有砸手机的冲动。而且我们还想让模型的识别效果好，不能经常出错，就得用一些跑起来很慢的大模型来提升效果。那么问题来了，咋在这两者之间做个平衡呢，既让模型跑得又快又好，又要占用空间够少，还能效果爆表？那就不得不提模型压缩和最近兴起的自动机器学习技术了，在压缩模型的同时还能自动搜索出效果更优的结构，岂不美哉。

在NLP领域，预训练语言模型（如BERT等）在大量任务上都取得了不错的效果，但是它们有着过大的体积以及推理时间，这严重限制了实际的部署。为了解决这个问题，很多模型压缩方法被提出，例如知识蒸馏（Knowledge Distillation，KD）技术。虽然用于预训练模型压缩的知识蒸馏方向取得了很大的进展，但针对Transformer结构中的前馈神经网络（Feed-forward Network, FFN）压缩和改进方面的研究还很少，其计算代价是多头注意力（Multi-head Attention，MHA）模块的2-3倍；同时[1]也指出，前馈网络里的多层感知器（Multilayer Perceptron，MLP）结构可以防止自注意机制带来的秩崩塌（rank collapse）的问题，可以帮助模型的优化。因此，如何在最大限度压缩前馈网络的同时，寻找到具有更优秀非线性的前馈网络结构，进而提升模型效果，是个很值得关注的问题。

此外，自动机器学习技术，特别是模型架构自动化搜索（Neural Architecture Search，NAS），在计算机视觉和自然语言处理领域已经崭露头角，可以自动地从庞大的搜索空间中搜索出比人工设计的结构效果更优的模型。因此可以很自然地想到利用这类技术搜索出符合需求的前馈网络结构。但由于搜索空间往往极其庞大和复杂，如何高效地进行模型架构搜索也是亟待解决的问题。

基于网络架构搜索的模型压缩方法之前也有很多人进行研究，比如AdaBERT[2]基于卷积神经网络搜索任务自适应的小模型。然而，这些方法没有针对Transformer中前馈网络的非线性进行探究；搜索空间受限于人工设计，不够灵活；泛化性也值得考究。

为了解决上述问题，今天我们介绍一篇收录于Findings of EMNLP 2021的工作EfficientBERT。在搜索空间方面，该方法主要针对Transformer中前馈网络的数学表达式进行搜索，以发现具有更优非线性的模型结构；同时也将网络的堆叠层数和中间通道数加入进来，以综合平衡模型的参数量和推理时间。在搜索方法方面，该文提出了一种由粗到细的搜索策略，渐进地对搜索空间中的每个部分进行搜索，以提升搜索效率；同时提出了一种热启动知识蒸馏（Warm-up KD）策略，使每个搜索子结构在训练时可以更快地收敛，并提升模型的泛化性能。

论文标题
EfficientBERT: Progressively Searching Multilayer Perceptron via Warm-up Knowledge Distillation

论文链接
https://arxiv.org/abs/2109.07222

开源代码
https://github.com/cheneydon/efficient-bert

1 方法

本文提出的方法主要分成两个部分：搜索空间的设计，以及基于热启动知识蒸馏的渐进搜索方法。

1.1 搜索空间

搜索空间设计方面，首先对前馈网络的数学表达式进行搜索，以此反映网络的非线性能力，并为其定义了如下几种基本的数学操作符号：

同时该文也对前馈网络的堆叠层数和中间通道扩增比例进行搜索，以平衡模型的参数量和推理时间，分别从{1, 2, 3, 4}和{1, 1/2, 1/3, 1/4}中选取。在搜索时，每个前馈网络结构通过一个有向无环图构建，且上述操作被放入中间节点中来处理输入特征。

值得一提的是，正如上述框架图所示，该文对每层前馈网络整体的数学表达式进行搜索，包括激活函数以及一些复杂的嵌套或组合表达，且各层网络的表达式也各不相同。比如该文搜到的EfficientBERT模型第4, 5层的表达式为linear2_2(gelu(linear2_2(max(x,tanh(x)),wb1)),wb2)和linear1_2(relu(mul(linear1_2(x,wb1),x)),wb2)，其中linear x_y表示其所在的前馈网络的堆叠层数为x、中间通道扩增比例为1/y，wb i表示第i个线性层的权重和偏置。

1.2 搜索方法

首先是基础搜索模型（Base Model）结构的设计。该文综合采用了之前工作的一些压缩方法，比如嵌入层因子分解（embedding factorization），减少模型宽度、深度、中间通道扩增比例等。

接下来是搜索过程，整个搜索过程被分为三个阶段：

• 在第一阶段，对整个搜索空间进行联合搜索，其基于上述基础搜索模型结构。每个搜索子模型通过LaNAS[3]中提出的一种可学习的采样决策树进行采样，使得更具有潜力的子模型可以以更大的概率被采样到。为了避免不同数学表达式之间参数干扰的问题，在这个阶段不同的子模型被单独地进行训练。同时，为了使每个子模型的训练更快地收敛，该文提出一种热启动知识蒸馏方法。首先建立一个热启动超级网络，该网络有着最大的堆叠层数和中间通道扩增比例。将该超级网络通过知识蒸馏预训练好之后固定其权重，在搜索时每个子模型从对应位置继承其权重，之后利用知识蒸馏只需要少量的预训练和微调步数即可达到收敛，节省了大量搜索时间。

• 在第二阶段，对数学表达式进行单独搜索，以发掘其更多样化的表达并评估其对模型最终效果的影响，并将其它两个部分的结构固定住（即堆叠层数和中间通道扩增比例）。该阶段基于第一阶段搜索到的结构进行搜索，并采用和第一阶段相同的知识蒸馏方式。

• 在第三阶段，对堆叠层数和中间通道扩增比例进行联合搜索。该阶段基于第二阶段搜索到的结构进行搜索，其数学表达式被固定，但使用了不同的堆叠层数和中间通道扩增比例组合，以增加搜索多样性。由于无需对数学表达式进行搜索，该阶段采用权重共享方法加速搜索。在该阶段的热启动知识蒸馏过程中，将热启动超级网络预训练好之后，其权重不再被固定；在搜索时，每个搜索子模型被均匀采样，继承超级网络参数后进行训练，且不同子模型的参数可以共享。为了提高模型的泛化性，在训练时采用了多任务训练的方法，让每个子模型在不同种类的下游任务上进行微调。其中嵌入层和Transformer层的参数在所有任务上共享，但不对预测层参数进行共享。

2 实验

首先是在GLUE test和SQuAD dev数据集上与之前的模型压缩方法进行比较：

并在GLUE dev数据集上与之前相关的利用架构搜索进行模型压缩的方法进行比较：

可以看到，在相似的参数量下，本文搜索到的模型EfficientBERT效果远远超过了之前效果最好的模型压缩方法（TinyBERT[4], MobileBERT[5], MiniLM[6]等），也超过了之前利用架构搜索的模型压缩方法（AdaBERT[2]，NAS-BERT[7]等）。

接下来，本文也给出各阶段搜索到模型的最终性能，验证提出的渐进架构搜索方法、以及搜索空间中各个部分的有效性：

本文也给出各个阶段的搜索时间，验证提出的渐进搜索方法的加速能力：

可以看到，本文提出的渐进搜索方法不仅可以搜索出更具有潜力的模型结构，而且可以加速搜索进程。

为了测试搜到的模型结构的泛化能力，EfficientBERT结构随后被迁移到了TinyBERT上：

此外，本文也给出各个子模型在搜索阶段和最终测试阶段预测结果的相关性结果，验证所提出的热启动知识蒸馏方法对不同子模型效果的预测排序能力：

可以看到，使用热启动知识蒸馏方法，子模型在各个下游任务上都可以在搜索和最终测试阶段的预测结果之间保持很高的正相关性。

为了直观地评估模型的非线性能力，本文对搜索到的前馈网络结构进行了可视化：

其中(d)-(f)是随机采样出的效果更差的几个搜索子模型。从上图可以看出，(a)-(c)相比于(d)-(f)曲线更加流畅且异常突增区域更少；并且从(a)到(c)，曲线的复杂程度越来越低，显示了搜索到的模型EfficientBERT优秀的非线性能力。

最后，本文将各个阶段搜索到的模型整体结构进行了可视化：

从最终阶段搜到的EfficientBERT结构中可以发现一些有趣的规律：

• 由于多数前馈网络中间通道扩增比例为1/2，且多数前馈网络堆叠层数少于2，因此搜索得到的EfficientBERT非常轻量；

• 更低层有着更大的堆叠层数或中间通道扩增比例（如第1, 2层），以最大程度地丰富模型的语义表达能力；

• 更高层有着更复杂的数学表达式（如第4, 5层），以提高模型的非线性表达能力。

3 总结

本文主要针对前馈网络的尺寸和非线性能力进行压缩和提升，设计了一种非常灵活庞大的搜索空间进行自动化架构搜索。为了提高搜索效率，本文提出了一种由粗到细的架构搜索方法，并在每个搜索阶段采用了一种热启动知识蒸馏方法来加速搜索并提高模型的泛化能力。大量实验证明了搜索到的模型EfficientBERT的有效性。

后台回复关键词【入群】

加入卖萌屋NLP/IR/Rec与求职讨论群

后台回复关键词【顶会】

获取ACL、CIKM等各大顶会论文集！

 

[1] Yihe Dong, et al. Attention is not all you need: Pure attention loses rank doubly exponentially with depth. ICML 2021.

[2] Daoyuan Chen, et al. AdaBERT: Task-adaptive bert compression with differentiable neural architecture search. IJCAI 2020.

[3] Linnan Wang, et al. Sample-efficient neural architecture search by learning action space. TPAMI 2021.

[4] Xiaoqi Jiao, et al. TinyBERT: Distilling BERT for natural language understanding. Findings of EMNLP 2020.

[5] Zhiqing Sun, et al. MobileBERT: a compact task-agnostic BERT for resource-limited devices. ACL 2020.

[6] Wenhui Wang, et al. MiniLM: Deep self-attention distillation for task-agnostic compression of pre-trained transformers. NeurIPS 2020.

[7] Jin Xu, et al. NAS-BERT: Task-agnostic and adaptive-size bert compression with neural architecture search. KDD 2021.