推理速度提升29倍，参数少1/10，阿里提出AdaBERT压缩方法

2020 年 1 月 23 日 机器之心

机器之心发布

机器之心编辑部

作为当前最佳的自然语言处理模型，BERT 却存在规模大、成本高和实时性差等缺点。为了能在实际应用中部署这种技术，有必要对 BERT 进行压缩。此前机器之心就已经介绍了几种来自不同研究机构的压缩方案，参阅《内存用量 1/20，速度加快 80 倍，腾讯 QQ 提出全新 BERT 蒸馏框架，未来将开源》和《 AAAI 2020 | 超低精度量化 BERT，UC 伯克利提出用二阶信息压缩神经网络》。

今天，我们要介绍的是一项来自阿里巴巴的压缩方案 AdaBERT，这一方案的特点是能针对具体任务得到性能不会显著下降的小型模型。

论文链接： https://arxiv.org/pdf/2001.04246v1.pdf

BERT 等大型预训练语言模型研究已经表明，它们能够有效地处理多种不同的自然语言处理任务。但是，这些模型的参数规模庞大，难以部署到需要通过有限资源执行快速推理的实时应用中。

对于这一问题，现有的方法是将 BERT 压缩成小型模型，但这种压缩方法与任务无关，也就是说对于不同的下游任务而言，压缩方法是一样的。

面向任务的 BERT 压缩方法是有必要的而且很有用，为此阿里巴巴的研究者提出了一种全新的压缩方法 AdaBERT。 该方法利用了可微神经架构搜索来自动将 BERT 压缩成适应不同特定任务的小型模型。

研究者为 AdaBERT 提出了两种不同的损失函数。一是面向任务的知识注入损失，可为搜索过程提供提示；二是效率感知型损失，这能提供搜索约束。这两个损失能为任务适应型 BERT 压缩提供效率和有效性之间的平衡。

研究者在多个 NLP 任务上对 AdaBERT 进行了评估，结果表明 这些任务适应型压缩模型在保证表现相当的同时，推理速度比 BERT 快 12.7 到 29.3 倍，同时参数缩小至 BERT 的 11.5 到 17.0 之一的规模。

方法

如下图 1 所示，研究目标是针对一项具体任务，将给定的大型 BERT 模型压缩成一个有效且高效的任务适应型模型。 压缩后模型的结构是以可微分方式搜索到的， 这个过程需要用到来自大型 BERT 模型的面向任务的知识，同时也要重点考虑模型的效率。

图 1：AdaBERT 概况。通过考虑来自原始 BERT 的对任务有用的知识，同时兼顾模型效率，阿里巴巴提出的新方法能以可微分的方式为目标任务搜索合适的小型模型。

在数学形式上，研究者将在目标数据集 D_t 上进行微调后的 BERT 模型记为 BERT_t，将架构搜索空间记为 A。任务目标是找到一种最优架构 α∈A，方法是最小化以下损失函数：

其中 w_α 是架构 α 的可训练的网络权重（比如一个前馈层的权重）; L_CE、L_KD 和 L_E 分别是目标任务损失、面向任务的知识注入损失和效率损失。

具体来说，L_CE 是与目标数据 D_t 相关的交叉熵损失；L_KD 是面向任务的知识注入损失，能为该任务提供用于寻找适当结构的线索；L_E 是效率感知型损失，用于提供搜索约束，以便找到轻量级的高效结构。γ 和 β 是超参数，用以保持这些损失项之间的平衡。

面向任务的知识注入

找到对任务有用的知识的探针

为了从教师模型 BERT_t 中分层地分解出对任务有用的知识，研究者使用了一组探针分类器，然后将分解出的知识注入到压缩后的模型中。

具体来说，研究者先是冻结了 BERT_t 的参数，然后根据 groundtruth 任务标签为每个隐藏层训练一个 softmax 探针分类器。总体模型有 J 个分类器（BERT-base 的 J=12），第 j 个分类器的分类 logit 可被视为从第 j 层学习到的知识。给定一个输入实例 m，用

表示 BERT_t 的第 j 层的隐藏表征，用

表示压缩后学生模型的第 i 层上通过注意力机制求和的隐藏状态，注入对任务有用的知识（分类 logit）的方式如下：

注意力机制的分层迁移

如 Liu et al., 2019b 研究所示的那样，对于不同的任务，BERT 中每一层的作用都不一样。这里研究者通过注意机制来结合所有层的分解后的知识：

此外，为了丰富对任务有用的知识，研究者还在目标任务数据集上执行了数据增强，使用的增强过程与 Jiao et al., 2019 的一样，其中利用了 BERT 和 GloVe 来替换原始文本中的词。

效率感知型损失

因为该研究的目标是将原始 BERT 模型压缩成高效的压缩模型，所以研究者进一步从两个方面将模型效率纳入到了损失函数中，即参数规模和推理时间。具体来说，对于被搜索的架构 α_c 和 K，研究者如等式 (1) 那样定义了效率感知型的损失：

可微分架构搜索

新提出的方法与已有的 BERT 压缩方法有一个主要差别， 即新提出的 AdaBERT 方法的目标是为不同的任务寻找适应任务的结构 。研究者提出使用可微分的结构搜索方法来通过前述的损失项来搜索适应任务的结构。

搜索空间设置

首先看看该方法的搜索空间，如下图所示：

图 2：包括堆叠的层和单元的搜索空间

对于单元中的候选操作，研究者采用了基于轻量级 CNN 的操作，因为它们能够有效地处理 NLP 任务，而且已有研究表明基于 CNN 的操作的推理速度比基于 RNN 和基于自注意力的模型更快，因为基于 CNN 的操作对并行处理比较友好。

搜索算法

通过直接地暴力穷举所有候选操作来优化等式 (1) 中的整体损失是不可能办到的，因为组合操作具有巨大的搜索空间，在 α 上训练 w_α 需要消耗大量时间。为了解决这一问题，研究者在这里采用的方法是将被搜索的架构

建模为离散的变量，并使其服从离散概率分布

和

。

因此，K 和 o_{i,j} 都被建模为了 one-hot 变量，并且分别是从层范围 [1, K_max] 和候选操作集 O 采样的。但是，由于这个离散采样过程让梯度无法传递回可学习的参数 P_K 和 P_o，因此 L 是不可微分的。研究者受之前的研究（Xie et al., 2019; Wu et al., 2019）的启发，使用了 Gumbel Softmax 技术来将分类样本弛豫化到连续样本向量

1和

中：

通过这种方式，y^K 和 y^o 就成了离散样本的可微分代理变量，然后就可以直接使用基于梯度的优化器来有效地优化 L 了。

注意，整体损失 L 的优化可被视为学习一个由搜索空间 α 所定义的父图，其候选操作的权重 w_α 和架构分布 P_α是同时训练的。在训练阶段，P_α 所引入的随机性可以增强对合适的上下文编码器的探索，这些编码器可以在资源限制下模仿特定于任务的教师 BERT_t。在完成父图（parent graph）的训练之后，可以通过在 P_α 上使用 argmax 来推导得到一个高效的任务适应型子图，并以此作为压缩后的模型。

前面已经提到，知识注入损失 L_KD 能为 P_α 上的架构采样提供正则化，而效率感知型损失 L_E 则有助于得到稀疏结构，使模型紧凑和高效。

实验

实验采用了 6 个数据集进行评估，下表总结了压缩结果， 包括参数规模、推理速度和分类准确度。

表 1：模型压缩结果。

下表总结了 AdaBERT 方法在不同任务上的详细结果。

表 2：通过 AdaBERT 为不同任务搜索到的结构的层数、参数数量和推理速度提升情况。

整体上讲，AdaBERT 方法在所有评估数据集上都实现了显著的效率提升，同时还保持着相当的性能。相比于 BERT_12-T，压缩后的模型的参数规模小了 11.5 到 17.0 倍，推理速度则提升了 12.7 到 29.3 倍，同时平均准确度表现仅有 2.79% 的降低。这说明 AdaBERT 可以将 BERT 压缩成特定于任务的小型模型。

适应性研究

为了进一步研究通过 AdaBERT 搜索到的架构的适应能力，研究者将搜索得到的压缩模型结构应用到了不同的下游任务上。下表 3 给出了跨任务验证的结果。