如何针对指针网络批量计算交叉熵？

我之前在做指针网络的批次训练时也碰到过一模一样的问题——填充后的logits计算损失很容易把padding部分也算进去，直接导致模型训练跑偏。结合你给出的代码片段，我来补全并解释正确的处理方案：

首先明确核心问题：指针网络输出logits维度为[batch, max_length]，但每个样本的实际有效输入长度不一致，padding位置完全不应该参与损失计算，否则会干扰模型对有效位置的学习。

1. 补全概率计算的完整逻辑 #

你已经完成了前几步的数值稳定处理，这里补全概率归一化的关键部分，确保padding位置的概率不参与归一化：

logits = stabilize(logits(inputs))  # [batch, max_length]. 减去logits的max值做数值稳定，防止exp溢出
masks = masks(inputs)  # [batch, max_length]. 有效输入位置标记为1，padding位置标记为0
exp_logits = torch.exp(logits)
exp_logits_masked = exp_logits * masks  # 直接把padding位置的exp结果置为0

# 计算归一化分母，仅对有效位置的exp结果求和
sum_exp = torch.sum(exp_logits_masked, dim=1, keepdim=True)
probs = exp_logits_masked / sum_exp  # padding位置的概率会自动变成0（分子为0）

2. 计算带Mask的交叉熵损失 #

接下来要针对目标标签做同样的mask处理，只计算有效位置的损失。这里提供两种常用实现方式：

方式一：手动计算负对数似然损失 #

假设你的目标标签targets维度为[batch]（每个样本对应输入中的一个有效索引），步骤如下：

# 生成目标的one-hot编码，维度匹配logits：[batch, max_length]
target_onehot = torch.zeros_like(logits).scatter_(1, targets.unsqueeze(1), 1)
# 仅保留有效位置的目标one-hot
target_onehot_masked = target_onehot * masks

# 计算负对数似然损失，加小epsilon避免log(0)报错
log_probs = torch.log(probs + 1e-10)
loss_per_sample = -torch.sum(target_onehot_masked * log_probs, dim=1)
# 对批次内有效样本的损失取平均
mean_loss = torch.mean(loss_per_sample)

方式二：利用PyTorch内置函数简化实现 #

用nll_loss结合mask更简洁，只需把padding位置的log概率设为极小值，同时指定忽略的目标标签：

log_probs = torch.log(probs + 1e-10)
# 把padding位置的log_probs设为-1e9，这样softmax后几乎不贡献损失
log_probs_masked = log_probs * masks + (1 - masks) * (-1e9)
# ignore_index设为你定义的padding标签值（比如-1），自动跳过无效目标的损失计算
loss = torch.nn.functional.nll_loss(log_probs_masked, targets, ignore_index=-1)

3. 必须注意的细节 #

• 数值稳定性不能省：你做的stabilize（减去logits的max值）是关键，否则exp计算很容易出现数值溢出，导致后续概率计算完全失效。
• Mask要前后一致：输入的mask标记必须和目标标签的padding标记对应，比如输入padding用0，目标中无效样本的标签就要设为ignore_index指定的值（比如-1）。
• 避免分母为0：计算probs时要确保每个样本至少有一个有效输入位置，这一步要在数据预处理阶段就做好校验。

我当初踩过的坑就是一开始没处理mask，直接计算损失，结果模型一直在学预测padding位置，训练曲线完全乱掉，加上mask后立刻就恢复正常了。

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