CS336 - BPE Tokenizer

本文是我学习 Stanford CS336 BPE tokenizer 的学习笔记。

Tokenizer 负责文本和 tokens 的转换，一个典型的大模型运行流程是：

        encode()                                   decode()
 text ──────────▶ token IDs ──▶ LLM ──▶ token IDs ──────────▶ text
        └── Tokenizer ──┘                      └── Tokenizer ──┘

为什么一定要 tokenization，不能直接用文本呢？Token 需要被映射为一个高维向量（embedding）才能被 LLM 使用。

又引伸出一个问题：文本应该以何种方式被切分为一系列的 tokens？这是 tokenizer 核心问题。切分必须平衡好向量空间尺寸和输入量。两种极端情况：（1）输入太长，计算量爆炸。一个 byte 占用一个 token，单词表的长度略大于 256（包含一些特殊 token），模型输入长度比文本长度还大，如果输入的是中文那结果更糟，因为中文字符长度三倍于英文字符；（2）语义信息太稀疏。一个单词对一个 token，输入长度能很好的被压缩，但是语义空间爆炸，LM 很难学习到这句话所表示的含义。

UTF-8

UTF-8 采用了一种可变编码结构，如下所示。

首 bits 表示长度，直到遇到第一个 0。比如前 128 个字符仅占用一个字节，首 bit 必须为 0，大于 2 个字节，第一个字节需要表示需要有几个 bytes，比如 3 个字节是 1110，接下来的字节首 bits 比如是 10 开头。

在 Python 中，bytes 很多操作与 str 保持一致，比如 b"".join([b"a", b"b"]) 等效于 b"a" + b"b"，除了前面必须有个前缀 b。但是不能混合操作，比如 "".join([b"a", b"b"])。

一个字符串可以方便的与 bytes 相互转换，利用 encode() 和 decode()。

>>> test_str = "hi"
>>> print(test_str)
hi
>>> test_bytes = "hi".encode("utf-8")
>>> print(test_bytes)
b'hi'
>>> test_str_from_bytes = test_bytes.decode("utf-8")
>>> print(test_str_from_bytes)
hi

BPE Tokenzier

生成单词表

BPE tokenizer 解决压缩的思路来源于哈夫曼编码：高频短码，低频长码。高频出现通常意味着具有更明确的意义，比如“人工智能”这四个词大概率会同时出现。

每一轮循环寻找一个在语料库范围频率最高的 byte-pair，合并 byte-pair 为一个 bytes。以下是一个具体的例子：

首先将字符组合出 byte-pairs，比如“智能”（UTF-8 字节是 \xe6\x99\xba\xe8\x83\xbd）的 byte-pairs 是 [[\xe6,\x99], [\x99,\xba], [\xba,\xe8], [\xe8,\x83], [\x83,\xbd]]，比如 "lower" 的 byte-pairs 是 [[b'l',b'o'], [b'o',b'w'], [b'w',b'e'], [b'e',b'r']]。

在整个语料库的范围内统计出现频率最高的 pairs，假设 b'er' 出现了最多次，更新 l,o,w,e,r 为 l,o,w,er，下一轮计算 byte-pairs 变为了 [b'l',b'o'], [b'o',b'w'], [b'w',b'er']。

进入下一轮寻找，直到满足词汇表的尺寸限制……

为什么要大费周章的通过合并 bytes 构建词汇表呢？我们能不能直接统计高频词汇构建词表呢？答案是不能。直接统计词汇的频率就像是直接用词汇，但是不知道词汇的来源，导致泛化性不足。场景 1：英文词汇有很多时态变种，比如 happy, happiness, unhappy，直接学 3 个词，还是学 "happ" 与 "y"、"iness", "un" 的组合，肯定是后者泛化性更好。

Byte-level 具有更好的表示性，遇到没见过的字符退化为按字节理解，所有的字符都可以被表示为 bytes。

BPE 天然对英文友好，英文很多单词可以通过词根拼装，因此 byte-level 可以很好的学到这些特征，但是中文的偏旁部首并没有体现在 UTF-8 编码中，大模型就没办法学习到这种特征了。

BPE 说到底是一个 byte-pair 频率统计器，语料的分布也可能会影响性能，比如如果在大部分是英文的语料上训练可能无法充分学习到中文词汇特征。根据 Gemini 提供的信息，模型词汇表从早期的 3-5 万已经拓展到 10-25 万。

实验

CS336 课后的第一个编程任务是自己实现一个 BPE tokenizer，使用 TinyStories 训练 tokenizer 并序列化保存到本地磁盘，时间限制 30m，内存限制 30GB（我不知道限制为什么大的离谱），我的实现时间 115.7s（1.9m），峰值内存占用 0.24G，细节参见 GitHub。

文本处理需要利用正则表达式将文本切分为单词，首先用 re.split() 切分 special tokens 和小段文本，然后再用 re.finditer() 处理小段文本为单词列表，并构建 vocab_map 维护单词 bytes 到 count 的映射。这部分可以将 dataset 切分成不同的小块，统计工作之间不干扰，因此我选择用多进程（Python multi-process）加速处理进程。

接下来就是全局维度合并 byte-pair，一次迭代寻找一个最高频的 byte-pair，直到达到词汇表的上限。第一版利用了一种 naive 的方案，每次迭代重新统计全局最高频的 byte-pair，但是速度很慢，一个性能单元测试限制 1.5s，实际用时 3s 多。第二版利用三个 dicts，缓存全局统计信息，增量更新，性能有了比较大的提升。我维护了三个 dicts：

• vocab bytes tuple -> count（a.k.a. vocab_map）: vocab bytes tuple 将单词拆成 bytes，并在合并 byte-pair 时更新；
• byte-pair -> (count, set(vocab bytes tuple): count 是 byte-pair 全局频率，后面包含 byte-pair 的全部单词。
• count -> set(byte-pairs): byte-pair 频率到 byte 反向映射。

最后，用 TinyStories 的全部训练数据训练了 10,000 个词汇，并用 pickle 序列化词汇表到磁盘持久化保存。统计最长的词汇有三个 [b' accomplishment', b' disappointment', b' responsibility']，结果符合预期：（1）GPT2 tokenizer 会学习到前一个空格；（2）生成的词汇是有意义的。