| 87 | return kl_loss |
| 88 | |
| 89 | def re_loss(self, labels, student_logits, teacher_logits): |
| 90 | mask = (labels != -100) |
| 91 | |
| 92 | # Compute the log probabilities of the teacher's logits (true distribution) |
| 93 | teacher_output_log_prob = F.log_softmax(teacher_logits, dim=2) |
| 94 | |
| 95 | # Compute the softmax of the student's logits (approximate distribution) |
| 96 | student_output_soft = F.softmax(student_logits, dim=2) |
| 97 | |
| 98 | # Calculate the reverse KL Divergence (KL(teacher_logits || student_logits)) |
| 99 | kl_loss = F.kl_div(teacher_output_log_prob, student_output_soft, reduction="none") |
| 100 | kl_loss = kl_loss.sum(-1) * mask |
| 101 | kl_loss = kl_loss.sum(-1).mean() |
| 102 | return kl_loss |
| 103 | |
| 104 | def TLSD_loss(self, labels, student_logits, teacher_logits): |
| 105 | shift_logits = student_logits[..., :-1, :].contiguous() |