亚马逊:我们提取了BERT的一个最优子架构,只有Bert-large的16%,CPU推理速度提升7倍

自然语言处理领域,BERT 是一个里程碑式的进展。只需要添加一个单层线性分类器和一个简单的微调策略,它就能在多项任务中达到优异的性能。但另一方面,BERT 的应用也面临很多问题,如规模大、推理速度慢、预训练过程复杂。研究人员已经做了许多尝试来提取一个更简单的子架构,希望这个子架构能够保持原始 BERT 的优异性能,同时简化预训练过程,缩短推理时间。这些研究取得了不同程度的成功。然而,他们提取的这些子架构在准确率方面仍然低于原始架构的实现,而且架构参数集的选择往往看起来很随意。

虽然这个问题在计算上很难解决,但 de Wynter 最近的一项研究表明:存在一种近似算法——更具体地说,一种完全多项式时间近似模式(FPTAS)——在一定条件下能够有效地提取出具有最优保证的此类集合。

在本文中,来自 Amazon Alexa 团队的研究者将提取 BERT 最优子架构参数集这一问题细化为三个指标:推断延迟、参数大小和误差率。该研究证明:BERT 具备 strong AB^nC 属性,可满足这些条件组合,使上述算法表现得像 FPTAS。然后,研究者从一个高性能的 BERT 变体中提取了一个最优的子架构,称为 Bort,其大小是 BERT-large 的 16%,在 CPU 上的推理速度提升到原来的 8 倍。

  • 论文地址:https://arxiv.org/pdf/2010.10499.pdf

  • GitHub地址:https://github.com/alexa/bort/

尽管 FPTAS 可以确保找到表现最优的架构,但它返回的是在上述三个指标上表现最优的架构参数集,而不会输出一个训练到收敛的架构。因此,研究者对 Bort 进行了预训练,发现与原先的训练相比,预训练速度有了明显的提高:在相同的 GPU、数据集大小也相当的情况下,Bort 训练了 288 小时,BERT-large 训练了 1153 小时,而 RoBERTa-large 训练了 24,576 小时。

研究者还在 GLUE、SuperGLUE 以及 RACE 公共 NLU 基准上对 Bort 进行了评估。结果表明,与 BERT-large 相比,Bort 在所有这些基准上都获得了显著提高,提升幅度从 0.3% 到 31% 不等。

研究者在 GitHub 上开源了训练模型以及代码:https://github.com/alexa/bort/

Bort:BERT 的「最优」子架构

Bert 是一种基于 transformer 的双向全连接架构,它包括一个依赖于词汇量的嵌入层(BERT 的词汇量 V = 28,996 tokens)、包含 Transformer 的 D 编码器层,以及一个输出层。BERT 架构刚推出时有两个变体:

  • BERT-large(D = 24 编码器层、A = 16 注意力头、H = 1,024 隐藏层大小、I = 4,096 中间层大小);

  • BERT-base(D =12、A = 12、H = 768、I = 3072)。

在形式上,令 Ξ 表示包含四元组 <D, A, H, I>(即架构参数)数值有效组合的有限集合。与 de Wynter (2020b) 的研究一致,该研究将 BERT 架构族描述为某个函数的陪域(codomain),如下公式 1 所示:

算法

该研究想要找出一个架构参数集 ξ = <D, A, H, I>,对推理速度 i(b(X; ?))、参数量 p(b(?; W) 和误差率 e(b(X; W^∗ ), Y ) 这三个度量指标进行优化。

de Wynter (2020b) 表明, 对于任意架构而言,这都是一个 NP-Hard 问题。de Wynter (2020b) 中的 FPTAS 是一种近似算法,该算法依赖于对 i(?)、p(?) 和 e(?, ?) 代理函数的优化,这三个代理函数分别表示为 iˆ(?)、pˆ(?) 和 eˆ(?, ?)。执行过程中将它们表示为 Ξ 的函数,并通过选择一个参数最多、推理时间最长的架构 T∈B(T 被称为极大点,maximum point)和 W - 系数指标对它们进行标量化处理,具体如下公式 2 所示:

求取 i(?) 和 p(?) 的代理相对简单,事实上,的代理必须通过损失函数来获得。同样地,保证运行时和可逼近性依赖于两个额外的输入参数:选定的最大训练步数量 n > 0,预期间隔大小 1 ≤  ≤ |Ξ|。? 的的选择直接影响该近似算法求得解的质量。

用知识蒸馏进行预训练

尽管 FPTAS 能够确保我们获得描述最优子架构的架构参数集,但如何高效预训练参数化模型仍是一个待解决问题。

根据以往的研究(详见论文第二章)可以得出结论,使用知识蒸馏(KD)来预训练上述语言模型可以在前述评估指标上达到良好的性能。鉴于代理误差函数 eˆ(?, ?) 是关于极大点的交叉熵,因此将上述评估通过 KD 进行扩展是很自然的事情。

该研究还比较了 Bort 架构的自监督预训练和基于 KD 的预训练,发现与另一种方法相比,使用学生模型的最后一层和教师模型之间的一个简单交叉熵就足以找到一个优秀模型,该模型可以获得更高的遮蔽语言模型(MLM)准确率和更快的预训练速度。

评估

为了验证在经过最佳子架构提取过程之后,Bort 是否保留了 BERT 和 RoBERTa 的强大泛化性能,研究者在 GLUE、SuperGLUE 基准以及 RACE 数据集上微调了 Bort。结果表明,Bort 比其他类似 BERT 的压缩模型要好得多,在多个任务上的性能都大大优于其他模型

GLUE

GLUE(Generalized Language Evaluation benchmark)包含一组常见的自然语言任务。主要是侧重于自然语言推理(NLI),它由十个数据集组成。

研究者通过在所有任务中添加单层线性分类器来微调 Bort,但 CoLA 除外。在 CoLA 中,研究者发现在 Bort 和分类器之间添加额外的线性层可以提高收敛速度。研究者使用 Agora 对所有任务进行了微调。

结果如表 4.15 所示。除了 QQP 和 QNLI 以外,Bort 几乎在所有任务上表现优异,它的性能比其他基于 BERT 的同等模型要好得多。相比于 BERT-large,Bort 的性能提高了 0.3%-31%。研究者将这种提升归功于 Agora 的微调,因为它允许模型更好地学习每个任务的目标分布。

SuperGLUE

SuperGLUE 包含一组常见的自然语言任务,它由十个数据集组成。

研究者通过添加单层线性分类器来微调 Bort,并在所有任务中运行 Agora 至收敛。结果如表 5 所示:Bort 在除了 ReCoRD 以外的所有任务中获得了超越或媲美 BERT-large 的结果

RACE

RACE 数据集是一个文本阅读的多项选择题数据集,经过专业注释,被分成两个数据集:RACE-H(从高中考试题目中提取的数据集)和 RACE-M(从初中考试题目中提取的数据集)。

与之前实验一样,研究者通过添加单层线性分类器来微调 Bort,并运行 Agora 进行收敛。

结果如表 6 所示。总体而言,Bort 获得了良好的结果,在两个任务上都比 BERT-large 的性能高出 9-10%