📚 课程目录第4章蒸馏损失函数

4、蒸馏损失函数：KL散度详解、蒸馏损失与学生损失的加权融合

损失函数，说白了就是告诉模型「你学得怎么样」的那把尺子。在知识蒸馏里，这把尺子有两头——一头量学生和老师的差距，另一头量学生和真实标签的差距。怎么把这两头拧在一起，就是咱们这章要聊的核心。

我个人习惯把蒸馏损失拆成三块来理解：KL散度怎么算、蒸馏损失怎么搭、学生损失怎么加权。咱们一块一块啃。

4.1 KL散度：衡量两个分布的距离

KL散度，全称Kullback-Leibler散度，也叫相对熵。它衡量的是两个概率分布之间的「不匹配程度」。在蒸馏里，就是老师输出的软标签和学生输出的软标签之间的距离。

公式长这样：

KL(P || Q) = Σ P(x) * log(P(x) / Q(x))

其中P是老师分布，Q是学生分布。值越小，说明学生越像老师。

几个关键点：

KL散度不对称：KL(P||Q) ≠ KL(Q||P)。蒸馏里我们固定老师，优化学生，所以用KL(老师||学生)。
取值范围：[0, +∞)。完全相同时为0。
当P(x)=0时，该项为0；当Q(x)=0而P(x)>0时，该项无穷大——所以学生不能把老师认为可能的类别直接判为0概率。

我的经验： 实际项目中，我遇到过学生输出层用softmax后出现极小值（比如1e-10），导致KL散度爆炸。后来我习惯在log里加个epsilon（1e-7），稳定得多。

在PyTorch里，KL散度可以直接用torch.nn.KLDivLoss。但有个坑——它要求输入是log概率，目标（老师）是概率。代码长这样：

import torch
import torch.nn.functional as F

# 学生输出 logits
student_logits = torch.randn(4, 10)  # batch=4, 10类
# 老师输出 logits
teacher_logits = torch.randn(4, 10)

# 温度参数
T = 4.0

# 软化后的概率
student_soft = F.log_softmax(student_logits / T, dim=1)
teacher_soft = F.softmax(teacher_logits / T, dim=1)

# KL散度损失
kl_loss = F.kl_div(student_soft, teacher_soft, reduction='batchmean') * (T * T)

注意最后乘了T²。为什么？因为软化后梯度会缩小T倍，乘回来保持梯度量级。我刚开始做蒸馏时忘了这步，结果温度调高了损失就变小，调低了又变大，折腾了两天才发现是这里的问题。

4.2 蒸馏损失（Distillation Loss）的构成

蒸馏损失，就是让学生去模仿老师的软标签。它由三部分组成：

温度缩放：把logits除以温度T，让分布变「软」。
Softmax软化：得到软化后的概率分布。
KL散度计算：衡量学生和老师的差距。

写成公式就是：

L_distill = KL( softmax(teacher_logits / T) || softmax(student_logits / T) ) * T²

温度T是个超参数。T越大，分布越平滑，类别间的相对关系保留得越多。T越小，越接近硬标签。

温度T	分布特点	适用场景
T=1	接近原始softmax	类别少、任务简单
T=2~5	中等软化	大部分分类任务
T=6~10	高度平滑	类别多、需要保留细粒度关系
T>10	接近均匀分布	极少用，信息稀释严重

避坑： 我曾经在一个1000类的细粒度分类任务里，把T设成20，结果学生学出来跟随机差不多。后来发现温度太高，老师输出的信息被稀释成白噪声了。一般建议从T=4开始调，观察验证集效果。

4.3 学生损失（Student Loss）的加权融合

学生损失，就是标准的交叉熵损失，让学生去拟合真实标签：

L_student = CrossEntropy(student_logits, true_labels)

最终的总损失是两者的加权和：

L_total = α * L_distill + (1 - α) * L_student

α是蒸馏损失的权重，取值范围[0,1]。α=0时退化为普通训练，α=1时完全依赖老师。

α怎么设？ 我个人的经验是：

老师模型很强（准确率95%+），α可以设大一点，比如0.7~0.9。
老师模型一般（准确率80%~90%），α设在0.3~0.5。
老师模型弱（准确率低于80%），α设0.1~0.3，甚至考虑换老师。

完整损失函数的代码实现：

def distillation_loss(student_logits, teacher_logits, labels, T=4.0, alpha=0.7):
    # 蒸馏损失
    student_soft = F.log_softmax(student_logits / T, dim=1)
    teacher_soft = F.softmax(teacher_logits / T, dim=1)
    kl_loss = F.kl_div(student_soft, teacher_soft, reduction='batchmean') * (T * T)
    
    # 学生损失
    ce_loss = F.cross_entropy(student_logits, labels)
    
    # 加权融合
    total_loss = alpha * kl_loss + (1 - alpha) * ce_loss
    return total_loss, kl_loss, ce_loss

  核心思路： 蒸馏损失让学生学「老师怎么想」，学生损失让学生学「正确答案是什么」。两者互补，缺一不可。

4.4 知识蒸馏损失函数全景图

下面这张图把整个损失函数的构成和流程串起来了：

从图上可以清楚看到：老师和学生分别经过温度缩放和softmax，然后算KL散度得到蒸馏损失；学生同时跟真实标签算交叉熵得到学生损失；最后加权融合。

调参建议： 我一般先固定T=4，调α。α找到合适的值后，再调T。两个一起调容易陷入局部最优，浪费时间。

嗯，到这里损失函数这块就讲透了。KL散度怎么算、蒸馏损失怎么搭、学生损失怎么加权，每一步都有代码和避坑经验。你动手写代码的时候，记得先跑通单batch，再上完整训练循环。