ForwardKLLoss

class torchtune.modules.loss.ForwardKLLoss(ignore_index: int = - 100)

有效索引的 Kullback-Leibler 散度损失。 https://github.com/jongwooko/distillm/blob/17c0f98bc263b1861a02d5df578c84aea652ee65/distillm/losses.py 的实现

参数:

ignore_index (int) – 指定一个被忽略且不影响输入梯度的目标值。损失在非忽略目标上进行划分。默认值：-100。

forward(student_logits: Tensor, teacher_logits: Tensor, labels: Tensor, normalize: bool = True) → Tensor

参数:

• student_logits (torch.Tensor) – 来自学生模型的 logits，形状为 (batch_size*num_tokens, vocab_size)。

• teacher_logits (torch.Tensor) – 来自教师模型的 logits，形状为 (batch_size*num_tokens, vocab_size)。

• labels (torch.Tensor) – 真实标签，形状为 (batch_size, vocab_size)。

• normalize (bool) – 是否通过未屏蔽元素的数量来标准化损失。

返回:

形状为 (1,) 的 KL 散度损失。

返回类型:

torch.Tensor