修改Transformer结构，设计一个更快更好的MLM模型

2020 年 8 月 11 日 PaperWeekly

©PaperWeekly 原创 · 作者｜苏剑林

单位｜追一科技

研究方向｜NLP、神经网络

大家都知道，MLM (Masked Language Model) 是 BERT、RoBERTa 的预训练方式，顾名思义，就是 mask 掉原始序列的一些 token，然后让模型去预测这些被 mask 掉的 token。

随着研究的深入，大家发现 MLM 不单单可以作为预训练方式，还能有很丰富的应用价值，比如笔者之前就发现直接加载 BERT 的 MLM 权重就可以当作 UniLM 来做 Seq2Seq 任务（参考这里），又比如发表在 ACL 2020 的 Spelling Error Correction with Soft-Masked BERT [1] 将 MLM 模型用于文本纠错。

然而，仔细读过 BERT 的论文或者亲自尝试过的读者应该都知道，原始的 MLM 的训练效率是比较低的，因为每次只能 mask 掉一小部分的 token 来训练。ACL 2020 的论文 Fast and Accurate Deep Bidirectional Language Representations for Unsupervised Learning 也思考了这个问题，并且提出了一种新的 MLM 模型设计，能够有更高的训练效率和更好的效果。

论文标题：Fast and Accurate Deep Bidirectional Language Representations for Unsupervised Learning

论文来源：ACL 2020

论文链接：https://arxiv.org/abs/2004.08097

代码链接：https://github.com/joongbo/tta

MLM模型

假设原始序列为，表示将第 i 个token 替换为后的序列，那么 MLM 模型就是建模：

我们说它效率低，是因为每次只能选择一小部分 token 来 mask，比如 15%，那么也就是说每个样本只有 15% 的 token 被训练到了，所以同一个样本需要反复训练多次。在 BERT 里边，每个样本都被 mask 了多遍然后存为 tfrecord，训练效率低的同时还增加了硬盘空间占用。

如果训练的时候每个样本的所有 token 都可以作为预测目标，那么训练效率自然就能提升了。像 GPT 这样的单向语言模型是可以做到的，但是 MLM 是双向的模型，并不能直接做到这一点。为了达到这个目标，我们需要简化一下上式，假设每次只 mask 掉一个 token，也就是要构建的分布为：

然后我们希望通过单个模型一次预测就同时得到：

怎么做到这一点呢？这就来到本文要介绍的论文结果了，它提出了一种称之为 T-TA (Transformer-based Text Autoencoder) 的设计，能让我们一并预测所有 token 的分布。

T-TA介绍

首先，我们知道 Transformer 的核心运算是，在 BERT 里边都是同一个，也就是 Self Attention。而在 MLM 中，我们既然要建模，那么第 i 个输出肯定是不能包含第 i 个 token 的信息的。

为此，第一步要做出的改动是：去掉里边的 token 输入，也就是说第一层的 Attention 的不能包含 token 信息，只能包含位置向量。这是因为我们是通过把的信息聚合起来的，如果本身就有 token 信息，那么就会造成信息泄漏了。

改了之后，我们就可以防止来自的信息泄漏，接下来我们就要防止来自的信息泄漏了，这需要修改 Attention Mask，把对角线部分的 Attention（也就是自身的）给 Mask 掉，如图所示：

但是，这种防泄漏的 Attention Mask 只能维持一层！也就是说即便这样做之后，已经融入了第 i 个 token 的信息了，所以从第二层开始，如果你还是以第一层的输出为，即便配合了上述 Attention Mask，也会出现信息泄漏了。

原论文的解决很粗暴，但貌似也只能这样解决了：每一层 Attention 都共用输入为！所以，设为 token 的 embedding 序列，为对应的位置向量，那么 T-TA 与 BERT 的计算过程可以简写为：

当然残差、FFN 等细节已经省略掉了，只保留了核心运算部分，预训练阶段 T-TA 的 Attention 是进行了对角线形式的 Attention Mask 的，如果是下游任务的微调，则可以把它去掉。

实验结果

基于上述设计，T-TA 它能一次性预测所有的 token，并且不需要额外的符号，所以训练效率高，并且实现了预训练和微调之间的一致性（没有）。但是不难理解，T-TA 实则是对标准 Transformer 的一种简化，所以理论上它的拟合能力是变弱了。

这样一收一放之下，具体结果还有没有提升呢？当然，论文的实验结果是有的。原论文做了多个实验，结果显示 T-TA 这种设计在同样的参数情况下基本都能媲美甚至超过标准的 MLM 训练出来的模型。作者还很慷慨地开源了代码，以便大家复现结果。

说到修改 Transformer 结构，大家可能联想到大量的 GPU、TPU 在并行运算。但事实上，虽然作者没有具体列出自己的实验设备，但从论文可以看到设备阵容应该不算“豪华”。为此，作者只训练了 3 层的 T-TA，并且按照同样的模式复现了 3 层的 MLM 和 GPT（也就是单向语言模型），然后对比了效果。

没错，论文中所有 T-TA 的结果都只是 3 层的模型，而其中有些都超过了 Base 版本的 BERT。所以作者生动地给我们上了一课： 没有土豪的设备，也可以做修改 Transformer 的工作，也可以发ACL，关键是你有真正有效的 idea。

个人分析

最后，再来简单谈谈 T-TA 为什么有效。读者可能会质疑，既然作者只做了 3 层的实验，那么如何保证在更多层的时候也能有效呢？那好，我们来从另外一个角度看这个模型。

从设计上看，对于 T-TA 来说，当输入给定后，在所有 Attention 层中的保持不变，变化的只有，所以读者质疑它效果也不意外。

但是别忘了，前段时候 Google 才提出了个 Synthesizer（参考 Google 新作Synthesizer：我们还不够了解自注意力），里边探索了几种 Attention 变种，其中一种简称为“R”的，相当于固定为常数，结果居然也能 work 得不错！要注意，“R”里边的是彻彻底底的常数，跟输入都没关系。