摘要

人类每天阅读和撰写数千亿条消息。得益于大规模数据集、高性能计算系统和更优的神经网络模型，自然语言处理（NLP）技术在理解、校对和组织这些消息方面取得了显著进展。因此，将 NLP 部署于各类应用中，以帮助网页用户、社交网络和企业具有重要价值。特别是，我们认为智能手机和其他移动设备是大规模部署 NLP 模型的关键平台。然而，当今高精度的 NLP 神经网络模型（如 BERT 和 RoBERTa）计算开销极大，例如在 Pixel 3 手机上运行 BERT-base 对一个文本片段进行分类需要约 1.7 秒。本文观察到，在计算机视觉（CV）网络中，分组卷积等方法已实现显著加速，但许多此类技术尚未被 NLP 模型设计者采用。我们展示了如何在自注意力层中用分组卷积替换若干操作，并基于此提出了新架构 SqueezeBERT。在 Pixel 3 上，SqueezeBERT 比 BERT-base 快 4.3 倍，同时在 GLUE 测试集上保持了具有竞争力的准确率。SqueezeBERT 代码即将开源。

1 引言与动机

人类每天撰写超过 3000 亿条消息 [1–4]。其中，半数以上的电子邮件在移动设备上阅读，近一半的 Facebook 用户仅通过移动设备访问 Facebook [5,6]。NLP 技术有望在多方面帮助这些用户和社区：当用户撰写消息时，NLP 模型可辅助拼写和语法检查及句子补全；当内容发布到社交网络时，NLP 可在其出现在其他用户的动态前进行内容审核；当用户阅读消息时，NLP 模型可以帮助将消息分类到不同文件夹、生成新闻推送、优先排序消息以及识别重复消息。

近年来，注意力神经网络的发展和应用使得 NLP 几乎所有领域都取得了巨大进步。2017 年，Vaswani 等人提出了多头自注意力模块，在英德机器翻译任务上超过了循环神经网络的性能 [7]。此后，GPT [8] 和 BERT [9] 将其用于句子分类，GPT-2 [10] 和 CTRL [11] 则用于句子补全与生成。近期的 ELECTRA [12] 和 RoBERTa [13] 等研究表明，更大规模的数据集和更复杂的训练方案能进一步提升自注意力网络的准确率。

考虑到人们在移动设备上产生的海量文本数据，一个自然思路是将 NLP 模型直接部署在移动设备上，嵌入到常用的阅读、撰写和分享文本的应用中。不幸的是，许多当下最先进的 NLP 模型计算开销巨大，令移动部署变得不切实际。例如，我们观察到在 Google Pixel 3 手机上运行 BERT-base 进行一次文本分类大约需要 1.7 秒。尽管过去一年中已有不少研究开始关注高效自注意力网络在 NLP 中的应用，但自 SqueezeNet [14] 问世以来，移动计算机视觉社区已在过去四年里不断优化面向移动设备的神经网络。直观来看，移动 CV 领域丰富的优化经验理应能为加速移动端 NLP 提供良好借鉴。下面，我们回顾了已被采用的技术，并提出两种附加的 CV 技术用于加速 NLP。

1.1 计算机视觉研究已为 NLP 在高效网络设计方面提供了哪些启示？

近年来，为实现更快推理，涌现了多种新型自注意力网络。目前，MobileBERT 网络在移动设备上的低延迟文本分类领域定义了最先进水平 [15]。在 Google Pixel 3 手机上，MobileBERT 对一段文本的分类耗时约 0.6 秒。在包含九个自然语言理解数据集的 GLUE 基准上 [16]，MobileBERT 的准确率也高于其他高效网络，如 DistilBERT [17]、PKD [18] 及多种方法 [19–22]。为此，MobileBERT 在其 NLP 自注意力网络中引入了两项 CV 网络中已广泛使用的概念：

1. 瓶颈层。在 ResNet [23] 中，3×3 卷积计算量较大，因此在每个 3×3 卷积层前使用 1×1 “瓶颈”卷积，以减少通道数。同理，MobileBERT 在每个自注意力层前使用瓶颈层，降低了自注意力层的计算开销。
2. 高信息流残差连接。在 BERT‑base 中，残差连接仅存在于低通道数层（768 通道）之间，而高通道数层（3072 通道）则无残差连接。ResNet 和 Residual‑SqueezeNet [14] 等 CV 网络在高通道数层之间也添加了残差连接，以增强信息流。MobileBERT 同样在高通道数层间加入了残差连接。

1.2 计算机视觉研究还能为 NLP 在高效网络设计上提供哪些启示？

我们对 MobileBERT 利用 CV 文献中流行思路来加速 NLP 的进展表示鼓舞。但我们注意到 CV 中还有两项技术尚未被 MobileBERT 采用，也可用于加速 NLP：

1. 卷积操作。自 1980 年代以来，CV 神经网络大量依赖卷积层 [24,25]。卷积既灵活又软件优化成熟，可实现从一维全连接层到三维扩张卷积（用于上/下采样）等多种功能。
2. 分组卷积。这是现代移动端优化网络中的常用技术（见第 3 节）。该方法最早由 Krizhevsky 等人在 2012 年 ImageNet 冠军模型中提出 [26–28]，虽曾一度淡出文献，约在 2016 年重新出现 [29,30]，并广泛用于 MobileNet [31]、ShuffleNet [32]、EfficientNet [33] 等高效 CV 网络中。尽管在 CV 领域常见，我们尚未发现将分组卷积应用于 NLP 的研究。

1.3 SqueezeBERT：将 CV 经验应用于 NLP

本文介绍了如何在新型自注意力网络 SqueezeBERT 的设计中引入卷积，尤其是分组卷积。实验证明，SqueezeBERT 在智能手机上的延迟低于 BERT‑base、MobileBERT 及多种其他高效 NLP 模型，同时保持了具有竞争力的准确率。

2 实现将自注意力用卷积替代

在本节中，我们首先回顾自注意力网络的基本结构；接着指出其最大计算瓶颈在于逐位置全连接（PFC）层；然后展示这些 PFC 层等价于核大小为 1 的一维卷积。

2.1 自注意力网络

在大多数基于 BERT 的网络中，通常包含三个阶段：嵌入层、编码器和分类器[9,13,12,15,19]。嵌入层将预处理后的词（以整数 token 表示）转换为浮点特征向量；编码器由一系列自注意力及其他层组成；分类器输出网络的最终结果。正如表 1 所示，嵌入层和分类器的运行时间占比不到 1%，因此我们将重点讨论编码器。

下面我们描述 BERT‑base[9] 中使用的编码器结构。该编码器由多个块依次堆叠而成。每个块包含：一个自注意力模块，随后是三个逐位置全连接层（即前馈网络，FFN）。在自注意力模块内部，又包含三个单独的逐位置全连接（PFC）层，用于为特征嵌入中每个位置生成查询（Q）、键（K）和值（V）激活向量。这些 PFC 层对嵌入序列中每个位置独立地执行相同的线性变换。

传统神经网络通常将权重与激活相乘，而注意力网络的显著特点在于相乘的是激活与激活，这使得网络能够根据输入动态地加权张量元素。此外，注意力网络可以建模任意位置间的依赖关系，而不受它们在输入或输出序列中距离的限制[7]。

Vaswani 等人提出的自注意力模块（同样被 GPT[8]、BERT[9]、RoBERTa[13]、ELECTRA[12] 等采用）按照公式$\begin{array}{r}softmax(\frac{Q{K}^T}{\sqrt{d_k}})V\end{array}$ ，计算得到输出，其中 $d_k$ 是单个注意力头的通道数。

2.2 针对移动端推理的基准测试

为了识别 BERT 中耗时最多的部分，我们在智能手机上对其进行了性能剖析。具体而言，我们使用 PyTorch 和 TorchScript，在 Google Pixel 3 手机上测量网络延迟，输入序列长度为 128，批量大小为 1。在表 1 中，我们展示了 BERT 网络主要组件的 FLOPs 和延迟占比，发现自注意力模块中的 $\begin{array}{r}softmax(\frac{Q{K}^T}{\sqrt{d_k}})V\end{array}$ 计算仅占总延迟的 11.3%。然而，自注意力模块中的逐位置全连接（PFC）层占 18.9%，前馈网络（FFN）模块中的 PFC 层占 69.4%，合计 PFC 层占据了 88.3% 的延迟。鉴于 PFC 层几乎主导了整个网络的延迟，我们接下来将重点关注如何降低 PFC 层的计算开销。

2.3 用卷积替换逐位置全连接层（PFC）

为了解决这一问题，我们计划用分组卷积替换 PFC 层——该方法已被证明可显著加速计算机视觉网络。首先，我们展示 BERT 编码器中使用的全连接层是无分组一维卷积的特例。下文中，f 表示输入特征向量，w 表示权重。给定输入特征维度为$ (P, C_{in})$，其中P为位置数、$C_{in}$ 为通道数，输出特征维度为$ (P, C_{out})$，则逐位置全连接层的运算可定义如下：
$$PositionwiseFullyConnecte{d}_{p,c_{out}}(f,w)=\sum _i^{C_{in}}{w}_{c_{out},i}\ast f_{p,i}(1)$$
那么，如果我们考虑一个核大小为 K，且输入和输出维度相同的一维卷积的定义：
$$Convolutio{n}_{p,c_{out}}(f,w)=\sum _i^{C_{in}}\sum _k^K{w}_{c_{out},i,k}\ast f_{(p-\frac{K-1}{2}+k),i}(2)$$
我们发现逐位置全连接操作等价于核大小为1的一维卷积。因此，Vaswani 等人[7]、GPT、BERT 及类似自注意力网络中的 PFC 层可直接用卷积实现，而不改变网络的数值属性或行为。

3 将分组卷积引入自注意力

既然我们已展示了如何使用卷积实现自注意力网络中计算量大的 PFC 层，现在可以将高效的分组卷积整合进自注意力网络。分组卷积定义如下：给定输入特征维度为 $(P,C_{in})$，其中 P 为位置数、$C_{in}$ 为通道数，输出特征维度为 $(P,C_{out})$，核大小为 K、分组数为 G 的一维卷积可定义为：
$$GroupedConvolutio{n}_{p,c_{out}}(f,w)=\sum _i^{\frac{C_{in}}{G}}\sum _k^K{w}_{c_{out},i,k}\ast f_{(p-\frac{K-1}{2}+k),(i+\lfloor \frac{(i)(G)}{G_{out}}\rfloor \frac{C_{in}}{G})}(3)$$
这等价于将输入向量沿通道维度分成 G 个大小为$(P,\frac{C_{in}}{G})$的子向量，并对每个子向量分别使用独立权重执行卷积，生成大小为$(P,\frac{C_{out}}{G})$ 的输出。然而，分组卷积所需的浮点运算量和权重数量仅为普通卷积的 $1/G$（不计通常在卷积层中包含的少量通道偏置操作）。

3.1 SqueezeBERT

下面介绍我们提出的 SqueezeBERT 网络架构，该架构使用分组卷积。SqueezeBERT 与 BERT‑base 非常相似，但将 PFC 层实现为卷积，并在多层中采用分组卷积。回顾第 2 节，BERT‑base 编码器的每个块包含一个自注意力模块（含 3 个 PFC 层）以及 3 个前馈网络层$（FF{N}_1、FF{N}_2、FF{N}_3）$。其维度分别为：$FF{N}_1$ 中 $C_{in}=C_{out}=768$，$FF{N}_2$ 中 $C_{in}=768$、$C_{out}=3072$，$FF{N}_3$ 中 $C_{in}=3072$、$C_{out}=768$。在自注意力模块的所有 PFC 层及 $FF{N}_2$、$FF{N}_3$ 层中，我们使用分组数 G=4 的分组卷积；为了在不同组的通道间保持信息混合，在计算量较小的 $FF{N}_1$ 层中使用 G=1。注意，在 BERT‑base 中，$FF{N}_2$ 和 $FF{N}_3$ 的运算量均为 $FF{N}_1$ 的 4 倍，但采用 G=4 后，所有 FFN 层的运算量一致。

SqueezeBERT 的嵌入维度（768）、编码器块数（12）、每个自注意力模块头数（12）、分词器（WordPiece [35,36]）等均沿用 BERT‑base。除卷积实现和分组卷积设计外，SqueezeBERT 架构与 BERT‑base 完全一致。

4 实验方法论

4.1 数据集

预训练数据。我们使用维基百科和 BooksCorpus [37] 的组合作为预训练数据，并将合并数据集的 3% 留作测试集。遵循 ALBERT 论文做法，我们采用掩码语言模型（MLM）和句子顺序预测（SOP）作为预训练任务 [19]。

微调数据。我们在通用语言理解评估（GLUE）任务集上对 SqueezeBERT（及其它基线模型）进行微调和评估。GLUE 包含九个多样化的自然语言理解任务；凭借其任务结构和广度（详见补充材料中的任务级信息），GLUE 已成为 NLP 研究的标准评测基准。模型在 GLUE 各任务上的综合表现可较好地反映其泛化能力（尤其是对其他文本分类任务）。

4.2 训练方法论

近期有关高效 NLP 网络的论文往往在训练中使用蒸馏、对抗训练和/或跨 GLUE 任务迁移学习等“花哨”技巧，但各论文之间对这些方案缺乏统一标准，难以区分模型设计与训练方案对最终准确率的贡献。为此，我们首先使用简单训练方案训练 SqueezeBERT（第 4.2.1 节，结果见第 5.1 节），然后再加入蒸馏等技术进行训练（第 4.2.2 节，结果见第 5.2 节）。

温馨提示：
阅读全文请访问"AI深语解构" SqueezeBERT：计算机视觉能为自然语言处理在高效神经网络方面带来哪些启示？