构建更好的BERT：使用GroupBERT实现训练速度两倍提升

2022/02/07

构建更好的BERT：使用GroupBERT实现训练速度两倍提升

作者:

GroupBERT组件的深度增加和算术强度降低[2]突出了存储访问。与密集操作相比，分组操作对给定的输入激活张量执行更少的FLOP。为了利用计算，这些低算术强度操作需要比密集计算更快的数据访问。由于IPU允许将所有权重和激活存储在具有极高47.5 TB/秒带宽的片上SRAM中，因此它可以将GroupBERT的理论效率优势转化为缩短跨各种模型的实际训练时间。

分组全连接层

原始的Transformer编码器层由两个模块组成：多头注意力（MHA）和全连接网络（FFN）。许多工作都集中在提高其效率上。论文 Tay et al. (2020) [3] 概述了各种不同的方法，其中大多数修改都集中在减少MHA模块对序列长度的二次计算依赖性。然而，在BERT中，大部分计算是在相对适中的128序列长度下执行的。到目前为止FFN模块消耗的资源最多，在模型执行期间贡献了近三分之二的FLOP。

FFN模块的结构非常简单：它由两个矩阵和一个非线性组成。第一个矩阵将表示投影到更高的维度，通常是模型隐藏表示的四倍。这个维度扩展操作之后是一个非饱和激活函数，该函数执行表示的非线性转换并将其稀疏化。最后，稀疏扩展表示通过向下投影矩阵收缩到模型维度。

GroupBERT引入了一种新颖的方案，使FFN计算更便宜、更快。一个关键的观察结果是，分组引起的稀疏性应用于接收稀疏输入的矩阵是最好的。因此，我们只将分组矩阵乘法引入第二个向下投影矩阵，使其成为块对角线。然而，分组操作会对传入的隐藏表示产生局部约束，从而限制了转换过程中的通道信息传播。为了解决这个问题，我们将组与输出投影混合在一起，类似于MHA块，它也为多个头划分隐藏表示。

总体而言，这种Grouped FFN (GFFN)方案允许GroupBERT将FFN层的参数数量减少25%，而任务性能的降低最小。这与所有先前尝试使用分组转换、但将其应用于两个矩阵的方法形成对比。这导致了不相交的隐藏表示，并引起性能显著下降。

组卷积作为注意力的补充

计算BERT中使用的序列长度的多对多注意力不会导致显著的计算开销。然而，最近的一项研究(Cordonnier et al. 2020) [4]表明，对于语言模型而言，首先只使用多头注意力可能是多余的。Transformer层中的注意力头子集折叠为卷积模式，以仅对局部标记交互进行建模。

为了对序列内的局部交互进行建模，需要计算密集注意力图。为了减少这一计算的冗余，GroupBERT引入了一个专用的卷积模块。组卷积充当滑动窗口，在彼此接近的单词标记之间混合信息。然后，我们用一个额外的GFFN扩展编码器跟随卷积模块，确保每个标记处理模块都与一个特征处理GFFN模块相耦合。

通过这些附加模块，局部标记交互具有专用的轻量级模型元素。这反过来又使MHA注意力能够在仅对远程交互进行建模时表现更好，因为在局部标记交互上花费的注意力容量更少。图3显示了BERT和GroupBERT在预训练验证集上的注意力图。这是一个清晰的可视化，展示了利用卷积如何让注意力在远程交互中更有效，因为模式更平滑、更分散。

图3：BERT（左）和GroupBERT（右）从验证集数据生成的注意力图。GroupBERT注意力较少依赖于相对位置，使其能够更好地关注长期交互。

让模型参数有意义

归一化

许多工作都考虑了在Transformer中应用归一化的最佳方法。虽然层归一化始终是归一化的首选方法，但它的应用方式有两种变体：PreNorm（前归一）和PostNorm（后归一）。

PostNorm（后归一）对跳过连接和残差之和的输出进行归一化，而PreNorm（前归一）在应用任何投影之前对残差分支的表示进行归一化，如图 4 所示。

图4：PostNorm（后归一）（左）和PreNorm（前归一）（右）将层归一化 (LN) 函数定位在相对于残差和跳过连接的不同位置。

BERT的标准实现使用PostNorm（后归一），它的好处是在使用默认学习率的设置下，可以提供比PreNorm（前归一）更强的任务性能。但是，我们发现PreNorm（前归一）配置明显更稳定，并且可以使用更高的学习率，而PostNorm（后归一）模型不可用。

GroupBERT使用PreNorm（前归一），并实现了卓越的任务性能，因为与在如此高的学习率下发散的PostNorm（后归一）基线相比，该模型现在可以稳定地适应4倍的学习率增长。更高的学习率对于增加模型泛化和实现更好的收敛速度很重要。

更高的学习率不会直接导致计算节省。然而，实现更高的任务性能需要更大的模型，从而导致更高的计算成本。因此，通过提高模型稳定性来实现更高的学习率会提高模型参数的利用效率。

Dropout

许多基于Transformer架构的语言模型都受益于使用dropout，因为它减少了对训练数据集的过度拟合并有助于泛化。但是，BERT在非常大的数据集上进行了预训练，在这种情况下，过度拟合通常不是问题。

出于这个原因，GroupBERT在预训练期间消除了dropout。与dropout掩码相乘所涉及的FLOP可以忽略不计，因此也可以将其视为FLOP中性优化。但是，生成dropout掩码可能会产生很大的吞吐量开销，因此在BERT中去除这种正则化方法有利于更快的模型执行和更好的任务性能。

虽然在没有dropout的情况下对维基百科数据集执行预训练是有利的，但dropout在GroupBERT微调期间仍然是一个必不可少的工具，因为微调数据集比预训练语料库小几个数量级。

GroupBERT：一个IPU transformer

在将所有概述的修改组合到一个模型中之前，我们必须验证每个组件的有效性。然而，分组FFN减少了模型FLOP的数量，而添加卷积模块会增加FLOP。出于这个原因，直接与BERT进行比较是不公平的，因为模型消耗的资源的数量是不尽相同的。为了确定我们各种尺寸的模型增强的质量，我们必须将它们与不同尺寸的BERT模型之间的对数插值进行比较，这些尺寸是：Medium（中等）、Base（基础）和Large（大型）。

图5展示了对GroupBERT引入的每个增强功能的零碎消融研究。与基线BERT 相比，模型的每一次添加都会提高帕累托效率。组合后，GroupBERT Base实现了与BERT Large模型相同的MLM验证损失，尽管它只有不到50%的参数。

图5：GroupBERT消融研究。GroupBERT中引入的所有增强都提高了BERT的帕累托效率，因为它们实现了较低的验证MLM损失。图片由作者提供。

为了明确得出GroupBERT优于BERT的结论，我们展示了GroupBERT Base相对于BERT Large的优势在广泛的模型规模上都普遍存在。我们用两个额外的模型尺寸对GroupBERT基础模型进行了补充，以创建一个连续的帕累托前部。这种方法（在图6中进行了可视化）使我们能够展示 GroupBERT相对于BERT以及其他尝试在Transformer中使用分组转换的模型的持续优势。

由于GroupBERT是IPU本地模型，在FLOP方面实现的理论节省也转化为达到给定验证准确性水平所需的更好的端到端预训练时间。

图7：IPU-POD₁₆上不同优化程序的GroupBERT Base与BERT Large的训练时间。所有模型都达到了同等的预训练验证准确性。

关键要点

我们的结果表明：

分组转换可以创建更高效的模型
即使对于需要远程交互的任务，将卷积与密集的多对多注意力相结合也会产生有益的影响
IPU将理论上的FLOP节省转化为计算时间的实际节省
通过架构更改而不是简单的模型缩放，可以实现显著的性能提升

综上所述，GroupBERT在预训练性能方面可以比原始BERT实现高达2倍的提升，从而实现更快、更高效的计算。

阅读论文

访问GitHub上的代码

鸣谢

感谢Daniel Justus、Douglas Orr、Anastasia Dietrich、Frithjof Gressmann、Alexandros Koliousis、Carlo Luschi和David Svantesson，他们也为这项研究做出了贡献，并感谢Graphcore的其他同事给予的支持和见解。

参考文献

[1] I. Chelombiev, D. Justus, D. Orr, A. Dietrich, F. Gressmann, A. Koliousis, C. Luschi, GroupBERT: Enhanced Transformer Architecture with Efficient Grouped Structures (2021), arXiv
[2] D. Masters, Delving Deep into Modern Computer Vision Models (2020), graphcore.ai
[3] Y. Tay, M. Dehghani, D. Bahri, D. Metzler, Efficient Transformers: A Survey (2020), arXiv
[4] J.B. Cordonnier, A. Loukas, M. Jaggi, On the Relationship between Self-Attention and Convolutional Layers (2020), arXiv