🐻

トークナイザー構築のナレッジとチームの取り組み紹介【Team Kuma】

2024/09/10に公開

はじめに

はじめまして。GENIAC 松尾研LLM開発プロジェクト Team Kuma の宮澤と申します。

本プロジェクトについては様々な記事で紹介されているため詳細は割愛しますが、本記事は5月下旬までの「Phase1」の取り組みをもとに書いたものとなります。

本記事を執筆している現在は「Phase2」の開発期間もすでに終了していますが、「Phase1」での各チームの取り組みや結果発表は以下から視聴することができます。

本記事の概要

私はPhase1の開発にて、トークナイザーの構築・事後学習・評価周りを担当していました。今回は特に時間を使っていた「トークナイザーの構築」についての知見を共有したいと思います。

ただし、トークナイザーの工夫が直接モデルの精度に寄与していたかを測定することは難しい(事前学習を何度も実施することはできない)ため、ここでは「このように作ればよい!」というベストプラクティスを提案するのではなく 1. トークナイザー構築の取り組みを通じて学んだこと 2. チームでの取り組み紹介 という2点を軸に書かせていただきます。

それを踏まえて、本記事の想定読者としては、LLMを事前学習からゴリゴリ経験したことのある方というよりも、LLMを事前学習から開発したことはないがLLMを扱う機会がある方、もしくはあまり使ったことはないが興味がある方とさせていただきます。

1. トークナイザー構築の取り組みを通じて学んだこと

1-1. 【基礎知識】トークナイザーについて

トークナイザーとは

そもそもトークナイザーとは何かについて説明しておきます。
トークナイザーとは端的に言えば「テキストをトークン単位で分割するもの」です。例えば、本記事のタイトルをトークナイズすると、以下のようになります。

from transformers import AutoTokenizer
tokenizer = AutoTokenizer.from_pretrained("elyza/ELYZA-japanese-Llama-2-7b-fast")
print(tokenizer.tokenize("トークナイザー構築のナレッジとチームの取り組み紹介"))
# 出力
['▁', 'ト', 'ーク', 'ナ', 'イ', 'ザー', '構', '築', 'の', 'ナ', 'レッ', 'ジ', 'と', 'チーム', 'の', '取り', '組み', '紹介']

ここで上の分割結果を見るとわかるように、「トークン」とは、単語でもなく文字でもない単位で区切られています。したがって、トークナイザーの学習とは「どのようにテキストを分割するかのルール決め」であると言い換えることができます。

「語彙」とは

トークナイザーを理解する際に必ず出てくるのが「語彙」です。これは、上で示した「トークン」とそれに対するIDを紐付けた辞書のようなものを意味します。例えば、elyza/ELYZA-japanese-Llama-2-7b-fastのtokenizer.jsonを見てみると、以下のように語彙が設定されています。

"<unk>": 0,
"<s>": 1,
"</s>": 2,
…
"Price": 13026,
"antin": 13027,
"emento": 13028,
"may": 13029,
…
"構": 36116,
…
"紹介": 42858,
…

先ほどトークナイザーの学習は「どのようにテキストを分割するかのルール決め」と表現しましたが、その過程の一つとして上のような「語彙」を作ることも含まれます。

ここで重要な基礎知識としては、トークナイザーでは語彙に設定されているトークンの粒度でテキストが分割されるということです。

実際に先ほどテキストを分割したのと同じトークナイザーを使って、今度はテキストをトークン単位に区切ってIDにする「エンコード」をしてみます。

print(tokenizer.encode("トークナイザー構築のナレッジとチームの取り組み紹介"))
出力:
[1, 29871, 30279, 43004, 30576, 30260, 43349, 36116, 37450, 30199, 30576, 44313, 30391, 30364, 43978, 30199, 42850, 43451, 42858]

テキストが分割されてIDになりました。「紹介」というトークンに注目すると、語彙の中ではID:42858で登録されており、実際にトークナイズおよびエンコードをした時に、[紹,介]ではなく[紹介]として分割されていることがわかります。

代表的なアルゴリズム

トークナイザーで重要な「語彙」について理解したところで、実際にどのようにその語彙を作っていくのか(学習していくのか)について簡単に説明していきたいと思います。

よく使われているアルゴリズムとして、BPE(Byte-Pair Encoding)Unigramがあります。

BPE

こちらは、文字単位での分割をした状態から隣接する文字同士を結合していくことによって行われるデータ圧縮の手法の一つです。BPEでは、テキストをすべて1文字単位に分割してから、隣り合う文字の頻度が高いものを結合する過程を通してテキストの結合ルールを作ります。

例えば、「トークナイザーとアナライザーは語尾が似ている。」というテキストから結合ルールを作る場合、以下のように考えます。

文字単位に分割
ト ー ク ナ イ ザ ー と ア ナ ラ イ ザ ー は 語 尾 が 似 て い る 。

隣り合う頻度が高いものを結合する
ト ー ク ナ イザー と ア ナ ラ イザー は 語 尾 が 似 て い る 。

ここでは隣り合う頻度が高い「イザー」が一つのトークンとして作られました。

BPEではこのように、隣接する頻度が高い文字同士を繋げることでトークンを作っていきます。これによって語彙を形成し、結合ルールも学習していきます。また、どのくらいまで結合するかはハイパーパラメータで設定することができます。

先ほど例に挙げたelyza/ELYZA-japanese-Llama-2-7b-fastのtokenizer.jsonを再度見ると、"vocab"(語彙)の後に"merges"という結合ルールの設定が書かれていることがわかります。

"merges": [
      "▁ ▁",
      "▁ t",
      "e r",
      "i n",
      "▁ a",
      "e n",
      "o n",
      "▁ th",
      "▁t h",
      "e s",
        …

Unigram

Unigramの仕組みはBPEと比べると少し複雑であるため、簡潔な説明にとどめさせていただきます。Unigramによる学習プロセスはざっくりいうと「初めに大きなサイズの語彙を用意してそれぞれのトークンの生起確率を計算し、EMアルゴリズムを用いてテキスト分割の対数尤度が高くなるように、指定の語彙の大きさになるまでトークンを絞っていく手法」と言えるかと思います。

例えば、ある分割パターンにおいて、あるテキストをトークナイズしたときのトークンのリストを \mathbf{x} = [x_1, \dots, x_M] とすると、x_iのテキスト内での生起確率は p(x_i) = \frac{c(x_i)}{N} で計算できます。(c(x_i)w_iの出現回数。)

この時のテキスト分割(トークナイズ)における対数尤度は \log P(\mathbf{x}) = \sum_{i=1}^{M} \log p(x_i) となります。

この対数尤度が最大化されるように学習してトークンとその生起確率をもとに語彙を作っていくのがUnigramの仕組みです。

Unigramを用いて構築されたトークナイザーを使っているllm-jp/llm-jp-13b-v1.0のtokenizer.jsonを見ると、以下のように語彙にトークンと対数尤度が保存されていることがわかります。

…
[
"read",
-9.904170036315918
],
[
"ctrl",
-11.401485443115234
],
[
"▁returns",
-11.236441612243652
],
…

なぜ目的に適したトークナイザーを作るのか

トークナイザーの基礎知識編の最後に、「LLMを開発する上でなぜトークナイザーを作ることが求められるのか」を説明しておきます。トークナイザーはテキストを分割するだけのシンプルな役割に見えますが、単純な単語区切りや文字区切りにはない優れた特徴があります。

サブワードトークン化による効率性

ここまでに説明していませんでしたが、トークナイザーに関して押さえるべき概念として「サブワード」というものがあります。これはトークナイザーで分割されるトークンのことを表します。例えば、unbelievableという単語は['un', 'bel', 'iev', 'able']と分割されます。unは否定の意味を持っていたり、ableは「できる」という意味を持っていたりします。このように区切られたトークンを「サブワード」と呼ぶことがあります。

このような分割によって、語彙に登録するトークン数を節約することができます。例えば、先ほどの'un', 'bel', 'iev', 'able'が語彙に含まれていれば、unableというトークンが直接語彙に含まれていなくても、['un', 'able']とトークン化することができます。

これは言い換えると、語彙に含まれていない未知語への適応性が高いとも言い換えられます。日本語の例をあげると、上で見た例で['構','築']と分割されていましたが、他にも'成','造','える'などのトークンが語彙にあれば、直接その単語が語彙になくても'構成', '構造', '構える'といった語をトークナイズしたりエンコード・デコードできます。

ここで、「それならいっそ文字単位に区切ればいいじゃないか」と思われるかもしれませんが、この場合はある問題が発生します。それについては後述します。

多様な言語への適用性

テキストを分割する方法として考えられるのは、英語の場合は単語区切り、日本語の場合は形態素解析などが挙げられます。この時点で、言語によって異なる分割方法が候補となっているということがわかります。しかし、LLMを開発する上でこれは望ましくありません。理由は、汎用的なLLMを使っていてわかる通り、LLMで解きたいタスクには翻訳やコード生成といった言語横断的なものが多くあるためです。そこで、上に挙げたようなBPEやUnigramを用いたトークナイザーでは生のテキストから学習を行うため、言語依存性の問題を解決することができます。

また英語と日本語の大きな違いとしては「単語が空白で区切られているかどうか」であると言えますが、これらの言語の扱いを近づけるために、空白もトークンとして扱います。

例えば以下のようなトークナイズを行います。(_は空白を意味する。)
"Hello World.” → [ 'Hello', '_Wor', 'ld', '.' ]
“こんにちは世界。” → [ 'こんにちは', '世界', '。' ]

これによって、トークンの並びをデコードした際に、Hello World.こんにちは世界。を正しく元に戻すことができます。つまり、空白をトークンとして扱ってトークナイザーの学習を行うことで、空白区切りの言語とそうでない言語を一緒に扱うことができるということです。

LLMの計算効率の向上

先ほど、文字単位の分割には課題があると述べましたが、その理由について説明します。まず前提としてLLMの事前学習には大量のコーパスと計算リソースが必要です。事前学習は簡単に言えば「次に来るトークンの予測」によって行われるため、トークンの数が多いほど必要な計算量は多くなります。したがって、1トークンで表現される文字数が多いほど計算効率は高いということです。

一般的な日本語LLMでは、日本語は1トークン=1.5文字〜2文字程度で表現されるため、1文字単位で区切る場合に比べて、計算効率が1.5倍〜2倍ほど高くなります。多くのLLM開発では計算量に比例して膨大な経済的コストがかかるため、このような計算効率の向上は非常に重要なポイントであると言えます。

後述しますが日本語をバイト列で表記すると1文字あたり3バイトで表現されることも多いため、1文字=3トークンとなることもあります。よって、1文字=1トークンとなることで2倍以上の効率化につながると考えられます。

また、トークナイズの効率が上がることは事前学習だけではなく推論速度にも寄与します。少し前に公開されたGPT-4oを使われている方は感じたかもしれませんが、こちらは従来のGPT-4と比べて高速にレスポンスを返すようになっています。主な要因は、語彙にほとんど含まれていなかった日本語も語彙に追加され、同じテキストでもより少ないトークンで入力・出力できるようになったことであると考えられます。

1-2. 【先行研究調査】トークナイザーの評価指標

ここからは、GENIACプロジェクトにてトークナイザーを構築するにあたって調べたことについて紹介していきます。私自身、上のような基礎知識はあったもののトークナイザーを作るという経験はなかったため、まず「どのようなトークナイザーがよいとされるのか?」という疑問がありました。そこで評価指標について調べてみました。

一般的な評価としては以下のようなものがあることがわかりました。

  • Fertility
    • 文章を表現するのに必要なトークンの平均数。スコアが高いほどトークナイザーの圧縮率が低いことを意味する。
  • Parity
    • 異なる言語に対してトークナイズしたときのトークン数の長さのバランスを示す指標。スコアが高いほど多言語でのトークナイズのバランスがいい。FLORES-200のような多言語翻訳コーパスを使って、言語間のトークン数によって均等性を比較することができる。言語Aの文のトークン数 / 言語Bの文のトークン数が1に近いほどParityが高い。
  • LPT(Length Per Token)
    • 1トークンあたりの文字数を意味する。テキストの圧縮比の逆数で、文字数 / トークン数で求める。大きいほど圧縮率が高い。LPTは語彙数と正の相関をもつ。

1-3. 【先行研究調査】アルゴリズムの違いと下流タスクの精度

代表的なアルゴリズムとしてBPEとUnigramを上げましたが、多くのLLMのトークナイザーではこれらのどちらかが用いられていたため、チームの開発としてもどちらかを用いることを検討していました。ただし、どちらが今回の日本語LLM開発に適しているかを決定づけることが難しく、最終的には参考にしていた日本語モデルと使用したいライブラリとの互換性を踏まえて決定しました。

しかし、念のためどのような先行研究があるかは調べていました。上に挙げたTokenizer Choice For LLM Training: Negligible or Crucial? (2023)では、BPEとUnigramのどちらがLLMの下流タスクの性能が高いかを実験していました。本論文では英語およびヨーロッパ圏の多言語を扱っていたため、日本語ではありません。結論としては、多言語の場合はUnigramを使った場合の下流タスクの性能が高いということが述べられていました

1-4. 【先行事例調査】代表的な日本語モデルのトークナイザー

トークナイザーを構築するにあたって、他の日本語を扱うのに適したLLMのトークナイザーについて調査をしていました。

モデル 学習コーパスの言語(*1) 語彙数 アルゴリズム 参考
stabilityai/japanese-stablelm-base-alpha-7b 英語(*2), 日本語 65,536(*3) BPE blog
elyza/ELYZA-japanese-Llama-2-7b 英語(*2), 日本語 45,043(*3) BPE blog
llm-jp/llm-jp-13b-v1.0 英語, 日本語 50,688 Unigram paper
stabilityai/japanese-stablelm-base-ja_vocab-beta-7b 英語(*2), 日本語 49,408(*3) BPE blog
tokyotech-llm/Swallow-7b-hf 英語(*2), 日本語 43,176(*3) BPE paper
rinna/nekomata-14b 英語(*4), 中国語(*4), 日本語 151,936 BPE blog

*1) 主として構成されている言語のみ記載しています。構成比が小さい言語やプログラムは省略しています。
*2) Llama 2を継続事前学習したモデルであるため、Llama 2の事前学習時で使用されているデータにおいて英語が主であることを示しています。
*3) Llama 2の語彙を拡張しているため、これらのうち32,000トークンはLlama 2ですでに設定されている語彙になります。
*4) Qwen 7Bを継続事前学習したモデルであるため、Llama 2の事前学習時で使用されているデータにおいて英語と中国語が主であることを示しています。

これをみると、多くのモデルがBPEを用いていることがわかりますが、これはベースとなっているLlama 2のトークナイザーがBPEを用いているためです。ゼロから事前学習しているllm-jp/llm-jp-13b-v1.0ではUnigramが用いられており、このモデルは事前学習したデータの言語ドメインの割合として英語と日本語をおおおよそ半分ずつ学習していました。今回のLLM開発で私たちのチームも同様に英語と日本語を1:1程度の割合で学習することを検討していたため、llm-jp/llm-jp-13b-v1.0を参考とさせていただきました。

1-5. 【実装方法】松尾研標準コードの理解

ここからは実装に関する内容に入っていきます。
本プロジェクトでは、運営である松尾・岩澤研究室から「LLM開発の標準コード・手順」として、サンプルプログラムをご提供いただいていました。私はトークナイザーの構築を経験したことがなかったので、実装についてはまずこちらのコードの理解から始めました。

https://github.com/matsuolab/ucllm_nedo_prod/tree/main

トークナイザーの学習コード

コード理解のメモ

  • spm.SentencePieceTrainer.trainという記載からわかるように、トークナイザーの構築にはSentencePieceが用いられている。
  • 引数が様々あるがそれぞれ以下のような意味である。(全ての引数について知りたい場合はドキュメントを確認。)
input=args.input, # インプット(学習データ)
model_prefix=args.model_prefix, # モデルの接頭辞
vocab_size=args.vocab_size, # 語彙サイズ
input_sentence_size=args.input_sentence_size, # 学習に使う最大のセンテンス数
shuffle_input_sentence=args.shuffle_input_sentence, # センテンスをシャッフルするかどうか
num_threads=args.num_threads, # 学習のスレッド数
character_coverage=args.character_coverage, # 文字のカバー割合
model_type=args.model_type, # アルゴリズム
train_extremely_large_corpus=args.train_extremely_large_corpus, # Unigramのビット深度を増加させるかどうか(大規模コーパスを使う場合に有効にする)
user_defined_symbols=[BOS_TOKEN,] # ユーザー定義トークンの設定
byte_fallback=True, # バイトフォールバック
split_digits=True, # 数字を分割するかどうか
allow_whitespace_only_pieces=True, # 空白をトークンとして扱うかどうか
remove_extra_whitespaces=False, # 連続する空白を削除するかどうか
  • ユーザー定義トークンとしては、special_token_list.pyにて以下のように設定されている。
UNK_TOKEN = "<unk>"
BOS_TOKEN = "<s>"
EOS_TOKEN = "</s>"
PAD_TOKEN = "<pad>"
CLS_TOKEN = "<CLS>"
SEP_TOKEN = "<SEP>"
EOD_TOKEN = "<EOD>"
MASK_TOKEN = "<MASK>"
NEWLINE_TOKEN = "\n"
  • 以下のようなスクリプトでargsを設定しつつトークナイザーの学習を実行する。
python train_sentencepiece_tokenizer.py
 --input /path/to/dataset1.jsonl,/path/to/dataset2.jsonl,/path/to/dataset3.jsonl \
    --model_prefix ${YOUR_TOKENIZER_NAME} \
    --vocab_size 32000 \
    --input_sentence_size 3000000 \
    --shuffle_input_sentence True \
    --num_threads 16
  • 学習した後、SentencePieceからHuggingFace型へ変換する。変換コードは以下。

https://github.com/matsuolab/ucllm_nedo_prod/blob/16c3b3b9cbc93f43b68d9616a6976b527bfe8cd2/train/scripts/step3_upload_pretrained_model/convert_tokenizer_from_sentencepiece_to_huggingface_transformers.py

  • from transformers import T5TokenizerでT5Tokenzierというクラスを取得して、モデルファイルや各種設定を渡してHuggingFaceのtransformers`ライブラリで扱えるトークナイザーに変換していると読み取れる。

この時点で各種引数の意味などを全て理解することはできませんでしたが、全体的な実装の流れは掴むことができました。これをもとに意図通りの挙動をするかなどを実験的に確かめながらトークナイザー構築することとしました。

前半パートの「1. トークナイザー構築の取り組みを通じて学んだこと」はここまでとします。聞いたことも使ったこともあるが仕組みをよく知らなかったという方も、ここまでの内容である程度の基礎知識を抑えることができていれば幸いです。

2. チームの取り組み紹介

ここからはチームのLLM開発で実際に取り組んだトークナイザー構築について紹介したいと思います。

2-1. どのようなトークナイザーを作るか

まずどのようなトークナイザーを作るかを検討しました。特に大きなポイントとしては以下が挙げられました。

語彙数

上述した通り、トークナイザーは語彙数が大きなポイントになります。先行事例調査で見たように当時の日本語LLMの多くが約50,000語の語彙を設定していました。そのため、私たちのチームでも同等サイズにすることとしました。ただし、当時はQwen(約15万語)やGemma(約25万語)といった高性能なモデルがちょうど登場し始めた頃であり、自分の中では「語彙サイズが大きい方がLLMの性能が上がるのか?」という仮説が出てきていました。そこで、最終的には50,000語より少しだけ多く、56,320語に設定しました。

また、語彙の内訳についてですが、上に挙げたLlama 2を継続事前学習したモデルでは32,000語に約20,000語のトークンを追加しているため、英語と日本語の比率はおおよそ3:2となっています。ただし、これらのモデルはあくまでもLlama 2の事前学習がメインであったため、英語のトークンの割合が多いことは自然であるように思います。

私たちのチームでは、英語と日本語のテキストを同じくらいの量で事前学習を行う予定であったことから、トークナイザーの語彙も1:1程度を想定してトークナイザーを学習しようと考えました。しかし、日本語テキストで獲得した語彙にも多くのアルファベットが含まれているため、結果的には英語と日本語の語彙は3:2に近い割合になりました

アルゴリズム

こちらは上述した通り、Unigramを用いています。先行研究の「ヨーロッパ圏での多言語の場合にはUnigramの方が下流タスクの性能が高かった」という示唆や、llm-jp/llm-jp-13b-v1.0を参考にしていたという理由のほか、llm-jp-tokenizerというリポジトリが複数言語の語彙の割合を調整しながらトークナイザーを構築するのに役立ちそうであり、こちらがUnigramベースの方法となっていたことから、Unigramを使うこととしました。

学習データ

トークナイザーの学習にはできるだけ整頓されたテキストであることが望ましいため、主にはwiki40bを用いました。日本語と英語のどちらも整備されていたため、両言語ともこちらのテキストデータをメインで学習に使いました。また、今回のLLM開発の評価タスクにはプログラム生成タスクや数学タスクも含まれることがわかっていたため、語彙にもコード特有の英語や数学記号を含めることを目的として、事後学習に利用しようと考えていたMBPPGSM8Kも語彙獲得に使うこととしました。

2-2. 実装について

全体像

全体的な流れとしては以下のように、各ドメインコーパスで別々にトークナイザーを学習して語彙を獲得し、任意の割合で統合して最終的な語彙とモデルを作るという手順で進めました。

スクリプト

チームのリポジトリにてスクリプトを公開しています。
基本的には松尾研LLM開発標準コードllm-jp-tokenizerを参考にさせていただいています。

https://github.com/geniacllm/tokenizer/tree/main

モデル

チームのHuggingFaceにてトークナイザーを公開しています。
https://huggingface.co/geniacllm/ja-en-tokenizer-unigram-v5

実装手順

学習用データが準備できている前提で説明します。

1. 日本語テキストの形態素解析

日本語については事前に形態素解析の処理を含めることで、形態素の大きさ以上でトークンが区切られないようにしました。実装方法としては、fugashiを用いて形態素ごとに||||という文字列で分割を行いました。こちらの実装はSentencepiece の分割を MeCab っぽくするという記事を参考にさせていただきました。

形態素解析を行ったテキストに対して以下のようにSentencePieceでpretokenization_delimiter="||||"と設定することで、||||という区切りを使った分割を適用することができます。

spm.SentencePieceTrainer.train(
        input=args.input, # インプットファイル
        …
        pretokenization_delimiter="||||"
    )

2. 言語ごとにトークナイザーを学習

上の通り、英語・日本語・算数・コードのそれぞれでトークナイザーを学習して語彙を獲得しました。
英語(算数およびコード)と日本語の学習の設定として変えた部分としては、(先ほどのpretokenization_delimiter="||||"の設定以外では)1トークンに対する最大の文字数の設定があります。

spm.SentencePieceTrainer.train(
        …
        max_sentencepiece_length=16, # 最大の長さ)

日本語では8文字英語では16文字として設定しました。

そのほか、設定としては以下のようにしていました。(一部のみ抜粋)

vocab_size=56,320 # 語彙サイズ
character_coverage=0.9995 # 文字のカバー率99.95%
model_type="unigram" # アルゴリズム
normalization_rule_name="identity" # 正規化なし
byte_fallback=True # バイト変換あり
split_digits=True # 数字分割あり
allow_whitespace_only_pieces=True # 空白のトークンを許可する
remove_extra_whitespaces=True # 余分な空白の削除あり

ここまでで説明していませんでしたが、重要である引数を3つ取り上げて触れておきます。

normalization_rule_name(正規化)

こちらは文字通り「正規化」の設定です。identitynmt_nfkcを設定でき、前者は正規化なし、後者はNFKC正規化ありを意味します。NFKC正規化とは「互換等価」で文字を分解し「再合成」する正規化処理です。(参考:[Unicode 正規化])(https://ja.wikipedia.org/wiki/Unicode正規化)
正規化ありの場合、例えばThisthisテストテスト1が同じトークンとして扱われます。

llm-jp/llm-jp-13b-v1.0tokyotech-llm/Swallow-7b-hfのトークナイザーで試したところ、上の例は全て別のトークンとして扱われていたため、おそらく正規化なしの設定であると考えられました。今回の私たちの開発目的を踏まえても、正規化なしによるメリットの方が多いと考えられたため、normalization_rule_name="identity"としました。

byte_fallback(バイトフォールバック)

基礎編で見てきたように、トークナイザーは語彙に含まれているトークンをベースとして分割
エンコード・デコードを行います。逆に言えば、語彙に含まれていないトークンを扱うことができず、<unk>といったトークンに置き換えたりします。この問題を解決するために、バイトフォールバックという機能があります。これは、扱いたい文字が語彙に含まれていないとき、その文字をバイト列(通常はUTF-8のバイト列)として変換するものです。

例えば、語彙に漢字のトークンをほとんど含んでいないmeta-llama/Llama-2-7b-hfのトークナイザーに日本語をトークナイズさせてみると以下のようになります。

tokenizer = AutoTokenizer.from_pretrained("meta-llama/Llama-2-7b-hf")
print(tokenizer.tokenize("トークナイザー構築のナレッジとチームの取り組み紹介"))
# 出力
['▁', 'ト', 'ー', 'ク', 'ナ', 'イ', 'ザ', 'ー', '<0xE6>', '<0xA7>', '<0x8B>', '<0xE7>', '<0xAF>', '<0x89>', 'の', 'ナ', 'レ', 'ッ', 'ジ', 'と', 'チ', 'ー', 'ム', 'の', '取', 'り', '<0xE7>', '<0xB5>', '<0x84>', 'み', '<0xE7>', '<0xB4>', '<0xB9>', '介']

これをみると、構築'<0xE6>', '<0xA7>', '<0x8B>', '<0xE7>', '<0xAF>', '<0x89>'というバイト列に変換されていることがわかります。語彙にはといった漢字が含まれていないため、バイト列に変換することで未知語に対応しています。逆にテキストをデコードするときにはこのバイト列を用いることで、構築という漢字の文字列を出力することができます。このように、漢字2文字を6トークンで表現しているため効率は悪いですが、バイトフォールバックを使うことで、未知語を扱うことができます。

split_digits(数字の分割)

こちらはドキュメントに"split all digits (0-9) into separate pieces"と記載があるように、数字を分割するかどうかの設定です。今回はTrueに設定しました。理由としては、数を表現する際に桁数が揃わないことがLLMの性能に影響を及ぼす可能性があるといった知見が見られたためです。

3. 語彙の追加・整備

ここまでで、各言語に対する語彙が獲得できました。実際に中身を見てみると、日本語についてはカタカナ語が多く、一般的に利用する語や表現があまり含まれていないことがわかりました。(上述した反省点の通り、多様なテキストを使うことで防ぐことができたかもしれません。)

そこで、日本語および記号等を中心に手動で語彙を追加することとしました。

  • 一般語
    • wikitionary 「日本語の基本語彙1000」を抽出して追加。
    • wikitinary を参考に名詞・形容詞など品詞ごとによく使うものを定性的に選別して追加。
    • 文化庁 「常用漢字一覧表」の例から一部をサンプリングして追加。
    • 時間・季節・方角に関する語を追加。
    • 都道府県・観光地・東京23区などの土地の名称を追加。
    • 頻出する日本の苗字を追加。
  • 定型表現
    • 「こんにちは」「よろしく」「ございます」などの定型的な表現を追加。
  • 数字および記号
    • 漢数字
    • 半角数字0~9
    • 全角数字0〜9
    • 上付き数字0〜9
    • 数学に出てくるギリシャ文字
    • 「(」や「*」などよく使われる記号
    • 様々な長さの空白の連続(プログラム生成などで役に立ちます。)

記憶が定かではありませんが、合計で数千〜1万ほどの語彙を追加したかと思います。

語彙を追加した後に重複を削除し、スコアの低いトークンを削除することで、任意の語彙数に調整しました。

4. スコアの再推定

各言語で語彙を獲得してマージした後は、Unigramによって各トークンに対してスコア付けを行います。詳細はスクリプトをご参照いただければと思いますが、語彙を設定してテキストからスコアを再学習することでスコア推定ができます。ここでは各言語のテキストを全て使いました。

5. 語彙からトークナイザー構築

上のステップまでで、語彙と各トークンに対するスコア付与ができました。これを使ってトークナイザーのモデルファイルを作る必要がありますが、こちらはllm-jp-tokenizerこちらのスクリプトを使用させていただきました。

6. HuggingFaceのトークナイザー型へ変換

松尾研LLM開発標準コードにあったように、HuggingFaceのクラスに適用するトークナイザーとする必要があります。標準コードではT5Tokenizerのクラスが使われていました。しかし、私たちのチームではdMoEというMoEアーキテクチャのLLMを構築する方針としており、Mixtralベースのモデルを作るという流れになっていました。ここで、本家mistralai/Mixtral-8x7B-v0.1LlamaTokenizer形式を使用していることがわかったため、チームにおいてもT5TokenizerではなくLlamaTokenizerでトークナイザーを作ることとしました

実装はシンプルで、標準コードのT5TokenizerLlamaTokenizerにするだけでしたが、標準コードで1点注意が必要なことがありました。それは、クラスの設定の一つであるsplit_special_tokensTrueになっていたことです。こちらは特殊トークンを分割するかどうかの設定ですが、Trueの場合、特殊トークンが分割されてしまうため、文頭トークンや文末トークンが意図していた一つのトークンIDでは扱われないことになります。

tokenizer.split_special_tokens=True
tokenizer.tokenize("これはトークナイザーのテストです。 </s>")

tokenizer.split_special_tokens=False
tokenizer.tokenize("これはトークナイザーのテストです。 </s>")
# 出力
['▁', 'これは', 'ト', 'ーク', 'ナ', 'イ', 'ザー', 'の', 'テスト', 'です', '。', '▁</', 's', '>']
['▁', 'これは', 'ト', 'ーク', 'ナ', 'イ', 'ザー', 'の', 'テスト', 'です', '。', '▁', '</s>']

上の文章をみると、文末トークン</s>が分割されてしまっていることがわかります。本来1トークンで処理したいのに対して3トークン使うことになってしまう上に、文末としての意味を持たせるという目的も達成できない恐れがあるため、注意が必要です。

以上が、トークナイザーの構築手順でした。

2-3. 苦労した点

事前・事後学習ライブラリとの整合性の調査

上でT5TokenizerLlamaTokenizerを出しましたが、このようなトークナイザーのクラスごとの大きな違いの一つとして「特殊トークンの設定」があります。それは文頭トークンが<s>であるかといった文字列そのものの違いもありますが、これはさほど大きな問題ではありません。重要なのは「テキストをエンコードしたときにデフォルトで設定される特殊トークンは何か」という点です。

例えば、meta-llama/Llama-2-7b-hfllm-jp/llm-jp-13b-v1.0のトークナイザーの違いを見てみます。

from transformers import AutoTokenizer

tokenizer_A = AutoTokenizer.from_pretrained("meta-llama/Llama-2-7b-hf")
tokenizer_B = AutoTokenizer.from_pretrained("llm-jp/llm-jp-13b-v1.0")

print("="*30, "Llama 2", "="*30, )
print(tokenizer_A.tokenize("This is a pen.", add_special_tokens=True))
print(tokenizer_A.encode("This is a pen.", add_special_tokens=True))

print("="*30, "LLM-jp", "="*30, )
print(tokenizer_B.tokenize("This is a pen.", add_special_tokens=True))
print(tokenizer_B.encode("This is a pen.", add_special_tokens=True))
# 出力
============================== Llama 2 ==============================
['<s>', '▁This', '▁is', '▁a', '▁pen', '.']
[1, 910, 338, 263, 6584, 29889]
============================== LLM-jp ==============================
['▁This', '▁is', '▁a', '▁', 'pen', '.', '<EOD|LLM-jp>']
[406, 316, 311, 31, 7605, 10013, 7]

これをみると、同じコードを実行していますが、meta-llama/Llama-2-7b-hfの方は文頭トークンである<s>が付与されているのに対して、llm-jp/llm-jp-13b-v1.0の方は文末トークンである<EOD|LLM-jp>が付与されていることがわかります。トークナイザーの設定はtokenizer_config.jsontokenizer_classから確認できますが、これを見るとそれぞれLlamaTokenizerPreTrainedTokenizerFastというクラスが使われていることがわかります。

したがって、使われているトークナイザーのクラスによって、エンコード時の特殊トークンの付与の設定は異なるということです。

これがなぜ問題であり苦労につながったかというと、このような特殊トークン付与の設定は、事前学習や事後学習のライブラリにおけるスクリプトと深く関係するためです

例えば、私たちのチームで最終的に使用した学習ライブラリであるLlaMA2-Accessoryでは、エンコード処理が以下のように定義されていました。

def encode(self, s: str, bos: bool, eos: bool) -> List[int]:
        assert type(s) is str
        if self.tokenizer_type == "transformers":
            t = self.tokenizer.encode(s, truncation=False, add_special_tokens=False)
        else:
            t = self.tokenizer.encode(s)
        if bos:
            t = [self.bos_id] + t
        if eos:
            t = t + [self.eos_id]
        return t

こちらを見ると、tokenizer.encode(s, truncation=False, add_special_tokens=False)となっていることから、エンコード時に特殊トークンは付与されないことがわかります。ただし、encodeメソッドにboseosの引数を指定することができ、それぞれTrueの場合に付与されるように書かれています。実際にこれらの引数がどのように設定されて実行されているかを確認するためには、学習スクリプトを見に行く必要がありますが、結論としてはinput_data = torch.tensor(self.tokenizer.encode(ann, bos=True, eos=True), dtype=torch.int64)として、テキストに対して文頭トークンと文末トークンを両方とも付与して処理した後に事前学習に進んでいることがわかりました

このような処理を確認しない場合、<s><s>というように意図せず連続で付与していたり、そもそも特殊トークンが付与されていなかったりといった事象が発生する恐れがあります。 私たちのチームでは事前学習の本番が開始するまで、いくつかのライブラリで検討をしていたため、全てに対してのトークナイザーの処理を読み解きに行く必要があり、大変苦労しました。

同じく事後学習ライブラリについてもいくつか検討していたため同じような苦労がありました。事後学習では、入力トークンをマスクして出力部分のみを損失計算するといった処理をするのが一般的であるため、その区切り位置での文末トークンの処理がどうなっているかといった点の読み解きも必要でした。

何が最適かわからないという難しさ

苦労したこととは若干異なりますが、精神的にきつかったこととしては、「自分の作っているトークナイザーが本当にLLMの性能向上に寄与しているのか」という疑問が最後まで拭えなかったことです。もちろん事前学習を何度も行ってトークナイザーの違いによる精度を確かめるということはできないため、「先行研究や他の事例を参考にすると、おそらくこうするのがよさそう」という程度の根拠で進んでいたことが、個人的にはメンタルが削られる部分でした。

最近ではScaling Laws with Vocabulary: Larger Models Deserve Larger Vocabularies (2024)で最適な語彙サイズについて研究されているように、LLMの性能に寄与するトークナイザーのベストプラクティスは研究され続けていくと思いますので、引き続き学んでいこうと思います。

2-4. その他 実装を通して学んだこと

Chat Template

HuggingFaceのトークナイザーにはChat templateというものが設定できます。これは、指示形式や対話形式でファインチューニングする際の形式の統一や、その形式で推論するためにプロンプトに指定のテンプレートを自動適用するための機能だと思われます。

例えば、mistralai/Mistral-7B-Instruct-v0.1では以下のように動作させることができます。

from transformers import AutoTokenizer
tokenizer = AutoTokenizer.from_pretrained("mistralai/Mistral-7B-Instruct-v0.1")

chat = [
  {"role": "user", "content": "Hello, how are you?"},
  {"role": "assistant", "content": "I'm doing great. How can I help you today?"},
  {"role": "user", "content": "I'd like to show off how chat templating works!"},
]

print(tokenizer.apply_chat_template(chat, tokenize=False))
# 出力
<s>[INST] Hello, how are you? [/INST]I'm doing great. How can I help you today?</s> [INST] I'd like to show off how chat templating works! [/INST]

この例では、chatとして辞書が入ったリストを作っています。rolecontentを与えてapply_chat_templateメソッドを使うと、特殊トークンを付与した状態で対話形式のようにテキストを変換してくれます。

このテンプレートはtokenizer_config.jsonで設定することができます。実際にjsonファイルを見てみると、以下のように定義されています。

"chat_template": "{%- if messages[0]['role'] == 'system' %}\n    {%- set system_message = messages[0]['content'] %}\n    {%- set loop_messages = messages[1:] %}\n{%- else %}\n    {%- set loop_messages = messages %}\n{%- endif %}\n\n{{- bos_token }}\n{%- for message in loop_messages %}\n    {%- if (message['role'] == 'user') != (loop.index0 % 2 == 0) %}\n        {{- raise_exception('After the optional system message, conversation roles must alternate user/assistant/user/assistant/...') }}\n    {%- endif %}\n    {%- if message['role'] == 'user' %}\n        {%- if loop.first and system_message is defined %}\n            {{- ' [INST] ' + system_message + '\\n\\n' + message['content'] + ' [/INST]' }}\n        {%- else %}\n            {{- ' [INST] ' + message['content'] + ' [/INST]' }}\n        {%- endif %}\n    {%- elif message['role'] == 'assistant' %}\n        {{- ' ' + message['content'] + eos_token}}\n    {%- else %}\n        {{- raise_exception('Only user and assistant roles are supported, with the exception of an initial optional system message!') }}\n    {%- endif %}\n{%- endfor %}\n",

見慣れない書き方だと思いますが、こちらはJinjaと呼ばれる記法で書かれているようです。(詳細はこちらのドキュメントをご参照ください。)

実際に任意のテンプレートを作ってみます。
ここでは、よくインストラクションチューニングで用いられるalpaca形式の日本語版をテンプレとして作ってみたいと思います。

イメージとしては以下のような形式を作れればOKです。

以下は、タスクを説明する指示です。要求を適切に満たす応答を書きなさい。

### 指示:
{text}

### 入力:
{text}

### 応答:

これをjinjaテンプレートで作ると、以下のように書けます。これをtokenizer_config.jsonに追記します。(この記法への変換はChatGPTなどを使うと楽です。)

"chat_template":"{% for message in messages %}{% if message['role'] == 'system' %}{{ message['content'] + '\n\n'}}{% elif message['role'] == 'instruction' %}{{ '### 指示:\n' + message['content'] + '\n\n'}}{% elif message['role'] == 'input' %}{{ '### 入力:\n' + message['content'] + '\n\n'}}{% endif %}{% endfor %}### 応答:\n\n"

テンプレの挙動確認をしてみます。

from transformers import AutoTokenizer
tokenizer = AutoTokenizer.from_pretrained("you/your_tokenizer_name")

messages = [
    {"role": "system", "content": "以下は、タスクを説明する指示です。要求を適切に満たす応答を書きなさい。"},
    {"role": "instruction", "content": "GENIACプロジェクトの一つである東京大学松尾研究室のLLM開発プロジェクトについて説明してください。"},
    {"role": "input", "content": "GENIACとは経済産業省による日本の生成AI開発力向上のための取り組みです。"},
]

print(tokenizer.apply_chat_template(messages, tokenize=False))
以下は、タスクを説明する指示です。要求を適切に満たす応答を書きなさい。

### 指示:
GENIACプロジェクトの一つである東京大学松尾研究室のLLM開発プロジェクトについて説明してください。

### 入力:
GENIACとは経済産業省による日本の生成AI開発力向上のための取り組みです。

### 応答:

テンプレート通りにテキストが作られることが確認できました。

基本的なトークナイザーのメソッド

最後に、トークナイザーの開発を通して、HuggingFaceのtokenizerクラスにおいてよく使うメソッドを覚えることができましたのでそちらを紹介します。(LlamaTokenizerT5Tokenizerなどクラスによって使えるメソッドが異なる場合があるので注意してください。)

以下はelyza/ELYZA-japanese-Llama-2-7b-fastのLlamaTokenizerを使った例となっています。

tokenizer.tokenize

こちらは本記事でも使っていた通り、テキストをトークン単位に分割するメソッドです。
add_special_tokens=Trueとすることで特殊トークンが付与されます。

tokenizer.tokenize("これはトークナイザーのテストです。")
tokenizer.tokenize("これはトークナイザーのテストです。", add_special_tokens=True)
# 出力
['▁', 'これは', 'ト', 'ーク', 'ナ', 'イ', 'ザー', 'の', 'テスト', 'です', '。']
['<s>', '▁', 'これは', 'ト', 'ーク', 'ナ', 'イ', 'ザー', 'の', 'テスト', 'です', '。']
tokenizer.encode

こちらも本記事で使っていた通り、テキストをトークン単位に分割してID変換するメソッドです。
tokenizeメソッドと異なる点としては、add_special_tokens=Trueがデフォルトとなっています。そのため、add_special_tokens=Falseとすることで特殊トークンなしでのエンコードがされます。以下の例では文頭トークンの<s>がIDが1なので、それが消えていることがわかります。さらにreturn_tensors="pt"と設定することで、テンソル形式で返すことができます。

tokenizer.encode("これはトークナイザーのテストです。")
tokenizer.encode("これはトークナイザーのテストです。", add_special_tokens=False)
tokenizer.encode("これはトークナイザーのテストです。", return_tensors="pt")
# 出力
[1, 29871, 43061, 30279, 43004, 30576, 30260, 43349, 30199, 44698, 42669, 30267]
[29871, 43061, 30279, 43004, 30576, 30260, 43349, 30199, 44698, 42669, 30267]
tensor([[    1, 29871, 43061, 30279, 43004, 30576, 30260, 43349, 30199, 44698, 42669, 30267]])
tokenizer.encode_plus

こちらはテキストをエンコードしたものを辞書型で返すことができ、input_idsattention_maskがキーとなります。実はtokenizer()でそのままテキストを渡すこともできますが、そちらの出力とencode_plusの出力は同じものとなりました。

tokenizer.encode_plus("これはトークナイザーのテストです。")
tokenizer("これはトークナイザーのテストです。")
# 出力
{'input_ids': [1, 29871, 43061, 30279, 43004, 30576, 30260, 43349, 30199, 44698, 42669, 30267], 'attention_mask': [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1]}
{'input_ids': [1, 29871, 43061, 30279, 43004, 30576, 30260, 43349, 30199, 44698, 42669, 30267], 'attention_mask': [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1]}
tokenizer.batch_encode_plus

こちらは複数のテキストをまとめてエンコードすることができます。

tokenizer.batch_encode_plus([
    "これはトークナイザーのテストです。", 
    "これはトークナイザーのテストではありません。"
], padding='max_length', max_length=16, return_tensors='pt')
# 出力
{'input_ids': tensor([[    1, 29871, 43061, 30279, 43004, 30576, 30260, 43349, 30199, 44698,
         42669, 30267,     0,     0,     0,     0],
        [    1, 29871, 43061, 30279, 43004, 30576, 30260, 43349, 30199, 44698,
         43418, 30267,     0,     0,     0,     0]]), 'attention_mask': tensor([[1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 0, 0, 0, 0],
        [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 0, 0, 0, 0]])}

こちらをみるとmax_lengthで長さを揃えて、長さを満たさない部分は0でパディングされていることがわかります。また、パディング部分は学習などの処理に使わないのが一般的なので、attention_mask0になっています。

逆にトークンを指定の長さで切り詰めたい場合はtruncation=Trueと設定します。切り詰めの方向はtokenizer.truncation_sideで設定されているため、rightleftで変更できます。以下の例では最大8トークンになるように左側のトークンから削除されています。

tokenizer.truncation_side = "left"
tokenizer.batch_encode_plus([
    "これはトークナイザーのテストです。", 
    "これはトークナイザーのテストではありません。"
], padding='max_length', max_length=8, return_tensors='pt', truncation=True)
# 出力
{'input_ids': tensor([[    1, 30576, 30260, 43349, 30199, 44698, 42669, 30267],
        [    1, 30576, 30260, 43349, 30199, 44698, 43418, 30267]]), 'attention_mask': tensor([[1, 1, 1, 1, 1, 1, 1, 1],
        [1, 1, 1, 1, 1, 1, 1, 1]])}
tokenizer.decode

こちらはエンコードと逆の処理をするデコードです。トークンIDのリストを渡すことで文字列に戻すことができます。エンコードと同じように特殊トークン付与の設定がありますが、引数はエンコードで使ったadd_special_tokensではなくskip_special_tokens=Trueとすることで、特殊トークンを無視してデコードします。

tokenizer.decode([1, 29871, 43061, 30279, 43004, 30576, 30260, 43349, 30199, 44698, 42669, 30267])
tokenizer.decode([1, 29871, 43061, 30279, 43004, 30576, 30260, 43349, 30199, 44698, 42669, 30267], skip_special_tokens=True)
# 出力
"<s> これはトークナイザーのテストです。"
"これはトークナイザーのテストです。"

おわりに

本記事ではGENIAC 松尾研LLM開発プロジェクトのPhase1で取り組んだトークナイザーの構築方法と、その過程で得た学びについて紹介しました。LLM開発というと事前学習が本丸であり最も重要な部分であるのは間違いないですが、今回の取り組みを通じて、トークナイザーは事前学習・事後学習・推論のほぼ全てに関連する大事な要素であることがわかりました。
また、私自身はこのプロジェクトに参加するまでLLM開発やトークナイザー構築には直接携わったことがなかったため、非常に多くの学びがありました。
改めて、経済産業省, 東京大学松尾・岩澤研究室, 開発メンバー, コミュニティメンバーを含む関係者の皆様に、この場をお借りして感謝申し上げたいと思います。ありがとうございました。

東大松尾・岩澤研究室 | LLM開発 プロジェクト[GENIAC]

Discussion