Skip to content

LLM Sampling

実装: llm-sampling/ / 実行: go test ./llm-sampling/

LLM が次の1トークンを選ぶところを、モデルが吐く生スコア(logits)から自分で計算する。softmax で確率にして、temperature で尖らせたり平らにしたり、top-p などで候補を絞って、最後に1つ抽選する。ChatGPT の応答の「性格」は、実はモデル本体と同じくらいこの選び方で決まっている。

「サンプリング」の同名別分野に注意

この章は LLM がテキストを生成するときのトークンサンプリングの話。 分散トレーシングの head-based / tail-based sampling(「全リクエストを記録できない中で どのトレースを残すか」)は別問題なので、Trace Sampling で独立して扱う。

この章で作るもの

LLM の出力層で行われているサンプリングを、logits から自分で計算する。

  1. softmax — logits を確率分布に変換する(数値安定化つき)
  2. temperature — 分布を尖らせる / 平らにする
  3. top-k / top-p / min-p — 分布の尻尾を切り落とす3つの流儀
  4. greedy / sample — 分布から1トークン決める

各節にはその場で動かせるデモを置いてある。どれも Go 版と同じ手順を JS に移しただけのもので、 ブラウザの中だけで計算している。

先に押さえることが3つある。

  • LLM の「生成」とは、logits を確率分布に変えて1トークン抽選する行為の繰り返しである
  • temperature は softmax 直前の割り算1つ。特別なパラメータではない
  • top-k / top-p / min-p はどれも「尻尾を -Inf に落として renormalize する」フィルタで、 切る基準(順位 / 累積確率 / 最大確率との比)だけが違う

前提: トークン・語彙・logits

  • トークン: LLM が文章を扱う最小単位。単語より少し細かい(「食べました」が 「食べ」「まし」「た」のように割れる)。この割り方を決める仕組みは前章で作った
  • 語彙(vocabulary): モデルが知っている全トークンの一覧。GPT 系で5万〜25万種類
  • logits: 語彙の各トークンに対する「次に来る確からしさ」の生スコア(実数の列)。 確率ではない(合計1でも0〜1でもない)。これを確率に直すのが softmax で、この章の出発点

モデルの出力は「次のトークンの点数表」

モデルが1ステップで出すのは、語彙の全トークンに対する logits(実数の点数の列)だけ。 「明日の天気は」まで読んだモデルなら、「晴れ」に 4.0、「曇り」に 3.3、「猫」に -1.0、 のような点数を全語彙ぶん並べる。文章が流れるように出てくるのは、 この点数表から1個選んでコンテキストに足し、また点数表を出す、を繰り返しているだけ。

つまり ChatGPT の応答の「性格」は、モデル本体と同じくらい 点数表から1個選ぶ方法(=サンプリング)に支配されている。

softmax: logits を確率にする

go
// Softmax は logits を確率分布に変換する。
// exp の overflow を防ぐため、全要素から最大値を引いてから計算する(数値安定化)。
// 定数を引いても exp の比は変わらないので、結果の分布は同じ。
func Softmax(logits []float64) ([]float64, error) {
	if len(logits) == 0 {
		return nil, errors.New("sampling: logits must not be empty")
	}
	maxLogit := math.Inf(-1)
	for _, l := range logits {
		maxLogit = math.Max(maxLogit, l)
	}

	exps := make([]float64, len(logits))
	sum := 0.0
	for i, l := range logits {
		exps[i] = math.Exp(l - maxLogit)
		sum += exps[i]
	}

	probs := make([]float64, len(logits))
	for i, e := range exps {
		probs[i] = e / sum
	}
	return probs, nil
}

見どころは maxLogit を引いている部分。logits が 1000 を超えると exp(1000) が overflow して全部 +Inf になるが、softmax は「差」しか見ないので、 全要素から最大値を引いてから exp しても結果は変わらない。 テストでは logits = [1000, 1000] を入れてこの安定化を検証している。

もう1つの仕掛けが -Infexp(-Inf) = 0 なので、 logit を -Inf にする = そのトークンの確率を0にして renormalize することになる。 この後の top-k / top-p / min-p はすべてこの性質を使ったフィルタとして書ける。

temperature: 分布の尖りを操作する

go
// ApplyTemperature は logits を t で割る。softmax の直前に掛かる唯一のスカラー操作で、
// t < 1 は logits の差を拡大して分布を尖らせ(確定的寄り)、
// t > 1 は差を縮小して分布を平らにする(多様性寄り)。
// t = 0 は0除算になるため受け付けない。決定的にしたければ Greedy を使う。
func ApplyTemperature(logits []float64, t float64) ([]float64, error) {
	if t <= 0 {
		return nil, errors.New("sampling: temperature must be positive (use Greedy for t=0)")
	}
	out := make([]float64, len(logits))
	for i, l := range logits {
		out[i] = l / t
	}
	return out, nil
}

やっていることは割り算1つだが、softmax が exp(指数関数)を通すため効果は大きく、

  • t → 0: 差が無限に拡大され、最大 logit のトークンが確率1に近づく(greedy と同じ)
  • t = 1: モデルが出した分布そのまま
  • t → 大: 差が消えて一様分布に近づく(でたらめ)

試してみる: スライダーを 0.1 まで下げると「晴れ」がほぼ100%になり、 3.0 まで上げると「猫」や「無」にも現実的な確率がつく。 何回か抽選して、低温では同じ結果ばかり、高温ではばらつくことを確認してほしい。

デモ次トークンのサンプリング

「明日の天気は◯◯」の次に来るトークンの確率分布

晴れ51.5%
曇り25.6%
17.2%
3.1%
台風1.4%
0.7%
0.3%
0.1%

なお t = 0 は0除算なので、実装では受け付けずに greedy(argmax)を使う。

go
// Greedy は最も logit の大きいトークンを選ぶ。temperature → 0 の極限と同じ。
// 常に同じ入力から同じ出力が出る(決定的)。
func Greedy(logits []float64) (int, error) {
	if len(logits) == 0 {
		return 0, errors.New("sampling: logits must not be empty")
	}
	best := 0
	for i, l := range logits {
		if l > logits[best] {
			best = i
		}
	}
	return best, nil
}

top-k: 上位k個しか見ない

温度を上げて多様性を出すと、副作用として「猫」のような明らかにおかしいトークンにも 確率が漏れる。そこで抽選の前に候補を絞るのがフィルタ系の役割。 最も素朴な top-k は、logit の上位 k 個だけ残して残りを -Inf に落とす。

go
// FilterTopK は logit の大きい上位 k 個だけを残し、他を -Inf にする。
func FilterTopK(logits []float64, k int) ([]float64, error) {
	if k <= 0 {
		return nil, errors.New("sampling: k must be positive")
	}
	if k >= len(logits) {
		return append([]float64(nil), logits...), nil
	}

	order := argsortDesc(logits)
	keep := make([]bool, len(logits))
	for _, idx := range order[:k] {
		keep[idx] = true
	}
	return filterByKeep(logits, keep), nil
}

試してみる: k を下げていくと下位のトークンから順に消え、 残った候補だけで確率が再分配(renormalize)される。k=1 は greedy と同じ。

デモ次トークンのサンプリング

「明日の天気は◯◯」の次に来るトークンの確率分布

晴れ51.5%
曇り25.6%
17.2%
3.1%
台風1.4%
0.7%
0.3%
0.1%

top-k の弱点は k が固定なこと。分布が尖っている(答えがほぼ決まっている)場面では k=40 は緩すぎ、分布が平らな(どう続けてもいい)場面では k=40 は厳しすぎる、 ということが同じ文章の中で起きる。

top-p: 累積確率で切る(nucleus sampling)

top-k の「個数固定」問題への答えが top-p。確率の大きい順に足していき、 累積が p に達するまでの最小の集合(nucleus)だけを残す。

go
// FilterTopP(nucleus sampling)は確率の大きい順に足していき、
// 累積が p 以上になる最小の集合だけを残す。
// top-k と違い「残る個数」が分布の形に応じて変わるのが特徴。
func FilterTopP(logits []float64, p float64) ([]float64, error) {
	if p <= 0 || p > 1 {
		return nil, errors.New("sampling: p must be in (0, 1]")
	}
	probs, err := Softmax(logits)
	if err != nil {
		return nil, err
	}

	order := argsortDesc(probs)
	keep := make([]bool, len(logits))
	cum := 0.0
	for _, idx := range order {
		keep[idx] = true // 累積が p を超える境界のトークンまでは含める
		cum += probs[idx]
		if cum >= p {
			break
		}
	}
	return filterByKeep(logits, keep), nil
}

試してみる: p を下げると尻尾から消えていくのは top-k と似ているが、 残る個数が分布次第で変わるのが違い。temperature と組み合わせたとき、 尖った分布では1〜2個、平らな分布では多数が残る「適応的な」切り方になる。

デモ次トークンのサンプリング

「明日の天気は◯◯」の次に来るトークンの確率分布

晴れ51.5%
曇り25.6%
17.2%
3.1%
台風1.4%
0.7%
0.3%
0.1%

min-p: 最大確率との比で切る

比較的新しい流儀。「最大確率の minP 倍未満のトークンは切る」という相対閾値で、 top-p の累積計算より単純なのに、分布の尖り具合に適応する性質は保たれる。

go
// FilterMinP は「最大確率の minP 倍」を閾値として、それ未満のトークンを切る。
// 分布が尖っているときは厳しく、平らなときは緩く切れる相対的なフィルタで、
// top-p の「累積」ではなく「個々の確率」で判定するのが違い。
func FilterMinP(logits []float64, minP float64) ([]float64, error) {
	if minP < 0 || minP > 1 {
		return nil, errors.New("sampling: minP must be in [0, 1]")
	}
	probs, err := Softmax(logits)
	if err != nil {
		return nil, err
	}

	maxProb := 0.0
	for _, p := range probs {
		maxProb = math.Max(maxProb, p)
	}
	threshold := minP * maxProb

	keep := make([]bool, len(logits))
	for i, p := range probs {
		keep[i] = p >= threshold
	}
	return filterByKeep(logits, keep), nil
}

試してみる: min-p = 0.1 なら「1位の10分の1未満の泡沫候補は切る」という意味になる。 top-p と見比べると、切れ方が「累積」ではなく「個々の確率」で決まることがわかる。

デモ次トークンのサンプリング

「明日の天気は◯◯」の次に来るトークンの確率分布

晴れ51.5%
曇り25.6%
17.2%
3.1%
台風1.4%
0.7%
0.3%
0.1%

抽選: 確率分布から1個選ぶ

フィルタを通した最終分布から実際に1トークン選ぶのが inverse CDF 法。 [0,1) の乱数を引き、累積確率が乱数を超えた位置のトークンを返す。

go
// Sample は確率分布から1トークン抽選する(inverse CDF法)。
// rng は [0,1) の乱数を返す関数。累積確率が乱数を超えた位置のトークンを返す。
// 浮動小数点誤差で累積が1に届かない場合に備え、最後のトークンにフォールバックする。
func Sample(probs []float64, rng func() float64) (int, error) {
	if len(probs) == 0 {
		return 0, errors.New("sampling: probs must not be empty")
	}
	if rng == nil {
		return 0, errors.New("sampling: rng must not be nil")
	}

	r := rng()
	cum := 0.0
	for i, p := range probs {
		cum += p
		if r < cum {
			return i, nil
		}
	}
	return len(probs) - 1, nil
}

テストでは乱数生成器を「固定値を返す関数」に差し替えて、 境界(累積0.2ちょうど、など)で正しいトークンが選ばれることを確認している。 時計を外から渡したレートリミッタと同じで、非決定的なものを注入してテスト可能にする定石になる。

全部つなげる

実際の推論エンジン(llama.cpp など)では、これらが1本のパイプラインになっている:

logits → ÷temperature → top-k → top-p → min-p → softmax → 抽選

試してみる: 全部のつまみを同時に動かして、 「temperature で形を作り、フィルタで尻尾を切り、最後に抽選する」流れを確認してほしい。

デモ次トークンのサンプリング

「明日の天気は◯◯」の次に来るトークンの確率分布

晴れ51.5%
曇り25.6%
17.2%
3.1%
台風1.4%
0.7%
0.3%
0.1%

設計の観点

  • モデルと選び方を分ける: 同じ logits でも、選び方を変えれば出てくる文章の性格が変わる。モデルを再学習せずに振る舞いを変えられる層が、ここに1枚あることになる
  • フィルタを1つの形にそろえる: top-k も top-p も min-p も「切って正規化し直す」に還元できる。同じ形に落とすと、組み合わせても壊れないし、増やすのも簡単になる
  • 切る基準が適応するかどうか: 個数で切ると分布の形に追随できない。累積確率や最大値との比で切ると、尖った場面では狭く、平らな場面では広く、自動で変わる
  • 数値の安定化は仕様の一部: 最大値を引く一手を省くと、大きな logits で全部が無限大になる。理屈が正しくても計算が壊れるところを、実装で塞いでおく
  • 決定的にできる口を残す: 温度 0 と乱数の注入で、同じ入力から同じ出力を出せる。再現性が要る用途と、多様性が要る用途を同じコードで賄える
  • 順番に意味がある: 温度で形を作ってから切るのか、切ってから温度をかけるのかで結果が変わる。パイプラインの並びそのものが設計になっている

つまみの比較

方式何で切るか特徴実例
greedy切らない(argmax)決定的。同じ入力から常に同じ出力temperature=0 指定。コード生成・分類など再現性重視の用途
temperature切らない(形を変える)尖り/多様性のトレードオフの主役ほぼすべての LLM API の基本パラメータ
top-k順位(固定k個)単純。分布の形に適応しないHugging Face transformers の既定(k=50)
top-p累積確率分布の形に適応する。長年の実務標準OpenAI / Anthropic API の top_p
min-p最大確率との比実装が単純で適応的。近年の推論エンジンで採用増llama.cpp・vLLM 等の OSS 推論エンジン

裏どり:

  • nucleus sampling の動機: top-p の原論文は、greedy や beam search が繰り返しに陥る現象を観察するところから始まっている。尤度を最大にすると、人間の書く文章から遠ざかるという逆説がこの分野の出発点になった
  • 温度 0 でも完全な再現にはならない: 同じ入力でも、バッチの組み方や GPU の並列加算の順序で浮動小数の丸めが変わり、argmax が入れ替わることがある。API が決定性を保証しないのはこのため
  • 繰り返しへの対処は別のつまみ: 同じ語が続くのを抑える repetition penalty や frequency penalty は、過去に出したトークンの logit を下げる処理になる。この章のフィルタと同じ層に、履歴を見る係が並ぶ
  • min-p は新しい: 提案が 2024 年と比較的新しく、実装が単純なわりに分布の尖りへ適応する。llama.cpp や vLLM といったローカル・自前運用の推論エンジンから広まった
  • 既定値は用途で違う: コード生成は温度を低く、創作は高く。API の既定が 1.0 前後なのは中庸を取っているだけで、用途ごとに調整する前提の値になっている

簡略化したこと

  • logits は手作りの8語彙: 実モデルでは数万〜数十万次元。ただし計算は完全に同じ
  • 履歴依存の補正なし: repetition penalty / frequency penalty は「過去に出したトークンの logit を下げる」処理で、本質は同じ引き算・割り算
  • 1トークンずつの抽選のみ: beam search のような複数候補の探索は扱わない
  • デモの乱数は Math.random: 実務では seed 固定で再現性を作る(Go 版は rng 注入済み)

参考資料