LLM はトークンを1つずつ予測している — 自己回帰ループを手元で覗く

ChatGPT の返事が、左から少しずつ流れるように出てくるのを見たことがあると思います。あれは演出ではありません。モデルが本当に、文章をトークンという小さなかたまりに分け、それを 1 つずつ予測して繋げているからです。

この連載は、富士通が発表した新アーキテクチャ PHOTON¹（Transformer 比で GPU あたり最大 475 倍というニュースで話題になりました）を、最終回で手元の Mac に擬似再現することを目的にしています。ただ、その「475 倍」を魔法ではなく理解できるトレードオフとして受け取るには、まず言語モデルの中身を知っておく必要があります。Part 1 の今回は、いちばん土台にある「LLM は次のトークンを予測しているだけ」という事実を、自分のコードで覗いていきます。

この連載の現在地 — LLM は「次トークン予測器」

連載の各回は、同じかたちの「現在地」ボックスから始めます。どこまで分かっていて、今回どこを進めるのかを見失わないためです。

• ここまで分かっていること — まだありません。ここから始めます。
• 今回新しく説明できること — トークン化のしくみ、自己回帰ループ、そして生成が逐次である理由
• 今回はやらないこと — Attention の中身は Part 2、KV キャッシュのメモリ代償は Part 3、サンプリング戦略の比較は Part 4 へ送ります
• PHOTON のどの主張につながるか — 「生成は 1 トークンずつ・逐次」という性質が、最終回の「decode が重い場面で速くする」という話の出発点になります

言語モデルの本体は、ものすごく単純化すると「次に来るトークンの確率分布を返す関数」です²。文章を生成するというのは、その関数を繰り返し呼び、出てきたトークンを末尾に足し、また呼ぶ、というループにすぎません。難しさは関数の中身（次の Part 以降）にありますが、外側のループはとても素直です。

トークンは文字でも単語でもない

最初につまずきやすいのが「トークン」という単位です。トークンは文字でもなければ単語でもありません。多くの LLM は Byte-Pair Encoding（BPE）という方法で、よく一緒に出てくる文字の並びを 1 つの記号にまとめた語彙を持っています²。結果として、語彙にはありふれた単語、単語の断片、語幹などが混ざります。マージの手順そのものを 1 から実装して追うこともできます³。

面白いのは、同じ単語でもモデルによって切れ方が違うところです。strawberry は、ある語彙では straw というかたまりが存在せず 3 片に割れますが、別のモデルの語彙には straw があるので 2 片で済みます⁴。LLM が「strawberry に r はいくつ？」のような問いを苦手にしがちなのは、こうしてよくある単語をかたまりで覚えていて、個々の文字をそもそも見ていないことが多いからです⁴。

これは説明されるより、自分で叩いてみるほうが腹に落ちます。companion の token_loop.py --demo-tokenize を Llama 3.2 1B で走らせると、いくつかの文字列がこう分解されました。

# tokenize-demo（Llama-3.2-1B-Instruct-4bit / 特殊トークンは除外）
  'strawberry'   -> 3 tokens: ['str', 'aw', 'berry']
  '苺がすきです'  -> 6 tokens: ['�', '�', 'が', 'す', 'き', 'です']
  '    spaces'   -> 2 tokens: ['   ', ' spaces']
  '1234567890'   -> 4 tokens: ['123', '456', '789', '0']
  'tokenization' -> 2 tokens: ['token', 'ization']

strawberry は str / aw / berry の 3 トークンで、r という文字はどこにも単独で現れません。モデルが「r の数」を直接は見ていないことが、ここから分かります。苺 はこの語彙に無いので 2 個の文字化け片（�）に割れています。これは UTF-8 のバイトに落として処理する byte fallback の様子で、漢字 1 文字が 2 トークンを食っています。数字は 123 456 789 0 と 3 桁ごとに切られ、4 つの連続空白は「3 個ぶんの空白」と「空白 + spaces」に畳まれました。

トークン数は文字数とも単語数とも一致しません。この事実だけ持ち帰れば十分です。これが後で「料金はトークン課金」「文脈長はトークンで数える」という話に効いてきます。

自己回帰ループを最小実装する

生成のループそのものは、拍子抜けするほど短く書けます。流れは「文字列をトークンに直す → モデルに通して次トークンの分布を得る → そこから 1 つ選ぶ → 末尾に足す → 繰り返す」です。

MLX の mlx_lm には、この 1 ステップを回す低レベルの API があります⁵。generate_step はジェネレータで、呼ぶたびに（選ばれたトークン, 全語彙の対数確率）を返します。これを使うと、ループの中身を自分の目で確認できます。

from mlx_lm import load
from mlx_lm.sample_utils import make_sampler
from mlx_lm.generate import generate_step
import mlx.core as mx

model, tokenizer = load("mlx-community/Llama-3.2-1B-Instruct-4bit")
sampler = make_sampler(temp=0.0)            # greedy: 分布の最大を選ぶ
prompt = mx.array(tokenizer.encode("The capital of France is"))

for i, (token, logprobs) in enumerate(generate_step(prompt, model, sampler=sampler)):
    if i >= 32:
        break
    print(tokenizer.decode([int(token)]), end="", flush=True)

ここで大事なのは、generate_step が内部で KV キャッシュを使っている点です。つまり毎ステップでプロンプト全体を計算し直しているのではなく、過去の計算結果を取っておいて、新しい 1 トークンぶんだけを足しています⁶。「なぜそんなことができるのか」「その代償は何か」は Part 3 の主題なので、今回は「cached decode で動いている」という事実だけ押さえておきます。最初のプロンプトをまとめて処理する prefill と、そのあと 1 つずつ吐く decode が別フェーズだ、という区別もここで芽出ししておきます⁷。

生成にかかる時間はトークン数に比例する

自己回帰が「1 つずつ・逐次」だと、素直な予想が立ちます。1 トークンあたりの時間がだいたい一定なら、生成トークン数 N と、かかった累積時間は直線に乗るはずです。これは手元で測れます。

token_loop.py --bench は、最初のトークン（prefill を含むので別物です）を除いた decode 区間で、トークン数ごとの累積秒数を記録します。warmup を 1 回入れてから複数回測り、トークン数ごとに中央値を取ります。出力 results.jsonl には実行環境（モデル・mlx-lm のバージョン）も一緒に残します。

予想どおり、累積時間は生成トークン数にほぼ完璧な直線で乗りました。決定係数 R² は 0.99993、生成速度は約 303 tok/s でした（Apple M5 Pro 48GB・4bit 量子化の Llama 3.2 1B・cached decode・5 回の median）。この直線こそが、後半の話の出発点です。生成を「速くする」とは、結局この 1 トークンの単価を下げる戦いになります。

ひとつ注意があります。これは decode 区間の話で、プロンプトが長い場合は prefill の時間が別に乗ります⁷。また tok/s はモデル・量子化・機種に強く依存し、別の環境にそのまま一般化はできません。無料で動くことと速いことは別です。モデルのダウンロードを含めると、最初の 1 回は数分かかります。

logits を覗く — 次トークンは「分布」

もう 1 つ、generate_step が返してくれるものがあります。各ステップで、選ばれた 1 トークンだけでなく、全語彙ぶんの対数確率です⁸。token_loop.py はここから上位 k 個を取り出して表示します。

# loop  prompt='The capital of France is'（確率はこのモデル・量子化での値）
  step  0  pick=' Paris'    top5=[' Paris':0.70 | '...':0.03 | ':':0.02 | ' what':0.02 | ':\n':0.02]
  step  4  pick=' of'       top5=[' of':1.00 | ' is':0.00 | ' city':0.00 | ' o':0.00 | ' was':0.00]
  step  5  pick=' France'   top5=[' France':0.36 | ' Germany':0.19 | ' the':0.13 | ' Italy':0.04 | ' Spain':0.03]

確率の立ち方がステップごとに違うのが面白いところです。of の次はほぼ 1.00 で確定していますが、The capital of の次は France が 0.36、Germany が 0.19 と割れていて、モデルが迷っているのが見えます。ここで分かるのは、モデルが返すのは 1 個の答えではなく、候補に確率を割り振った分布だということです。今回は最大のものをそのまま選ぶ greedy で回しましたが、分布から引く選び方にすると出力は変わります。その「どう選ぶか」と、さらに踏み込んで「複数引いて束ねると何が起きるか」は Part 4 で扱います。

手元に残るもの と 次回予告

今回手を動かして残ったものを確認します。

• 動くコード — token_loop.py（自己回帰ループ・top-k 表示・累積時間の実測）
• 測った数字 — 生成速度（tok/s）と「累積時間はトークン数に線形」という観察
• メンタルモデル — LLM は次トークン予測器で、生成はその関数を回すループ

次回からは、この素直なループの「代償」と「活かし方」に入ります。1 トークンずつ・逐次ということは、過去のトークンを毎回参照しているということです。それを賢く貯めるのが KV キャッシュで、速度とメモリの本当のボトルネックになります（Part 3）。逆に、逐次生成が安いからこそ、複数の答えを束ねて品質を買い戻す、という発想も生まれます（Part 4）。最終回では、これらが PHOTON の「475 倍」（正確には実スループット約 44 倍 × メモリ約 11 倍の積で、品質はマルチクエリで補償する条件付きの数字です）の分解にそのまま繋がります¹。

コードは companion リポジトリの part-01 タグで再現できます（token_loop.py と、MLX 不要で通る selftest.py）。

参考文献