次トークン予測とは？logitsから文章の続きを1つ選ぶまで
LLMに「今日は」と入力したあと、モデルが一度に完成した文を決めているわけではない。文脈から次に続くトークンの候補ごとに点数を付け、1つ選び、選んだトークンを文脈へ加えてまた予測する。この反復が自己回帰的な文章生成の基本だ。
「次の単語を予測する」と説明されることもあるが、予測単位は単語とは限らない。トークナイザーによって、文字の一部や単語の一部、句読点などもトークンになりうる。
モデルが最初に出すのは確率ではない
Transformerの出力表現は、最終的に語彙サイズ分の点数へ写される。この正規化前の点数をlogitsと呼ぶ。語彙に3トークンしかない説明用モデルなら、次のような値を想像できる。
| トークン | logit |
|---|---|
| 猫 | 2.0 |
| 犬 | 1.0 |
| 。 | 0.0 |
logitの2.0は確率2.0ではない。0.0の「。」も候補から消えたわけではない。logitsの差を確率の配分へ変えるのがsoftmaxだ。
は候補 のlogit、 は変換後の確率を表す。結果は0から1の範囲になり、候補全体の和が1になる。
3候補をPythonで計算する
次のコードは標準ライブラリだけで動く。大きいlogitから指数を計算すると数値があふれやすいため、すべてのlogitから最大値を引いている。softmaxの結果は変わらない。
from math import exp
vocab = ["猫", "犬", "。"]logits = [2.0, 1.0, 0.0]
shifted = [score - max(logits) for score in logits]unnormalized = [exp(score) for score in shifted]total = sum(unnormalized)probs = [value / total for value in unnormalized]
for token, score, probability in zip(vocab, logits, probs): print(f"{token}: logit={score:.1f}, probability={probability:.3f}")
best = max(range(len(probs)), key=probs.__getitem__)print("選択:", vocab[best])print("合計:", round(sum(probs), 3))実行結果:
猫: logit=2.0, probability=0.665犬: logit=1.0, probability=0.245。: logit=0.0, probability=0.090選択: 猫合計: 1.0最も確率の高い候補を選ぶ方法を貪欲法と呼ぶ。この例では「猫」が選ばれる。しかし0.665は「猫が事実として正しい確率」ではない。この説明用の文脈と語彙、モデルの点数を前提とした、次トークン候補間の配分だ。例のlogitsは理解のために手で決めた値であり、実在するLLMの出力ではない。
1つ選んだあとに何が起こるか
選んだトークンを入力列の末尾へ追加し、更新された文脈で次のlogitsを計算する。たとえば「今日は」から「猫」が選ばれたら、次回は「今日は猫」に続く候補を採点する。先ほどの [2.0, 1.0, 0.0] をそのまま使い回すわけではない。
実際の語彙は3候補よりずっと多く、トークンIDと表示される文字列の対応もトークナイザーで決まる。生成の終わりを示す特殊トークンが選ばれたり、指定された最大長へ達したりすると、生成は停止する。選択方法も貪欲法だけではなく、確率に従って抽選するサンプリングなどがある。
学習時のsoftmaxと生成時の選択を混同しない
学習では、文脈に続く正解トークンへ高い確率を割り当てられるよう、交差エントロピーなどで損失を計算する。PyTorchの CrossEntropyLoss は確率ではなくlogitsを直接受け取るので、先にsoftmaxした値を渡さない。
生成では、学習済みモデルのlogitsから次に加えるトークンを決める。ここでは正解ラベルを見ながら損失を下げているのではない。毎回最も高い候補を選ぶか、分布から抽選するかによって、同じlogitsでも出力は変わりうる。
この小さな例で見たのは、モデルの内部表現が語彙上の候補へ変わり、次の1個が入力へ戻る接点だ。次はTemperatureを使うと候補の確率配分がどう変わるかを調べる。