如何获取单个Token的Perplexity而非序列平均Perplexity？

问题场景 #

我已经能通过以下代码计算整句的Perplexity：

device = "cuda"
from transformers import GPT2LMHeadModel, GPT2TokenizerFast

device = "cuda"
model_id = "gpt2"
model = GPT2LMHeadModel.from_pretrained(model_id).to(device)
tokenizer = GPT2TokenizerFast.from_pretrained(model_id)
sent = 'Happy Birthday!'
input_ids = tokenizer(sent, return_tensors='pt')['input_ids']
target_ids = input_ids.clone()
outputs = model(input_ids.to(device), labels=target_ids)
ppl = torch.exp(outputs.loss)
print(ppl)
>>>tensor(1499.6934, device='cuda:0', grad_fn=<ExpBackward0>)

但我需要获取每个Token的Perplexity值，而非整句的平均结果。示例中的'Happy Birthday!'被tokenizer拆分为3个Token，按照Perplexity公式，应该得到三个值：

• 第一个Token的对数概率（基于空上下文）
• 第二个Token的对数概率（基于第一个Token）
• 第三个Token的对数概率（基于前两个Token）
  每个值取指数后就是对应Token的Perplexity。

我的尝试及问题 #

我尝试了以下代码，但结果有问题：第一个Token的Perplexity为nan，最后一个Token的结果和整句Perplexity完全相同。

import torch
from transformers import GPT2LMHeadModel, GPT2TokenizerFast

device = "cuda"
model_id = "gpt2"
model = GPT2LMHeadModel.from_pretrained(model_id).to(device)
tokenizer = GPT2TokenizerFast.from_pretrained(model_id)

sent = 'Happy Birthday!'
input_ids = tokenizer(sent, return_tensors='pt')['input_ids'].to(device)
target_ids = input_ids.clone()

# 存储每个Token的Perplexity
perplexities = []

# 逐个计算Token的Perplexity
for i in range(input_ids.shape[1]):
    output = model(input_ids[:, :i+1], labels=target_ids[:, :i+1])
    log_prob = output.loss.item()
    perplexity = torch.exp(torch.tensor(log_prob))
    perplexities.append(perplexity.item())

# 打印结果
for i, token in enumerate([tokenizer.decode(i) for i in input_ids[0]]):
    print(f"Token: {token}, Perplexity: {perplexities[i]}")
>>> Token: Happy, Perplexity: nan
Token:  Birthday, Perplexity: 54192.46484375
Token: !, Perplexity: 1499.693359375

问题原因及解决方案 #

问题根源 #

1. 第一个Token出现nan：GPT2LMHeadModel的loss计算逻辑是自动移位的——用前N-1个Token的logits预测后N-1个Token。当输入序列长度为1时，没有有效的预测目标，导致loss为nan。
2. 最后一个Token与整句结果相同：当输入完整序列时，模型计算的loss是后N-1个Token的平均负对数概率，exp后得到的是整句的平均Perplexity，并非第三个Token单独的结果。

正确实现代码 #

要计算单个Token的Perplexity，需要手动计算每个Token的负对数概率，再取指数：

import torch
from transformers import GPT2LMHeadModel, GPT2TokenizerFast

device = "cuda" if torch.cuda.is_available() else "cpu"
model_id = "gpt2"
model = GPT2LMHeadModel.from_pretrained(model_id).to(device)
tokenizer = GPT2TokenizerFast.from_pretrained(model_id)

sent = 'Happy Birthday!'
input_ids = tokenizer(sent, return_tensors='pt')['input_ids'].to(device)
target_ids = input_ids.clone()

perplexities = []
seq_len = input_ids.shape[1]

with torch.no_grad():
    for t in range(seq_len):
        if t == 0:
            # 计算空上下文下第一个Token的概率（用bos_token作为初始输入）
            bos_input = torch.tensor([[tokenizer.bos_token_id]]).to(device)
            bos_logits = model(bos_input).logits
            # 取bos之后预测第一个Token的logits
            log_prob = -torch.nn.functional.log_softmax(bos_logits[0, 0, :], dim=-1)[target_ids[0, t]]
        else:
            # 基于前t个Token的前缀，预测当前Token的概率
            current_input = input_ids[:, :t]
            current_logits = model(current_input).logits
            # 最后一个位置的logits对应预测下一个Token（即当前目标Token）
            log_prob = -torch.nn.functional.log_softmax(current_logits[0, -1, :], dim=-1)[target_ids[0, t]]
        
        # 单个Token的Perplexity是exp(log_prob)
        ppl = torch.exp(log_prob).item()
        perplexities.append(ppl)

# 打印每个Token的结果
tokens = [tokenizer.decode(tok) for tok in input_ids[0]]
for token, ppl in zip(tokens, perplexities):
    print(f"Token: {token}, Perplexity: {ppl:.4f}")

代码说明 #

• 第一个Token：用GPT2的bos_token作为初始输入，计算模型预测该Token的概率
• 后续Token：输入前t个Token的前缀，取模型最后一个位置的logits，计算预测当前Token的概率
• 每个Token的Perplexity为对应负对数概率的指数值，完全符合Perplexity的定义

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