📘
BPEのOSS「tiktoken」を探る

2025/03/10に公開
自然言語処理
LLM
生成 AI
tech
 はじめに「つくりながら学ぶ！LLM自作入門」を読み進めています。良書ですね。
https://www.amazon.co.jp/dp/4839987807
第2章にBPE (Byte Pair Encoding)でOpenAI社のOSS「tiktoken」が紹介されているのですが、複雑とのことで本書での解説が無く、詳しく知りたいと思いまとめてみます。

 tiktokenhttps://github.com/openai/tiktoken

 全体アーキテクチャ大きく3つの構成から成り立っている。
Python(./tiktoken): ユーザが使用するAPIを提供
Rust(./src): 高速なトークン化処理を実装
Encodding定義: 各モデルのトークン辞書と特殊トークンを定義
これらを、PyO3を通じて連携し、高速なトークン化処理を実現している。

https://github.com/PyO3/pyo3
./tiktoken/setup.py
from setuptools import setup
from setuptools_rust import Binding, RustExtension

setup(
    name="tiktoken",
    rust_extensions=[
        RustExtension(
            "tiktoken._tiktoken",
            binding=Binding.PyO3,
            # Between our use of editable installs and wanting to use Rust for performance sensitive
            # code, it makes sense to just always use --release
            debug=False,
            features=["python"],
        )
    ],
    package_data={"tiktoken": ["py.typed"]},
    packages=["tiktoken", "tiktoken_ext"],
    zip_safe=False,
)

 主要なコンポーネント
 Encodingクラス(./tiktoken/core.py)ユーザが直接操作するPythonクラス
encode(): テキストをトークンIDに変換
decode(): トークンIDをテキストに戻す
encode_ordinary(): 特殊トークンを無視してエンコード
encode_batch(): 複数のテキストを並列でエンコード
このクラスはRustで実装されたCoreBPEのラッパーとして機能

 CoreBPEクラス(./src/lib.rs および ./src/py.rs)トークン化の中核ロジックをRustで実装したクラス
正規表現による文字列分割
バイトペアエンコーディングアルゴリズムの実装
高速な辞書検索
byte_pair_mergeとbyte_pair_encodeがBPEアルゴリズムの実態

 Encoding定義(./tiktoken_ext/openai_public.py)各モデルで使用されるエンコーディングの定義
gpt2: GPT-2モデル用エンコーディング
r50k_base: 旧GPTモデル用の50k語彙
p50k_base: 少し新しいモデル用の50k語彙
cl100k_base: GPT-3.5およびGPT-4用の100k語彙
o200k_base: GPT-4o用の200k語彙
各エンコーディングはそれぞれ異なる正規表現パターンとマージ可能なトークンのランク、特殊トークンを持つ

 ModelとEncodingのマッピング(./tiktoken/model.py)モデル名から適切なエンコーディングを選択するためのマッピングが定義されている
./tiktoken/model.py
...
# TODO: these will likely be replaced by an API endpoint
MODEL_PREFIX_TO_ENCODING: dict[str, str] = {
    "o1-": "o200k_base",
    "o3-": "o200k_base",
    # chat
    "chatgpt-4o-": "o200k_base",
    "gpt-4o-": "o200k_base",  # e.g., gpt-4o-2024-05-13
    "gpt-4-": "cl100k_base",  # e.g., gpt-4-0314, etc., plus gpt-4-32k
    "gpt-3.5-turbo-": "cl100k_base",  # e.g, gpt-3.5-turbo-0301, -0401, etc.
    "gpt-35-turbo-": "cl100k_base",  # Azure deployment name
    # fine-tuned
    "ft:gpt-4o": "o200k_base",
    "ft:gpt-4": "cl100k_base",
    "ft:gpt-3.5-turbo": "cl100k_base",
    "ft:davinci-002": "cl100k_base",
    "ft:babbage-002": "cl100k_base",
}
...

 BPEアルゴリズムの実装tiktokenの中核
./src/lib.rs
fn _byte_pair_merge(ranks: &HashMap<Vec<u8>, Rank>, piece: &[u8]) -> Vec<(usize, Rank)> {
    // 各位置でのランク（優先度）を初期化
    let mut parts = Vec::with_capacity(piece.len() + 1);
    
    // 最小ランクを探す（最も優先度の高いマージ）
    let mut min_rank: (Rank, usize) = (Rank::MAX, usize::MAX);
    for i in 0..piece.len() - 1 {
        let rank = *ranks.get(&piece[i..i + 2]).unwrap_or(&Rank::MAX);
        if rank < min_rank.0 {
            min_rank = (rank, i);
        }
        parts.push((i, rank));
    }
    parts.push((piece.len() - 1, Rank::MAX));
    parts.push((piece.len(), Rank::MAX));

    // メインループ: 最も優先度の高いペアを繰り返しマージする
    while min_rank.0 != Rank::MAX {
        // トークンをマージし、影響を受ける隣接ペアを更新
        // （詳細は省略...）
        
        // 次のマージ対象を見つける
        min_rank = (Rank::MAX, usize::MAX);
        for (i, &(_, rank)) in parts[..parts.len() - 1].iter().enumerate() {
            if rank < min_rank.0 {
                min_rank = (rank, i);
            }
        }
    }
    
    parts
}
このアルゴリズムを読み解くと、
バイトシーケンスを連続した2バイトのペアに分解
最も優先度の高い(ランクの低い)ペアを見つける
そのペアをマージして1つのトークンにする
新しく形成されたペアのランクを計算
マージできるペアがなくなるまで繰り返す

 パフォーマンス最適化高速化のために以下の最適化を行っている
Rustの使用: 計算処理をRustで実装
スレッド処理: マルチスレッドによる並列エンコーディング
効率的なハッシュマップ: 標準のHashMapの代わりにFxHashMapを使用
正規表現最適化: スレッドローカルな正規表現エンジンを使用してロックの競合を回避
メモリ最適化: Python配列のコピーを避けるためのバッファインターフェースース

 拡張性tiktoken_extからも分かる通り、拡張可能

これにより新しいエンコーディングを追加することができる

ENCODING_CONSTRUCTORS辞書に追加し、これがランタイムにregistry.pyによってロードされる
./tiktoken/registry.py
def get_encoding(encoding_name: str) -> Encoding:
    if not isinstance(encoding_name, str):
        raise ValueError(f"Expected a string in get_encoding, got {type(encoding_name)}")

    if encoding_name in ENCODINGS:
        return ENCODINGS[encoding_name]

    with _lock:
        if encoding_name in ENCODINGS:
            return ENCODINGS[encoding_name]

        if ENCODING_CONSTRUCTORS is None:
            _find_constructors()
            assert ENCODING_CONSTRUCTORS is not None

        if encoding_name not in ENCODING_CONSTRUCTORS:
            raise ValueError(
                f"Unknown encoding {encoding_name}.\n"
                f"Plugins found: {_available_plugin_modules()}\n"
                f"tiktoken version: {tiktoken.__version__} (are you on latest?)"
            )

        constructor = ENCODING_CONSTRUCTORS[encoding_name]
        // ./tiktoken/core.pyのEncodingクラス
        enc = Encoding(**constructor())
        ENCODINGS[encoding_name] = enc
        return en

 まとめtiktokenはBPEトークン化を高速に実行するための工夫がなされている
トークン辞書: 各モデルは事前に学習された語彙とマージランクを持つ
正規表現分割: テキストはまず正規表現で分割される
BPEアルゴリズム: バイト単位で動作し、優先順位に従ってマージされる
特殊トークン処理: 特殊トークンは通常のBPEプロセスを迂回して直接マッピングされる
Pythonの使いやすさとRustの性能をうまく組み合わせることで、高速な処理が可能となっている

Pythonだけで実現するよりも3-6倍速い処理が可能とのこと
はじめに

tiktoken

全体アーキテクチャ

主要なコンポーネント

Encodingクラス(./tiktoken/core.py)

CoreBPEクラス(./src/lib.rs および ./src/py.rs)

Encoding定義(./tiktoken_ext/openai_public.py)

ModelとEncodingのマッピング(./tiktoken/model.py)

BPEアルゴリズムの実装

パフォーマンス最適化

拡張性

まとめ

Discussion
