1
# 标准 softmax（T=1）
2
probs = softmax(logits / 1.0)  # 分布尖锐
3

4
# 高温 softmax（T>1）
5
probs = softmax(logits / 5.0)  # 分布平滑，暗知识更明显
6

7
# 蒸馏损失
8
# 让 Student 的 softened output 匹配 Teacher 的 softened output
9
loss_distill = KL_divergence(
10
    softmax(student_logits / T),
11
    softmax(teacher_logits / T)
12
) * T * T  # 乘以 T² 保持梯度量级

1
def distillation_loss(student_logits, teacher_logits, labels, T, alpha):
2
    """
3
    知识蒸馏损失函数
4

5
    T: 温度参数
6
    alpha: 蒸馏损失权重（通常 0.5-0.7）
7
    """
8
    # 1. 蒸馏损失：Student 匹配 Teacher 的软标签
9
    soft_teacher = F.softmax(teacher_logits / T, dim=-1)
10
    soft_student = F.log_softmax(student_logits / T, dim=-1)
11
    loss_distill = F.kl_div(soft_student, soft_teacher, reduction='batchmean') * (T * T)
12

13
    # 2. 标准损失：Student 匹配真实标签
14
    loss_hard = F.cross_entropy(student_logits, labels)
15

16
    # 加权组合
17
    total_loss = alpha * loss_distill + (1 - alpha) * loss_hard
18
    return total_loss

1
def distilbert_loss(student, teacher, masked_tokens, T=5):
2
    """DistilBERT 的三重损失"""
3

4
    # 1. MLM 损失（标准 BERT 目标）
5
    mlm_loss = cross_entropy(
6
        student.predict_masked(masked_tokens),
7
        masked_tokens.labels
8
    )
9

10
    # 2. 蒸馏损失（软标签匹配）
11
    student_logits = student.forward(masked_tokens)
12
    teacher_logits = teacher.forward(masked_tokens)
13

14
    soft_teacher = softmax(teacher_logits / T, dim=-1)
15
    log_soft_student = log_softmax(student_logits / T, dim=-1)
16
    distill_loss = kl_div(log_soft_student, soft_teacher) * (T * T)
17

18
    # 3. 余弦嵌入损失（隐藏状态对齐）
19
    student_hidden = student.get_hidden_state()
20
    teacher_hidden = teacher.get_hidden_state()
21
    cosine_loss = 1 - cosine_similarity(student_hidden, teacher_hidden).mean()
22

23
    # 总损失
24
    total_loss = mlm_loss + distill_loss + cosine_loss
25
    return total_loss

1
# TinyBERT 还蒸馏了注意力矩阵和隐藏状态
2
def tinybert_loss(student, teacher, input_ids, labels):
3
    # 1. 蒸馏隐藏层输出
4
    hidden_loss = mse_loss(student.hidden, teacher.hidden)
5

6
    # 2. 蒸馏注意力矩阵（每层每个头）
7
    attn_loss = 0
8
    for s_attn, t_attn in zip(student.attentions, teacher.attentions):
9
        attn_loss += kl_div(s_attn, t_attn)
10

11
    # 3. 蒸馏预测层
12
    pred_loss = kl_div(
13
        softmax(student.logits / T),
14
        softmax(teacher.logits / T)
15
    )
16

17
    # 4. 真实标签损失
18
    label_loss = cross_entropy(student.logits, labels)
19

20
    return hidden_loss + attn_loss + pred_loss + label_loss

1
from transformers import DistilBertTokenizer, DistilBertForSequenceClassification
2

3
# 加载 DistilBERT
4
tokenizer = DistilBertTokenizer.from_pretrained('distilbert-base-uncased')
5
model = DistilBertForSequenceClassification.from_pretrained('distilbert-base-uncased-finetuned-sst-2-english')
6

7
# 推理
8
inputs = tokenizer("This movie is great!", return_tensors="pt")
9
outputs = model(**inputs)
10
predicted_class = outputs.logits.argmax().item()
11
print(f"Sentiment: {'Positive' if predicted_class == 1 else 'Negative'}")

1
import torch
2
import torch.nn.functional as F
3
from transformers import BertForSequenceClassification, DistilBertForSequenceClassification
4

5
# Teacher 和 Student
6
teacher = BertForSequenceClassification.from_pretrained('bert-base-uncased')
7
student = DistilBertForSequenceClassification.from_pretrained('distilbert-base-uncased')
8

9
teacher.eval()  # Teacher 冻结
10

11
def distillation_train_step(student, teacher, batch, T=5, alpha=0.5):
12
    """一步蒸馏训练"""
13
    with torch.no_grad():
14
        teacher_outputs = teacher(**batch)
15
        teacher_logits = teacher_outputs.logits
16

17
    student_outputs = student(**batch)
18
    student_logits = student_outputs.logits
19

20
    # 蒸馏损失
21
    loss_distill = F.kl_div(
22
        F.log_softmax(student_logits / T, dim=-1),
23
        F.softmax(teacher_logits / T, dim=-1),
24
        reduction='batchmean'
25
    ) * (T * T)
26

27
    # 硬标签损失
28
    loss_hard = student_outputs.loss
29

30
    # 总损失
31
    loss = alpha * loss_distill + (1 - alpha) * loss_hard
32
    return loss