技術

Top-kとは

Top-k Sampling

確率が高い上位 k 個のトークンのみを候補に絞ってサンプリングする手法

ひとことで言うと

候補の中で確率の高い上位いくつかだけからAIに答えを選ばせる方法。

概要

Top-k(Top-kサンプリング)とは、LLMが次のトークンを生成する際に、確率の高い順に並べた上位 k 個のトークンのみを候補として残し、それ以外を除外したうえで確率に応じてサンプリングする手法。候補数 k を固定するシンプルな仕組みで、確率が非常に低く不自然になりがちなトークンを排除しつつ、貪欲法よりは出力に多様性を持たせられる。候補数 k を1にするとほぼ常に最も確率の高いトークンだけが選ばれる貪欲法(Greedy Decoding)と同等になり、k を大きくするほど出力の多様性が増す。 kの値が小さすぎると出力が単調になり、大きすぎると確率の低い不自然なトークンが選ばれやすくなるため、TemperatureやTop-pと組み合わせて調整されることが多い。

背景

貪欲法(常に最も確率の高いトークンを選ぶ)は出力が単調になりやすく、単純な確率的サンプリングでは確率の低い不自然なトークンが選ばれてしまうことがある。 Top-kは、候補をあらかじめ上位 k 個に絞ることで、この不自然な低確率トークンの混入を防ぎながら多様性を確保する手法として考案された。

歴史

2018年: Fanらが論文「Hierarchical Neural Story Generation」でTop-kサンプリングを用いた物語生成を提案し、生成テキストの品質向上を示す。

アーキテクチャ

各トークン候補を確率の高い順に並べ、上位 k 個のみを残す。残った候補集合内で確率を再正規化し、その分布に従ってサンプリングする。 k が小さいほど出力は決定的になり、大きいほど多様性が増す。

コード例

Hugging Face Transformersでtop_kを指定して生成する

from transformers import AutoModelForCausalLM, AutoTokenizer

tokenizer = AutoTokenizer.from_pretrained("gpt2")
model = AutoModelForCausalLM.from_pretrained("gpt2")

inputs = tokenizer("The future of AI is", return_tensors="pt")
output = model.generate(
    **inputs,
    do_sample=True,
    top_k=50,  # 確率上位50個のトークンのみを候補にする
    max_new_tokens=30,
)

print(tokenizer.decode(output[0], skip_special_tokens=True))

利点

実装がシンプルで、確率が極端に低い不自然なトークンの混入を防げる
kの値によって出力の多様性を直接制御できる
計算コストが低く、他のサンプリング手法と組み合わせやすい

欠点

候補数 k が固定のため、確率分布が特定のトークンに集中している場面でも不要に多くの候補を残してしまう場合がある
最適な k の値はタスクやモデルによって異なる
確率分布の形状を考慮しないため、Top-pと比べて柔軟性に欠ける

比較

Top-p — Top-kは固定数の候補に絞るのに対し、Top-pは累積確率に基づき動的に候補数を決める
Temperature — TemperatureとTop-kは、出力候補の選び方を調整する目的で組み合わせて使われることが多い

よくある質問

kにはどんな値を設定する？

モデルやAPIのデフォルトは数十程度に設定されていることが多いが、決定的な出力にしたい場合は小さい値、多様性を重視する場合は大きい値を設定する。

k=1にするとどうなる？

常に最も確率の高いトークンのみが選ばれるため、貪欲法(Greedy Decoding)とほぼ同等の決定的な出力になる。

参考文献

Research PaperHierarchical Neural Story Generation