3. 损失函数基石
聊到知识蒸馏,绕不开两个核心概念:KL散度和交叉熵。很多初学者一上来就被这两个公式吓住了。其实没那么玄乎,说白了它们就是衡量「两个概率分布之间差多远」的尺子。
我在项目中见过不少同学,把蒸馏损失直接写成交叉熵,结果学生模型学得跟老师一模一样,反而丢了泛化能力。嗯,这里有个关键点——蒸馏损失的标准形式,其实是在交叉熵的基础上加了一个温度系数。我们一步步拆开看。
3.1 从信息论说起:熵、交叉熵、KL散度
先回顾一下基础概念。信息论里,熵衡量的是一个分布的不确定性。交叉熵则衡量用模型分布Q去拟合真实分布P时,所需的平均编码长度。KL散度是两者之差:
KL(P || Q) = H(P, Q) - H(P)
换句话说,KL散度 = 交叉熵 - 熵。当P是真实分布时,最小化交叉熵等价于最小化KL散度,因为H(P)是常数。
但在蒸馏场景里,我们的目标分布是老师模型的软标签,不是真实硬标签。这时候直接用交叉熵,其实是在强迫学生模仿老师的「确定性」——这恰恰是我们要避免的。
核心理解:蒸馏损失的本质,是让学生分布去匹配老师分布的「软性结构」,而不是硬生生复制老师的预测结果。
3.2 温度系数:软化概率分布的关键
为什么需要温度?你想想看,老师模型输出的概率分布,通常是一个接近one-hot的尖锐分布。比如猫、狗、车三类,老师可能输出[0.98, 0.01, 0.01]。这种分布里,类别间的关系信息几乎丢失了。
Hinton在2015年的那篇经典论文里提出:在softmax之前除以一个温度系数T,把分布「软化」:
q_i = exp(z_i / T) / Σ_j exp(z_j / T)
T越大,分布越平滑。T=1就是标准softmax。T趋于无穷时,分布趋于均匀。T趋于0时,分布趋于one-hot。
我个人习惯把T设置在2-8之间。T太小,学生学不到类别间的相似关系;T太大,分布太平坦,梯度信号太弱,训练效率低。
| 温度T | 分布特点 | 适用场景 |
|---|---|---|
| T=1 | 标准softmax,分布尖锐 | 硬标签训练 |
| T=2~4 | 适度软化,保留类别间关系 | 大多数蒸馏任务 |
| T=5~8 | 分布较平滑,强调结构信息 | 类别数多、相似类别多的任务 |
| T>10 | 接近均匀分布,梯度微弱 | 极少使用,除非特殊设计 |
经验之谈:我曾经在ImageNet子集上做蒸馏,T=3时学生准确率比T=1高了2.3个点。但换到CIFAR-100上,T=5效果更好。建议你从T=4开始调,观察验证集上的KL散度变化。
3.3 蒸馏损失的标准形式
蒸馏损失由两部分组成:
- 软损失(蒸馏损失):学生和老师在高温下的KL散度
- 硬损失(学生损失):学生预测与真实标签的交叉熵
标准公式如下:
L_distill = α * T² * KL(q_s^T || q_t^T) + (1-α) * CE(q_s^1, y_true)
其中:
q_s^T和q_t^T分别是学生和老师在温度T下的软化分布q_s^1是学生在T=1下的原始分布α是平衡系数,控制软硬损失的权重T²是梯度缩放因子,保证不同温度下的梯度量级一致
为什么要乘T²?因为KL散度对logits的梯度与1/T²成正比。乘上T²后,梯度量级与T无关,方便调参。
注意:很多开源代码里漏掉了T²这个缩放因子。如果你发现蒸馏效果不稳定,先检查这里。我曾经踩过这个坑,调了两天参数才发现是缩放因子的问题。
3.4 代码实现:从零搭建蒸馏损失
下面是一个PyTorch风格的蒸馏损失实现。我习惯把它封装成一个类,方便复用:
class DistillationLoss(nn.Module):
def __init__(self, T=4.0, alpha=0.7):
super().__init__()
self.T = T
self.alpha = alpha
self.kl_loss = nn.KLDivLoss(reduction='batchmean')
self.ce_loss = nn.CrossEntropyLoss()
def forward(self, student_logits, teacher_logits, labels):
# 软损失:高温下的KL散度
s_soft = F.log_softmax(student_logits / self.T, dim=1)
t_soft = F.softmax(teacher_logits / self.T, dim=1)
loss_soft = self.kl_loss(s_soft, t_soft) * (self.T ** 2)
# 硬损失:标准交叉熵
loss_hard = self.ce_loss(student_logits, labels)
# 加权组合
return self.alpha * loss_soft + (1 - self.alpha) * loss_hard
这段代码里,我用了KLDivLoss配合log_softmax和softmax的组合。注意reduction='batchmean',它会按batch大小取平均,而不是求和。这样损失值不会随batch size变化。
3.5 蒸馏损失的知识结构图
下面这张图梳理了蒸馏损失的核心逻辑。你可以看到,损失函数由两条路径汇合而成:
3.6 常见问题与避坑指南
我在实际项目中遇到过几个典型问题,分享出来帮你少走弯路:
- 梯度爆炸:当T很小时,KL散度的梯度会非常大。建议T不要低于2,同时配合梯度裁剪。
- α选择:α=0.7是个不错的起点。如果学生模型容量小,适当增大α(比如0.9),让老师多带一带。如果学生容量大,α可以降到0.5。
- 温度与α的耦合:T和α不是独立的。T越大,软损失越平滑,α可以适当调大。我一般先固定T=4,调好α,再微调T。
- 老师logits的预处理:如果老师logits的数值范围很大(比如几百),建议先做LayerNorm或标准化,否则KL散度会不稳定。
调试技巧:训练过程中,同时监控软损失和硬损失的变化。如果软损失下降很快但硬损失不动,说明学生过度模仿老师,忽略了真实标签。这时可以降低α或增大T。
3.7 小结
蒸馏损失的核心就三件事:KL散度衡量分布差异、温度系数软化分布、加权组合平衡软硬损失。理解了这三者的关系,你就能灵活设计自己的蒸馏损失函数了。