Mask-Predict: Parallel Decoding of Conditional Masked Language Models

来源：EMNLP2019 链接：https://arxiv.org/pdf/1904.09324.pdf 代码：https://github.com/facebookresearch/Mask-Predict

动机

大部分的神经机器翻译模型(seq2seq)都是以自回归的方式(autoregressive)进行的,decode阶段从左到右依次生成token，因此，decoder阶段的时间会随着输入句子的长度增加而增加。虽然有一些工作探索了非自回归的方式的seq2seq模型，decoder阶段同时生成target端的所有token，但是与自回归的seq2seq模型相比，效果相差较大。

解决方法

论文提出了一种新的decoder方法：mask-predict

在decoder过程中，首先用非自回归的方式生成target端全部token（遮住全部token），然后再遮住模型不确定的一些token，依据target端其他未遮盖token文以及源文本重新进行预测，迭代进行。

模型及训练

Conditional Masked Language Models (CMLM)

$P\left(y|X,Y_{obs}\right)$

给定源文$X$和部分的目标文本$Y_{obs}$，来预测剩下的目标文本$Y_{\text{mask}}$

$N$即是目标端文本的长度$N=\left|Y_{mask}\right|+\left|Y_{obs}\right|$

论文采用标准的transformer的模型，不同之处在于decode端的self-attention不再使用attention mask防止看到预测单词之后的词。

训练的时候，首先确定需要遮盖的token个数（均匀分布在0到target端文本长度），然后随机选取target端的token进行遮盖，并进行预测，训练目标是交叉熵损失函数，只计算遮盖位置的loss。

预测target端文本的长度

自回归的方法中，decode从左到右依次生成token，直到预测出EOS结束符，或者超过的给定的最大长度。论文提出的方法并行的生成生成target端全部token，因此必须事先知道target端文本的长度。

仿照Bert原论文，添加一个 [CLS] token来处理分类任务，论文在encoder端添加了一个 [LENGTH] token来预测target端文本的长度，计算的loss和CMLM的loss相加。

inference

使用Mask-Predict解码

主要可以分为两步：

Mask

$\begin{array}{rl}{Y}_{\text{mask}}^{(t)}& =\arg \underset{i}{\min }\left(p_i,n\right)\\ Y_{obs}^{(t)}& =Y\backslash Y_{\text{mask}}^{(t)}\end{array}$

$n=N\cdot \frac{T-t}{T}$`

第t次迭代，遮住的单词数n，T为预先设定的迭代次数。第0次迭代，遮住target端所有token；随着迭代次数的增加，逐步减少遮盖的单词数。
$p_i$是当前迭代次数的预测概率

Predict

$\begin{array}{rl}{y}_i^{(t)}& =\arg \underset{w}{\max }P\left(y_i=w|X,Y_{obs}^{(t)}\right)\\ p_i^{(t)}& =\underset{w}{\max }P\left(y_i=w|X,Y_{obs}^{(t)}\right)\end{array}$

给定源文本$X$和部分的目标文本$Y_{obs}$，来预测mask的目标文本$Y_{\text{mask}}$,并更新概率

没有被mask的token $Y_{obs}$，概率不变

$\begin{array}{rl}{y}_i^{(t)}& =y_i^{(t-1)}\\ p_i^{(t)}& =p_i^{(t-1)}\end{array}$

两个步骤迭代进行，下面是一个例子

第0步，遮盖整个target端所有tokens，获得target端全部输出，可以看到有很多重复的单词，整个句子语法也不合理。

第1步，根据第0步得到的概率，从12个单词中选出8个低概率的单词(标黄的部分)，使用CMLM重新预测。此时的结果比第一步更加准确合理。

第2步，根据第1步得到的概率（被重新预测的8个单词的概率需要更新），从12个单词中选出4个低概率的单词，再使用CMLM重新预测。值得注意的是第1步没有被重新预测的词也可能被选中。

预测target端文本的长度

获取encoder部分 [LENGTH] token处预测到的长度分布之后，选在top l个长度作为候选，并行decode，然后选择平均概率 $\frac{1}{N}\sum \log p_i^{(T)}$ 最好的结果

实验结果

在机器翻译上面的结果：

1. 论文提出的方法(CMLM with Mask-Predict) 超过了其他的非自回归的方法。
2. 相比自回归的方法，结果差了 0.5-1.2 BLEU, 但是速度加快了。
3. 在中英(两种语言相差比较大)翻译的数据集上结果也比较好

速度与效果的trade-off

1. 牺牲2个BLEU的情况下，可以提升3倍多的速度（T=4, l= 2）
2. 达到相当的效果（27.03 vs 27.74），速度提升30%

模型简化测试

分析decoder是否必须迭代多次

decoder迭代多次主要是为了改善，非自回归方法共有的 token repetitions 问题。 随着token重复率的减少，bleu也有所上升

分析长句是否需要迭代更多的次数

长句迭代多次有更好的效果。

分析使用更多的长度候选是否能改善效果

不能

分析知识蒸馏是否必要

这里知识蒸馏的意思是，训练非自回归模型时，训练集不使用原始训练集中的target端数据，而使用teacher模型预测的target结果。

必要

总结

1. 在机器翻译任务上，与其他非自回归seq2seq模型相比，效果达到最好；
2. 在WMT en-de任务上，比baseline自回归模型低0.7 BLEU，速度提升了30%； bleu损失两个点，inference的速度可以达到原来的三倍多。
3. target端文本长度的预测也比较重要，如果完全预测对，bleu还有提升空间
4. 知识蒸馏的teacher网络使用了强模型，这也会使得论文的结果强于一般模型。同时，论文也和先前其他论文一样没有解释清楚为什么需要知识蒸馏