理論的に言えば、LLMが返す最良の回答は、常に最も確率の高いトークンを選び続けることで得られる。この手法ではLLMの出力が決定的になってしまうが、多くのアプリケーションにとってそれは望ましい性質ではない。そのため、文脈への忠実さを保ちつつLLMの創造性のバランスを取るために、近年さまざまなサンプリングアルゴリズムが提案されてきた。

今日、最もよく使われているもの、ほぼデフォルトと言えるのがtop-pと呼ばれる手法だ。これはnucleus sampling（核サンプリング）の一種で、スコアの高いトークンから順に合計確率が「p」に達するまで集め、そこから重み付きランダムサンプリングを行うものである。

このブログ記事では、なぜnucleus samplingが最良のアプローチではないと考えるのかを検討し、その問題を回避するためのシンプルで理解しやすい代替手法を示したい。このアルゴリズムはまだ開発途上だが、現時点で公開することで議論やハッキングのきっかけになればと思っている。

ロジットの中に金脈がある

LLMのロジットは、LLM内部の動作の中でも数少ない完全に理解可能な部分であるにもかかわらず、その特徴を研究したり、より高度なサンプリング手法を探求したり、不確実性やハルシネーションの可能性を検出してユーザーに知らせたりといったことへの関心は、総じてあまり見られない。連続するトークンの確率分布を可視化することは、洞察を得るためのシンプルで実践的な手段だ。

画像では、上位32個の候補トークンが確率に応じて色付けされているのがわかる（白＝0、青＝1）。また、選択されたトークンと、そのトークンのランク（最も確率が高いものが0、その次が1、といった具合）も示されている。

上記の例では、Mistralのベースモデルはウンベルト・エーコの生没年を把握しているため、最も確からしいトークンに確率の大部分を割り当て、自信を持って示している。一方で、文章の続きを表現する方法が複数ある場合や、ある事実について確信が持てない場合など、モデルがより困惑するケースもある。学習時に誕生日を覚えていない人物の生年月日を尋ねると、トークン確率の分布は異なる様相を示す。

しかし一般的に避けたいのは、最適とは言えないトークンを選択し、LLMの生成を最小パープレキシティかつハルシネーションが最も少ない経路から外してしまうことだ。nucleus samplingはこの点で失敗する。というのも、「p」に達するまでトークンを集める方式では、分布によっては最有力候補より極端に弱いトークンまで含まれてしまうからだ。例えば上の画像にある「writer」というトークンを考えてみよう。ウンベルト・エーコは確かに作家である。このトークンはそれほど突出して高いスコアではないものの、2番手の候補よりはるかに高い。それでもpを0.3に設定すると、0.01のような低い値の2番手のトークンまで集めてしまい、数パーセントの確率でそれが選ばれることになり、生成を誤った方向へ導くリスクがある。

一方、その次のトークンである「who」では、比較的近いスコアを持つ複数の選択肢が存在し、それらを使って別の生成パターンを作り出すことが可能だ。

雪崩効果を活用する

ここで重要な観察は、多様性を得るために弱すぎるトークンを選ぶのは割に合わないということだ。出力の多様化のために、少数の有力な候補が存在する生成タイミングだけを利用したとしても、雪崩効果が助けになる。入力コンテキストが変わればLLMの出力も揺らぎ、その結果としてテキストの別バージョンが生成されやすくなるのだ。

First Token Cutoff（FTC）アルゴリズム

注意：ここで説明するアルゴリズムは、いかなる科学的な検証も受けていない。私は「最悪トークンの質を制限する」特性を持つさまざまなサンプリングアルゴリズムを2日ほど実験してみたが、その中でも本手法が適用しやすさ、結果、理解しやすさのバランスが最も良いように思われる。

ここで提案するアルゴリズムは、非形式的には次のように述べることができる。

• LLMがある特定の候補に強く偏っているときは、それを選択する。
• 複数の有力な候補があるときは、代替案を生成する。
• 選択されうる最も悪いトークンの質は、一定の範囲に制限されるべきである。

Tail Free Samplingのような先行研究でも、LLMが出力する高品質なトークンの小さな集合の中から選択すべきだと指摘されている。しかしTFSでは、その集合は微分を行うことで特定され、トークンの質が急激に低下する手前の緩やかなカーブに対応するクラスターを選ぶ形になっている。

一方、ここで提案するアルゴリズムでは、最高スコアのトークンT0を基準とした、より制限が明確で理解しやすいカットオフに従って選択を行いたい。T0には他のすべてのトークンとは異なる特別な意味を持たせる。すなわち、そのトークンを出力する際のLLMの確信度（基本的にパープレキシティの代理指標）である。したがって本アルゴリズムでは、T0と比較して一定の割合以上に悪いトークンはすべて除外する。

このカットオフの割合は0から1の値を取り、「co」と呼ぶ。有効な「co」の一例としては0.5が挙げられる。

アルゴリズムの動作は次のとおりである。

1. logitsに対してsoftmax()を計算する。
2. トークンを確率順にソートする。
3. 最良トークンの確率T0を基準に、他のすべてのトークンについて比率 r = 1 - (T[i] / T0) を計算する。
4. r <= co を満たすトークンのみを選択する。
5. 選択されたトークンの中から重み付きランダム選択を行う。

このようにすることで、トークンの値が単調に滑らかに減少していく場合であっても、可能性の集合に含められるトークンにはハードな上限が設けられることになる。一方、高スコアのクラスターを特定しようとする他の手法では、このような保証はない。

実例

nucleus top-pサンプリングが実用上破綻しているように見えない理由の一つは、しばしば最初のトークンの確率が非常に高く、パープレキシティが低いときや、たまたま低品質なトークンを集めて選んでしまわないときには、生成が妥当な経路をたどり続けるからだ。

次のような確率が連続する場合には、より問題が顕著になる。

0.25, 0.14, 0.01

p=0.4の場合、3番目の低品質なトークンまで集めてしまい、約3%の確率でそれが選ばれることになる。

ここで、coを0.5としたFirst Token Cutoffを考えてみよう（トークンは最上位より最大50%まで悪くても許容される）。

2番目のトークンのr値は次のとおりである。

r[t1] = 1-(0.14/0.25) = 0.44 # 0.44 <= 0.5、このトークンは採用される

r[t2] = 1-(0.01/0.25) = 0.96 # 0.96 > 0.5、このトークンは除外される

出力例

プロンプト「Sorted sets are」に対して、co=0.7でMistralベースモデル（instructなし）を実行した出力例である。以下は品質で選別することなく得られた、3つの成功例だ。

1. Sorted sets are a powerful data structure in Redis. They can be used to store sorted lists, to store unique values, to store scores for ranking, and to store a sorted list of sorted sets.
2. Sorted sets are a very powerful data structure. They allow you to store data in a way that makes it easy to find the highest or lowest values in the set, and they also allow you to sort the data. This can be useful for many different tasks, such as ranking users by their score, or finding the most popular items in a database.
3. Sorted sets are a very powerful data structure that can be used to solve many different problems. The most common use case is to store a list of unique elements, each of which has an associated value. For example, you could use a sorted set to store the names of all the people in your family, with their ages as the associated values.

なぜ理解しやすさが重要なのか

サンプリングパラメータは、LLMのエンドユーザーやAPI利用者が調整しなければならない数少ない要素の一つである。試行錯誤が必要になることも多いが、現実世界での意味や直感がすぐに理解できる単一の調整可能なパラメータがあることは大きな助けになる。さらに「co」は線形なパラメータであるため、temperatureのようなパラメータや、線形ではあるもののロジットの分布形状に強く依存するtop_pの「p」と比べて、推論がとりわけシンプルだ。

今後の課題

私はまだ調査を始めたばかりなので、今後このアルゴリズムをより深く研究し評価していくつもりだ。何よりも、サンプリングアルゴリズムへの関心が高まり、top-p的な手法から一歩先へ進もうという機運が高まることを願っている。

ロジットの分布からは、おそらく興味深い情報を収集できる。例えば、LLMの隠れ層が事実情報を扱っているタイミングを、線形プローブが学習できる可能性は高い。これをトークンのパープレキシティと組み合わせることで、出力の一部が誤っている可能性が高いことをLLMのユーザーに示すために利用できるかもしれない。一般的に、トークンの確率分布を可視化することは非常に有益であり、ある意味でLLMがどのように動作し、各ステップでどのような候補が存在するのかを、手で触れるように体感させてくれる。

リファレンス実装

logits = mx.softmax(logits)
np_logits = np.array(logits) # MX -> NumPy
np_logits = np_logits.flatten()
sorted_indices = np.argsort(np_logits)
sorted_indices = sorted_indices[::-1]

co = 0.7
j = 1
t0 = np_logits[sorted_indices[0]]
while 1 - (np_logits[sorted_indices[j]] / t0) < co and j < len(np_logits):
    j += 1
accepted_logits = []
for i in range(0,j):
    accepted_logits.append(float(np_logits[sorted_indices[j]]))
accepted_logits = mx.array(accepted_logits)

idx = mx.random.categorical(accepted_logits)
idx = int(np.array(idx)) # Convert zero-dim array to scalar
token_id = sorted_indices[idx]

謝辞

これらの実験は、AppleのMLXライブラリと、精力的に開発を続けるクールで優秀な開発者たちのおかげで、非常に簡単に行うことができた。MLXはあるべき姿のとおり極めてアクセスしやすい。結局のところ、LLM自体は不可解だが、推論そのものはシンプルなプロセスなのだから。