LLMの中身を理解するため、tiny-jp-llmという小さな日本語LLMを一つずつ作っています。

最初に作ったのは、文章を細かく分けてトークンIDへ変換するByte-level BPEトークナイザーです。次に、直前のトークンから次のトークンを選ぶBigramモデルを作りました。

コードと実行方法は、GitHubのNXSW/tiny-jp-llmへ置いています。

前回のBigramモデルでは、Next Developers Blogにある388件の記事タイトルを使い、タイトルらしき文字列を生成しました。

実際に出てきたのがこちらです。

ChatGPTがかしてみた

何となく技術ブログらしい単語は入っていますが、意味はよくわかりません。

前回のBigramモデルが覚えたのは、隣り合うトークンが何回出てきたかという回数だけです。トークンの意味を見て選んだ結果ではありません。

そこで次は、LLMの構成を調べる中で見かけたEmbeddingを作ることにしました。

Embeddingについて調べると、意味が近い言葉は近い数値になる、という説明をよく見かけます。最初は、Embeddingへ通せば言葉の意味が数字になるのだと思っていました。

LLMが文章を返すまで

調べた範囲では、LLMが文章を返すまでの流れは、ざっくり次のようになるようです。

LLMが文章を返すプロセス

文章
  ↓ トークナイザーでトークンに分割し、それぞれをトークンIDへ置き換える
トークンIDの並び
  ↓ 【今回】Embeddingで各トークンIDに対応するベクトルを取り出す
トークンごとのベクトル
  ↓ Transformerで、それまでのトークンの情報を反映する
計算後のベクトル
  ↓ 出力層で、次の候補となる各トークンに点数を付ける
トークンごとの点数
  ↓ 点数をもとに1つ選ぶ
次のトークンID
  ↓ 選んだトークンIDを末尾に加え、次のトークンを再び予測する
トークンIDの並び
  ↓ トークナイザーでトークンIDを文字へ戻してつなげる
文章

今回は、この中にあるEmbeddingを作ります。

ただし、まだTransformerは作りません。Embeddingから次のトークンを予測するため、今回は小さな出力層も一緒に用意します。

Embeddingとは?

前回は、388件の記事タイトルでByte-level BPEの結合ルールを作りtokenizer-800.jsonへ保存しました。

このJSONには、どのトークン同士を結合するかが記録されています。読み込むと800種類のトークンとトークンIDの対応が再現されます。この対応をまとめたものが語彙ということのようです。

実際にAIを変換した結果がこちらです。

{
  "text": "AI",
  "utf8_bytes": 2,
  "token_count": 1,
  "token_ids": [379],
  "pieces": ["AI"],
  "decoded": "AI"
}

piecesには分割後のトークン、token_ids`には対応するトークンIDが同じ順番で入っています。

今回の語彙では、AIというトークンに対応するトークンIDは379です。

Embeddingは、このトークンIDを受け取ります。ここで、なぜトークンIDを使うのかわからなくなりました。

調べてみると、Embeddingはトークンごとに複数の数値を用意する仕組みのようです。どのトークンの数値を取り出すのか指定するために、トークンIDを使います。

AIのトークンIDが379なら、先頭を0行目として、Embeddingの379行目を取り出します。

トークンID 379
  ↓ Embeddingの379行目を取り出す
[AIに対応する複数の数値]

PyTorchが379を見て、AIだと判断しているわけではありません。トークナイザーが出したトークンIDを、そのままEmbeddingの行に対応させています。

複数の数値を並べたものを、ここではベクトルと呼ぶようです。

つまり、トークンIDをベクトルへ変換するという説明は、トークンIDに対応する1行分の数値を取り出す、という意味でした。

文章だけではまだ自信がないので、まずは小さなEmbeddingを動かして確かめます。

実行の準備

今回のプログラムではPyTorchを使います。

自分の環境はApple SiliconのMacとPython 3.12.9です。

まだリポジトリを取得していない場合は、次のコマンドを実行します。

git clone https://github.com/NXSW/tiny-jp-llm.git
cd tiny-jp-llm

このプロジェクト専用のPython環境を作り、PyTorchを含む必要なライブラリをインストールします。

python3 -m venv .venv
source .venv/bin/activate
python3 -m pip install -e .

source .venv/bin/activateを実行すると、ターミナルの行頭に(.venv)と表示されます。以降のコマンドは、この表示がある状態で実行します。

まずEmbeddingだけを動かす

PyTorchにはtorch.nn.Embeddingが用意されています。

いきなり800種類のトークンを扱うと結果を追いにくいため、まずは4種類だけを扱う確認用のEmbeddingを作りました。

該当ファイルはscripts/inspect_embedding.pyです。

torch.manual_seed(42)

embedding = nn.Embedding(
    num_embeddings=4,
    embedding_dim=3,
)

token_ids = torch.tensor([0, 2])

num_embeddings=4は、トークンID0から3までの4種類を用意する設定です。

embedding_dim=3は、1種類につき3個の数値を持たせる設定です。つまり、4行3列の2次元配列ができます。

この確認用のトークンIDには、AIなどの実際のトークンを対応させていません。今回は、トークンIDと同じ行が返るかだけを確認します。

結果を比べやすいように、4行3列の配列全体と、トークンID02を渡した結果をJSONで表示します。

python3 scripts/inspect_embedding.py

実行結果はこちらです。

{
  "token_ids": [0, 2],
  "embedding_table": [
    [0.3367, 0.1288, 0.2345],
    [0.2303, -1.1229, -0.1863],
    [2.2082, -0.6380, 0.4617],
    [0.2674, 0.5349, 0.8094]
  ],
  "output": [
    [0.3367, 0.1288, 0.2345],
    [2.2082, -0.6380, 0.4617]
  ]
}

embedding_tableが4行3列の配列全体です。

outputには、そのうち0行目と2行目が入っています。トークンID02を使って何かを計算したのではなく、同じ位置の行を取り出していました。

1行には3個の数値があります。この1行が、1つのトークンに対応する3次元のベクトルです。

なお、この時点の数値はPyTorchが最初に入れた値です。まだ記事タイトルを使った予測は行っていません。

ここまで動かして、Embeddingからベクトルを取り出す処理は理解できました。

なぜトークンIDのままでは駄目なのか

次にわからなかったのは、なぜトークンIDとは別に、複数の数値を持つのかです。

今回の語彙では、AIのトークンIDは379AWSのトークンIDは285です。

この2つは、各トークンを区別するために割り当てられています。379285の差を計算しても、AIAWSの関係を表すことにはなりません。

さらに、トークンIDそのものを変更すると、トークナイザーとの対応が崩れてしまいます。

一方、Embeddingのベクトルは、トークンIDとの対応を保ったまま、中にある数値だけを変更できます。

ただし、Embeddingを作っただけでは数値は変わりません。どのような結果に近づけるのか、練習問題が必要です。

そこで今回は、直前のトークンから次のトークンを予測する問題を使い、予測を外すたびにEmbeddingの数値を変更してみます。

今回試すこと

学習に使うのは、前回と同じ388件の記事タイトルです。

前回のBigramモデルは、直前のトークンの次に何が何回出たかを数え、その回数から次のトークンを選びました。

今回は見る範囲を同じく直前の1トークンに限定し、回数の代わりにEmbeddingと小さな出力層を使います。

現在のトークンID
  ↓ Embeddingから対応するベクトルを取り出す
32個の数値
  ↓ 出力層で次の候補に点数を付ける
候補ごとの点数
  ↓ 実際に続いたトークンIDと比べる
予測がどのくらい外れたか
  ↓ 外れ方が小さくなるよう数値を変更する
Embeddingと出力層を更新

この実装は、まだTransformerではありません。まずはEmbeddingの数値が予測結果を受けて変わるところまでを確認します。

記事タイトルから予測問題を作る

記事タイトルが次の3トークンに分かれた場合を考えました。

AWS → で → 試してみた

現在のトークンから次のトークンを当てる組にすると、次のようになります。

文頭       → AWS
AWS        → で
で         → 試してみた
試してみた → 文末

最初のAWSを予測するときには、直前のトークンがありません。また、試してみたの次で文章が終わることも予測させたいです。

そこで、文章の始まりを示す文頭と、終わりを示す文末を追加しました。

実装では、通常のトークンID0から799に続けて、文頭800文末801を割り当てています。これは今回の予測モデル内だけで使うトークンIDです。

current = [bos_token_id, *token_ids]
following = [*token_ids, eos_token_id]

inputs.extend(current)
targets.extend(following)

388件の記事タイトルから、10,102組の予測問題ができました。

Embeddingの配列を作る

最初の確認ではnn.Embeddingをそのまま使いました。

実際の予測モデルでは、トークンIDで行を選んでいることがコードからもわかるよう、PyTorchの2次元配列を直接用意しました。

PyTorchでは、数値を並べたこのようなデータをTensorと呼ぶようです。今回は行と列があるため、2次元のTensorです。

該当ファイルはsrc/tiny_jp_llm/embedding.pyです。

self.embedding_table = nn.Parameter(
    torch.randn(self.total_tokens, embedding_dim) * 0.02
)

扱うのは、通常の800種類に文頭文末を加えた802種類です。1種類につき1行を使います。

各行に置く数値は32個にしました。したがって、Embedding全体は802行32列です。

Embedding全体:802行 × 32列

1行分:1トークンに対応する32次元のベクトル

802次元のベクトルではありません。32次元のベクトルが802行あります。

ベクトルの数値は8個や16個でも作れます。今回はMacで短時間に試せる小ささを優先し、32個にしました。最適な個数の比較は行っていません。

nn.Parameterで用意すると、PyTorchが変更対象として扱うようです。

トークンIDに対応する行を取り出す部分はこちらです。

vectors = self.embedding_table[token_ids]

ベクトルから次のトークンへ点数を付ける

Embeddingから32個の数値を取り出しただけでは、次のトークンを選べません。

そこで、32個の数値を受け取り、次の候補802種類それぞれに点数を付ける出力層を追加しました。

return vectors @ self.output_weight + self.output_bias

この計算結果には、候補ごとの点数が802個入ります。

記事全体の流れで示した出力層に当たる部分です。今回はTransformerを挟まず、Embeddingの直後に出力層を置いています。

予測の外れ方を一つの数値にする

次は、802個の点数と、実際に続いたトークンIDを比べます。

ただ、どのくらい外れたのかをPyTorchへどう伝えるのかわかりませんでした。

調べてみると、複数の候補から正解を選ぶ問題では、Cross Entropyという計算が使われるようです。今回はPyTorchのcross_entropyを使いました。

logits = model(inputs)
loss = F.cross_entropy(logits, targets)

logitsには、次の候補802種類の点数が入っています。targetsには、実際に続いたトークンIDが入っています。

cross_entropyが返すlossは、予測の外れ方をまとめた値として使います。正解率ではありません。

今回は、学習前よりlossが下がるかを確認します。

Embeddingの数値を更新する

lossは計算できましたが、Embeddingのどの数値を、どちらへ動かせばよいのかはわかりません。

ここではPyTorchの自動微分を使いました。今回の計算を逆向きにたどり、lossを下げるには各数値をどちらへ動かせばよいかを求める仕組みのようです。

for step in range(1, steps + 1):
    logits = model(inputs)
    loss = F.cross_entropy(logits, targets)

    optimizer.zero_grad()
    loss.backward()
    optimizer.step()

backward()で動かす方向を計算し、step()でEmbeddingと出力層の数値を更新します。

数値の更新には、PyTorchのAdamWを使いました。今回は更新方法の比較までは行わず、用意されている方法を一つ使って先へ進みます。

記事タイトルを使って実行する

取得した記事タイトルと学習結果は、GitHubのリポジトリへ含めていません。

初めて実行する場合は、記事タイトルを取得し、語彙数800のBPEトークナイザーを作ります。

tiny-jp-llm corpus fetch-next-titles

tiny-jp-llm tokenizer train \
  --input data/next_titles.txt \
  --model artifacts/tokenizer-800.json \
  --vocab-size 800

続けて、記事と同じ条件で予測モデルを学習します。

export PYTHONPATH="$PWD/src"
python3 scripts/run_embedding_experiment.py

今回はCPUで実行しました。

結果は次の2ファイルへ保存されます。

artifacts/
├── embedding-bigram.pt
└── embedding-summary.json

Embeddingの数値は変わったのか

32次元のベクトルを持つEmbeddingを500回更新しました。

{
  "documents": 388,
  "vocab_size": 800,
  "embedding_dim": 32,
  "training_pairs": 10102,
  "steps": 500,
  "initial_loss": 6.687054,
  "final_loss": 2.434602,
  "training_seconds": 9.807,
  "mean_embedding_change": 4.73714
}

loss6.687054から2.434602まで下がりました。

step   0: 6.687054
step  50: 2.512620
step 100: 2.440894
step 200: 2.435926
step 300: 2.435008
step 500: 2.434602

最初の100回で大きく下がり、その後はほとんど変わっていません。

また、802行それぞれについて、32個の数値が学習前からどのくらい移動したかを距離として計算しました。その平均は4.73714でした。

これは正解率や品質を表す点数ではありません。0ではなかったため、Embeddingの数値が実際に変わったことを確認するために使いました。

これで、Embeddingの数値が、次トークン予測の結果を受けて変わったことを確認できました。

ただし、このlossは学習に使った10,102組を、もう一度予測した結果です。まだ見せていない記事タイトルへの予測性能を測った数字ではありません。

似たトークンが近くなるのか

Embeddingについて調べたときに見かけた、意味が近い言葉は近い数値になる、という説明も確かめてみます。

32個の数値をそのまま見比べるのは難しいため、今回はコサイン類似度を使いました。調べてみると、ベクトルの向きが近いほど1に近づく値のようです。今回は、この値が大きかったトークンから順に並べます。

してみたに近かったトークンはこちらです。

順位トークンコサイン類似度
1てみた0.748659
2調べてみた0.710783
3!0.681862
4]0.652217
5する方法0.635683

てみた調べてみたが上位に来ました。これは少しそれらしく見えます。

ほかのトークンも確認しました。

基準近かったトークン
AIEC2ール_B
AWS開発んなuxれる
生成de
Amazon×/Claude Code

こちらは、意味が近い言葉ばかりとは言えませんでした。

考えてみると、今回与えた問題は、直前の1トークンから次を当てることだけです。記事タイトルも388件しかなく、文章全体の意味は見ていません。

Embeddingを使えば、自動的に言葉の意味が入るわけではなさそうです。どのような問題を与え、どのようなデータで数値を更新したかによって、ベクトルが表すものも変わるようです。

今回できたのは、言葉の意味を表す万能なベクトルではなく、10,102組の次トークン予測に合わせて更新されたベクトル、と考えるのがよさそうです。

テストする

今回追加したテストでは、次の動きを確認しました。

  • 指定したトークンIDに対応する行を取得できる
  • 文頭と文末を含む予測問題を作れる
  • 更新後にlossが下がる
  • Embeddingの数値が学習前から変わる
  • 保存したモデルを読み直せる
  • CLIから学習と近いトークンの確認ができる

前回までのテストと合わせて実行します。

python3 -m unittest discover -s tests -v

16件とも通りました。

Ran 16 tests in 1.096s

OK

作ってみて

前回のBigramは言葉の回数表でしたが、今回の embedding.py から初めて 学習で数値が変わるニューラルネットワーク が登場しました。

実際に小さく動かしてみると、最初に行っていたのは、トークンIDに対応するベクトルを取り出す処理でした。そのベクトルは最初から意味を持つわけではなく、与えた予測問題に合わせて数値が変わります。

今回は直前の1トークンしか見ていないため、意味が近いトークンばかりが近くに並ぶ結果にはなりませんでした。

次は、AttentionとCausal Maskを調べます。

今のモデルは、それより前に出てきたトークンを見ていません。複数のトークンを使って次を予測する処理が、どのような計算になるのか、小さく動かすところから始めます。

参考

ここまで読んでいただき、ありがとうございます。もしこの記事の技術や考え方に少しでも興味を持っていただけたら、ネクストのエンジニアと気軽に話してみませんか。

  • 選考ではありません
  • 履歴書不要
  • 技術の話が中心
  • 所要時間30分程度
  • オンラインOK

エンジニアと話してみる

関連リンク

AI・クラウド・データ分析のご相談はネクスト株式会社までお問い合わせください。