从字到词，大词典中文BERT模型的探索之旅

最近有一些工作尝试将词的信息引入到BERT中去。百度的ERNIE通过遮罩策略将词的信息引入到模型之中，但是其本质仍然是基于字的模型[2]。另外一种方式是对语料进行分词，直接训练基于词的中文BERT模型。不过由于中文词典很大，在谷歌BERT上使用大词典会导致显存溢出。使用小词典的话会导致大量的OOV。为了减缓OOV的影响，我们需要扩大词典。

因此本文引入了一些语言模型训练的技术，提出了训练大词典BERT模型的方法。大词典能提升词语的覆盖度，减缓未登录词（OOV）问题[3]。不过虽然大词典能够减缓OOV，但是相比于字模型，仍然存在一定比例的OOV，会受到OOV的影响。此外，训练基于词的BERT模型能够得到高质量的上下文相关词向量。这些向量可以用于各种应用和下游任务之中。

在谷歌原来的BERT模型中，embedding 层和softmax 层的参数数量随着词典大小呈线性变化。如果在BERT-base模型上使用50万大小的词典，那么embedding层和softmax层会包含500000*768+768*500000=7.68亿个参数。要知道BERT-base的编码器层（12层transformer）也仅仅包含8500万个参数。因此，直接在BERT上使用大词典会造成显存溢出。实际上，对于BERT-base模型（在P40型号的GPU上，batch size为32），当句子长度为128的时候，最多支持16万的词典大小，当句子长度为192的时候，只能支持8万的词典大小。

在下面的章节中，我们针对embedding 层和softmax 层进行优化。对softmax层优化后词典可以达到五十万大小，对softmax层和embedding层同时优化可以把词典扩大到100万大小。

对softmax层的优化一直是自然语言处理领域研究的重点。这个方向的研究很多，常用的技术包括基于采样的方法[4]，基于树的方法[5]等等。这里我们使用facebook提出的adaptive softmax[6]进行优化。之所以选择adaptive softmax，一方面是因为其是针对GPU进行的优化。我们在多机多GPU上进行预训练，使用adaptive softmax能帮助我们在效率上得到显著的提升。另一方面，adaptive softmax可以节约显存，这是我们引入大词典的关键。下面简要介绍adaptive softmax的原理。

因为高频词占据了语料中绝大部分的频数（Zipf’s Law），所以在大多数情况下，我们只需要用head这个规模较小的前向神经网络去做预测，这很大程度上减少了计算量。当遇到低频词的时候（比如上图中的“二氧化碳”），模型先使用head去预测，发现“tail2”的概率值最大，这表明需要继续使用tail2前向神经网络去预测。然后我们得到tail2前向神经网络对于“二氧化碳”的预测值之后，让其与head的“tail2”位置的概率值相乘，就能得到“二氧化碳”这个词的预测概率。在具体实现中，tail1和tail2一般是两层的前向神经网络，中间隐层维度设为一个较小的值，从而实减少模型的参数。下图是PyTorch实现的一个adaptive softmax的示例。词典大小为50万。

可以看到，中间频率的词语的向量维度是192；低频词语的向量维度是48。相对于原来的参数量，adaptive softmax参数减少了500000*768-37502*768-62500*192-400000*48-768*768-192*768-48*768 = 323224320个，相当于原来参数量的84.2%。在使用adaptive softmax之后，词典的规模可以从8万扩展到50万。50万词典已经可以覆盖绝大部分常见词，能有效的减少OOV的情况。

Softmax层一定程度上可以看作embedding层的镜像。因此可以把softmax上的优化操作搬到embedding层。在今年的ICLR上facebook提出了adaptive input方法[7]。下面简要介绍adaptive input的原理。

使用上面介绍的方式进行以词搜词，相对于传统的上下文无关词向量，在多义词的检索效果上有着明显的提升。下面给出几个以词搜词的示例。给定一句话以及这句话中要检索的词语，返回这个词语的最近邻词语以及相似度。

候选词语：“吉利”

候选句子：主要演员 有 扎克· 布拉 夫 、 萨拉 · 朝克 、 唐纳德 · 费森 、 尼尔 · 弗林 、 肯 · 詹金斯 、 约翰 · 麦 吉利 、 朱迪 · 雷耶斯 、 迈克尔 · 莫斯利 等 。

在这个例子中，“吉利”指一个人名。可以看到基于字和基于词的模型均通过上下文捕捉到了这里的“吉利”是人名的信息。基于字的模型返回的人名均包含“利”和“吉”。基于词的模型返回的单词则更加的多样。

候选句子：他 出生 时 身上附有 胎膜 ， 其母 认为 这是 一个 吉利 的 征兆 。

在这个例子中，“吉利”指这个单词的本义。可以看到基于字和基于词的模型均表现不错。基于字的模型返回的单词大多包含“吉”、“利”两字。基于词的模型除了会返回和“吉利”意思相近的词，比如“仁德”、“光彩”、“愉快”，还会返回“食用油”、“玉”这样的名词。这些名字放在原句中比较合适。

在这个例子中，“苹果”指水果。两种模型效果均不错。基于字的模型倾向返回包含“果”的词。

在这个例子中，“苹果”指科技公司。基于词的模型效果不错。但是基于字的模型效果一般，虽然也在topn列表中返回了“谷歌”、“微软”，但是前面几个单词均是包含“果”的词语。

候选词语：“腾讯”

候选句子：马化腾于 1984 年 随 家人 从 海南 移居 深圳 ， 14 年后 的 1998 年 11 月 ， 他 和 其他 四位 联合创始人一起 创立 腾讯 。

在这个例子中，我们可以看到基于词的模型返回的均是互联网公司或者相关的内容。基于字的模型也会返回和互联网公司相关的单词，比如“微信”、“百度”，但是也会返回一些包含“腾”的词语。这些词语和互联网公司关系较弱。

通过上面的定性分析，我们可以看到使用基于字的BERT模型做词的检索有一定的局限性。基于字的模型的词向量会受到字向量的影响，倾向于返回包含相同字的词语。基于词的模型效果较好，能对多义词进行高质量的检索。

上面使用的检索方法效率较低。假如词典大小为2万，那么进行一次检索，需要通过BERT编码器2万次。如果不加以优化，检索一次的时间需要十几秒以上。目前我们正在进行优化。我们发现合理的减少编码器的层数和attention head的个数，可以在大幅度提升效率的同时，对检索结果质量产生较小的影响。目前，我们对2万词语进行检索，时间可以控制在3秒以内，并能得到较高质量的检索结果。在这个方向上还有很多的优化点，后续我们会进一步探索如何提升检索效率。

和Google BERT一样，我们使用中文维基百科作为语料进行预训练。对于语料中的中文，我们使用jieba默认模式进行分词；对于语料中的英文等其它语言，我们使用和Google相同的分词策略。所有的模型均训练50万步；学习率是1e-4（warmuplinear策略，warmup设为0.01）；优化器为BERT默认的AdamW；batchsize为256；序列长度为192。我们使用一个百科词条作为一个文档。文档中根据“。”、“！”、“？”进行分句。目前我们有8万词典和50万词典的模型，100万词典的模型正在训练当中。我们和Google BERT预训练使用的策略基本一样。只是受限于资源，我们的seq_length和训练步数小于Google的预训练模型。