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

大家都知道，MLM（Masked Language Model）是BERT、RoBERTa的预训练方式，顾名思义，就是mask掉原始序列的一些token，然后让模型去预测这些被mask掉的token。随着研究的深入，大家发现MLM不单单可以作为预训练方式，还能有很丰富的应用价值，比如笔者之前就发现直接加载BERT的MLM权重就可以当作UniLM来做Seq2Seq任务（参考这里），又比如发表在ACL 2020的《Spelling Error Correction with Soft-Masked BERT》将MLM模型用于文本纠错。

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

MLM模型 #

假设原始序列为$\boldsymbol{x}=[x_1,x_2,\dots,x_T]$，$\boldsymbol{x}\backslash \{x_i\}$表示将第i个token替换为$\text{[MASK]}$后的序列，那么MLM模型就是建模
\begin{equation}p\big(x_i, x_j, x_k, \cdots\big|\,\boldsymbol{x}\backslash \{x_i,x_j,x_k,\cdots\}\big)\end{equation}
我们说它效率低，是因为每次只能选择一小部分token来mask，比如15%，那么也就是说每个样本只有15%的token被训练到了，所以同一个样本需要反复训练多次。在BERT里边，每个样本都被mask了多遍然后存为tfrecord，训练效率低的同时还增加了硬盘空间占用。

如果训练的时候每个样本的所有token都可以作为预测目标，那么训练效率自然就能提升了。像GPT这样的单向语言模型是可以做到的，但是MLM是双向的模型，并不能直接做到这一点。为了达到这个目标，我们需要简化一下上式，假设每次只mask掉一个token，也就是要构建的分布为
\begin{equation}p\big(x_i\big|\,\boldsymbol{x}\backslash \{x_i\}\big),\,i=1,2,\dots,T\end{equation}
然后我们希望通过单个模型一次预测就同时得到$p(x_1|\,\boldsymbol{x}\backslash \{x_1\}),p(x_2|\,\boldsymbol{x}\backslash \{x_2\}),\dots,p(x_T|\,\boldsymbol{x}\backslash \{x_T\})$。怎么做到这一点呢？这就来到本文要介绍的论文结果了，它提出了一种称之为T-TA（Transformer-based Text Autoencoder）的设计，能让我们一并预测所有token的分布。

T-TA介绍 #

首先，我们知道Transformer的核心运算是$Attention(\boldsymbol{Q},\boldsymbol{K},\boldsymbol{V})$，在BERT里边$\boldsymbol{Q},\boldsymbol{K},\boldsymbol{V}$都是同一个，也就是Self Attention。而在MLM中，我们既然要建模$p(x_i|\,\boldsymbol{x}\backslash \{x_i\})$，那么第$i$个输出肯定是不能包含第$i$个token的信息的，为此，第一步要做出的改动是：去掉$\boldsymbol{Q}$里边的token输入，也就是说第一层的Attention的$\boldsymbol{Q}$不能包含token信息，只能包含位置向量。这是因为我们是通过$\boldsymbol{Q}$把$\boldsymbol{K},\boldsymbol{V}$的信息聚合起来的，如果$\boldsymbol{Q}$本身就有token信息，那么就会造成信息泄漏了。然后，我们要防止$\boldsymbol{K},\boldsymbol{V}$的信息泄漏，这需要修改Attention Mask，把对角线部分的Attention（也就是自身的）给Mask掉，如图所示。

如果还不理解这一点，我们可以从Attention的一般形式来理解：Attention的一般定义为
\begin{equation}Attention(\boldsymbol{Q},\boldsymbol{K},\boldsymbol{V})_i = \frac{\sum\limits_{j=1}^n \text{sim}(\boldsymbol{q}_i, \boldsymbol{k}_j)\boldsymbol{v}_j}{\sum\limits_{j=1}^n \text{sim}(\boldsymbol{q}_i, \boldsymbol{k}_j)}\label{eq:gen-att}\end{equation}
所以很明显，$Attention(\boldsymbol{Q},\boldsymbol{K},\boldsymbol{V})_i$一定跟$\boldsymbol{q}_i$有联系，所以$\boldsymbol{q}_i$绝对不能包含第$i$个token的信息；但它不一定跟$\boldsymbol{k}_i,\boldsymbol{v}_i$有联系，因为只需要当$\text{sim}(\boldsymbol{q}_i, \boldsymbol{k}_i)=0$时$\boldsymbol{k}_i,\boldsymbol{v}_i$就相当于不存在了，因此需要Mask掉对角线部分的Attention。

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

原论文的解决很粗暴，但貌似也只能这样解决了：每一层Attention都共用原始输入为$\boldsymbol{K},\boldsymbol{V}$！所以，设$\boldsymbol{E}$为token的embedding序列，$\boldsymbol{P}$为对应的位置向量，那么T-TA与BERT的计算过程可以简写为：
\begin{equation}
\begin{array}{c}\bbox[border: 1px dashed red; padding: 5px]{\begin{aligned}&\boldsymbol{Q}_0 = \boldsymbol{E}+\boldsymbol{P}\\
&\boldsymbol{Q}_1 = Attention(\boldsymbol{Q}_0,\boldsymbol{Q}_0,\boldsymbol{Q}_0)
\\
&\boldsymbol{Q}_2 = Attention(\boldsymbol{Q}_1,\boldsymbol{Q}_1,\boldsymbol{Q}_1)
\\
&\qquad\vdots\\
&\boldsymbol{Q}_n = Attention(\boldsymbol{Q}_{n-1},\boldsymbol{Q}_{n-1},\boldsymbol{Q}_{n-1})
\end{aligned}} \\ \text{BERT运算示意图}\quad\end{array}\qquad
\begin{array}{c}\bbox[border: 1px dashed red; padding: 5px]{\begin{aligned}&\boldsymbol{Q}_0 = \boldsymbol{P}\\
&\boldsymbol{Q}_1 = Attention(\boldsymbol{Q}_0,\boldsymbol{E}+\boldsymbol{P},\boldsymbol{E}+\boldsymbol{P})
\\
&\boldsymbol{Q}_2 = Attention(\boldsymbol{Q}_1,\boldsymbol{E}+\boldsymbol{P},\boldsymbol{E}+\boldsymbol{P})
\\
&\qquad\vdots\\
&\boldsymbol{Q}_n = Attention(\boldsymbol{Q}_{n-1},\boldsymbol{E}+\boldsymbol{P},\boldsymbol{E}+\boldsymbol{P})
\end{aligned}} \\ \text{T-TA运算示意图}\quad\end{array}\end{equation}
当然残差、FFN等细节已经省略掉了，只保留了核心运算部分，预训练阶段T-TA的Attention是进行了对角线形式的Attention Mask的，如果是下游任务的微调，则可以把它去掉。

实验结果 #

基于上述设计，T-TA它能一次性预测所有的token，所以训练效率高，并且不需要额外的$\text{[MASK]}$符号，所以实现了预训练和微调之间的一致性。但是不难理解，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来说，当输入给定后，$\boldsymbol{K},\boldsymbol{V}$在所有Attention层中的保持不变，变化的只有$\boldsymbol{Q}$，所以读者质疑它效果也不意外。但是别忘了，前段时候Google才提出了个Synthesizer（参考《Google新作Synthesizer：我们还不够了解自注意力》），里边探索了几种Attention变种，其中一种简称为“R”的，相当于$\boldsymbol{Q},\boldsymbol{K}$固定为常数，结果居然也能work得不错！要注意，“R”里边的$\boldsymbol{Q},\boldsymbol{K}$是彻彻底底的常数，跟输入都没关系。

所以，既然$\boldsymbol{Q},\boldsymbol{K}$为常数效果都还可以，那么$\boldsymbol{K},\boldsymbol{V}$为什么不能为常数呢？更何况T-TA的$\boldsymbol{K},\boldsymbol{V}$动态依赖于输入的，只是输入确定后它才算是常数，因此理论上来讲T-TA的拟合能力比Synthesizer的“R”模型要强，既然“R”都能好了，T-TA能好应该也是不奇怪。

当然，还是期望后续会有更深的实验结果出现。

转载到请包括本文地址：https://kexue.fm/archives/7661

更详细的转载事宜请参考：《科学空间FAQ》