大词表语言模型在续写任务上的一个问题及对策

对于LLM来说，通过增大Tokenizer的词表来提高压缩率，从而缩短序列长度、降低解码成本，是大家都喜闻乐见的事情。毕竟增大词表只需要增大Embedding层和输出的Dense层，这部分增加的计算量几乎不可感知，但缩短序列长度之后带来的解码速度提升却是实打实的。当然，增加词表大小也可能会对模型效果带来一些负面影响，所以也不能无节制地增加词表大小。本文就来分析增大词表后语言模型在续写任务上会出现的一个问题，并提出参考的解决方案。

优劣分析 #

增加词表大小的好处是显而易见的。一方面，由于LLM是自回归的，它的解码会越来越慢，而“增大词表 → 提高压缩率 → 缩短序列长度”，换言之相同文本对应的tokens数变少了，也就是解码步数变少了，从而解码速度提升了；另一方面，语言模型的训练方式是Teacher Forcing，缩短序列长度能够缓解Teacher Forcing带来的Exposure Bias问题，从而可能提升模型效果。

不过增大词表的缺点也很明显，最直接的就是会割裂token与token之间在字符层面之间的联系，从而可能会影响泛化，甚至会损失做某些任务的能力。比如“太阳能”和“太阳”都是词表中的一个词的话，模型是不知道“太阳能”是由“太阳”和“能”组成，也不知道“太阳”是“太”和“阳”，这样如果要做一些子词相关的任务就会比较艰难，比如最经典的问“‘太阳能’反过来怎么读？”，期望回答时“能阳太”，但由于模型不知道它是“太”、“阳”、“能”三个字组成，从而很难回答正确。

续写问题 #

近日 @Armen Aghajanyan 分享了另一个问题。他们在训练代码模型时使用了超大词表，结果就是常见的命令如“import numpy as np”都变成了一个token，然后发现当用户输入“import numpy”时，模型无法续写出“ as np”。原因很简单，“import numpy as np”被当作了一个token，于是当“import numpy”单独出现时，模型会发现它后面永远不会接“ as np”（接“ as np”的都被合并成单独的“import numpy as np”了），自然也无法完成续写。

这个现象确实很经典，其实不单是代码模型，常见的自然语言模型也会出现。比如当“太阳能”和“太阳”都成为了一个独立的token时，用户输入“太阳”后，接下来续写的字就基本不会是“能”了，这可能不符合用户的分布期望；又比如“白云”、“白云山”、“白云机场”都是一个独立的token时，用户输入“广州的白云”后，接下来也几乎不会续写出“广州的白云机场”、“广州的白云山”，等等。

参考对策 #

然而，笔者认为 Armen Aghajanyan 所提的现象，并不能构成增大词表的缺点，反而稍微处理一下之后，它还有可能成为增大词表的优点。其实这个问题很简单，以前没有LLM的时候，基于“词表+前缀搜索”我们也能做一定的补全任务，现在有了LLM，难道我们就一定要囿于LLM，不能将基于LLM的续写和基于词表的续写结合起来吗？

还是刚才的例子，假设用户输入了“广州的白云”，Tokenizer将它分为“广州/的/白云”，现在如果将这三个词直接转为id输入到模型中，就会无法续写出“广州/的/白云机场”等结果。这本质上是因为Tokenizer无法提前预估未来的文本，从而导致分词结果出错（当然，也可以考虑在训练阶段就使用带有随机性的tokenize算法，这种情况下“白云机场”可能作为一个词出现，也可能作为“白云/机场”出现，此时分词结果不至于严重影响后续效果，甚至能增强泛化能力，参考《Subword Regularization: Improving Neural Network Translation Models with Multiple Subword Candidates》）。

那么，我们是否可以预估一下未来的文本呢？假设分词为“广州/的/白云”后，我们回退一步，拿“白云”去词表做前缀搜索，不妨再假设搜索结果为“白云”、“白云机场”、“白云山”、“白云路”四个词，这步搜索是纯粹基于词表做的，相比LLM的计算量可以忽略不计。有了搜索结果后，我们用LLM计算：
\begin{equation}\begin{aligned}
p(\text{白云}|\text{广州},\text{的}) \\p(\text{白云机场}|\text{广州},\text{的}) \\
p(\text{白云山}|\text{广州},\text{的}) \\
p(\text{白云路}|\text{广州},\text{的}) \\
\end{aligned}\end{equation}
由于输入都是相同的，所以计算这四个条件概率只需要运行一次LLM。有了这四个条件概率后，我们将它们重新归一化然后进行采样。假如采样结果是“白云”，那么我们就按照“广州/的/白云”来做续写；如果采样到“白云机场”，那么就可以输出“机场”，并按照“广州/的/白云机场”来做续写；依此类推。这就轻松解决了Armen Aghajanyan所提到的问题，并且将缺点转化为优点了（压缩率高时，即便回退了一步，但是前缀搜索出来的词可能很长，可以一次性生成更多的字）。特别地，回退操作只需要在采样第一步进行，它只是为了避免输入不完整导致的分词错误，从第二步开始就不需要回退操作了，因此新增的计算量是非常少的。

值得一提的是，微软有一个名为“guidance”的库，也提出了同样的技巧（参考这里）。此外，考虑更一般的场景，有时候回退一步也不够，比如“import numpy as np”的例子，单输入“import numpy”时，可能被分为“import/ numpy”了，这时候起码要回退两步才能完整合理的序列。但这没有本质的区别，只是细节上稍微复杂一些，这里就不展开了，读者部署推理模型的时候自行构造就好。

文章小结 #

本文介绍了超大词表的LLM在做文本续写任务时可能出现的一个问题，并分享了参考的解决方案。

转载到请包括本文地址：https://kexue.fm/archives/9762

更详细的转载事宜请参考：《科学空间FAQ》