GuidanceのReadmeを読んでみると、色々ある中の一つにToken healingについて触れられています。この概念自体にあまり馴染みがなく、GPT-4にも聞きながら理解に努めましたが、よく分かりませんでした。

ただリポジトリをよく見ていくと、このToken healingの理解に繋がる興味深いnotebookを発見したので、本日はそちらについて日本語でまとめていきたいと思います。

プロンプトデザインの技術：プロンプトの境界とトークンヒーリング

これは、 Marco Tulio Ribeiro と共同で執筆されたプロンプトデザインの技術に関するシリーズの第2部です（第1部はこちら）。この記事では、大規模言語モデル（LLM）を指導する方法について説明します。

この投稿では、言語モデルが使用する貪欲なトークン化方法が、プロンプトに意図しないトークン分割を導入し、不可解な生成物が現れることがあることを説明します。

言語モデルは、生のテキストではなく、トークン（よく一緒に現れるテキストの塊）で訓練されます。これは言語モデルがテキストをどのように「見る」か、プロンプトも含めてどのように理解するかに影響を与えます。GPTスタイルのモデルは、Byte Pair Encoding（BPE）などのトークン化方法を使っていますが、これはすべての入力バイトを貪欲にトークンIDにマップする方法です。これは訓練には問題ありませんが、推論時に微妙な問題が生じることがあります。

以下に、プロンプトの境界問題の例を挙げます。

次の例では、HTTP URL 文字列を生成しようとしています。

import guidance

# 例としてStableLMを使用しますが、これらの問題はすべてのモデルにさまざまな程度で影響します。
guidance.llm = guidance.llms.Transformers("stabilityai/stablelm-base-alpha-3b", device=0)

# トークンヒーリング（後で説明）をオフにすることで、guidanceが通常のプロンプトライブラリのように動作するようにします
program = guidance('''The link is <a href="http:{{gen max_tokens=10 token_healing=False}}''')
program()
The link is <a href="http: //www.google.com/search?q

LLMが生成した出力は、URLを完了するための明らかな次の文字列（2つのスラッシュ）ではなく、URLの途中にスペースが含まれる無効なURLを生成します。これは驚くべきことですが、「http:」の後には「//」が続くことが明らかであるにもかかわらず、URLが生成されません。

なぜこれが起こるのかを理解するために、まず、プロンプトの境界を変更して、コロン文字を含まないようにしてみましょう。

guidance('''The link is <a href="http{{gen max_tokens=10 token_healing=False}}''')()
The link is <a href="http://www.youtube.com/v/s

今度は、言語モデルが期待通りに有効なURL文字列を生成しています。

コロンを含むプロンプトとコロンを含まないプロンプトの違いを理解するために、プロンプトのトークン化された表現を見てみましょう。以下のように、ここではコロンで終わるプロンプトのトークン化表現を示します（コロンがないプロンントは、最後のトークンを除いて同じトークン化を持っています）：

def print_tokens(tokens):
    print("len = " + str(len(tokens)))
    for i in tokens:
        print(str(i) + "\t`" + guidance.llm.decode([i]) + "`")

print_tokens(guidance.llm.encode('The link is <a href="http:'))
len = 9
510	  `The`
3048  ` link`
310	  ` is`
654	  ` <`
66	  `a`
3860  ` href`
568	  `="`
2413  `http`
27	  `:`

有効なURLのトークン化がどのように見えるかに注意して、特に「http:」の後にあるトークン1358に注意してください。

print_tokens(guidance.llm.encode('The link is <a href="http://www.google.com/search?q'))
len = 18
510	  `The`
3048  ` link`
310	  ` is`
654	  ` <`
66	  `a`
3860  ` href`
568	  `="`
2413  `http`
1358  `://`
2700  `www`
15	  `.`
9906  `google`
15	  `.`
681	  `com`
16	  `/`
8716  `search`
32	  `?`
82	  `q`

この特定のLLMでは、常に可能な限り最も長いトークンを優先する貪欲なトークン化方法が使われています。つまり、:// が常にテキスト全体の : に優先されます。

しかし、訓練中のURLがトークン1358（://）でエンコードされる一方で、プロンプトはLLMにトークン27（:）を見せています。このため、:// を分割することで、予測が妨げられます。実際、モデルは、トークン27（:）を見た時に、それに続くものが「より長いトークン」（例えば://）でエンコードできたものではないことをかなり確信しています。なぜなら、モデルの訓練データではそれらの文字がコロンと一緒にエンコードされていたからです。

プロンプトの境界が重要であることを忘れやすくなりますが、そのことを気に留めておくことは重要です。

モデルの語彙にあるすべてのトークンの文字列表現を検索し、コロンで始まるものを見つけると、次のようになります。

print_tokens(guidance.llm.prefix_matches(":"))
len = 34
27	`:`
21610	`:／`
1358	`://`
1450	`::`
41210	`::::`
5136	`:"`
46064	`:")`
18031	`:"){`
49777	`:",`
27506	`:*`
6098	`:**`
48471	`:**]{}`
8048	`:\`
10477	`:(`
13522	`:＝`
25942	`:=\`
18459	`:#`
19282	`:</`
21382	`:[`
22314	`:-`
42841	`:--`
22426	`:'`
23338	`:_`
25731	`:@"`
27976	`:%`
30337	`:``
34417	`:]`
35490	`:$`
47279	`:$$\`
37731	`:)`
41924	`:{`
46186	`:{\`
43118	`:.`
44662	`:&`

コロンで始まるトークンは34個あります。そのため、プロンプトの最後にコロンがある場合、モデルはこれら34のトークン文字列のいずれも生成しないことが期待されます。この微妙で強力なバイアスは、意図しない結果をもたらす可能性があります。そして、これは単にコロン（:）に限定されたことではなく、より長いトークンに拡張される可能性のあるあらゆる文字列に適用されます。「http」で終わる「修正された」プロンプトもまた、バイアスが組み込まれています。なぜなら、それはモデルに「http」の後に続くものが「s」である可能性が低いことを伝えているためです。

print_tokens(guidance.llm.prefix_matches("http"))
len = 2
2413	`http`
3614	`https`

これがトークン境界の問題が影響する URL に関連する限定的な問題だという誤解が生じやすいですが、ほとんどのトークン化器は、空白、句読点、引用符などで始まるトークンとその他のトークンとを異なる方法で扱います。その結果、プロンプトの終わりにこれらの要素を含めると、トークン境界が間違った場所に入り込み、生成物が壊れることがあります。

# 誤ってスペースを追加すると、生成がおかしくなります。
program = guidance('''I read a book about {{gen max_tokens=5 token_healing=False temperature=0}}''')
program()
I read a book about ~~the~~ the history
# スペースがない場合、予想通りに機能します。
program = guidance('''I read a book about{{gen max_tokens=5 token_healing=False temperature=0}}''')
program()
I read a book about the history of the New

また、「[」文字を含むプロンプトと生成物を見てみると、次のようになります。

guidance('''An example ["like this"] and another example [{{gen max_tokens=10 token_healing=False}}''', caching=False)()
An example ["like this"] and another example [like this] are shown in FIG. 1.

なぜ2つ目の文字列が引用符で囲まれていないのですか？プロンプトの最後が「[」トークンで終わったため、モデルは次に生成される 27 の長いトークンのいずれも、生成物と一致しないものと考えます。「15640」のようなトークンは引用符を追加しますが、そこでカットされます。

print_tokens(guidance.llm.prefix_matches(" ["))
len = 27
544	    ` [`
1008	` [@`
3921	` [*`
4299	` [**`
23734	` [****,`
8168	` []`
24345	` [],`
26991	` [];`
27501	` []{`
8605	` [[`
44965	` [[*`
14412	` ['`
15640	` ["`
16731	` [$`
20629	` [$\`
21810	` [(`
49824	` [(\[`
21938	` […]`
24430	` [\`
27075	` [^`
28591	` [-`
31789	` [...]`
33440	` [{`
42989	` [_`
43521	` [<`
44308	` [``
49193	` [#`

トークンの境界バイアスは、どこにでも存在します。上記の StableLM モデルの 10k の最も一般的なトークンのうち、70% 以上が長いトークン化可能なトークンの接頭辞であり、プロンプトの最後のトークンとして表示された場合、トークン境界バイアスを引き起こします。プロンプトデザインの際にすべての可能な拡張バイアスを管理するのは現実的ではないので、ほとんどの人はそれらを無視しています。

# より長い拡張子があるトークンの数をカウントします
count = 0
for i in range(10000):
    m = guidance.llm.prefix_matches(guidance.llm.decode([i]))
    if len(m) > 1:
        count += 1
print(str(100 * count / 10000) + "%")
70.23%

意図しないバイアスを「トークンヒーリング」で修正

これらの意図しないバイアスを回避するにはどうすればいいでしょうか？提案の一つは、プロンプトを必ず拡張できないトークンで終わらせることですが、これは大幅な制限となります。

その代わりに、guidance には「トークンヒーリング」という機能があります。これは、プロンプトの終わりの1つ前のトークンに逆戻りして生成プロセスを開始し、生成された最初のトークンがプロンプトの最後のトークンと一致する接頭辞を持つように制約します。URLの例で言えば、コロンを削除し、最初のトークンの生成にコロンの接頭辞を持つように強制することです。

トークンヒーリングを使用すると、ユーザーはトークンの境界を心配することなく、プロンプトを思い通りに表現できます。

例えば、上記のURLの例をトークンヒーリングが有効になった状態（Transformerモデルではデフォルトで有効になっているため、token_healing=Falseを削除）で再度実行してみましょう。

# トークンヒーリングを使用すると、プロンプトがコロンで終わっても有効なURLが生成されます。
guidance('''The link is <a href="http:{{gen max_tokens=10}}''')()
The link is <a href="http://www.youtube.com/v/s_

同様に、余分なスペースがあっても問題ありません。

# 誤ってスペースを追加しても生成に影響を与えません。
program = guidance('''I read a book about {{gen max_tokens=5 temperature=0}}''')
program()
I read a book about the history of the New Orleans
# これは上記と同じテキストを生成します。
program = guidance('''I read a book about{{gen max_tokens=6 temperature=0}}''')
program()
I read a book about the history of the New Orleans

また、「[」文字でプロンプトが終わっている場合でも、引用符で囲まれた文字列が得られます。

guidance('''An example ["like this"] and another example [{{gen max_tokens=10}}''', caching=False)()
An example ["like this"] and another example ["like that"] are used to illustrate the invention.

部分単語の正規化はどうですか？

言語モデルの訓練方法に精通している場合は、部分単語の正規化がどのように関与しているか疑問に思うかもしれません。部分単語の正規化は、訓練中に最適でないトークン化をランダムに導入することで、モデルがトークン境界の問題に対してより頑健になるようにする手法です。

これは、モデルが常に最適な貪欲なトークン化を見るわけではないことを意味します。部分単語の正規化は、モデルがトークン境界に対してより頑健になるのに役立ちますが、モデルが標準的な貪欲なトークン化に対するバイアスを除去するわけではありません。これは、訓練中の部分単語の正規化の量によって、モデルがトークン境界バイアスを多かれ少なかれ示すことを意味しますが、すべてのモデルがこのバイアスを持ちます。そして、上記で示したように、これはモデルの出力に強力で予期しない影響を与えることがあります。

結論

プロンプトを作成するときは、貪欲なトークン化が言語モデルがプロンプトに対する解釈に大きな影響を与えることを覚えておいてください。特にプロンプトが、それ自体がより長いトークンに拡張される可能性のあるトークンで終わる場合です。この容易に見落とされがちなバイアスの源は、予想外かつ意図しない方法で結果に影響を与えることがあります。

これを解決するために、プロンプトを拡張できないトークンで終わらせるか、guidanceの「トークンヒーリング」という機能を使って、トークン境界の破片を気にせずにプロンプトを思い通りに表現できるようにしてください。

付録：リンクの例は運が悪かっただけ？

いいえ、それは運が悪かったわけではありません。ランダムサンプリングがそれを検証しています。

# コロンがあると、ほとんど常に無効なリンクが生成されます。
program = guidance('''The link is <a href="http:{{gen 'completions' max_tokens=10 token_healing=False temperature=1.0 n=5}}''')
program()["completions"]
The link is <a href="http:
- Bookmark for Bookmark Settings </a
['\n- Bookmark for Bookmark Settings </a',
 '\\\\www.substance.gov.uk\\\\',
 ' //www.facebook.com/pages/P',
 ' //www.parentrepublic.org/releases',
 ' //www.example.com/en/reg']

# コロンがないと、常に有効なリンクが生成されます。
program = guidance('''The link is <a href="http{{gen 'completions' max_tokens=10 token_healing=False temperature=1.0 n=5}}''')
program()["completions"]
The link is <a href="http://www.cpmindia.com"
['://www.cpmindia.com"',
 '://www.gizmag.com/author',
 '://www.instagram.com/oops',
 '://www.yellow-bellied-knight',
 '://www.gizmag.com/p']

所感

主にChat.Completionを利用していた自分にとってはあまり意識することのない概念でしたが、確かにCompletionのような文の途中から生成するようなケースの場合は、プロンプトの境界問題について向き合う必要がありそうだなと感じました。

ただ Guidance にはデフォルトでトークンヒーリングがTrueになっていることもあり、意識せずともこの問題についてはある程度カバーができていると思って良いのかもしれません。

ひとまずトークンヒーリングが何をしているのか、その意義についてがある程度イメージできるようになったため良かったです。