ALBERT：轻量级BERT语言模型 ICLR2020

论文链接：https://arxiv.org/pdf/1909.11942.pdf
代码链接：https://github.com/google-research/ALBERT

导读

今天阅读的是 Google 同学 2020 年的ICLR论文《ALBERT: A LITE BERT FOR SELF-SUPERVISED LEARNING OF LANGUAGE REPRESENTATIONS》。

我们知道模型效果会随着模型深度的增加得到提升，然而模型深度的增加也会使得训练变得更困难，为了解决这个问题，Google 的同学提出了一个轻量级的 BERT：ALBERT，参数比 BERT-large 更少，且效果更好。

摘要

预训练自然语言表征时，增加模型大小一般是可以提升模型在下游任务中的性能。但是这种纯粹依赖模型尺寸进而期望大力出奇迹的想法在未来会越发困难。进一步增加模型大小将带来以下困难：

• (1)GPU/TPU内存不足
• (2)训练时间会更长
• (3)模型退化。

所以，为了解决上述这些问题，本文提出通过两种参数精简技术来降低内存消耗，并加快BERT的训练速度。此外，本文还引入一个自监督损失(self-supervised loss)，用于对句子连贯性(inter-sentence coherence)建模，并证明该损失函数能够提升多句子作为输入的下游任务的性能。本文所提出的模型ALBERT在 GLUE、RACE 和 SQuAD 这3个基准上都取得了新的SOTA结果，且参数量还少于 BERT-large。

1、引言

通常来说，模型深度与模型效果成正比，但是模型越深也会带来动则数亿甚至数十亿的参数量，这就对计算内存有了一定的要求。而在分布式训练中，通信开销与参数也成正比，所以其对训练速度也产生了显著的影响。

目前已有的解决方案要么是并行化，要么是内存管理，但是都没有解决通信开销，即降低模型本身的参数。

研究者还观察到，仅仅增加 BERT-large 等模型的隐藏层大小也会导致性能下降。如下表 1 和图 1 所示，研究者将 BERT-large 的隐藏层大小增加一倍，该模型（BERT-xlarge）在 RACE 基准测试上的准确率显著降低。

在本文，作者设计一种轻量级的 Bert，并取名为 ALBERT（A Lite BERT），ALBERT 的参数比传统的 BERT 要少得多，有效的解决了模型通信开销的问题。

ALBERT 提出了两种降低参数的技术，包括 factorized embedding parameterization 和 cross-layer parameter sharing。前者是将 Embedding 参数矩阵分解为两个小矩阵，后者是跨层的参数共享。

除了降低参数外，为了提高性能，作者还引入了一种自监督的句子预测目标（sentence-order prediction，SOP），主要用来关注句子间的连贯性，旨在解决 BERT 中 next sentence prediction（NSP）的无效性问题。

最终 ALBERT 比 BERT-large 参数更少，效果更好，并在多个 NLP 任务中刷榜。

2、ALBERT模型

ALBERT 架构的主干和 BERT 类似，都使用了基于 GELU 的非线性激活函数的 Transformer。但是其分别在两个地方减少了参数量。

以下图为例可以看到模型的参数主要集中在两块，一块是 Token embedding projection block，另一块是 Attention feed-forward block，前者占有 20% 的参数量，后者占有 80% 的参数量。

2.1、Factorized embedding parameterization

我们先来看 Token embedding projection 模块。

在 BERT 中，Token Embedding 的参数矩阵大小为$(V\times H)$$ ，其中 $V$ 表示词汇表长度，$H$ 为隐藏层大小。即：

而 ALBERT 为了减少参数数量，在映射中间加入一个大小为 $E$ 的隐藏层，这样矩阵的参数大小就从 $O(V\times H)$ 降低为 $O(V\times E+E\times H)$ ，而 $E<<h$ 。

之所以可以这样做是因为每次反向传播时都只会更新一个 Token 相关参数，其他参数都不会变。而且在第一次投影的过程中，词与词之间是不会进行交互的，只有在后面的 Attention 过程中才会做交互，我们称为 Sparsely updated。如果词不做交互的话，完全没有必要用一个很高维度的向量去表示，所以就引入一个小的隐藏层。（这个过程是不是很像矩阵分解）

这里便完成了参数降低的第一个模块。

2.2、Cross-layer parameter sharing

ALBERT 的参数共享主要是针对所有子模块内部进行的，这样便可以把 Attention feed-forward 模块参数量从 $O(12\times L\times H\times H)$ 降低到 $O(12\times H\times H)$ ，其中 $L$ 为层数，$H$ 为隐藏层大小。

参数共享能显著减少参数。共享可以分为全连接层、注意力层的参数共享；注意力层的参数对效果的减弱影响小一点。

ALBERT 之所以这样做是因为，考虑到每层其实学习到内容非常相似，所以尝试了将其进行参数共享。下图为不同层 Attention 学到的东西：

这样变完成了参数降低的第二个模块。

2.3、Sentence order prediction

谷歌自己把它换成了 SOP。这个在百度 ERNIE 2.0 里也有，叫 Sentence Reordering Task，而且 SRT 比 SOP 更强，因为需要预测更多种句子片段顺序排列。ERNIE 2.0 中还有一些别的东西可挖，比如大小写预测 Captialization Prediction Task、句子距离 Sentence Distance Task。

• NOP：下一句预测， 正样本=上下相邻的2个句子，负样本=随机2个句子
• SOP：句子顺序预测，正样本=正常顺序的2个相邻句子，负样本=调换顺序的2个相邻句子
• NOP任务过于简单，只要模型发现两个句子的主题不一样就行了，所以SOP预测任务能够让模型学习到更多的信息

BERT使用的NSP损失，是预测两个片段在原文本中是否连续出现的二分类损失。目标是为了提高如NLI等下游任务的性能，但是最近的研究都表示 NSP 的作用不可靠，都选择了不使用NSP。

作者推测，NSP效果不佳的原因是其难度较小。将主题预测和连贯性预测结合在了一起，但主题预测比连贯性预测简单得多，并且它与LM损失学到的内容是有重合的。

SOP的正例选取方式与BERT一致（来自同一文档的两个连续段），而负例不同于BERT中的sample，同样是来自同一文档的两个连续段，但交换两段的顺序，从而避免了主题预测，只关注建模句子之间的连贯性。

具体的损失函数表达式读者可以查阅原论文，但研究者表示，在使用了该损失函数后，ALBERT 能显著提升下游多句子编码任务的性能。

• 使用段落连续性任务。正例，使用从一个文档中连续的两个文本段落；负例，使用从一个文档中连续的两个文本段落，但位置调换了。
• 避免使用原有的NSP任务，原有的任务包含隐含了预测主题这类过于简单的任务。

3、实验

我们简单看下实验。

3.1、实验设置

为了进行更公平的对比，本文一方面使用与原始 BERT相同的配置训练试验模型，另一方面采用 BOOKCORPUS 和 English Wikipedia 共计 16GB 的纯文本作为预训练任务的数据。与BERT一样，使用的词典大小是30000；此外还借鉴了XLNet中使用的SentencePiece。在MLM目标函数上使用n−gram的masking，随机选用n−gram遮蔽输入。预测生成的n-gram的概率：$$p(n)=\frac{1/n}{{\sum }_{k=1}^{N}1/k}$$本文设置的n-gram最大长度为3，即MLM目标最多由3个全词组成。

BERT的MLM目标是随机MASK15%的词来预测，ALBERT预测的是N-gram片段，包含更多的语义信息，每个片段长度n（最大为3），根据概率公式计算得到。比如1-gram、2-gram、3-gram的的概率分别为6/11、3/11、2/11.越长概率越小。

下图为 BERT 和 ALBERT 的参数对比：

下图为模型的参数、精度、速度的对比，可以看到模型的参数变小了，但是速度下降很多。从Table 2看出，ALBERT-xxlarge 的参数量只有 BERT-Large 70% ，但性能却能够显著超越BERT-large。具体表现在SQuAD v1.1上提升1.7%，在SQuAD v2.0上提升4.2%，MNLI上提升2.2%，在SST-2上提升3.0%，在RACE上提升8.5%。此外，还观察到BERT-xlarge在全部的指标上全面溃败于BERT-base。这说明形如BERT-large的大参数模型相较于更小参数量的模型是更难训练的。另一个有趣的发现是，相同训练配置下(相同TPUs数量)训练一次，数据的吞吐速度差异。由于ALBERT模型通信更少，计算量更低，所以相比于BERT有更高的数据吞吐量。这里选用最慢的BERT-large作为速度参照基准。

下图对比了模型共享参数下的精度和参数大小：

下图为 SOP 相对 NSP 的提升：本文这里对比了3种策略：没有句子间损失(比如XLNet和RoBERTa)、NSP(比如BERT)、SOP(ALBERT)。这里采用的ALBERT也是ALBERT-base。对比过程，一方面对比自身任务中的准确率，另一方面是下游任务的性能表现。在自身任务这一维度，可以看出NSP损失对于SOP几乎是没有任何益处，NSP训练后，在SOP上的表现只有52%，这跟瞎猜差不了多少。据此，可以得出结论：NSP建模止步于主题识别。反观SOP损失，确实一定程度上能够解决NSP任务，其准确率为78.9%，而自身的准确率为86.5%。更为重要的是，在下游任务上SOP损失统统起到促进作用，具体表现在SQuAD1.1提升1%，SQuAD 2.0提升2%，RACE提升1.7%。

下图为 ALBERT 的刷榜结果：


下图为去掉 dropout 前后的精度。这里解释一下：dropout 对于防止过拟合来说具有非常好的效果，但是对于 Masked Language Model 来说，其学习本身就比较困难，所以不用担心过拟合）：

另外加入 dropout 后会增加很多临时变量，删掉 dropout 后可以提升内存的利用率。此外 ALBERT 还加了十倍的数据量= =：

4、结论

ALBERT 设计了 factorized embedding parameterization 和 cross-layer parameter sharing 两种方式分别对模型的两个部分降低参数，并利用 Sentence order prediction 来改进 BERT 的 NSP。通过设计一个更深的网络，移除 dropout，并在十倍的数据集上进行训练，ALBERT 在多项 NLP 任务重刷榜，但由于其结构更大，计算代价昂贵，所以其训练速度下降。

ALBERT 更像是一个工程优化加上一些小 Trick，所以看来意外简单。虽然其减少了一定内存，但计算量增加了呀，该用不起的还是用不起的。

本文的贡献主要在3个方面：

• (1)对嵌入矩阵分解，解除词嵌入和隐含层大小的关系，便于隐含层大小的扩展而不剧增模型参数
• (2)跨层参数共享，虽然一定程度上微微降低性能，但是可以大大地降低模型参数，收益很高
• (3)放弃NSP，引入SOP，更为有力地学习句子间的连贯性