温度・top-k・top-p を手で作る

目標

ロジット1式から次のトークンを1つ選ぶまでに、分布が通る位置を、標準ライブラリだけですべて作ります。温度で割ったあとの安定したソフトマックス、エントロピー、top-k、top-p、この4つを決まった順序でつなぐ関数、逆累積分布による抽出、greedyと温度のスイープまでです。最後に、同じシードで2回選んで同じ数列が出るかと、温度を下げたほうがgreedyと同じになるかを、記録に残します。

なぜ重要なのか

Transformerブロックの仕事は、語彙のサイズと同じ長さの実数のリストを1つ出力するところまでです。実際に使うトークンを選ぶのはモデルの外側のルールで、同じモデルが毎回違うことを言う理由も、毎回同じ言葉を繰り返す理由も、すべてそこにあります。このラボは、モデルを呼び出しません。このPodにはtransformersもtorchもなく、numpyは/opt/onnx-lab/bin/pythonの中にしかないので、システムのPythonではimport numpyができません。代わりに、ロジット1式を自分で書いておき、その後ろを手で作ります。そのため、ここに出てくる数字はすべて、自分で書いたロジットから測ったものです。難しいのは細部です。温度を割る位置がソフトマックスの前か後か、最大値をなぜ引くのか、top-kの同点をどう解消するか、top-pがpに届く項目を含めるか、温度とtop-pとtop-kをどの順序で適用するか。1つだけ違っても、同じ設定が別のシステムになります。採点ツールは、書いておいた説明を信用しません。自分で作ったモジュールを実際に呼び出し、毎回異なるロジットと異なる温度で関数を直接叩いて、採点ツールが別に計算した確率ベクトルと照合します。無作為な抽出そのものは判定しません。シードを固定した番号の数列と確率ベクトルだけを見ます。

ステップ

/root/work/tf-sample/sample.pyにVOCAB・LOGITSと、softmax_t(logits, temperature)を作成してください。ロジットを温度で割ってから最大値を引き、ソフトマックスにかけます。
entropy_bits(probs)を追加して、分布の広がり具合を、底が2のエントロピーで測るようにしてください。確率が0の位置は、0として扱います。
top_k_filter(probs, k)を追加して、確率の大きいk個だけを残し、残りを正確に0.0にして、再び正規化するようにしてください。同点なら、番号が小さいほうが残ります。
top_p_filter(probs, p)を追加して、降順の累積和がpに初めて届くその項目までを残すようにしてください。その項目を除くと、合計がpに届きません。
filtered_probs(logits, temperature, top_k=None, top_p=None)を作成して、温度 → ソフトマックス → top-k → top-pの順序でつなぐようにしてください。Noneなら、そのステップを飛ばします。
sample_index(probs, u)とsample_sequence(logits, temperature, top_k, top_p, n, seed)を作成してください。逆累積分布で選び、random.Random(seed)からrandom()をn回呼び出して、同じシードなら同じ数列が出るようにします。
greedy(logits)とsweep(logits, temps, p)を作成してください。greedyは最も大きいロジットの番号(同点なら小さい番号)で、sweepは、温度ごとに(온도, 엔트로피, top-p 가 남기는 개수)(プレースホルダーは温度、エントロピー、top-pが残す個数です)を測ります。
決められた設定で測って、/root/work/tf-sample/sample_report.jsonと、/root/work/tf-sample/sample_report.mdに結果を記録してください。

参考

実行の契約: 採点ツールは、/root/work/tf-sample/sample.pyをPythonモジュールとして読み込み、VOCAB・LOGITS・softmax_t・entropy_bits・top_k_filter・top_p_filter・filtered_probs・sample_index・sample_sequence・greedy・sweepを直接使います。スクリプトとして実行しないので、if __name__ == "__main__"はなくてかまいません。
VOCABは、互いに異なる短い文字列12個以上のリストです。内容は自由に決めてください。
LOGITSは、VOCABと長さが同じ実数のリストで、値が互いに異なる必要があります。確率ではなく、正規化されていないスコアです。1位と2位の差が0.5以上、最大値と最小値の差が3.0以上であれば、温度を変えたときに分布が動くのが見えます。
softmax_t(logits, temperature)は、まずロジットをtemperatureで割り、割った値から最大値を引いてから指数をとって、正規化します。最大値を引かないと、温度が小さいときにmath.expがあふれます。temperatureが0以下なら、例外を出してください。
entropy_bits(probs)は、-sum(p * log2(p))です。pが0の位置は飛ばします。均等に広がったn個なら、log2(n)が出ます。
top_k_filter(probs, k)は、捨てた位置を正確に0.0にして、残ったものだけを、合計が1になるように正規化します。同点なら、番号が小さいほうが残ります。kがリストの長さ以上なら、何も捨てずに、正規化だけを行います。kが1より小さければ、例外を出してください。
top_p_filter(probs, p)は、確率を降順(同点なら番号が小さいほうが先)に並べて累積和を測り、累積和がp以上になる最初の項目までを含めて止まります。残りは正確に0.0で、残ったものを正規化します。pは0より大きく1以下である必要があり、それ以外なら、例外を出してください。
filtered_probsの順序は、温度 → ソフトマックス → top-k → top-pです。この順序がこのラボの契約で、順序を変えると、残るトークンが変わります。
sample_index(probs, u)は、番号順に確率を足していって、uを初めて超える位置を返します(u < 누적합(プレースホルダーは累積和です)が真になる最初の位置)。uは0以上1未満です。確率が0.0の位置は、絶対に返しません。
sample_sequence(logits, temperature, top_k, top_p, n, seed)は、分布を1回だけ作り、random.Random(seed)のrandom()をn回呼び出して、番号n個のリストを返します。シードが同じなら、結果も同じである必要があります。
sweep(logits, temps, p)は、tempsの温度ごとに、(온도, 엔트로피, top-p 가 남기는 개수)(プレースホルダーは温度、エントロピー、top-pが残す個数です)の3つ組のリストを返します。個数は、top_p_filterを通したあとで0でない位置の数です。
ステップ8の設定は固定です。シードは20260917、選ぶ回数は24、top_kは5、top_pは0.9です。高温側は温度1.0、低温側は温度0.2で、同じシードから選びます。
ステップ8のsweepの温度は0.25・0.5・1.0・2.0・4.0で、そのときのpは0.9です。nucleusは、温度1.0でpを0.5・0.8・0.9・0.95と変えながら測った[p, 남은 개수](プレースホルダーはpと残った個数です)で、topk_massは、温度1.0のもとの確率で、上位k個が占める合計を、kが1・3・5・10の場合について測った[k, 합](プレースホルダーはkと合計です)です。
sample_report.jsonのキー: vocab_size・greedy_index・greedy_token・top_prob_t1・entropy_t1・sweep・nucleus・topk_mass・seed・draw_count・hot_draws・cold_draws・cold_is_greedy・hot_distinct・cold_distinct・repeat_matches。
sample_report.mdは、## 무엇을 쟀나 ## 온도가 분포를 어떻게 바꾸나 ## top-k 와 top-p 가 남기는 것 ## 같은 시드는 같은 수열을 준다の4つの節で書きます(韓国語の見出しは順に、「何を測ったか」「温度は分布をどう変えるのか」「top-kとtop-pが残すもの」「同じシードは同じ数列を返す」という意味です)。
このPodにはインターネットがありません。pip installはできず、transformersもtorchもありません。numpyは/opt/onnx-lab/bin/pythonの中にしかないので、システムのPythonではimport numpyができません。mathとrandomで十分です。
公式ドキュメント: nucleus samplingの原論文・top-kの原論文・Hugging Face — Generation strategies・Python — math
よくある間違い: 確率に温度をかける、最大値を引くのを省略して、低い温度であふれる、切り落としたあとの正規化を忘れる、top-pでpを超える最初の項目を除いてしまう、同点のルールを決めない、top-pをtop-kより先に適用する、選ぶたびに分布を作り直す。

温度で割ってからソフトマックスにかける

/root/work/tf-sample/sample.pyに、VOCAB(互いに異なる短い文字列12個以上)とLOGITS(同じ長さ、値が互いに異なる、1位と2位の差が0.5以上、最大値と最小値の差が3.0以上)を置き、softmax_t(logits, temperature)を作成してください。ロジットをtemperatureで割ってから最大値を引き、指数をとって正規化します。temperatureが0以下なら、例外を出してください。

割る位置が要点です。確率を先に求めておいて、それに温度をかけると、まったく違う結果が出ます。最大値を引く処理は、温度が小さいときに効いてきます。0.01で割ると、ロジット5.2が520になり、math.exp(520)はそのままあふれます。同じ値から同じ数を引くだけなので、確率は変わりません。温度0は割り算ができないので、ValueErrorを出すほうがよいでしょう。greedyは「温度0」ではなく、別のルールです。

広がり具合をビットで測る

entropy_bits(probs)を追加してください。底が2のエントロピー-sum(p * log2(p))です。確率が0の位置は、0として扱います。均等に広がったn個なら、log2(n)が出ます。

math.log2(0)は例外を出します。そのため、0の位置は飛ばす必要がありますが、これは便宜的な処置ではなく、正しい処理です。p * log2(p)のpが0に近づくときの極限が0だからです。温度を上げながらこの値を測ってみると、大きくなるのが見えます。分布がどれだけ広がっているかを、数字1つにまとめる箇所です。

上位k個だけを残す

top_k_filter(probs, k)を追加してください。確率の大きいk個だけを残し、残りは正確に0.0にして、残ったものを合計が1になるように正規化します。同点なら、番号が小さいほうが残ります。kがリストの長さ以上なら、何も捨てずに正規化だけを行い、kが1より小さければ、例外を出してください。

番号を確率の降順に並べ替えて、先頭のk個を集合として持っておけば、残りはその集合を見るだけです。並べ替えのキーを(-확률, 번호)(プレースホルダーは確率と番号です)にすると、同点のルールまで1行に収まります。正規化を忘れると、合計が1にならず、あとで選ぶときに最後の位置に偏ります。捨てた位置を非常に小さい値にしてはいけません。低い確率で復活するからです。

累積和がpに届くその項目までを残す

top_p_filter(probs, p)を追加してください。確率を降順(同点なら番号が小さいほうが先)に並べて累積和を測り、累積和がp以上になる最初の項目までを含めて止まります。残りは正確に0.0で、残ったものを正規化します。pが0以下か、1より大きければ、例外を出してください。

含めるか除くかが、このステップのすべてです。pを超えさせたその項目を除いてしまうと、残ったものの合計がpに届きません。そうなると、「p分の質量を残す」という言葉そのものが成り立ちません。累積和がpに届いた瞬間に、その番号をすでに入れたあとで止まれば、それでよいのです。pがどんなに小さくても、少なくとも1つは残ります。

4つを決まった順序でつなぐ

filtered_probs(logits, temperature, top_k=None, top_p=None)を作成してください。温度 → ソフトマックス → top-k → top-pの順序でつなぎ、Noneのステップは飛ばします。

順序が結果を変えます。先に温度をかけると、分布そのものが変わるので、同じpでも残る個数が変わります。先にtop-kを行うと、生き残ったものが再び正規化されて確率が膨らむので、top-pのほうが早く止まります。関数は4行で終わります。前に作った3つの関数を、この順序で呼ぶだけです。top_kやtop_pがNoneかどうかを見る条件を、忘れないでください。

同じシードなら同じ数列になる

sample_index(probs, u)とsample_sequence(logits, temperature, top_k, top_p, n, seed)を作成してください。sample_indexは、番号順に確率を足していって、uを初めて超える位置を返し、sample_sequenceは、分布を1回だけ作ってから、random.Random(seed)のrandom()をn回呼び出して、番号n個を返します。

u < 누적합(プレースホルダーは累積和です)が真になる最初の位置です。確率が0.0の位置は、累積和を増やさないので、絶対に選ばれません。切り落としたトークンが復活しないということで、これがこの構造の安全装置です。random.Random(seed)は、関数の中で1回だけ作ってください。選ぶたびに新しい生成器を作ると、同じ番号がn個出るだけになります。分布も、ループの外で1回だけ作ります。

greedyと温度のスイープを作る

greedy(logits)とsweep(logits, temps, p)を作成してください。greedyは、最も大きいロジットの番号で、同点なら小さい番号です。sweepは、tempsの温度ごとに、(온도, 엔트로피, top-p 가 남기는 개수)(プレースホルダーは温度、エントロピー、top-pが残す個数です)の3つ組のリストを返します。

greedyは、確率を求める必要さえありません。ソフトマックスは、順序を変えないからです。同点でmaxがどちらを返すのかが不安なら、自分で走査して、より大きいときだけ更新してください。そうすれば、小さい番号が残ります。sweepの個数は、top_p_filterを通したあとで0でない位置の数です。温度を上げるほど、エントロピーが大きくなり、残る個数も増えていくのを、目で見てください。

つまみを回した結果を記録に残す

固定された設定で測ってください。シードは20260917、選ぶ回数は24、top_kは5、top_pは0.9で、高温側は温度1.0、低温側は温度0.2です。結果を、/root/work/tf-sample/sample_report.jsonにはvocab_size・greedy_index・greedy_token・top_prob_t1・entropy_t1・sweep・nucleus・topk_mass・seed・draw_count・hot_draws・cold_draws・cold_is_greedy・hot_distinct・cold_distinct・repeat_matchesとして、/root/work/tf-sample/sample_report.mdには## 무엇을 쟀나 ## 온도가 분포를 어떻게 바꾸나 ## top-k 와 top-p 가 남기는 것 ## 같은 시드는 같은 수열을 준다の4つの節で書いてください(韓国語の見出しは順に、「何を測ったか」「温度は分布をどう変えるのか」「top-kとtop-pが残すもの」「同じシードは同じ数列を返す」という意味です)。

数字は手で書かず、自分のコードを実際に動かして得た値で埋めてください。sweepは、温度0.25・0.5・1.0・2.0・4.0に、pは0.9です。nucleusは、温度1.0でpを0.5・0.8・0.9・0.95と変えながら測った[p, 남은 개수](プレースホルダーはpと残った個数です)で、topk_massは、温度1.0のもとの確率で、上位k個が占める合計を、kが1・3・5・10の場合について測った[k, 합](プレースホルダーはkと合計です)です。切ったあとに正規化した値ではなく、切る前の合計です。cold_is_greedyは、低温側の数列がすべてgreedy(LOGITS)と同じかどうかで、repeat_matchesは、同じシードでもう一度選んだときに同じ数列が出るかどうかです。自分で書いたロジットによっては、cold_is_greedyが偽になることもあります。その場合は、そのまま書いて、本文でなぜそうなるのかを説明してください。