Llama-2-7b-chat-hf生成Token概率显示100%异常问题排查

问题背景 #

用户希望查看Llama-2-7b-chat-hf针对提示语Could you give me 3 cities located in Europe ?生成每个token的真实概率，以评估模型置信度，编写了如下代码：

from transformers import LlamaForCausalLM, LlamaTokenizer, AutoModelForCausalLM, AutoTokenizer
import torch
import numpy as np

device = "cuda:0" if torch.cuda.is_available() else "cpu"

model=LlamaForCausalLM.from_pretrained("Llama-2-7b-chat-hf").to(device)
tokenizer= LlamaTokenizer.from_pretrained("Llama-2-7b-chat-hf")

prompt = "Could you give me 3 cities located in Europe ?"

inputs = tokenizer([prompt], return_tensors="pt").to(device)

outputs=model.generate(**inputs,return_dict_in_generate=True, output_scores=True,max_new_tokens=75)

transition_scores = model.compute_transition_scores(outputs.sequences, outputs.scores, normalize_logits=True)

input_length = 1 if model.config.is_encoder_decoder else inputs.input_ids.shape[1]
generated_tokens = outputs.sequences[:,input_length:]

for tok, score in zip(generated_tokens[0], transition_scores[0]):
        # | token | token string | logits | probability
            print(f"| {tok:5d} | {tokenizer.decode(tok):8s} | {score.numpy(force=True):.4f} | {np.exp(score.numpy(force=True)):.2%}")

运行后输出多数token的选中概率为100%，显然不符合逻辑（Llama-2词表含32000个token，合理概率不应为100%），因此判断代码存在问题，请求定位错误。

问题分析 #

认同该判断，代码核心错误在于compute_transition_scores的参数设置：

• 错误根源：当normalize_logits=True时，函数会返回归一化后的对数概率——即当前生成token的对数概率，减去所有token对数概率的logsumexp值。对于模型选出的概率最高的token，这个归一化后的值会趋近于0，指数化后就显示为100%，但这并非该token在原始概率分布中的真实占比，而是相对归一化的结果。

修正方案 #

将normalize_logits设为False，即可获取模型输出的原始对数概率，指数化后得到该token在整个词表中的真实概率占比。修正后的代码如下：

from transformers import LlamaForCausalLM, LlamaTokenizer
import torch
import numpy as np

device = "cuda:0" if torch.cuda.is_available() else "cpu"

model = LlamaForCausalLM.from_pretrained("Llama-2-7b-chat-hf").to(device)
tokenizer = LlamaTokenizer.from_pretrained("Llama-2-7b-chat-hf")

prompt = "Could you give me 3 cities located in Europe ?"

inputs = tokenizer([prompt], return_tensors="pt").to(device)

# 生成时保留输出分数
outputs = model.generate(**inputs, return_dict_in_generate=True, output_scores=True, max_new_tokens=75)

# 关闭归一化，获取原始对数概率
transition_scores = model.compute_transition_scores(outputs.sequences, outputs.scores, normalize_logits=False)

input_length = inputs.input_ids.shape[1]
generated_tokens = outputs.sequences[:, input_length:]

# 打印格式化结果
print("| token | token string | logits | probability")
for tok, score in zip(generated_tokens[0], transition_scores[0]):
    prob = np.exp(score.numpy(force=True))
    print(f"| {tok:5d} | {tokenizer.decode(tok):8s} | {score.numpy(force=True):.4f} | {prob:.2%}")

另一种直接从outputs.scores计算真实概率的方式：

# 遍历每个生成步骤的分数
for i, score_logits in enumerate(outputs.scores):
    # 对logits做softmax得到全词表概率分布
    full_probs = torch.softmax(score_logits, dim=-1)
    # 获取当前生成的token及其概率
    generated_tok = generated_tokens[0][i]
    tok_prob = full_probs[0, generated_tok].item()
    tok_str = tokenizer.decode(generated_tok)
    # 计算对数概率并打印
    log_prob = torch.log(full_probs[0, generated_tok]).item()
    print(f"| {generated_tok:5d} | {tok_str:8s} | {log_prob:.4f} | {tok_prob:.2%}")

修正效果：修正后输出的概率为该token在整个32000词表中的真实占比，不会出现100%的异常值，能准确反映模型生成该token的置信度。

内容的提问来源于stack exchange，提问作者Zerokerim