如何在HuggingFace GPT2LMHeadModel训练中加入样本权重计算损失？

核心思路 #

GPT2做下一个token预测时，每个序列的有效位置（除了最后一个token的输入）都会产生对应预测损失。我们需要先拿到每个位置的原始损失值，再将整个序列的所有位置损失统一乘以该样本的权重，最后对所有加权后的有效损失做归约（比如求平均）。

具体实现步骤 #

1. 手动计算逐位置损失
   不要依赖模型自动返回的归约后总损失，而是通过模型输出的logits自行计算：

   • 调用模型时仅传入input_ids和attention_mask，获取logits（跳过labels参数，避免模型自动计算均值损失）。
   • 构造移位后的目标labels：下一个token预测的目标是input_ids向右移位一位（即labels = input_ids[:, 1:]），对应的logits取前seq_len-1个位置的输出（logits = logits[:, :-1, :]）。
   • 使用CrossEntropyLoss并设置reduction='none'，得到每个位置的未归约损失。

2. 应用样本权重并过滤无效位置

   • 假设批次样本权重是形状为(batch_size,)的张量sample_weights，将其扩展为(batch_size, 1)以实现广播相乘，让样本权重作用于该序列的所有位置损失。
   • 结合attention_mask过滤padding位置：用attention_mask[:, 1:]（目标序列比输入短一位）屏蔽无效位置的损失，避免padding影响训练。

3. 计算最终加权损失
   将加权后的损失与有效掩码相乘，对所有有效损失值求和后除以有效位置总数，得到最终用于反向传播的损失。

代码示例 #

import torch
from torch.nn import CrossEntropyLoss
from transformers import GPT2LMHeadModel, GPT2Tokenizer

# 初始化模型与分词器
model = GPT2LMHeadModel.from_pretrained("gpt2")
tokenizer = GPT2Tokenizer.from_pretrained("gpt2")
tokenizer.pad_token = tokenizer.eos_token  # GPT2默认无pad token，需指定

# 定义损失函数，关闭自动归约
loss_fct = CrossEntropyLoss(reduction='none')

# 示例数据：2个样本，对应权重分别为0.8和1.2
texts = ["Hello world", "I love transformers"]
sample_weights = torch.tensor([0.8, 1.2])

# 数据预处理
inputs = tokenizer(texts, return_tensors="pt", padding=True)
input_ids = inputs["input_ids"]
attention_mask = inputs["attention_mask"]

# 模型前向传播获取logits
outputs = model(input_ids=input_ids, attention_mask=attention_mask)
logits = outputs.logits

# 构造移位后的logits、labels和掩码
shifted_logits = logits[:, :-1, :].contiguous()
shifted_labels = input_ids[:, 1:].contiguous()
shifted_attention_mask = attention_mask[:, 1:].contiguous()

# 计算逐位置损失
per_position_loss = loss_fct(shifted_logits.view(-1, shifted_logits.size(-1)), shifted_labels.view(-1))
per_position_loss = per_position_loss.view(shifted_labels.size())

# 应用样本权重
weighted_loss = per_position_loss * sample_weights.unsqueeze(1)

# 过滤padding位置损失
weighted_loss = weighted_loss * shifted_attention_mask

# 计算最终损失：有效损失的平均值
final_loss = weighted_loss.sum() / shifted_attention_mask.sum()

# 反向传播更新参数
final_loss.backward()

关键注意事项 #

• 样本权重必须与输入批次的样本顺序严格对应，确保权重正确匹配到对应样本的所有位置损失。
• 必须结合attention_mask处理padding，否则无效位置的损失会干扰模型训练。
• 若使用Dataset和DataLoader加载数据，只需在__getitem__中返回样本权重，训练循环中取出即可复用上述逻辑。

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