NLP Subword 之 WordPiece 算法原理

本文将介绍以下 内容：

• 1. WordPiece 与 BPE 的区别
• 2. WordPiece 算法原理
• 3. WordPiece 算法流程
• 4. WordPiece 算法源码实现Demo

WordPiece 算法的核心思想是： WordPiece是基于最大化语料似然来选择扩展 token，每次从词表中选出两个子词合并成新的子词加入词表。

1. WordPiece 与 BPE 的区别：

WordPiece 是 Google 在 BERT 系列里用的分词方法，本质上和 BPE 类似，区别在于：

• WordPiece：基于最大化语料似然来选择扩展 token；
• BPE：基于合并 token 对的频率来扩展词表；

2. 算法原理：

假设句子 $S=(t_1,t_2,\dots ,t_n)$ 由 n 个子词组成，$t_i$ 表示子词，且假设各个子词之间是独立存在的，则句子
的语言模型似然值等价于所有子词概率的乘积：
$$\log P(S)=\sum _{i=1}^n\log P(t_i)$$
假设把相邻位置的x和y两个子词进行合并，合并后产生的子词记为z，此时句子 $S$ 似然值的变化可表示为：

$$\log P(t_z)-(\log P(t_x)+\log P(t_y))=\log \left(\frac{P(t_z)}{P(t_x)P(t_y)}\right)$$
从上面的公式，很容易发现，似然值的变化就是两个子词之间的互信息。简而言之，WordPiece每次选择合并的两个子词，他们具有最大的互信息值，也就是两子词在语言模型上具有较强的关联性，它们经常在语料中以相邻方式同时出现。

3. 算法流程：

将单词拆分成多个前缀符号（比如BERT中的##）最小单元，再通过子词合并规则将最小单元进行合并为字词级别。

（1）计算初始词表：通过训练语料获得或者最初的英文中的26个字母加上各种符号以及常见中文字符，这些作为初始词表。
（2）计算合并分数：对训练语料拆分的多个子词单元通过合并规则计算合并分数。
（3）执行合并：选择分数最高的子词对，将它们合并成一个新的子词单元，并更新词表。
（4）重复合并步骤：不断进行重复步骤（2）和步骤（3），直到达到预定的词表大小。
（5）分词：使用训练得到的词汇表对文本进行分词。

4. 算法源码实现Demo：

import math
import collections
from typing import List, Dictclass WordPieceTokenizerMLE:def __init__(self, unk_token="[UNK]"):self.vocab = {unk_token: 0}self.unk_token = unk_tokendef tokenize_chars(self, text: str) -> List[str]:"""基础分词：逐字切分"""return list(text.strip())def corpus_to_tokenized(self, corpus: List[str]) -> List[List[str]]:return [self.tokenize_chars(sent) for sent in corpus]def compute_likelihood(self, corpus_tokenized: List[List[str]], vocab: Dict[str, int]):"""计算 log-likelihood"""token_counts = collections.Counter()for sent in corpus_tokenized:for tok in sent:token_counts[tok] += 1total = sum(token_counts.values())loglik = 0.0for tok, c in token_counts.items():p = c / totalloglik += c * math.log(p + 1e-12)return loglikdef train(self, corpus: List[str], vocab_size: int = 50000):"""基于最大似然的 WordPiece 训练"""# 初始语料corpus_tok = self.corpus_to_tokenized(corpus)# 初始 vocab = 单字符for sent in corpus_tok:for ch in sent:if ch not in self.vocab:self.vocab[ch] = len(self.vocab)while len(self.vocab) < vocab_size:# 1. 统计相邻 pair 频率pair_counts = collections.Counter()for sent in corpus_tok:for i in range(len(sent) - 1):pair = (sent[i], sent[i + 1])pair_counts[pair] += 1if not pair_counts:break# 2. 计算所有 pair 的 ΔLbase_likelihood = self.compute_likelihood(corpus_tok, self.vocab)best_pair, best_gain = None, -1e9for pair, _ in pair_counts.items():# 尝试合并new_token = "".join(pair)temp_corpus = []for sent in corpus_tok:new_sent = []i = 0while i < len(sent):if i < len(sent) - 1 and (sent[i], sent[i+1]) == pair:new_sent.append(new_token)i += 2else:new_sent.append(sent[i])i += 1temp_corpus.append(new_sent)# 计算新的 log-likelihoodnew_ll = self.compute_likelihood(temp_corpus, self.vocab)delta = new_ll - base_likelihoodif delta > best_gain:best_gain = deltabest_pair = pairbest_new_corpus = temp_corpusif best_pair is None:break# 3. 接受增益最大的合并new_token = "".join(best_pair)self.vocab[new_token] = len(self.vocab)corpus_tok = best_new_corpusprint(f"Add token: {new_token}, ΔL={best_gain:.4f}")self.corpus_tok = corpus_tok  # 保存最后的分词表示def tokenize(self, text: str) -> List[str]:"""贪心匹配 WordPiece (不强制加##，因为训练没加过)"""tokens = []chars = self.tokenize_chars(text)i = 0while i < len(chars):match = Nonefor j in range(len(chars), i, -1):  # 尝试最长匹配substr = "".join(chars[i:j])if substr in self.vocab:match = substri = jbreakif match is None:tokens.append(self.unk_token)i += 1else:tokens.append(match)return tokensif __name__ == "__main__":corpus = ["我爱中国","中国人友好","I love China",]tokenizer = WordPieceTokenizerMLE()tokenizer.train(corpus, vocab_size=30)print("Vocab:", tokenizer.vocab)print("Tokenize 示例:", tokenizer.tokenize("我爱中国人"))