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トークン決める

各節にはその場で動かせるデモを置いてある。今回のデモは Worker すら不要で、 すべてブラウザ内の計算(Go 版と同じロジックの JS ミラー)で動いている。

この章の肝は3つ。

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

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

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

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

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

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

softmax: logits を確率にする

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

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

// 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 まで上げると「猫」や「無」にも現実的な確率がつく。 何回か抽選して、低温では同じ結果ばかり、高温ではばらつくことを確認してほしい。

「明日の天気は◯◯」の次トークン分布

temperature: 1.00

晴れ51.5%

曇り25.6%

雨17.2%

雪3.1%

台風1.4%

虹0.7%

猫0.3%

無0.1%

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

// 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 に落とす。

// 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 と同じ。

「明日の天気は◯◯」の次トークン分布

top-k: 8

晴れ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)だけを残す。

// 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個、平らな分布では多数が残る「適応的な」切り方になる。

「明日の天気は◯◯」の次トークン分布

top-p: 1.00

晴れ51.5%

曇り25.6%

雨17.2%

雪3.1%

台風1.4%

虹0.7%

猫0.3%

無0.1%

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

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

// 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 と見比べると、切れ方が「累積」ではなく「個々の確率」で決まることがわかる。

「明日の天気は◯◯」の次トークン分布

min-p: 0.00

晴れ51.5%

曇り25.6%

雨17.2%

雪3.1%

台風1.4%

虹0.7%

猫0.3%

無0.1%

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

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

// 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ちょうど、など)で正しいトークンが選ばれることを確認している。 時計を注入した rate-limiter 編と同じ、非決定的なものを外から注入してテスト可能にする定石。

全部つなげる

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

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

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

「明日の天気は◯◯」の次トークン分布

temperature: 1.00top-k: 8top-p: 1.00min-p: 0.00

晴れ51.5%

曇り25.6%

雨17.2%

雪3.1%

台風1.4%

虹0.7%

猫0.3%

無0.1%

4方式の比較

方式	何で切るか	特徴	実例
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 推論エンジン

簡略化したこと

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

参考資料

• The Curious Case of Neural Text Degeneration — top-p(nucleus sampling)の原典。greedy/beam が繰り返しに陥る観察も面白い
• Min-p Sampling — min-p の提案論文
• llama.cpp sampling — 実務のパイプラインでフィルタがどう直列されているかが読める