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

©PaperWeekly 原创 · 作者｜苏剑林

单位｜追一科技

研究方向｜NLP、神经网络

大家都知道，MLM (Masked Language Model) 是 BERT、RoBERTa 的预训练方式，顾名思义，就是 mask 掉原始序列的一些 token，然后让模型去预测这些被 mask 掉的 token。

随着研究的深入，大家发现 MLM 不单单可以作为预训练方式，还能有很丰富的应用价值，比如笔者之前就发现直接加载 BERT 的 MLM 权重就可以当作 UniLM 来做 Seq2Seq 任务（参考这里），又比如发表在 ACL 2020 的 Spelling Error Correction with Soft-Masked BERT [1] 将 MLM 模型用于文本纠错。

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

论文标题：Fast and Accurate Deep Bidirectional Language Representations for Unsupervised Learning

论文来源：ACL 2020

论文链接：https://arxiv.org/abs/2004.08097

代码链接：https://github.com/joongbo/tta

MLM模型

假设原始序列为 ， 表示将第 i 个token 替换为 后的序列，那么 MLM 模型就是建模：

我们说它效率低，是因为每次只能选择一小部分 token 来 mask，比如 15%，那么也就是说每个样本只有 15% 的 token 被训练到了，所以同一个样本需要反复训练多次。在 BERT 里边，每个样本都被 mask 了多遍然后存为 tfrecord，训练效率低的同时还增加了硬盘空间占用。

▲ MLM任务示意图

如果训练的时候每个样本的所有 token 都可以作为预测目标，那么训练效率自然就能提升了。像 GPT 这样的单向语言模型是可以做到的，但是 MLM 是双向的模型，并不能直接做到这一点。为了达到这个目标，我们需要简化一下上式，假设每次只 mask 掉一个 token，也就是要构建的分布为：

然后我们希望通过单个模型一次预测就同时得到：

怎么做到这一点呢？这就来到本文要介绍的论文结果了，它提出了一种称之为 T-TA (Transformer-based Text Autoencoder) 的设计，能让我们一并预测所有 token 的分布。

T-TA介绍

首先，我们知道 Transformer 的核心运算是