随机分词浅探：从Viterbi Sampling到完美采样算法

原创苏剑林 PaperWeekly

©PaperWeekly 原创 · 作者 | 苏剑林

单位 | 月之暗面

研究方向 | NLP、神经网络

从Viterbi Decoding到Viterbi Sampling

上一篇文章《大词表语言模型在续写任务上的一个问题及对策》发布后，很快就有读者指出可以在训练阶段引入带有随机性的分词结果来解决同样的问题，并且已经有论文和实现。

经过进一步查阅学习，笔者发现这是一个名为 Subword Regularization [1] 的技巧，最早应用在 NMT（机器翻译）中，目前 SentencePiece 也有相应的实现。看起来这个技巧确实能缓解前述问题，甚至有助于增强语言模型的容错能力，所以就有了将它加进去 BytePiece 的想法。

那么问题来了，如何将确定性分词改为随机性分词呢？BytePiece 是基于 Unigram 模型的，它通过 Viterbi 算法找最大概率的分词方案，既然有概率，是否就可以自然地导出随机采样？本文来讨论这个问题，并分享自己的解决方案。

1.1 要点分析

现阶段，Unigram 分词是直接输出最大概率的切分方案，通常这是一个确定性的输出。具体来说，假设代表一个切分方案，对应的打分为

，代表句子所有可能的切分方案的集合，那么分词算法可以描述为

这可以通过 Viterbi 算法在线性时间内来完成，所以这个过程我们也称之为 “Viterbi Decoding”。看起来，Unigram 模型天然带有概率，所以似乎并不难将它改为依概率采样的形式，但细想之下才发现这并非一个平凡的问题，有很多细节上的困难需要克服。

笔者设想是模仿自回归语言模型设计一个递归采样流程，但这里最困难的地方是如何尽量保持原来的候选切分方案的排序不变，或者就算不能保持所有的排序不变，也至少满足最大概率不变，即 Viterbi 解码的最大概率路径应该对应所设计的递归采样算法的最大概率采样结果。

由于所有切分方案构成一个有向无环图（DAG，Directed Acyclic Graph），笔者一开始以为直接在有向无环图上随机游走是一个可行方案，但再思考后发现很难设计适当的转移概率来保证最大概率路径不变（因为同一起点的不同边不是平权的，不能简单按照边的频率为权重做采样）。

1.2 已有方案

由于一时半会没有新想法，所以笔者决定去翻翻“参考答案”——看看 Subword Regularization 的原始论文《Subword Regularization: Improving Neural Network Translation Models with Multiple Subword Candidates》[1] 是怎么做的。

然而，这个“标准答案”却让笔者有点哭笑不得。原来Subword Regularization的思路非常简单直接：先搜索出最大的个分词方案（n-best segmentations），然后构建如下分布

对这个方案进行依概率采样，其中是一个超参数。该算法已经集成在 SentencePiece 中，读者可以自行测试（使用方法参考这里 [2]）。

问题是，“简单直接”不代表着“高效”，尽管搜索 top- 个分词方案最优方案的复杂度也是线性的（有兴趣的读者可以自行找找 N-best Viterbi 的资料），但明显比只找 top1 的 Viterbi Decoding 要大很多（理论上是倍复杂度），所以直接的后果是开启了随机采样后，会比确定性的分词要慢很多，所以这并非是笔者心中的理想采样方法。

1.3 个人思路

思路再度陷入了僵局。一筹莫展之际，笔者决定把思路再捋一捋：我们的目标是想要找到复杂度类似 Viterbi Decoding 的随机采样算法，既然如此，Viterbi Decoding [3] 本身应该是一个不错的突破点。于是，笔者再次翻开了分词代码，当时的分词函数长这个样：

 1def _tokenize(self, bytes text):
 2    cdef int e, k, s
 3    cdef double v, score
 4    cdef list routes = [(0, None)] + [(-INFINITY, None) for _ in text]
 5    cdef list tokens = []
 6    for e, (k, v) in self._automaton.iter(text):
 7        s, e = e - k + 1, e + 1
 8        score = routes[s][0] + v
 9        if score > routes[e][0]:
10            routes[e] = score, s
11    while text:
12        s = routes[e][1]
13        tokens.append(text[s:e])
14        text, e = text[:s], s
15    return tokens[::-1]

反复读了几遍，总算有了灵感：Viterbi Decoding 的关键是 if score > routes[e][0]: 这一句，它代表保留截止到当前位置的最优切分方案，其中 score 是新切分方案分数（概率对数）， routes[e][0] 是历史最优分数，如果新方案更优则覆盖。这让笔者联想到了 MCMC 算法的接受率设计，如果在这里引入随机采样，不就可以将分词结果随机化了？

我们用表示接受/拒绝新方案，由于这一步只是一个二元选择，所以将它概率化也非常简单：

这里是均匀随机数，是超参数，是 Sigmoid 函数，分别是新旧方案的得分（概率对数）。不难发现，左端的确定性采样对应的随机性采样。

这样，在 Viterbi 解码的基础上我们得到了一个非常自然、非常轻量级的随机采样算法，这里称之为 “Viterbi Sampling”，实现它只需要将 if score > routes[e][0]: 这一判据换成带随机数的版本。

由于 Sigmoid 函数的单调性，当时，它自然会给新方案分配更大的概率，所以很明显原来的的最大概率切分在 Viterbi Sampling 之下也是最大概率结果，并且当越大，也越大，这意味着原本得分越大的方案被采样到的概率也越高，一定程度上保持了切分方案的排序不变（尽管还没有证明一定严格保序，但从应用角度看，近似保序就够了）。

1.4 简单测试

从 0.4.0 版本开始，Viterbi Sampling 就内置在 BytePiece 的分词函数中，只需要在 tokenizer.tokenize 或者 tokenizer.encode 时加入大于 0 的 alpha 参数，结果就是随机的：

 1import bytepiece
 2assert bytepiece.__version__ >= '0.4.0'
 3
 4tokenizer = bytepiece.Tokenizer('bytepiece_160k.model')
 5text = '今天天气不错'
 6print(tokenizer.tokenize(text))  # alpha默认值为-1，alpha≤0 都代表确定性分词
 7for i in range(5):
 8    print(tokenizer.tokenize(text, alpha=0.1))
 9
10# [b'\xe4\xbb\x8a\xe5\xa4\xa9', b'\xe5\xa4\xa9\xe6\xb0\x94', b'\xe4\xb8\x8d\xe9\x94\x99']
11# [b'\xe4\xbb\x8a\xe5\xa4\xa9', b'\xe5\xa4\xa9\xe6', b'\xb0\x94', b'\xe4\xb8\x8d\xe9\x94\x99']
12# [b'\xe4\xbb\x8a\xe5\xa4\xa9', b'\xe5\xa4\xa9\xe6\xb0\x94', b'\xe4\xb8\x8d\xe9\x94\x99']
13# [b'\xe4\xbb\x8a\xe5\xa4\xa9', b'\xe5\xa4\xa9\xe6\xb0\x94', b'\xe4\xb8', b'\x8d', b'\xe9\x94', b'\x99']
14# [b'\xe4\xbb\x8a\xe5\xa4\xa9', b'\xe5\xa4\xa9', b'\xe6\xb0\x94', b'\xe4\xb8\x8d\xe9\x94\x99']
15# [b'\xe4\xbb', b'\x8a\xe5\xa4\xa9', b'\xe5\xa4\xa9', b'\xe6\xb0\x94\xe4\xb8\x8d', b'\xe9\x94', b'\x99']

下面对比一下 SentencePiece 的 Subword Regularization 和 BytePiece 的 Viterbi Sampling 的速度（随机性分词时都设）：

可以看到，Subword Regularization（“SP-Unigram” 这一行）开启之后，分词速度不到原来的 1/4，这表明 Subword Regularization 的采样算法是相当低效的。

相比之下，本文提出的 Viterbi Sampling 只下降了 30% 左右，效率显然更高，下降的部分在于随机数的生成和 Sigmoid 函数的计算，如果能进一步优化这两部分，速度还能进一步提升。至于 BPE 模型，它的随机分词叫做 BPE Dropout [4]，这是专属于BPE模型的方法，有兴趣的读者自行了解，这里就不介绍了。

1.5 小结

上文主要探讨了将 Unigram 分词模型的确定性分词改为随机性分词的策略。尽管已有名为 “Subword Regularization” 的方法可以实现这一目标，但其效率相对较低。为此，笔者提出了一种更高效的采样算法 Viterbi Sampling，它仅需对确定性的 Viterbi Decoding 进行简单的修改，从而基本保持了原有的效率。实验证明，新的算法采样速度明显超越了 Subword Regularization。相应的实现已经内置在 BytePiece 最新版中。

从Viterbi Sampling到完美采样算法

在上一节中，笔者提出了一种名为 “Viterbi Sampling” 的随机分词算法，它只是在求最优解的 Viterbi Decoding 基础上进行小修改，保留了 Viterbi 算法的简单快速的特点，相比于已有的 Subword Regularization [1] 明显更加高效。不过，知乎上的读者 @鶴舞指出，当前的采样算法可能会在多次二选一“稀释”了部分方案的出现概率，直接后果是原本分数最高的切分并不是以最高概率出现。

经过仔细思考后，笔者发现相应的问题确实存在，当时为了尽快得到一种新的采样算法，在细节上的思考和处理确实比较粗糙。为此，本节将进一步完善 Viterbi Sampling 算法，并证明完善后的算法在效果上可以跟 Subword Regularization 等价的。

2.1 问题分析

首先，我们来看一下评论原话：

subword regularization 中可以保证按概率数据（具有 temperature 超参数）。提出的方法中对于每个 e，第一个算出的 route 会被多次 1v1 “挑战”，最终概率分布会不会和已有算法差蛮多的。

举个例子，watching 三种分法 watch ing，wat ching，和 watching 概率都是三分之一，提出的方案的采样概率概率就会变成，前两个的概率是四分之一，第三个的概率是二分之一，是这样的吗？

其实评论里边已经说得很清晰了，如果读者还不理解的话，这里笔者稍微再展开一下。假设有三种切分方案，每种方案的得分都一样，那么我们自然是期望采样过程中每种方案的出现概率都是。然而，Viterbi Sampling 是将多选一的采样过程转化为多步的二选一：

这样一来，前面的两种切分方案先二选一，概率都是；选出来一个结果之后，又跟第三种方案放一起来二选一，由于概率是按照各自得分来算的，所以这时候各自的概率还是。于是，在完整的采样过程中，前两种方案出现的概率是，后一种方案出现的概率是，越晚出现的方案相对来说越“占便宜”，而越早出现的方案概率被稀释得越严重。

而很不巧的是，按照 BytePiece 的 AC 自动机的返回顺序，越长的词（通常来说得分越高）出现的次序会越靠前，所以在 Viterbi Sampling 中，得分越高的方案反而更容易被稀释概率。

2.2 解决办法

现在看来，其实解决办法也很简单，每次进行二选一后，同时缓存累积概率就可以了，而从第二步开始，每次二选一时新进来的候选者不是跟已有候选者得分二选一，而是跟累积概率得分二选一，这就是俗称“水塘采样（Reservoir sampling）”[5] 的算法。

用前面的例子来说，先进来两种切分方案，按照的概率选出一种，然后它们总的累积概率是；接下来被选者跟新方案选一，新出现的方案被选到的概率应该是，也就是说要跟累积概率比，而不是跟被选者自己的概率比，这样完整的采样流程下来，每种切分方案出现的概率都是。

对于 Viterbi Sampling 来说，每个终点位置会有多个切分方案，我们要对其进行多选一采样，被选中的概率是由各自的得分构造出来的，是归一化因子。因为我们是递归处理的，所以我们不知道多选一的“多”是多少，也无法计算，不过这不重要，知道就够了，因为计算每一步的条件采样概率其实也用不到完整的，而是需要递归的：

实际计算时，由于指数爆炸的原因，直接缓存大概率会有溢出风险，所以我们一般缓存的是它的对数，并利用函数避免溢出：

相应的实现已经内置在 bytepiece>=0.5.0 中。

2.3 完美采样

总的来说，出现旧版 Viterbi Sampling 的缺陷，还是因为之前操之过急了，所以现在认真地给新版 Viterbi Sampling 补上数学证明。有意思的是，可以证明更新后的 Viterbi Sampling 跟 Subword Regularization 一样都是“完美采样”算法。

之前我们介绍过，Subword Regularization 的做法非常“粗暴”，直接找出得分最高的个切分方案，然后通过的方式计算被选中的概率，其中是第种方案的得分。

这种做法除了复杂度高外没有任何毛病，当不做限制（即找出全部切分方案）时，我们得到所有切分方案的一个随机采样，而每种方案被采样到的概率正比于 ——是得分的单调增函数，即采样概率与得分的大小排序都是一样的，满足这两个条件的，笔者称之为“完美采样”。

2.4 Decoding

为了证明新版 Viterbi Sampling 也是“完美采样”，我们先来回顾一下 Viterbi Decoding。设有一个长度为的字节串，用表示最优切分方案的得分，假设我们知道之间一定会分开，那么必然有

也就是说，最优切分方案的子串，一定也是对应的子字节串的最优切分方案，这是动态规划的根本依据。当然，事实上我们不能预知哪一处会被切开，所以只能用枚举的方式：

其中是指字节串作为一个 token 时的得分（如果它不是词表中的 token，那么记为）。这样一来，的计算就转化为

的计算，依此类推，的计算又可以转化为

的计算，等等，也就是的结果是可以复用的。所以，整个流程总结下来就是一句话：

扫描到每一个位置时，都记录到当前位置的最优切分方案及其得分。

当然，直接按照式（7）进行递归的话，理论上复杂度是，但事实上不可能每个子字节串都是词表中的一个 token，所以可以用 Trie 树、AC 自动机等方法根据词表提前扫描好所有可能出现的 token，那么复杂度就正比于搜索出来的候选 token 数，关于是线性的，如果非要估计一个数值，那么假设词表中 token 的最大长度为，那么长度为的字节串扫描出来的 token 数就不超过

2.5 Sampling

有了 Decoding 部分做铺垫后，理解 Sampling 就相对容易一些了。其实关键还是在式（7），我们用表示字节串的所有切分方案的归一化因子（完美采样），那么有

这个等式也表明，要实现从的所有切分方案中按的比重采样，可以从的所有切分方案中随机选一个然后接上 token 、从的所有切分方案中随机选一个然后接上 token 、从的所有切分方案中随机选一个然后接上 token 、...，得到这个采样结果后，分别再以权重、、

、... 从中选一个。

接下来跟 Decoding 情形一样，的计算又可以重用

的结果，的计算又可以重用

的结果，等等，以及采样结果也都是可以重用的。于是类似地，那么整个 Sampling 算法也可以总结为一句话：

扫描到每一个位置时，都对以当前位置为终点的所有切分方案按照权重进行采样，记录采样结果以及累积权重。

如果两边取对数，那么式（9）可以等价地改写成

跟 Viterbi Decoding 的式（7）区别就是代替了，代替了，而正好是的光滑近似，所以时能退化为 Viterbi Decoding。

另一方面，在实际计算时，同一终点的多个切分方案是逐一到达而不是一次性到达的，所以就需要将单步的“多选一”转化为多步的“二选一”，这就是“解决办法”一节所讨论的内容。至此，我们证明了（或者说从 Viterbi Decoding 出发重新推导了）修改后的 Viterbi Sampling 实际是跟 Subword Regularization 一样的完美采样算法。

2.6 小结

本文完善了之前提出的随机分词算法 Viterbi Sampling，并从数学上证明了它在效果上跟 Subword Regularization 一样都是“完美采样”算法，而在使用上有着比 Subword Regularization 明显更高的效率。