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位只有28=2562^8=25628=256种表示。优点是字典很小,但压缩率很低,恒定为1 byte/token,生成的token序列很长,计算代价很大。尤其是注意力机制的复杂度是O(L2)O(L^2)O(L2),和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

更多推荐