📚 课程目录第3章 Logits蒸馏实战

3、Logits蒸馏实战：基于KL散度的经典蒸馏损失实现

Logits蒸馏，说白了就是让「学生模型」去模仿「老师模型」的软输出。你想想看，老师教学生的时候，不光告诉学生正确答案是什么，还会说「这个选项虽然不对，但比那个选项靠谱一点」——这种信息，硬标签是给不了的。

我个人习惯把Logits蒸馏叫做「软标签学习」。它背后的数学工具就是KL散度。嗯，咱们今天就把这东西拆开揉碎了讲清楚。

3.1 蒸馏损失的核心公式

先看公式，别怕，就一行：

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

这里面有三个关键角色：

老师模型的Logits：未经softmax的原始输出，可以是任意实数
学生模型的Logits：同样格式，但数值分布通常更「平」
温度系数T：控制软化的程度，T越大，分布越均匀

为什么要乘T²？我在项目中遇到过这个问题——如果不乘，梯度会随着T增大而衰减，训练根本推不动。这个缩放因子是Hinton论文里明确写的，不是拍脑袋。

关键理解：蒸馏损失不是让学生输出和老师一模一样的logits，而是让学生学会老师「对不同类别的相对置信度」。

3.2 PyTorch代码逐行解析

直接上代码，我习惯把蒸馏损失封装成一个类：

import torch
import torch.nn as nn
import torch.nn.functional as F

class DistillationLoss(nn.Module):
    def __init__(self, T=4.0, alpha=0.7):
        super().__init__()
        self.T = T          # 温度系数
        self.alpha = alpha  # 蒸馏损失权重
        
    def forward(self, student_logits, teacher_logits, labels):
        # 1. 计算蒸馏损失（KL散度）
        # 用log_softmax对学生输出做对数概率
        s_log_prob = F.log_softmax(student_logits / self.T, dim=1)
        # 用softmax对老师输出做概率分布
        t_prob = F.softmax(teacher_logits / self.T, dim=1)
        # KL散度：sum(p * log(p/q))
        kl_loss = F.kl_div(s_log_prob, t_prob, reduction='batchmean')
        # 乘以T²恢复梯度尺度
        kl_loss = kl_loss * (self.T ** 2)
        
        # 2. 计算硬标签损失（交叉熵）
        ce_loss = F.cross_entropy(student_logits, labels)
        
        # 3. 加权融合
        total_loss = self.alpha * kl_loss + (1 - self.alpha) * ce_loss
        
        return total_loss, kl_loss, ce_loss

逐行拆解一下：

第10行：学生logits除以T后做log_softmax。注意是log_softmax，因为KL散度的第一个参数要求是对数概率
第12行：老师logits除以T后做普通softmax。老师不需要梯度，所以不用log
第14行：reduction='batchmean'会取batch的平均，这是KL散度的标准做法
第16行：T²缩放。如果不乘，T=8时梯度会缩小64倍，训练基本废了
第22行：alpha控制蒸馏损失和硬标签损失的混合比例

我的习惯：返回三个loss值方便调参时观察。如果kl_loss突然飙升，说明学生和老师的分布差异在变大，可能是温度设太高了。

3.3 温度系数与损失权重的协同调优

这两个参数不是独立的。我调过几十组参数，总结出一个规律：温度越高，蒸馏损失的权重应该越大。

为什么会这样？因为高温下softmax输出接近均匀分布，KL散度的值会变小。如果不加大alpha，蒸馏损失就被硬标签损失淹没了。

温度T	推荐alpha	适用场景	注意事项
1~3	0.3~0.5	分类任务、老师学生能力差距小	软标签信息量少，硬标签主导
4~8	0.5~0.7	通用场景、老师比学生强很多	最常用的区间，平衡性好
8~20	0.7~0.9	细粒度分类、老师极度自信	小心过平滑，学生可能学不到细节

我曾经在一个细粒度鸟类分类项目里踩过坑。老师模型的top-1准确率92%，学生只有78%。我一开始用T=3, alpha=0.5，训了50个epoch学生才涨到81%。后来改成T=8, alpha=0.8，20个epoch就到了85%。

原因很简单：老师对相似鸟类（比如「黑枕黄鹂」和「金黄鹂」）的logits分布差异非常细微，低温下这些信息被压缩掉了。高温才能把这些「相对关系」暴露给学生。

避坑指南：我曾经试过T=20, alpha=0.95，结果学生模型的logits变得太平坦，连硬标签都学不好。温度太高会让所有类别的概率都接近，学生反而失去了判别能力。T不要超过15，alpha不要超过0.9，这是我的经验红线。

3.4 知识体系结构图

下面这张图把Logits蒸馏的完整流程画清楚了：

3.5 实战调参建议

最后给几个实战中验证过的调参策略：

先固定T，再调alpha：我习惯先设T=4，跑几个epoch看kl_loss的量级，然后让alpha使得kl_loss和ce_loss大致相等
T和alpha可以动态变化：训练初期用高T(6~8)让学生快速学习分布，后期降低T(2~3)让学生精细化
监控学生logits的方差：如果方差太小（<0.1），说明温度太高或alpha太大，学生输出太「软」了
对比实验必做：至少跑三组——无蒸馏、T=4 alpha=0.5、T=8 alpha=0.7，看哪个组合在验证集上最好

小技巧：如果你发现学生模型在蒸馏后准确率反而下降，先检查老师logits是否经过了softmax。很多人把老师logits和老师概率搞混了——蒸馏用的是logits，不是概率！

Logits蒸馏是知识蒸馏里最基础也最实用的方法。你把它吃透了，后面再学特征蒸馏、关系蒸馏就会轻松很多。嗯，今天就到这儿，动手跑跑代码比看十遍文章都管用。