【実装】BPEトークナイザを自作する — マージ規則を学習し、日本語で殴られる

BPEの学習器とエンコーダを自分で書く。マージ規則が「順序付きの手順書」であること、素朴な数え直しをどう捨てるか、日本語では語彙の先頭数千枠が「文字を組み立てる」ことに消える理由、語彙サイズを振る実験の正しい測り方までを実装の手触りで追う。

Neural Machine Translation of Rare Words with Subword Units

使うだけでは見えないこと

既存のトークナイザは3行で呼べます。AutoTokenizer.from_pretrained して encode するだけです。だからこそ、そこで何が起きているのかを聞かれると答えに詰まります。なぜこの位置で切れたのか。なぜ同じ単語が文頭と文中で別のIDになるのか。

概念の地図はトークナイザを1からにあります。この記事はその続きで、学習器とエンコーダを自分で書く側に回ります。200行ほどで動きます。面白いのは、書いた瞬間に実務の落とし穴がほぼ全部顔を出すことです。ストリーミング出力の文字化けも、語彙を足したときの事故も、原因はこれから書くコードの中にあります。

比喩: 速記者が略記号を増やしていく

あなたが速記者だとします。手持ちの記号は五十音と英数字だけ。会議録を取っていると「よろしくお願いいたします」が1日に何十回も出てくるので、新しい記号「①」を作ってこの並びに割り当てる。次に多い「ということで」は「②」。こうして略記号を増やすほど、同じ会議録が短く書けるようになります。

あとで効いてくる性質が2つあります。1つは先に作った記号を材料にできること。「よろしく」に記号を振った後なら、「よろしくお願い」は〈その記号＋お願い〉として作れます。記号は雪だるま式に育ちます。

もう1つは作った順番を控えないと再現できないこと。同じ記号表を渡されても、当てる順番が違えば別の書き起こしになります。記号表だけでは足りない。順番こそが規則です。

直感: 学習の成果物は語彙ではなく「順序」

BPEの学習と聞くと「よく出る単語のリストを作る」と思いがちですが、実際に作られるのは順序付きのマージ規則の列です。1行目に「e と s をくっつける」、2行目に「es と t をくっつける」……と続く手順書で、語彙はその副産物にすぎません。

なぜ順序が本体なのか。エンコードのとき est を 〈e + st〉 と切るか 〈es + t〉 と切るかを決めているのは、どちらの規則が先に学習されたかだけだからです。頻度は学習が終われば用済みで、残るのは並び順です。

だから配布物としてのトークナイザは、語彙ファイル（トークン→ID）とマージファイル（規則を学習順に並べたもの）の2つ組になります。HuggingFace の tokenizer.json に vocab と merges が並ぶのはこの事情で、順序を1行入れ替えれば語彙が同じでも別のトークナイザになります。

仕組み: 1回のマージで系列はどれだけ縮むか

学習の手順は3語で言えます。数えて、選んで、くっつける。

式(1)は「いま数えた隣接ペアのうち、出現回数がいちばん多いものを選ぶ」と言っているだけです。$\arg\max$ は「最大にする引数を返す」記号、$\mathrm{count}(a,b)$ は $a$ の直後に $b$ が来た回数です。つまり、意味も文法も語の切れ目も一切考慮せず、コーパスがいちばん多く繰り返した並びが勝ちます。

この選び方が理にかなっているのは、回数がそのまま縮む量だからです。ペア $(a,b)$ を1個の新トークンに置き換えると、置き換えた箇所ごとにトークンが2個から1個に減ります。

式(2)は「1回のマージで、コーパス全体のトークン数 $L$ は出現回数のぶんだけ減り、語彙の大きさ $|V|$ は1つ増える」という意味です。要するに $L$ は「トークンにしたときのコーパス全体の長さ」、$|V|$ は「種類として何個のトークンがあるか」で、マージのたびにこの2つは必ず逆向きに動きます。語彙1枠という同じ値段で買える短縮量を最大にする貪欲法——BPEが圧縮アルゴリズム出身だと言われるのは、この目的関数を指しています。

なお aaa に (a,a) を当てる場合、左から処理すると1回しか適用できず、実際の短縮は式(2)よりわずかに小さくなります。自作したBPEが既存実装と一致しないときの犯人は、たいていこの手の端の処理です。

実装1: まず素朴に書く

式(1)をそのままコードにします。コーパスは「事前分割した断片 → 出現回数」の辞書として持ち、断片はバイトの列にしておきます。

from collections import Counter

def train(corpus, n_merges):          # corpus: Counter[tuple[bytes, ...]]
    merges = []
    for _ in range(n_merges):
        pairs = Counter()
        for word, f in corpus.items():
            for p in zip(word, word[1:]):
                pairs[p] += f          # 隣接ペアを全部数え直す
        if not pairs:
            break
        (a, b), _ = pairs.most_common(1)[0]
        merges.append((a, b))
        corpus = Counter({apply(w, a, b): f for w, f in corpus.items()})
    return merges                      # 学習順に並んだ規則の列