如何用 LLM 构建自己的决策模型
Build your own decision model
作者以 Qwen/Qwen3-1.7B 为例,演示通过约束解码把模型输出限制在固定选项(A-E),单次前向即可选出概率最高的答案;在 CommonsenseQA 随机样本上准确率为 0.5938,简单微调后提升到 0.6241。
原文给出从约束解码到温度校准的完整代码与评测,读者可以照搬这套方法把任意 LLM 改成固定选项的决策模型。
"系统一" 决策模型是一类能够以经过校准的概率为每个被允许的回答进行推断,并据此作答的模型。
以你日常使用的语言模型为例,如果你想从它那里获取结构化类型的输出(例如 JSON),可以使用 Structured Output 把模型的输出约束为保证合法的 JSON。虽然模型只需一次前向就能完成输入的预填充,但为了生成一个完整的有效响应,它仍然必须为每个 token 再各进行一次前向。
在本例中,生成最终输出需要进行 11 次前向。(此处尚未考虑推测解码以及其他推理优化技术。)
提示(预填充) 已生成(已预测) 正在预测
诸如 Jev 这类决策模型会做这样一个假设:存在一组固定的选项供我们挑选,而且只需一次前向即可迅速完成选择。在本例中,我们把可能的输出集合限制为选项 A、B、C、D、E。通过把词表中的其他条目屏蔽掉,模型就只能输出这些 token。选择概率最大的输出,即得到我们的答案。
提示(预填充) 已生成(已预测) 正在预测
由于输出被限制在一组固定的选项内,模型无法选择这些选项之外的任何内容。但这并不能保证输出结果就是正确的。在这种场景下,把输出 token 的概率视作置信度分数也是常见做法;不过,如果不进行额外的训练,这些分数反映的很可能是模型对“下一个 token 是什么”的把握,而非“该响应是正确答案”的真实概率。
动手搭建
我们可以通过使用大语言模型来约束输出 token,从而复现这种行为。这里我使用的是 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()
# load the tokenizer and the model
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(
model_name,
torch_dtype="auto",
device_map="auto"
)
# the token the model would emit for each option as the first assistant token
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]
# constrained decoding: only the option tokens are allowed
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%。这说明模型整体上对自己的预测过于自信。
既然我们的目标是让模型输出的分数能够真实反映它的准确率,那么一种可用的校准方法就是温度缩放。通过调节温度参数,我们可以压平其输出概率分布曲线,并将其缩放到近似其准确率的程度。
通过曲线拟合将温度参数拟合到模型的准确率,我得到 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 仓库,里面有脚本一步步带你完成数据集构建、评估、微调以及校准你自己的模型。建议你把它拉下来,在其他更大的模型上试试看。
来源:Hacker News · nishtahir.com