从字到词，大词典中文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的原理。

如上图所示，如果我们将词典分成三个部分，那么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 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的原理。

如上图所示，adaptive input对高频词使用较大的维度，对低频词使用较小的维度。更具体的，上图把50万词语按照词频排好序，并分成三个部分：第一部分包括37500个高频词语，词向量维度为768；第二部分包括62500个中间频率的词语，词向量维度为192；第三部分包括40万个低频词语，词向量维度为48。为了使维度可以统一到隐层维度，这里引入了三个不同的前向神经网络，把不同维度的词向量映射到相同的维度。经过计算，我们可以得到adaptive input 减少的参数量是500000*768-37500*768-62500*192-400000*48-768*768-192*768-48*768= 323225856，相当于原来的参数量的84.2%。当模型同时引入adaptive softmax和adaptive input的时候，词典可以进一步扩展到100万。

因此，对于基于词的BERT模型，无论是用新的语料进行叠加预训练，还是针对下游任务进行微调，我们都需要生成新的词典（动态词典），并根据新的词典去对预训练模型的embedding层和softmax层进行调整。调整的方式如下图所示（蓝色表示根据预训练模型的参数初始化，红色表示随机初始化）。如果使用固定词表，那么词语“人工智能”会被映射到UNK，这个词语无法在下游任务上进行训练。对于加入adaptive机制的模型，调整的过程会增加一些步骤。后面实验中基于词的BERT模型均默认使用动态词典。动态词典相对于固定词典在效果上有着显著的提升。

本文的另一项工作是尝试了一种上下文相关词向量的最近邻检索方法。在传统的词向量工作中，每个词对应一个固定的向量，我们称其为上下文无关词向量。对这类词向量进行最近邻检索只需要计算词向量之间的cos值，并根据cos值进行排序。对于上下文相关词向量，词向量不仅和词语本身有关，还和词语周围的上下文有关。也就是语料中的每一个token(词语的一次出现），都有一个独有的向量。ELMO[8]为每个token计算一个向量，并利用token向量去寻找其它的token向量。ELMO是基于整个语料寻找token向量的最近邻。这里我们尝试一种基于词典的上下文相关词向量检索方法，具体步骤如下：假如我们需要查找在句子“吉利汽车是中国汽车品牌”中词语“吉利”的最近邻（如下图），我们首先把这句话输入到BERT编码器中，取最上层的词语对应的隐层向量作为词向量。然后我们把词典中的所有词语依次替换到“吉利”的位置，这样就得到每个词语在特定上下文中的向量。我们计算这些向量之间的cos值，即可得到词语的最近邻。

我们还尝试使用基于字的BERT模型进行以词搜词作为基于词的模型的对比。具体步骤和基于词的模型相似，只是在基于字的模型中，每一个字对应一个向量，因此一个词会对应多个向量。我们对BERT编码器最上层的字的隐层向量取平均，从而得到每个词语在特定上下文中的向量。

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

在这个例子中，“吉利”指一种汽车品牌。可以看到基于词的模型的检索效果不错。其中“沃尔沃”、“金牛座”是汽车的品牌。“永利”、“天安”、“仁和”、“新光”均是公司名称，并且这些公司均以一个比较吉利的词语作为公司的名字。基于字的模型效果也不错，会倾向于返回包含“吉”、“利”这两个字的词语。

候选词语：“吉利”

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

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

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

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

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

在这个例子中，“苹果”指科技公司。基于词的模型效果不错。但是基于字的模型效果一般，虽然也在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的预训练模型。