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

4. 蒸馏损失函数

损失函数是蒸馏的灵魂。说白了，它决定了学生模型到底要学什么、学到什么程度。我见过不少团队，模型结构搭得漂漂亮亮，结果损失函数随便选了个 KL 散度就跑，最后效果还不如直接训练学生——嗯，这就是典型的「方向不对，努力白费」。

这一章，我把蒸馏损失拆成四大类来讲：软标签损失、硬标签损失、中间层特征匹配损失、关系蒸馏损失。每一类我都会结合项目里的实际踩坑经验，告诉你什么时候用、怎么配、有哪些坑。

核心观点：没有万能损失函数。好的蒸馏方案，往往是 2-3 种损失的加权组合。关键是要理解每种损失在「教什么」。

4.1 软标签损失：KL 散度与交叉熵

软标签损失是蒸馏最经典的玩法。教师模型输出一个概率分布，学生去拟合这个分布——注意，不是拟合 one-hot 标签，而是拟合那个「软」的分布。

为什么软标签有用？你想想看，一张猫的图片，教师模型可能输出「猫 0.85、狗 0.10、老虎 0.05」。这个分布里藏着类别之间的相似性信息。硬标签只告诉你「这是猫」，软标签告诉你「它有点像狗，但更像猫」。

具体实现上，最常用的是 KL 散度：

import torch
import torch.nn.functional as F

def kl_div_loss(student_logits, teacher_logits, temperature=4.0):
    """
    软标签蒸馏损失
    Args:
        student_logits: 学生模型 logits, shape [B, C]
        teacher_logits: 教师模型 logits, shape [B, C]
        temperature: 温度系数，越大分布越软
    """
    # 用温度软化分布
    soft_student = F.log_softmax(student_logits / temperature, dim=-1)
    soft_teacher = F.softmax(teacher_logits / temperature, dim=-1)
    
    # KL散度 + 温度缩放
    loss = F.kl_div(soft_student, soft_teacher, reduction='batchmean')
    loss = loss * (temperature ** 2)
    return loss

温度怎么选？我个人习惯先试 4.0。温度太低（< 2.0），软标签退化成硬标签，蒸馏效果不明显；温度太高（> 10.0），分布太平坦，学生学不到区分性信息。我在一个细粒度分类项目里试过，温度从 3 调到 8，准确率差了将近 2 个点。

也有团队直接用交叉熵代替 KL 散度。效果差不多，但梯度形式略有差异。KL 散度对教师输出的「相对关系」更敏感，我个人更推荐。

4.2 硬标签损失：别丢了真实答案

软标签虽好，但不能完全替代真实标签。为什么？因为教师模型也会犯错。如果学生只学教师，教师的错误就会被继承——这叫「错误传播」。

硬标签损失就是标准的交叉熵，让学生直接拟合 ground truth：

def hard_label_loss(student_logits, labels):
    """标准交叉熵损失"""
    return F.cross_entropy(student_logits, labels)

实际工程中，总损失通常是软标签损失和硬标签损失的加权和：

total_loss = alpha * kl_div_loss + (1 - alpha) * hard_label_loss

alpha 的坑：我曾经在一个 OCR 识别项目里把 alpha 设成 0.9，结果学生模型学了一堆教师的错误识别模式，准确率反而下降了。后来改成 0.5，效果才正常。建议 alpha 从 0.5 开始调，不要一上来就偏向软标签。

4.3 中间层特征匹配损失

只约束输出层，学生只能学到「结果」，学不到「推理过程」。中间层特征匹配就是为了解决这个问题。

这类方法有三个主流变体：

方法	核心思想	对齐方式	适用场景
FitNet	直接匹配中间层特征图	L2 距离 + 适配层	教师学生结构差异不大时
AT (Attention Transfer)	匹配注意力图（空间重要性）	注意力图上的 L2	需要空间语义对齐的任务
SP (Similarity Preserving)	匹配样本间相似性矩阵	Gram 矩阵上的 L2	教师学生宽度差异大时

FitNet 最直接，但有个问题：教师和学生中间层的通道数、空间尺寸往往不一样。需要加一个 1x1 卷积做适配层。我踩过这个坑——适配层不加 BatchNorm，训练特别不稳定。

AT 是我个人比较喜欢的方法。它不直接匹配特征值，而是匹配「注意力图」——说白了就是告诉学生「你应该关注图像的哪些区域」。这对检测、分割这类空间敏感的任务特别有效。

def attention_transfer_loss(student_feat, teacher_feat):
    """
    Attention Transfer 损失
    计算注意力图上的 L2 距离
    """
    # 计算空间注意力图: 对通道维度求均值
    student_attn = torch.mean(student_feat ** 2, dim=1)  # [B, H, W]
    teacher_attn = torch.mean(teacher_feat ** 2, dim=1)  # [B, H, W]
    
    # 归一化
    student_attn = student_attn / student_attn.view(student_attn.size(0), -1).sum(dim=1, keepdim=True)
    teacher_attn = teacher_attn / teacher_attn.view(teacher_attn.size(0), -1).sum(dim=1, keepdim=True)
    
    return F.mse_loss(student_attn, teacher_attn)

经验之谈：中间层损失不要加太多层。我见过有人把教师 5 个 stage 全对齐了，结果学生模型训练发散。建议只对齐 1-2 个关键层，比如 backbone 的倒数第二层和倒数第四层。

4.4 关系蒸馏损失：RKD 与 PKDT

前面几种方法都是「点对点」的对齐——每个样本独立计算损失。关系蒸馏换了个思路：它关注样本之间的关系。

举个例子，教师模型认为「样本 A 和样本 B 很相似，A 和 C 差异很大」。关系蒸馏让学生也学到这种「关系结构」。

RKD (Relational Knowledge Distillation) 是最有代表性的方法。它包含两个部分：

距离关系：样本对之间的欧氏距离要一致
角度关系：三元组之间的角度要一致

def rkd_loss(student_features, teacher_features):
    """
    RKD 关系蒸馏损失
    """
    batch_size = student_features.size(0)
    
    # 计算距离关系
    student_dist = torch.cdist(student_features, student_features)
    teacher_dist = torch.cdist(teacher_features, teacher_features)
    
    # 归一化到 [0, 1]
    student_dist = student_dist / student_dist.max()
    teacher_dist = teacher_dist / teacher_dist.max()
    
    # 距离损失
    dist_loss = F.mse_loss(student_dist, teacher_dist)
    
    # 角度关系（三元组）
    # 简化实现：取前 N 个三元组
    angle_loss = 0.0
    for i in range(batch_size):
        for j in range(i+1, batch_size):
            for k in range(j+1, batch_size):
                # 计算角度
                v1_s = student_features[i] - student_features[j]
                v2_s = student_features[k] - student_features[j]
                v1_t = teacher_features[i] - teacher_features[j]
                v2_t = teacher_features[k] - teacher_features[j]
                
                cos_s = F.cosine_similarity(v1_s.unsqueeze(0), v2_s.unsqueeze(0))
                cos_t = F.cosine_similarity(v1_t.unsqueeze(0), v2_t.unsqueeze(0))
                angle_loss += F.mse_loss(cos_s, cos_t)
    
    return dist_loss + angle_loss / (batch_size * (batch_size-1) * (batch_size-2) / 6)

PKDT (Probabilistic Knowledge Distillation via Transformers) 是更新的方法。它用 Transformer 建模样本间的全局关系，效果更好但计算量也更大。

什么时候用关系蒸馏？我个人的经验是：当你的任务对「结构」敏感时——比如人脸识别、行人重识别、小样本学习。这些任务里，样本间的相对关系比单个样本的绝对特征更重要。

4.5 损失组合策略：实战中的选择

说了这么多，到底怎么组合？我整理了一个速查表：

任务类型	推荐组合	权重建议	注意事项
图像分类	KL + CE	α=0.5~0.7	温度 3~6，不需要中间层
目标检测	KL + CE + AT	α=0.3, β=0.3, γ=0.4	AT 对齐 backbone 最后两层
语义分割	KL + AT + SP	α=0.2, β=0.4, γ=0.4	SP 保持空间一致性
人脸识别	KL + RKD	α=0.4, β=0.6	RKD 权重可以更高
小样本学习	CE + RKD + SP	α=0.3, β=0.4, γ=0.3	关系信息比特征更重要

权重调优的坑：不要一次性把所有损失都加上。我习惯的做法是：先跑一个 baseline（只有 KL+CE），然后逐个加入新损失，每次只调一个权重。这样出了问题，你知道是哪个损失搞的鬼。

最后说一句：损失函数的选择没有银弹。同一个任务，换一个教师模型、换一个学生架构，最优组合可能完全不同。多试、多对比、多记录——这才是工程化的正确姿势。