First Token Cutoff LLM sampling

2024-01-13

First Token CutoffによるLLMサンプリング

理論的には、LLMが返す最良の応答は、常に最も確率の高いトークンを選ぶことで得られます。この方法ではLLMの出力が決定論的になりますが、多くの用途において好ましい性質とはいえません。そこで近年、LLMの創造性と文脈への忠実さのバランスを取るため、さまざまなサンプリングアルゴリズムが提案されてきました。

現在、最も広く使われている、いわばデフォルトとなっている手法の一つがtop-pです。これはnucleus samplingの一種で、スコアの高いトークンを集め、確率の合計が「p」に達したところで止め、その中から重み付きランダムサンプリングを行います。

このブログ記事では、nucleus samplingが最善のアプローチではないと私が考える理由を検討し、nucleus samplingの問題を避けるための、シンプルで理解しやすい代替手法を紹介します。このアルゴリズムはまだ開発途中ですが、今の段階で公開することで、議論やハッキングを促せればと思っています。

logitには宝が眠っている

LLMのlogitは、LLMの内部動作の中でも数少ない、完全に理解可能な部分です。それにもかかわらず、その特徴を調べたり、より高度なサンプリング手法を研究したり、ユーザーの不確実性や幻覚の可能性を検出して知らせたりすることには、一般にほとんど関心が向けられていないように感じます。連続するトークンの確率分布を可視化するのは、理解を深めるためのシンプルで実用的な練習になります。

画像では、確率に応じて色付けした上位32個の候補トークン（白＝0、青＝1）、選択されたトークン、そしてそのトークンの順位（確率が最も高いもの＝0、次のもの＝1、以下同様）が確認できます。

上の例では、Mistralのベースモデルはウンベルト・エーコの生年月日と没年月日を知っているため、最も可能性の高いトークンに全確率の大部分を割り当て、自信を持って示しています。一方、文章の続きを表現する方法が複数ある場合や、特定の事実について確信がない場合には、モデルがより困惑していることもあります。学習中に誕生日を覚えなかった人物の名前について日付を尋ねると、トークン確率の分布は違ったものになります。

ただし一般に避けたいのは、最適ではないトークンを選んでしまい、LLMの生成を困惑度と幻覚が最小になる経路から外してしまうことです。nucleus samplingは、確率の合計が「p」に達するまでトークンを累積するため、分布によっては最初の候補より極端に弱いトークンまで含めてしまい、この問題を解決できません。たとえば、上の画像にある「writer」トークンを考えてみましょう。確かにウンベルト・エーコは作家でした。「writer」に対応するトークンのスコアはそれほど高くありませんが、それでも2番目の候補を大きく上回っています。ところがpを0.3にすると、0.01のような低い値の2番目のトークンまで累積され、1桁パーセントの確率で選ばれる可能性があります。その結果、生成が誤った経路に進む危険があります。

一方、次のトークンが「who」の場合は、スコアが比較的近い候補が複数あり、生成のバリエーションに利用できます。

雪崩効果を利用する

ここで重要な観察があります。多様性のために弱すぎるトークンを選ぶのは、割に合いません。出力を多様化するために、良い候補がいくつか存在する生成の瞬間だけを利用すれば、それで十分です。雪崩効果が働くからです。入力コンテキストが変化すれば、LLMの出力も揺さぶられ、その結果、文章の別のバージョンを生成しやすくなります。

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

免責事項：ここで説明するアルゴリズムは、科学的な検証を受けていません。「最悪のトークンに上限を設ける」性質を持つさまざまなサンプリングアルゴリズムを数日間試した結果、適用しやすさ、結果、理解しやすさのバランスでは、これが最も良さそうだと判断しました。

ここで説明するアルゴリズムは、くだけて言えば次のようになります。

LLMが特定の候補に強く偏っているときは、その候補を選ぶ。
実行可能な候補が複数あるときは、バリエーションを生成する。
選択される可能性がある最悪のトークンを、一定の範囲内に制限する。

Tail Free Samplingのような過去の研究でも、LLMが出力した質の高い少数のトークンから選択すべきだと指摘されています。ただしTFSでは、微分を使って、急なカーブの後にトークンの質が大きく低下しないトークン群に相当するクラスターを特定し、その集合を決めます。

これに対して、ここで提案するアルゴリズムでは、最も高いスコアを持つトークンT0を基準に、より上限が明確で理解しやすいカットオフに従って選択を行います。T0には他のすべてのトークンとは異なる特別な意味があります。それは、そのトークンを出力するときにLLMが持つ確信の度合い、つまり基本的には困惑度の代替指標です。したがってこのアルゴリズムでは、T0と比較して一定の割合を超えて悪いトークンをすべて拒否します。

0から1までの値を取るカットオフの割合を「co」と呼びます。実用的な「co」の例としては0.5が考えられます。

アルゴリズムは次のように動作します。

logitに対してsoftmax()を計算する。
トークンを確率順に並べ替える。
最良のトークンの確率をT0として、他のすべてのトークンについて次の比率を計算する：r = 1 - (T[i] / T0)
r <= coを満たすトークンだけを選ぶ。
選ばれたトークンの中から、重み付きランダム選択を行う。

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

実用例

nucleusのtop-pサンプリングが、実際には壊滅的な失敗をしているように見えない理由の一つは、最初のトークンの確率が非常に高いことが多いからです。つまり困惑度が低いときには、たまたま質の低いトークンを集めて選ぶこともなく、生成は妥当な経路に沿って続きます。

次のような確率のトークンが連続する場合には、問題がより明らかになります。

0.25, 0.14, 0.01

p=0.4の場合、3番目の質の低いトークンまで集められ、約3%の確率で選ばれる可能性があります。

ここでcoの値を0.5（トークンは最初のトークンより最大50%悪くてもよい）にしたFirst Token Cutoffを考えてみましょう。

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例です。

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.
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.
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のようなパラメータや、logitの分布形状に大きく依存するtop_pの「p」と比べても、特に考え方がシンプルです。

今後の課題

私はまだ調査を始めたばかりなので、このアルゴリズムをさらに研究し、評価していくつもりです。何よりも、サンプリングアルゴリズムへの関心が高まり、top-pに似たアプローチから先へ進むことに、もっと関心が集まってほしいと思っています。

logitの分布からは、興味深い情報を収集できる可能性があります。たとえば、線形プローブによって、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は不可解ですが、推論そのものはシンプルなプロセスなのです。