| 102 | return kl_loss |
| 103 | |
| 104 | def TLSD_loss(self, labels, student_logits, teacher_logits): |
| 105 | shift_logits = student_logits[..., :-1, :].contiguous() |
| 106 | tc_shift_logits = teacher_logits[..., :-1, :].contiguous() |
| 107 | |
| 108 | # Step 1. get per-token ce loss with teacher logits |
| 109 | tc_shift_labels = labels[..., 1:].contiguous().to(labels.device) |
| 110 | tc_loss_all = self.ce_loss_none(tc_shift_logits.view(-1,tc_shift_logits.size(-1)), tc_shift_labels.view(-1)) |
| 111 | |
| 112 | # Step 2. get token-scale with tc_loss_all and temperatured softmax function |
| 113 | tc_all = tc_loss_all.reshape(tc_shift_logits.shape[0], -1) |
| 114 | token_scale = torch.nn.functional.softmax(tc_all / 10, dim=-1).clone().detach() |
| 115 | |
| 116 | # Step 3. logit distillation with token-scale |
| 117 | student_likelihood = torch.nn.functional.log_softmax(shift_logits, dim=-1) |
| 118 | targets_prob = torch.nn.functional.softmax(tc_shift_logits, dim=-1) |
| 119 | tsld_loss = (torch.sum((- targets_prob * student_likelihood), dim=-1) * token_scale).sum() # SUM |
| 120 | |
| 121 | return tsld_loss |
| 122 | |
| 123 | def mse_loss(self, student_logits, teacher_logits): |
| 124 | return mse_loss(student_logits, teacher_logits) |