【CS336】lecture1 tokenization

原创已于 2026-07-15 16:57:38 修改 · 251 阅读

4 ·

本内容遵循CC 4.0 BY-SA版权协议

GEO检测

话题

#AtomGit「码动四季·开源同行」夏季征稿活动

于 2026-07-15 12:51:35 首次发布

1. 定义

语言模型处理的输入是字符串，但内部是对token预测概率，因此需要字符串到token的转换规则。也就是需要定义

encode：字符串转化到token
decode：token转化到字符

定义压缩率为每个 token 对应的 UTF-8 byte 数：

def get_compression_ratio(string: str, indices: list[int]) -> float:
    num_bytes = len(string.encode("utf-8"))
    num_tokens = len(indices)
    return num_bytes / num_tokens

压缩率越高越好，因为压缩率高能把相同一段字符串输入，变成更短的token序列，降低模型训练推理的计算开销。然而压缩率提升往往需要更大的词典，比如把一个多个词组成的短语也看成一个token，而不是每个词一个token，压缩率自然提高了，但代价是词典里必须保存这个短语，不能只保存单个词汇。而词典增大会导致稀疏（很多token根本用不上）的问题，因此tokenizer设计，需要在词典大小和压缩率之间取得平衡。

2. Character tokenize

每个字符视为一个token，这实现简单，只要找一个字符集就行。缺点是字典很大，例如unicode就有15万个字符，很多稀有字符基本不会用上。并且压缩率也不会太高，unicode设计时大部分字符的字节数都很小。

class CharacterTokenizer(Tokenizer):
    def encode(self, string: str) -> list[int]:
        return list(map(ord, string))

    def decode(self, indices: list[int]) -> str:
        return "".join(map(chr, indices))

assert ord("a") == 97
assert chr(97) == "a"

3. Byte tokenizer

按字节分词，也就是把输入数据重新解释为单字节类型，这样字典大小只有256，因为一个字节8位只有 $2^8=256$ 种表示。优点是字典很小，但压缩率很低，恒定为1 byte/token，生成的token序列很长，计算代价很大。尤其是注意力机制的复杂度是 $O(L^2)$ ，和token序列长度二次方成正比

class ByteTokenizer(Tokenizer):
    def encode(self, string: str) -> list[int]:
        return list(map(int, string.encode("utf-8")))

    def decode(self, indices: list[int]) -> str:
        return bytes(indices).decode("utf-8")

assert "🌍".encode("utf-8") == b"\xf0\x9f\x8c\x8d"

4. Word tokenizer

按真实的人类字典进行分词，这是LLM前传统NLP的做法，压缩率还可以，且本身就包含了人类的语义内容。缺点是词典还是可能太大，并且会有很多稀疏token词典里有，但输入文本里几乎不会出现。且固定字典，如果遇到没见过的新词，只能映射到UNK（未知token），相当于一句话里一堆星号，会影响loss的计算，影响模型学习。

import regex

string = "I'll say supercalifragilisticexpialidocious!"
chunks = regex.findall(r"\w+|.", string)

5. Byte-Pair Encoding（BPE）

目前LLM最常用的分词算法。BPE 最初由 Philip Gage 在 1994 年用于数据压缩，后来被用于 neural machine translation，并被 GPT-2 采用。

核心思想是，在当前文本上训练分词器，得到词典，而不是基于某个确定规则。让分词器根据数据中词汇出现的频率决定如何压缩，出现比较多的序列用一个token表示，出现较少的序列仍用多个token表示。

名称带Byte是因为最开始按字节分割输入，也就是把输入看成一个字节序列，然后在此基础上合并

统计所有相邻token pair的出现次数，在第一轮，这就是字节形成的pair，因此称为Byte-pair
找到出现最多的token pair，将其所有出现都合并
将这次合并结果作为一个新的token存入词典。

重复以上过程直到进行了一定轮数，或者说词典大小达到目标。可以看到这是经典的机器学习思想，不人为规定规则，而是让模型根据数据驱动自学规则，这样每个训练数据，对应的分词器都不一样，都是最适合当前数据的。

下面是一个朴素实现，首先来看token pair计数和合并

计数时，把原始序列和左移一位的序列放一起打包成pair，这样每次取出的都是一个长度为2的子数组，且不用显式处理下标内容。用一个哈希表对pair进行计数

合并时，采用追加式的合并，遍历整个序列，对于目标pair的位置合成一个token插入序列末尾，对于非目标不进行操作，逐个插入。

from collections import defaultdict

def count_adjacent_pairs(indices: list[int]) -> dict[tuple[int, int], int]:
    counts = defaultdict(int)
    for index1, index2 in zip(indices, indices[1:]):
        counts[(index1, index2)] += 1
    return counts

def merge(
    indices: list[int],
    pair: tuple[int, int],
    new_index: int,
) -> list[int]:
    new_indices = []
    i = 0
    while i < len(indices):
        if (
            i + 1 < len(indices)
            and indices[i] == pair[0]
            and indices[i + 1] == pair[1]
        ):
            new_indices.append(new_index)
            i += 2
        else:
            new_indices.append(indices[i])
            i += 1
    return new_indices

训练时，初始化按字节分割，0-255作为初始词典中的token，运行指定轮数，每轮先统计所有相邻对的出现次数，然后取出出现最多的pair传给merge函数，并把这个pair合并在合并记录merges和词典vocab中记录。

from dataclasses import dataclass

@dataclass(frozen=True)
class BPETokenizerParams:
    vocab: dict[int, bytes]                 # token index -> bytes
    merges: dict[tuple[int, int], int]      # pair -> new token index

def train_bpe(string: str, num_merges: int) -> BPETokenizerParams:
    indices = list(map(int, string.encode("utf-8")))
    merges = {}
    vocab = {x: bytes([x]) for x in range(256)}

    for i in range(num_merges):
        counts = count_adjacent_pairs(indices)
        pair = max(counts, key=counts.get)
        new_index = 256 + i
        merges[pair] = new_index
        vocab[new_index] = vocab[pair[0]] + vocab[pair[1]]
        indices = merge(indices, pair, new_index)

    return BPETokenizerParams(vocab=vocab, merges=merges)

利用训练好的分词器进行encode，decode。encode时遍历整个词典，对每个词分别进行合并，这显然是低效的实现，这里只是为了演示，后面会在assignment 1里实现一个更高效的BPE。

class BPETokenizer(Tokenizer):
    def __init__(self, params: BPETokenizerParams):
        self.params = params

    def encode(self, string: str) -> list[int]:
        indices = list(map(int, string.encode("utf-8")))
        for pair, new_index in self.params.merges.items():
            indices = merge(indices, pair, new_index)
        return indices

    def decode(self, indices: list[int]) -> str:
        pieces = map(self.params.vocab.get, indices)
        return b"".join(pieces).decode("utf-8")

params = train_bpe("the cat in the hat", num_merges=3)
tokenizer = BPETokenizer(params)
string = "the quick brown fox"
assert tokenizer.decode(tokenizer.encode(string)) == string