3.5 KiB

Raw Blame History

Lesson 20 笔记：蒸馏 Loss（KL + Temperature）详解

蒸馏 Loss 完整流程

def distillation_loss(student_logits, teacher_logits, temperature=1.0):
    with torch.no_grad():
        teacher_probs = F.softmax(teacher_logits / temperature, dim=-1).detach()

    student_log_probs = F.log_softmax(student_logits / temperature, dim=-1)

    kl = F.kl_div(
        student_log_probs,
        teacher_probs,
        reduction='batchmean'
    )
    return (temperature ** 2) * kl

Step 1：teacher 分布（soft target）

teacher_probs = softmax(teacher_logits / T).detach()

做了三件事：

步骤	操作	作用
①	`logits / T`	压缩 logits 差距
②	`softmax`	变成概率分布
③	`detach`	teacher 不参与反向传播

Step 2：student 分布

student_log_probs = log_softmax(student_logits / T)

做了：

logits / T（同样变软）
log_softmax（为了计算 KL）

Step 3：KL Loss

kl = F.kl_div(student_log_probs, teacher_probs)

PyTorch 的 F.kl_div 计算的是：

KL(P_teacher || P_student) = Σ P_teacher(x) * log(P_teacher(x) / P_student(x))

Q1：为什么除 T 会变"更均匀"

本质：softmax 对 logits 差值是指数放大的，除以 T 就是在"压缩差距"

T	logits 差距	分布
1	原始	很尖（winner takes all）
> 1	被压缩	更平

直觉记忆：T 越大 → "winner takes all" 越弱

T=1:  softmax([3.0, 1.0, 0.0]) → [0.90, 0.07, 0.03]   很尖
T=2:  softmax([1.5, 0.5, 0.0]) → [0.58, 0.31, 0.11]   更平
T=10: softmax([0.3, 0.1, 0.0]) → [0.35, 0.32, 0.33]   很平

Q2：KL 为什么可能负？为什么不取绝对值

KL 不是距离，而是"用 Q 编码 P 的额外信息成本"

性质	说明
❌ 不是距离	不对称，KL(P\|\|Q) ≠ KL(Q\|\|P)
✅ 整体 ≥ 0	Jensen 不等式保证
⚠️ 单项可能为负	如果 Q(x) > P(x)，则 log(P/Q) < 0

直觉理解：

就算某些 token "Q 猜得更好"，整体编码策略仍然不可能比真实分布更优。

Step 4：为什么要乘 T²

return (temperature ** 2) * kl

原因：softmax(z/T) 会让梯度整体缩小约 1/T²

发生了两次缩放：

步骤	梯度缩放
logits / T	scale ↓ 1/T
softmax gradient	再 ↓ 1/T
合计	≈ 1/T²

所以要乘 T² 补偿，保持训练强度。

一句话总结

概念	一句话
蒸馏本质	teacher 用 softmax(z/T) 生成"软标签"，student 用 log_softmax(z/T) 拟合 teacher 分布
T 的作用	T > 1 → logits 差距压缩 → 分布更平 → 信息更丰富
KL 本质	用 Q 逼近 P 时的信息损失（额外编码成本），不是距离
乘 T²	temperature 让梯度缩小约 T² 倍，需要放大回来

3.5 KiB Raw Blame History Unescape Escape