📚 课程目录第2章 Logits蒸馏

2. Logits蒸馏：KL散度损失函数、温度缩放策略、蒸馏权重α的调参经验

Logits蒸馏，说白了就是让“学生模型”去模仿“老师模型”的软输出。你想想看，老师教学生，不只是给个标准答案（硬标签），更重要的是教解题思路（软概率分布）。这个思路，就是Hinton在2015年那篇经典论文里提出的核心思想。

我个人觉得，Logits蒸馏是知识蒸馏里最基础也最实用的一招。它不需要动模型内部结构，只改损失函数和输出层。今天咱们就把KL散度、温度T、权重α这三个核心参数彻底聊透。

2.1 KL散度损失函数：到底在算啥？

KL散度，全称Kullback-Leibler divergence，用来衡量两个概率分布之间的差异。在蒸馏场景里，就是老师模型的软输出和学生模型的软输出之间的“距离”。

公式长这样：

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

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

为什么不用交叉熵？ 我刚开始做蒸馏时也困惑过。后来发现，KL散度天然适合做“模仿学习”——它不对称，惩罚的是学生分布偏离老师分布的程度。交叉熵虽然数学上等价，但KL散度在梯度回传时更稳定，尤其当温度较高时。

关键理解： KL散度不是对称的。KL(P||Q) ≠ KL(Q||P)。蒸馏时我们用老师当P，学生当Q，意思是“学生要尽量靠近老师”。反过来用就错了。

2.2 温度缩放策略：这个T到底怎么调？

温度T是Logits蒸馏的灵魂参数。它控制着软标签的“平滑程度”。

公式很简单：

softmax(z_i / T)

其中z_i是logits，T是温度。

温度T的作用：

T=1：标准softmax，输出就是硬概率分布
T>1：分布变平滑，小概率类别的信息被放大
T<1：分布变尖锐，接近one-hot

我在项目中遇到过一个问题：温度设太高，学生模型学到的全是噪声；温度设太低，又跟硬标签训练没区别。后来总结出一套经验：

温度范围	适用场景	我的经验
T=1~3	分类数少（10类以内）	图像分类常用，稳定
T=3~8	分类数中等（10~100类）	NLP任务常见，能学到类间关系
T=8~20	分类数多（100类以上）	大规模分类，需要平滑信息

调参技巧： 我习惯先固定T=4跑一轮，观察学生模型的收敛速度和验证集精度。如果loss下降太快但精度上不去，说明温度太高；如果loss下降慢且精度跟硬标签训练差不多，说明温度太低。

2.3 蒸馏权重α：软硬兼施的艺术

蒸馏损失函数通常由两部分组成：

L_total = α * L_soft + (1-α) * L_hard

其中L_soft是KL散度（学生vs老师），L_hard是交叉熵（学生vs真实标签）。α控制着软标签和硬标签的权重比例。

α的调参经验：

α=0.7~0.9：老师主导，适合老师模型很强、学生模型容量够大的情况
α=0.3~0.5：平衡策略，适合大多数通用场景
α=0.1~0.3：硬标签主导，适合学生模型容量小、容易过拟合的情况

我曾经踩过一个坑：在某个文本分类任务中，我把α设成0.9，结果学生模型完全跟着老师走，连老师的错误都学过来了。后来改成α=0.5，同时把温度从8降到4，效果反而更好。

避坑指南： 我曾经在CIFAR-100上做实验，发现α=0.9时学生模型精度反而比α=0.5低。原因是老师模型在少数类别上也有错误，学生学得太“像”老师，反而把错误也继承了。建议先跑一组α=[0.3, 0.5, 0.7]的对比实验，找到最优值。

2.4 核心流程：一张图看懂Logits蒸馏

下面这张图展示了Logits蒸馏的完整流程，从输入到损失计算：

2.5 完整训练流程与代码示例

下面是一个完整的Logits蒸馏训练代码片段，包含了温度缩放和KL散度计算：

import torch
import torch.nn.functional as F

def distillation_loss(student_logits, teacher_logits, labels, T=4.0, alpha=0.5):
    """
    Logits蒸馏损失函数
    Args:
        student_logits: 学生模型输出
        teacher_logits: 老师模型输出（不计算梯度）
        labels: 真实标签
        T: 温度
        alpha: 蒸馏权重
    """
    # 软标签损失：KL散度
    soft_student = F.log_softmax(student_logits / T, dim=1)
    soft_teacher = F.softmax(teacher_logits / T, dim=1)
    kl_loss = F.kl_div(soft_student, soft_teacher, reduction='batchmean')
    kl_loss = kl_loss * (T * T)  # 温度补偿
    
    # 硬标签损失：交叉熵
    ce_loss = F.cross_entropy(student_logits, labels)
    
    # 总损失
    total_loss = alpha * kl_loss + (1 - alpha) * ce_loss
    return total_loss

# 训练循环示例
for batch in dataloader:
    inputs, labels = batch
    inputs, labels = inputs.to(device), labels.to(device)
    
    # 老师模型不更新梯度
    with torch.no_grad():
        teacher_logits = teacher_model(inputs)
    
    # 学生模型前向
    student_logits = student_model(inputs)
    
    # 计算蒸馏损失
    loss = distillation_loss(
        student_logits, teacher_logits, labels,
        T=4.0, alpha=0.5
    )
    
    # 反向传播
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()

温度补偿： 代码中kl_loss乘以T²是因为梯度会随T缩放，补偿后梯度大小与T无关，方便调参。这个细节我一开始没注意，导致不同温度下学习率需要重新调，后来加上补偿就稳定多了。

2.6 调参经验总结

我把Logits蒸馏的三个核心参数总结成一张速查表，方便你快速定位问题：

现象	可能原因	我的建议
学生精度比硬标签训练还差	α太大，老师错误被继承	降低α到0.3~0.5，同时检查老师精度
训练loss下降快但验证集不涨	温度太高，学到噪声	降低T到2~4，或增加硬标签权重
学生模型收敛太慢	温度太低，软标签信息不足	提高T到6~8，或增大α
学生模型过拟合	α太小，硬标签主导	增大α到0.7~0.9，利用老师做正则化

嗯，这里要注意：没有万能参数组合。我每次接到新任务，都会先跑一组小规模对比实验，用网格搜索找T和α的最优值。虽然费点时间，但比凭感觉调参靠谱得多。

Logits蒸馏的核心就这些。说白了，KL散度告诉你“差多少”，温度T控制“看多细”，权重α决定“听谁的多”。三个参数配合好了，学生模型就能在保持轻量的同时，学到老师模型的精髓。