NovelAI の新しい LLM トークナイザーのご紹介

Finetune です。私が担当した新しいトークナイザーについて、少しお話しします。

まずは簡単なおさらいから。 ほとんどの場合、 私たちのモデルは文章を1文字ずつの並びとして見ているわけではありません。 文章はトークン、つまり単語や単語の断片へと分割されます。

モデルに使うトークナイザーを決めるときには、いくつもの判断基準があります。まず分かりやすいのが語彙サイズです。AI21 の Jurassic 系モデルが使うトークナイザーのように、どんな単語にも、場合によっては複数語のまとまりにも専用のトークンを割り当てられるよう、語彙サイズをとにかく大きく取りたくなるかもしれません。ただしその場合、モデルの汎化能力が落ちるという欠点があります。つまり、「-ize」で終わる単語どうしの共通点のような、綴りに潜む意味のあるパターンを活かせなくなるということです。綴り間違いにも弱くなります。とはいえ、語彙が小さすぎるのも良くありません。
よく使われる単語には専用のトークンを与えるべきですし、トークナイズ対象の文章に出てきそうな Unicode 文字も同様です。そうしないと、モデルはそれらをバイト単位で組み立てることになり、負担がかなり大きくなります。単一言語向けの語彙であれば、32000 トークン前後が語彙サイズの落としどころとして妥当です。この規模なら 16 ビットに余裕で収まるため、トークナイズ済みデータの扱いが多くの場面で楽になるという利点もあります。

トークナイザーの種類も重要な選択です。Unigram 系のトークナイザーは、単語をより意味の通る形に分割できることが示されています。一方で、これまで大規模言語モデル（LLM）で主流だったのは BPE（byte pair encoding）です。BPE の実装として最も広く知られているのは GPT2 のものでしょうが、Google の sentencepiece による BPE 実装には、バイトから組み立てずに Unicode 文字をそのままトークナイズできるという小さくない利点があります。バイトから組み立てる方式では、Unicode コードポイントの一部分を表す追加トークンを語彙に入れる必要があり、その分の領域が無駄になります。

たとえば「🙂」は「F0 9F 99 82」という4バイトで構成されます。従来の BPE では、まず F0 と 9F が結合して F09F になり、次に 99 と結合して F09F99 になり、さらに 82 と結合する、という手順を踏むため、中間段階のトークンを2つ余分に語彙へ追加しなければなりません。加えて sentencepiece は、バイトトークンを使って任意のバイナリデータをトークナイズすることにも対応しています。

最後に、トークナイザーが達成する圧縮率も見逃せません。同じ文章がより少ないトークン数に収まれば、LLM が一度に見られるコンテキスト長は決まっているので、その分だけ多くの文章を一度に扱えるようになります。これは LLM を使う人にとって重要です。また、たとえば一定量の学習データを用意したい場合に、目標のトークン数に届くまでどれだけの文章が必要になるかにも関わってきます。圧縮効率の低いトークナイザーであれば、同じトークン数のデータセットを作ること自体は簡単になりますが、そうして作られた同サイズのデータセットで学習したモデルは、圧縮率の高いトークナイザーで作った同サイズのデータセットで学習したモデルよりも学べる量が少なくなるはずです。実質的に、学習中に目にする情報量そのものが少なくなるからです。

こうした点を踏まえ、私たちは今後学習させるモデル向けに、 物語生成 といった自社の用途により最適化された、独自のトークナイザーを持つことを決めました。

トークナイザーはデータから学習させるものなので、まずはモデル学習用データセットのさまざまな部分集合から、小さなランダムサンプルを抽出し、それを使って利用可能なトークナイザー学習手法を評価しました。

Huggingface の tokenizers ライブラリも Google の sentencepiece も、異なる種類のトークナイザーの学習に対応しています。事前調査の結果、sentencepiece のトレーナーのほうがメモリ効率に優れていることが分かりました。それでも、十数ギビバイト規模の学習データを扱うには、1TB の RAM を積んだ計算ノードが必要でした。この点から、私たちは sentencepiece を採用しました。

当初は語彙サイズを 32000 にする予定でしたが、Genji V2 の学習の際に、既存のトークナイザーへ後から言語を追加するのは決して快適な作業ではないと痛感しました。今後も同じような言語間の転移学習を行う可能性が高いことから、トークナイザーには最初から英語と日本語の両方を収めることにしました。そのため語彙サイズを倍の 64000 に増やしたのですが、これは 16 ビットで表現できるトークン ID 空間をほぼ埋め尽くす規模だったので、思いきって 65535 トークンまで広げました。トークナイザーの学習にあたっては、2文字以上のラテン文字トークンと日本語トークンがおおよそ同じだけの領域を占めるよう、学習データの配分を丁寧に調整しました。語彙サイズを 65535 に引き上げたことで、絵文字などの Unicode 文字トークンもより多く収められるようになりました。日本語側の学習データには、既存の Genji 用データと、それより少量の日本語 Wikipedia を使いました。

複数の空白文字が連続する文字列については手作業でトークンを追加し、数字は1桁ずつトークナイズされるように設定しました。数字を1桁ずつ扱うと、数字の多い文章では圧縮率がわずかに下がる可能性がありますが、その代わりに LLM が数値の扱い方をより効果的に学べるようになります。

Unigram 系トークナイザーの利点を考慮して、私たちはまず Unigram のトークナイザーを学習させました。トークンの分布を狙いどおりにするため、言語間、そして主要データセット内の各部分集合間でデータのバランスを取り直しながら、何度も学習を回すことになりました。
Unigram の学習は1回あたり数時間かかります。比較のために BPE モデルも学習させましたが、こちらもデータの再調整のために複数回の実行が必要でした。BPE の学習ははるかに遅く、1回でほぼ丸一日を要しました。

こうして出来上がった2つのトークナイザーを、データセットから取り分けておいた部分で評価しました。狙いは、圧縮率が同程度、あるいは Unigram がわずかに劣る程度であれば、より自然な単語分割という利点を取って Unigram を採用する、というものでした。しかし実際には、英語データセットの主要部分において BPE のほうが 25〜29% 高い圧縮率を示しました。この予想外に大きな差から、Unigram ではなく BPE を採用することにしました。これは、LLM で BPE 系トークナイザーが使われ続けている理由の説明にもなっています。あわせて、32000 トークンの語彙を持つ sentencepiece ベースの BPE トークナイザーである LLaMa のトークナイザーとも圧縮率を比較しました。

LLaMa のトークナイザーと比べると、私たちのトークナイザーは英語データセットの主要部分で 7〜19% 高い圧縮率を達成しています。

最後に、トークンの分布についての数字をいくつか挙げておきます。

私たちのトークナイザーには、2文字以上のラテン文字から成るトークンが 28586 個含まれています。先頭に空白が付くトークンもここに含まれます。さらに、1文字を超える長さの日本語トークンが 18123 個、日本語と中国語の文字に対応するトークンが 9626 個あります。後者は Unicode の漢字統合の都合上、この統計では簡単には区別できません。そのほかのトークンが 9200 個あり、その多くは絵文字などの Unicode 文字が占めています。

比較として、LLaMa のトークナイザーには、ラテン文字だけから成るトークンが 23964 個、1文字を超える長さの日本語トークンは0個、日本語の文字が 836 個、そのほかのトークンが 7224 個含まれています。

新しいトークナイザーを実際に触ってみたい方は、こちらから sentencepiece モデルをダウンロードできます：https://github.com/NovelAI/novelai-tokenizer

golang-ベースの gpt_bpe トークナイザーに、BPE sentencepiece 対応と NovelAI の nerdstash 語彙を追加しました。1.5 兆トークンを現実的な時間でトークナイズするために使ったのが、この dataset_tokenizer ツールです。
作業中のプルリクエストはこちらです：https://github.com/wbrown/gpt_bpe/pull/30