Knowledge Distillation

import torch.nn.functional as F

def distillation_loss(
    student_logits,
    teacher_logits,
    labels,
    T: float = 4.0,
    alpha: float = 0.5,
):
    """Hàm loss distillation chuẩn Hinton 2015.

    Args:
        student_logits: output thô của student (chưa softmax), shape (B, C)
        teacher_logits: output thô của teacher, shape (B, C)
        labels:         ground truth hard labels, shape (B,)
        T:              temperature — càng cao, phân bố càng mềm
        alpha:          trọng số hard loss vs soft loss

    Trả về scalar loss để backprop.
    """
    # ---- Soft loss: KL divergence trên phân bố đã làm mềm ----
    soft_t = F.softmax(teacher_logits / T, dim=-1)
    log_soft_s = F.log_softmax(student_logits / T, dim=-1)
    soft_loss = F.kl_div(
        log_soft_s,
        soft_t,
        reduction="batchmean",
    ) * (T ** 2)  # nhân T^2 để bù gradient

    # ---- Hard loss: cross-entropy với ground truth ----
    hard_loss = F.cross_entropy(student_logits, labels)

    return alpha * hard_loss + (1 - alpha) * soft_loss


# ---- Vòng lặp huấn luyện ngắn gọn ----
for batch in loader:
    x, y = batch
    with torch.no_grad():
        t_logits = teacher(x)         # teacher đóng băng
    s_logits = student(x)
    loss = distillation_loss(s_logits, t_logits, y, T=4.0, alpha=0.5)

    optimizer.zero_grad()
    loss.backward()
    optimizer.step()

import torch
import torch.nn.functional as F

class PatternDistiller(torch.nn.Module):
    """Distill cả output cuối lẫn hidden state trung gian.

    Các tầng của student và teacher được ánh xạ qua student_to_teacher.
    """

    def __init__(self, student, teacher, student_to_teacher: dict[int, int]):
        super().__init__()
        self.student = student
        self.teacher = teacher
        self.teacher.eval()
        for p in self.teacher.parameters():
            p.requires_grad_(False)
        self.pattern_map = student_to_teacher
        # Projection nếu kích thước hidden khác nhau
        self.proj = torch.nn.Linear(
            student.hidden_size,
            teacher.hidden_size,
            bias=False,
        )

    def forward(self, x, labels, T=4.0, alpha=0.5, beta=0.1):
        s_out = self.student(x, output_hidden_states=True)
        with torch.no_grad():
            t_out = self.teacher(x, output_hidden_states=True)

        # 1) Logit loss (Hinton KD)
        logit_loss = (
            F.kl_div(
                F.log_softmax(s_out.logits / T, dim=-1),
                F.softmax(t_out.logits / T, dim=-1),
                reduction="batchmean",
            )
            * (T ** 2)
        )
        ce = F.cross_entropy(s_out.logits, labels)

        # 2) Hidden-state loss — MSE giữa các tầng tương ứng
        hidden_loss = 0.0
        for s_idx, t_idx in self.pattern_map.items():
            s_h = self.proj(s_out.hidden_states[s_idx])
            t_h = t_out.hidden_states[t_idx]
            hidden_loss = hidden_loss + F.mse_loss(s_h, t_h)

        return alpha * ce + (1 - alpha) * logit_loss + beta * hidden_loss