Dropout视角下的MLM和MAE

大家都知道，BERT的MLM（Masked Language Model）任务在预训练和微调时的不一致，也就是预训练出现了[MASK]而下游任务微调时没有[MASK]，是经常被吐槽的问题，很多工作都认为这是影响BERT微调性能的重要原因，并针对性地提出了很多改进，如XL-NET、ELECTRA、MacBERT等。本文我们将从Dropout的角度来分析MLM的这种不一致性，并且提出一种简单的操作来修正这种不一致性。

同样的分析还可以用于何凯明最近提出的比较热门的MAE（Masked Autoencoder）模型，结果是MAE相比MLM确实具有更好的一致性，由此我们可以引出一种可以能加快训练速度的正则化手段。

Dropout #

首先，我们重温一下Dropout。从数学上来看，Dropout是通过伯努利分布来为模型引入随机噪声的操作，所以我们也简单复习一下伯努利分布。

伯努利分布 #

伯努利分布（Bernoulli Distribution）算得上是最简单的概率分布了，它是一个二元分布，取值空间是{0,1}$\{0,1\}$，其中ε$\epsilon$取1的概率为p$p$，取0的概率为1−p$1-p$，记为

ε∼Bernoulli(p)(1)$$\begin{array}{}\text{(1)}& \epsilon \sim \text{Bernoulli}(p)\end{array}$$

伯努利分布的一个有趣的性质是它的任意阶矩都为p$p$，即

Eε[εn]=p×1n+(1−p)×0n=p(2)$$\begin{array}{}\text{(2)}& {\mathbb{E}}_{\epsilon }[{\epsilon }^n]=p\times 1^n+(1-p)\times 0^n=p\end{array}$$

所以我们知道它的均值为p$p$，以及方差为

Varε[ε]=Eε[ε2]−Eε[ε]2=p(1−p)(3)$$\begin{array}{}\text{(3)}& \mathbb{V}a{r}_{\epsilon }[\epsilon ]={\mathbb{E}}_{\epsilon }[{\epsilon }^2]-{\mathbb{E}}_{\epsilon }[\epsilon {]}^2=p(1-p)\end{array}$$

训练和预测 #

Dropout在训练阶段，将会以1−p$1-p$将某些值置零，而其余值则除以p$p$，所以Dropout事实上是引入了随机变量ε∼Bernoulli(p)$\epsilon \sim \text{Bernoulli}(p)$，使得模型从f(x)$f(x)$变成f(xε/p)$f(x\epsilon /p)$。其中ε$\epsilon$可以有多个分量，对应多个独立的伯努利分布，但大多数情况下其结果跟ε$\epsilon$是标量是没有本质区别，所以我们只需要针对ε$\epsilon$是标量时进行推导。

在《又是Dropout两次！这次它做到了有监督任务的SOTA》中我们证明过，如果损失函数是MSE，那么训练完成后的最佳预测模型应该是

Eε[f(xε/p)](4)$$\begin{array}{}\text{(4)}& {\mathbb{E}}_{\epsilon }[f(x\epsilon /p)]\end{array}$$

这意味着我们应该要不关闭Dropout地预测多次，然后将预测结果进行平均来作为最终的预测结果，即进行“模型平均”。但很显然这样做计算量很大，所以实际中我们很少会用这种做法，更多的是直接关闭Dropout，即将ε/p$\epsilon /p$改为1。而我们知道

f(x)=f(xEε[ε]/p)(5)$$\begin{array}{}\text{(5)}& f(x)=f(x{\mathbb{E}}_{\epsilon }[\epsilon ]/p)\end{array}$$

所以关闭Dropout事实上是一种“权重平均”（将ε$\epsilon$视为模型的随机权重）。也就是说，理论的最优解是“模型平均”，但由于计算量的原因，我们通常用“权重平均”来近似，它可以视为“模型平均”的一阶近似。

MLM模型 #

在这一节中，我们将MLM模型视为一种特殊的Dropout，由此可以清楚描述地预训练和微调的不一致之处，并且可以导出一个简单的修正策略，可以更好地缓解这种不一致性。

Dropout视角 #

简单起见，我们先来分析一个简化版本的MLM：假设在预训练阶段，每个token以p$p$的概率保持不变，以1−p$1-p$的概率被替换为[MASK]，并且第i$i$个token的Embedding记为xi$x_i$，[MASK]的Embedding记为m$m$，那么我们可以同样引入随机变量ε∼Bernoulli(p)$\epsilon \sim \text{Bernoulli}(p)$，将MLM的模型记为

f(⋯,xi,⋯)→f(⋯,xiε+m(1−ε),⋯)(6)$$\begin{array}{}\text{(6)}& f(\cdots ,x_i,\cdots )\to f(\cdots ,x_i\epsilon +m(1-\epsilon ),\cdots )\end{array}$$

这样，MLM跟Dropout本质是相同的，它们都是通过伯努利分布给模型引入了随机扰动。现在，按照Dropout的常规用法，它的预测模型应该是“权重平均”，即

f(⋯,Eε[xiε+m(1−ε)],⋯)=f(⋯,xip+m(1−p),⋯)(7)$$\begin{array}{}\text{(7)}& f(\cdots ,{\mathbb{E}}_{\epsilon }[x_i\epsilon +m(1-\epsilon )],\cdots )=f(\cdots ,x_{i}p+m(1-p),\cdots )\end{array}$$

此时，MLM在微调阶段的不一致性就体现出来了：我们将预训练的MLM视为一种特殊的Dropout，那么微调阶段对应的是“取消Dropout”，按照常规做法，此时我们应该将每个token的Embedding改为xip+m(1−p)$x_{i}p+m(1-p)$，但事实上我们没有，而是保留了原始的xi$x_i$。

修正Embedding #

按照BERT的默认设置，在训练MLM的时候，会有15%的token被选中来做MLM预测，而在这15%的token中，有80%的概率被替换为[MASK]，有10%的概率保持不变，剩下10%的概率则随机替换为一个随机token，这样根据上述分析，我们在MLM预训练完成之后，应该对Embedding进行如下调整：

Embedding[i]←0.85×Embedding[i]+0.15×⎛⎝⎜0.8×Embedding[m]+0.1×Embedding[i]+0.1×Avg[Embedding]⎞⎠⎟(8)$$\begin{array}{}\text{(8)}& \text{Embedding[i]}\leftarrow 0.85\times \text{Embedding[i]}+0.15\times \left(\begin{array}{l}0.8\times \text{Embedding[m]}+\\ 0.1\times \text{Embedding[i]}+\\ 0.1\times \text{Avg[Embedding]}\end{array}\right)\end{array}$$

其中Embedding[m]$\text{Embedding[m]}$是[MASK]的Embedding，而Avg[Embedding]$\text{Avg[Embedding]}$的全体token的平均Embedding。在bert4keras中，参考代码如下：

embeddings = model.get_weights()[0]  # 一般第一个权重就是Token Embedding
v1 = embeddings[tokenizer._token_mask_id][None]  # [MASK]的Embedding
v2 = embeddings.mean(0)[None]  # 平均Embedding
embeddings = 0.85 * embeddings + 0.15 * (0.8 * v1 + 0.1 * embeddings + 0.1 * v2)  # 加权平均
K.set_value(model.weights[0], embeddings)  # 重新赋值

那么，该修改是否跟我们期望的那样有所提升呢？笔者在CLUE上对比了BERT和RoBERTa修改前后的实验结果（baseline代码参考《bert4keras在手，baseline我有：CLUE基准代码》），结论是“没有显著变化”。

看到这里，读者也许会感到失望：敢情你前面说那么多都是白说了？笔者认为，上述操作确实是可以缓解预训练和微调的不一致性的（否则我们不是否定了Dropout？）；至于修改后的效果没有提升，意味着这种不一致性的问题并没有我们想象中那么严重，至少在CLUE的任务上是这样。一个类似的结果出现的MacBERT中，它在预训练阶段用近义词来代替[MASK]来修正这种不一致性，但笔者也在用同样的baseline代码测试过MacBERT，结果显示它跟RoBERTa也没显著差别。因此，也许只有在特定的任务或者更大的mask比例下，才能显示出修正这种不一致性的必要性。

MAE模型 #

不少读者可能已经听说过何凯明最近提出的MAE（Masked Autoencoder）模型，它以一种简单高效的方式将MLM任务引入到图像的预训练之中，并获得了有效的提升。在这一节中，我们将会看到，MAE同样可以作为一种特殊的Dropout来理解，从中我们可以得到一种防止过拟合的新方法。

Dropout视角 #

如下图所示，MAE将模型分为encoder和decoder两部分，并且具有“encoder深、decoder浅”的特点，然后它将[MASK]只放到decoder中，而encoder不处理[MASK]。这样一来，encoder要处理的序列就变短了，最关键的一步是，MAE使用了75%的mask比例，这意味着encoder的序列长度只有通常的1/4，加上“encoder深、decoder浅”的特点，总的来说模型的预训练速度快了3倍多！

MAE模型示意图

我们也可以从另一个角度来实现MAE模型：MAE把[MASK]从encoder中移除，这等价于剩下的token不与被mask掉的token交互，而对于Transformer模型来说，token之间的交互来源于Self Attention，所以我们依然可以保持原始输入，但在Attention矩阵中mask掉对应的列。如图所示，假设第i$i$个token被mask掉，事实上就相当于Attention矩阵的第i$i$列的所有元素被强制置0：

MAE的等价Attention Dropout示意图

当然，从实用的角度看，这种做法纯粹是浪费算力，但它有助于我们得到一个有意思的理论结果。我们设有n$n$的输入token，原始的Attention矩阵为A$A$（softmax后的），定义Mi$M_i$为一个n×n$n\times n$矩阵，它的第i$i$列为0、其余都为1，然后定义随机矩阵M~i${\tilde{M}}_i$，它以p$p$的概率为全1矩阵，以1−p$1-p$的概率为Mi$M_i$，那么MAE模型可以写成

f(⋯,A,⋯)→f(⋯,Norm(A⊗M~1⊗M~2⊗⋯⊗M~n),⋯)(9)$$\begin{array}{}\text{(9)}& f(\cdots ,A,\cdots )\to f(\cdots ,\text{Norm}(A\otimes {\tilde{M}}_1\otimes {\tilde{M}}_2\otimes \cdots \otimes {\tilde{M}}_n),\cdots )\end{array}$$

这里Norm$\text{Norm}$是指将矩阵重新按行归一化；⊗$\otimes$时逐个元素对应相乘；当有多个Attention层时，各个Attention层共用同一批M~1,M~2,⋯,M~n${\tilde{M}}_1,{\tilde{M}}_2,\cdots ,{\tilde{M}}_n$。

这样，我们将MAE转换为了一种特殊的Attention Dropout。那么同样按照微调阶段“取消Dropout”的做法，我们知道它对应的模型应该是

===f(⋯,Norm(A⊗E[M~1⊗M~2⊗⋯⊗M~n]),⋯)f(⋯,Norm(A⊗E[M~1]⊗E[M~2]⊗⋯⊗E[M~n]),⋯)f(⋯,Norm(Ap),⋯)f(⋯,A,⋯)(10)$$\begin{array}{}\text{(10)}& \begin{array}{rl}& f(\cdots ,\text{Norm}(A\otimes \mathbb{E}[{\tilde{M}}_1\otimes {\tilde{M}}_2\otimes \cdots \otimes {\tilde{M}}_n]),\cdots )\\ =& f(\cdots ,\text{Norm}(A\otimes \mathbb{E}[{\tilde{M}}_1]\otimes \mathbb{E}[{\tilde{M}}_2]\otimes \cdots \otimes \mathbb{E}[{\tilde{M}}_n]),\cdots )\\ =& f(\cdots ,\text{Norm}(Ap),\cdots )\\ =& f(\cdots ,A,\cdots )\end{array}\end{array}$$

其中第二个等号是因为E[M~i]$\mathbb{E}[{\tilde{M}}_i]$是一个第i$i$列为p$p$、其余为1的矩阵，那么E[M~1]⊗E[M~2]⊗⋯⊗E[M~n]$\mathbb{E}[{\tilde{M}}_1]\otimes \mathbb{E}[{\tilde{M}}_2]\otimes \cdots \otimes \mathbb{E}[{\tilde{M}}_n]$事实上就是一个全为p$p$的矩阵，所以与A$A$相乘的结果等价于A$A$直接乘以常数p$p$；第三个等号则是因为全体元素乘以同一个常数，不影响归一化结果。

从这个结果中看到，对于MAE来说，“取消Dropout”之后跟原模型一致，这说明了MAE相比原始的MLM模型，不仅仅是速度上的提升，还具有更好的预训练与微调的一致性。

防止过拟合 #

反过来想，既然MAE也可以视为一种Dropout，而Dropout有防止过拟合的作用，那么我们能不能将MAE的做法当作一种防止过拟合的正则化手段来使用呢？如下图所示，在训练阶段，我们可以随机扔掉一些token，但要保持剩余token的原始位置，我们暂且称之为“DropToken”：

DropToken示意图

之所以会这样想，是因为常规的Dropout虽然通常被直接地理解为采样一个子网络训练，但那纯粹是直观的想象，实际上Dropout的加入还会降低训练速度，而DropToken由于显式了缩短了序列长度，是可以提高训练速度的，如果有效那必然是一种非常实用的技巧。此外，有些读者可能已经试过删除某些字词的方式来进行数据扩增，它跟DropToken的区别在于DropToken虽然删除了一些Token，但依然保留了剩余token的原始位置，这个实现依赖于Transformer结构本身。

在CLUE上做的几个实验对比，基准模型为BERT base，下标的数字是drop比例，最终的效果参差不齐，除了IFLYTEK明确有效外，其他看缘分（其实很多防止过拟合手段都这样），最优drop比例在0.1～0.15之间：

CLUE分类任务对比实验（验证集）BERT0.00BERT0.10BERT0.15BERT0.25BERT0.50IFLYTEK60.0660.5660.1061.2961.45TNEWS56.8057.0056.6856.8857.02AFQMC72.4172.6172.5072.3469.76OCNLI73.9373.7674.5473.0970.68WSC78.6277.3077.3073.6869.41CSL83.9383.3383.3083.3782.56$$\begin{array}{c}\text{CLUE分类任务对比实验（验证集）}\\ \overline{)\begin{array}{ccccccc}& \text{IFLYTEK}& \text{TNEWS}& \text{AFQMC}& \text{OCNLI}& \text{WSC}& \text{CSL}\\ {\text{BERT}}_{\text{0.00}}& 60.06& 56.80& 72.41& 73.93& 78.62& 83.93\\ {\text{BERT}}_{\text{0.10}}& 60.56& 57.00& 72.61& 73.76& 77.30& 83.33\\ {\text{BERT}}_{\text{0.15}}& 60.10& 56.68& 72.50& 74.54& 77.30& 83.30\\ {\text{BERT}}_{\text{0.25}}& 61.29& 56.88& 72.34& 73.09& 73.68& 83.37\\ {\text{BERT}}_{\text{0.50}}& 61.45& 57.02& 69.76& 70.68& 69.41& 82.56\end{array}}\end{array}$$

本文小结 #

本文从Dropout的视角考察了MLM和MAE两个模型，它们均可视为特殊的Dropout，从这个视角中，我们可以得到了一种修正MLM的不一致性的技巧，以及得到一种类似MAE的防止过拟合技巧。

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

来自为知笔记(Wiz)