修改Transformer结构，设计一个更快更好的MLM模型
原文链接 | 2020-08-07 | 苏剑林
大家都知道，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能好应该也是不奇怪。
当然，还是期望后续会有更深的实验结果出现。