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

导语： BERT模型自诞生以来统治了各项NLP任务的榜单，谷歌也针对中文给出了基于字的模型。然而我们知道词是语言最重要的组成部分，因此，一个自然的想法就是做基于词的BERT模型。但是受限于显存，谷歌原生的BERT模型可以使用的词典比较小，只能覆盖一小部分的汉语的词。在本文中，我们提出了对BERT的一些改进，将中文BERT词模型的词典大小进行了大幅扩充，并在多个下游任务上测试了大词典BERT的表现。此外，我们尝试了一种针对上下文相关词向量的最近邻检索方法，基于BERT的词向量做以词搜词任务，相对于上下文无关词向量在效果上有着明显的提升。

1. 做大词典BERT词模型的意义

词是语言最重要的组成部分。谷歌发布基于字的BERT[1]模型之后，一个很自然的想法就是将词的信息引入到模型之中。词在汉语中扮演了非常重要的词是语言最重要的组成部分。谷歌发布基于字的BERT[1]模型之后，一个很自然的想法就是将词的信息引入到模型之中。词在汉语中扮演了非常重要的角色。很多词所表达的含义与其包括的字的含义没有直接联系，比如音译词“巧克力”、“圣代”。

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

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

2. 大词典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万大小。

3. Adaptive Softmax

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

如上图所示，如果我们将词典分成三个部分，那么adaptivesoftmax则由三个前向神经网络构成，我们分别称之为head、tail1、tail2。第一个词典会链到第二和第三个词典。我们根据词频对词语进行排序，令head前向神经网络去预测高频词语；tail1去预测中间频率的词语；tail2去预测低频词语。

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

可以看到，中间频率的词语的向量维度是192；低频词语的向量维度是48。相对于原来的参数量，adaptive soft