大家都知道,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,训练效率低的同时还增加了硬盘空间占用。

MLM任务示意图

MLM任务示意图

如果训练的时候每个样本的所有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介绍 #

T-TA的Attention Mask模式

T-TA的Attention Mask模式

首先,我们知道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在语义表达方面有它的独特优势。

原论文的实验表格之一。可以看到T-TA在语义表达方面有它的独特优势。

基于上述设计,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://spaces.ac.cn/archives/7661

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

如果您还有什么疑惑或建议,欢迎在下方评论区继续讨论。

如果您觉得本文还不错,欢迎分享/打赏本文。打赏并非要从中获得收益,而是希望知道科学空间获得了多少读者的真心关注。当然,如果你无视它,也不会影响你的阅读。再次表示欢迎和感谢!

如果您需要引用本文,请参考:

苏剑林. (Aug. 07, 2020). 《修改Transformer结构,设计一个更快更好的MLM模型 》[Blog post]. Retrieved from https://spaces.ac.cn/archives/7661

@online{kexuefm-7661,
        title={修改Transformer结构,设计一个更快更好的MLM模型},
        author={苏剑林},
        year={2020},
        month={Aug},
        url={\url{https://spaces.ac.cn/archives/7661}},
}