AI・機械学習
独自の意思決定モデルを構築する
Build your own decision model (nishtahir.com)
要約
この記事では、LLM(大規模言語モデル)を使用して、確率を推論して応答する「システム1」意思決定モデルを構築する方法を解説しています。出力トークンを固定の選択肢に制約することで、モデルの応答を制御し、その精度と信頼度を評価・調整する手法について、具体的なコード例と実験結果を交えて説明しています。
全文翻訳
「システム1」意思決定モデルとは、調整された確率または許容されるあらゆる回答で推論し応答するモデルのことです。
日常的に使用する言語モデルを考えてみましょう。モデルから(JSON形式で)型付けされた出力を得るには、Structured Outputを使用して、確実に有効なJSONに制約することができます。
モデルは入力のプリフィルを1回のパスで行いますが、有効な応答を生成するためには、トークンごとに1回のパスを実行する必要があります。この例では、最終的な出力を生成するために11回のパスが必要です。(投機的デコーディングやその他の推論最適化技術は考慮していません。)
ステップ プレイ リセット プロンプト(プリフィル) 生成済み(予測) 予測中
Jevのような意思決定モデルは、選択可能な固定オプションが存在するという仮定を立て、単一のパスで迅速にそれらを選択できるとします。この例では、可能な出力のセットをオプションA、B、C、D、Eに制約します。語彙内の他の項目をマスキングすることで、モデルはそのトークンしか出力できません。最も確率の高い出力を選択することで、回答が得られます。
ステップ プレイ リセット プロンプト(プリフィル) 生成済み(予測) 予測中
出力が固定のオプションセットに制約されているため、モデルはそれ以外のものを選択できません。しかし、これにより出力が正しいことが保証されるわけではありません。
この文脈で出力トークンの確率を信頼度スコアとして扱うことも一般的ですが、追加のトレーニングなしでは、これらのスコアは応答が正しい答えである真の確率よりも、次のトークンに対するモデルの信頼度を反映している可能性が高いです。
独自のモデルを構築する
LLMを使用して出力トークンを制約することで、この動作をエミュレートできます。ここではQwen/Qwen3-1.7Bを使用します。
import argparse
import json
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
model_name = "Qwen/Qwen3-1.7B"
options = ["A", "B", "C", "D", "E"]
parser = argparse.ArgumentParser()
parser.add_argument("--input", default="question.json")
args = parser.parse_args()
# トークナイザーとモデルをロード
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(
model_name,
torch_dtype="auto",
device_map="auto"
)
# 各オプションの最初の「アシスタント」トークンとしてモデルが出力するトークン
option_token_ids = [tokenizer.encode(opt, add_special_tokens=False)[0] for opt in options]
def format_prompt(item):
prompt = item["question"] + "\n"
for opt in options:
prompt += f"{opt}. {item[opt]}\n"
prompt += "Answer:"
messages = [
{"role": "user", "content": prompt}
]
return tokenizer.apply_chat_template(
messages,
tokenize=False,
add_generation_prompt=True,
enable_thinking=False
)
with open(args.input) as f:
item = json.load(f)
model_inputs = tokenizer(format_prompt(item), return_tensors="pt").to(model.device)
with torch.no_grad():
logits = model(**model_inputs).logits[0, -1] # 制約付きデコーディング: オプショントークンのみが許可される
probs = torch.softmax(logits[option_token_ids].float(), dim=-1)
print(f"prediction: {options[probs.argmax().item()]}")
for opt, prob in zip(options, probs.tolist()):
print(f"{opt}: {prob:.4f} {item[opt]}")
これを簡単な質問で実行すると、次の出力が得られます。
// input
{
"question": "What color is the sky?",
"A": "Red",
"B": "Blue",
"C": "Green",
"D": "Purple",
"E": "I don't know"
}
// output
prediction: B
A: 0.0000 Red
B: 0.9988 Blue
C: 0.0000 Green
D: 0.0000 Purple
E: 0.0012 I don't know
モデルは入力の意味を理解し、正しい答えに合理的に対応する予測を行うことができました。
公開データセットでモデルの精度をテストできます。CommonsenseQAのランダムサンプリングされたホールドアウトセットで実行しました。
precision recall f1 support
A 0.5733 0.7197 0.6382 239
B 0.5506 0.7686 0.6416 255
C 0.5372 0.6598 0.5922 241
D 0.7206 0.3904 0.5065 251
E 0.7519 0.4255 0.5435 235
accuracy: 725/1221 = 0.5938
macro f1: 0.5844
1.7Bモデルとしては悪くない結果です。簡単なファインチューニングを行うと、パフォーマンスがわずかに向上します。
precision recall f1 support
A 0.6475 0.6611 0.6542 239
B 0.6113 0.6784 0.6431 255
C 0.6234 0.5975 0.6102 241
D 0.6700 0.5418 0.5991 251
E 0.5808 0.6426 0.6101 235
accuracy: 762/1221 = 0.6241
macro f1: 0.6234
モデルのキャリブレーション
非常に曖昧な問題に対してモデルをテストすると、興味深い問題が明らかになります。
// input
{
"question": "Where would you most likely find a bat?",
"A": "Cave",
"B": "Baseball game",
"C": "Attic",
"D": "Zoo",
"E": "Sporting goods store"
}
// output
prediction: A
A: 0.9978 Cave
B: 0.0004 Baseball game
C: 0.0017 Attic
D: 0.0000 Zoo
E: 0.0001 Sporting goods store
ここには明確な答えがあるはずはありませんが、出力確率を疑似「信頼度」スコアとして扱うと、モデルはこの答えに極端に過信していることがわかります。
以前実行した評価で信頼度スコアの範囲をビン化すると、モデルの信頼度がその精度と一致していないことがわかります。これは、モデルがキャリブレーションされていないことを意味します。
bin count confidence accuracy
(0.00, 0.10] 0 0.0000 0.0000
(0.10, 0.20] 0 0.0000 0.0000
(0.20, 0.30] 3 0.2834 0.0000
(0.30, 0.40] 26 0.3761 0.2692
(0.40, 0.50] 41 0.4538 0.2683
(0.50, 0.60] 70 0.5490 0.3286
(0.60, 0.70] 74 0.6476 0.3649
(0.70, 0.80] 77 0.7495 0.4286
(0.80, 0.90] 121 0.8555 0.4711
(0.90, 1.00] 809 0.9855 0.7009
モデルは0.9〜1.0のビンで極端に過信している傾向がありますが、正解率は70%にすぎません。0.8〜0.9の信頼度で予測を行う場合、精度は約40%にすぎません。これは、モデルが一般的に予測に対して過信していることを意味します。
モデルの出力をその精度を反映するものにすることが目標であるため、キャリブレーションに使用できる方法の1つは、温度スケーリングです。温度値を変更することで、出力確率分布カーブを平坦化し、精度に近似するようにスケーリングできます。
temperature T Curve fitting
モデルの精度に温度パラメータを適合させると、3.797280788421631 という温度値が見つかりました。
bin count confidence accuracy
(0.00, 0.10] 0 0.0000 0.0000
(0.10, 0.20] 0 0.0000 0.0000
(0.20, 0.30] 82 0.2712 0.2317
(0.30, 0.40] 217 0.3507 0.3917
(0.40, 0.50] 199 0.4472 0.5126
(0.50, 0.60] 166 0.5475 0.5482
(0.60, 0.70] 139 0.6562 0.5827
(0.70, 0.80] 140 0.7492 0.7714
(0.80, 0.90] 169 0.8507 0.7988
(0.90, 1.00] 109 0.9333 0.9541
これにより、はるかに優れたキャリブレーションが得られます。
これを試したい場合は、データセットの構築、評価、ファインチューニング、キャリブレーションを行うスクリプトを含むGitHubリポジトリを作成しました。それをプルして、他のより大きなモデルで試すことをお勧めします。