3、特征迁移:中间层特征对齐、Hint Regression与Attention Transfer

特征迁移,说白了就是让教师模型把「思考过程」教给学生,而不仅仅是「最终答案」。我刚开始做蒸馏时,总觉得只拿 soft label 就够了,直到有一次在图像检索任务上,学生模型学到的特征分布完全走样了……嗯,从那以后我才意识到,中间层的特征对齐才是真正提升模型泛化能力的关键。

3.1 为什么需要特征迁移?

传统的 logit 蒸馏只关注输出层的概率分布,但教师模型在中间层已经提取了丰富的语义特征。你想想看,一个在 ImageNet 上预训练的 ResNet-50,它的浅层学的是边缘和纹理,中层学的是形状和部件,深层学的是语义概念。如果学生只模仿最终输出,它可能用完全不同的特征组合来达到同样的分类结果——这就像两个学生虽然都考了满分,但一个真正理解了知识,另一个只是背了答案。

核心观点:特征迁移的目标是让学生模型的中间层特征分布尽可能接近教师模型,从而继承教师强大的特征提取能力。

3.2 Hint Regression:直接回归教师特征

Hint Regression 是最直观的特征迁移方法。它的思路很简单:在学生模型的某个中间层后面接一个回归头,让它去拟合教师模型对应层的特征图。

3.2.1 基本流程

  1. 选择对齐层:确定教师和学生的哪一层需要对齐。通常选择学生网络的最后一个下采样层之前。
  2. 维度匹配:如果学生和教师的特征图通道数不同,需要加一个 1x1 卷积或全连接层做投影。
  3. 计算损失:使用 MSE Loss 或 L1 Loss 衡量特征图之间的差异。
import torch
import torch.nn as nn
import torch.nn.functional as F

class HintLoss(nn.Module):
    def __init__(self, student_channels, teacher_channels):
        super().__init__()
        # 投影层:将学生特征映射到教师特征空间
        self.regressor = nn.Conv2d(student_channels, teacher_channels, kernel_size=1)
    
    def forward(self, student_feat, teacher_feat):
        # student_feat: [B, C_s, H, W]
        # teacher_feat: [B, C_t, H, W]
        projected = self.regressor(student_feat)
        loss = F.mse_loss(projected, teacher_feat.detach())
        return loss

我的经验:在实际项目中,我建议对教师特征做 detach() 操作,避免教师模型的梯度回传。另外,MSE Loss 对异常值敏感,如果特征图有较大的激活值,可以先做 LayerNorm 再计算损失。

3.2.2 避坑指南

  • 层选择不是越深越好:我曾经尝试对齐教师模型的最后一层特征,结果学生模型反而过拟合了。浅层特征更通用,深层特征更任务相关,建议从中间层开始尝试。
  • 投影层不能太复杂:投影层只是维度匹配的工具,不要用多层 MLP,否则学生会依赖投影层来「作弊」。
  • 损失权重需要调参:Hint Loss 的权重通常设为 0.1~1.0,太大容易压制分类损失,太小又没效果。

3.3 Attention Transfer:让模型学会「看哪里」

Attention Transfer 是另一种特征迁移方式,它不直接对齐特征图的值,而是对齐特征图的注意力分布。说白了,就是让教师告诉学生:「你看,我关注的是图像的这些区域,你也应该关注这些区域。」

3.3.1 注意力图的生成

常用的注意力图生成方式有两种:

方法 公式 特点
Sum of Abs A = sum(|F|, dim=1) 简单高效,保留空间结构
Sum of Squares A = sum(F², dim=1) 对高激活区域更敏感
Channel-wise Max A = max(|F|, dim=1) 保留最显著的特征通道
def attention_map(feature, method='sum_squares'):
    # feature: [B, C, H, W]
    if method == 'sum_abs':
        return torch.sum(torch.abs(feature), dim=1)
    elif method == 'sum_squares':
        return torch.sum(feature ** 2, dim=1)
    elif method == 'channel_max':
        return torch.max(torch.abs(feature), dim=1)[0]
    else:
        raise ValueError(f"Unknown method: {method}")

3.3.2 AT Loss 的计算

得到注意力图后,AT Loss 通常使用 L2 距离来衡量学生和教师注意力图的差异:

class AttentionTransferLoss(nn.Module):
    def __init__(self, method='sum_squares', p=2):
        super().__init__()
        self.method = method
        self.p = p
    
    def forward(self, student_feat, teacher_feat):
        # 生成注意力图
        student_attn = attention_map(student_feat, self.method)
        teacher_attn = attention_map(teacher_feat, self.method)
        
        # 归一化,消除尺度差异
        student_attn = F.normalize(student_attn.view(student_attn.size(0), -1), p=self.p, dim=1)
        teacher_attn = F.normalize(teacher_attn.view(teacher_attn.size(0), -1), p=self.p, dim=1)
        
        # 计算损失
        loss = F.mse_loss(student_attn, teacher_attn.detach())
        return loss

注意:注意力图一定要做归一化!我踩过这个坑——如果不归一化,教师模型的注意力图值域可能比学生大很多,导致 Loss 被教师主导,学生根本学不到东西。

3.4 Hint Regression vs Attention Transfer 对比

维度 Hint Regression Attention Transfer
对齐目标 特征图的值 注意力分布
计算开销 较高(需要投影层) 较低(只需计算注意力图)
对维度敏感度 敏感,需要维度匹配 不敏感,注意力图是单通道
适用场景 教师和学生结构相似 结构差异大时更鲁棒
收敛速度 较快 较慢但更稳定

3.5 特征迁移的整体框架

下面这张图展示了特征迁移在蒸馏流程中的位置:

特征迁移蒸馏框架 教师模型 学生模型 中间层特征 F_t 中间层特征 F_s 特征迁移模块 Hint Regression / Attention Transfer 总损失 = 分类损失 + λ * 特征迁移损失 Hint Regression 直接对齐特征值;Attention Transfer 对齐注意力分布,对结构差异更鲁棒

3.6 实践建议与调参经验

我的调参策略:

  • 先单独用 Hint Loss 或 AT Loss 训练 10 个 epoch,观察学生模型的特征图是否开始模仿教师。
  • 如果特征图出现「模糊化」——所有区域激活值都差不多,说明损失权重太大,建议从 0.1 开始。
  • 如果特征图完全没有变化,说明权重太小或层选择不对,尝试换一个更深的层。

我记得有一次在目标检测任务中,学生模型用 AT Loss 后 mAP 提升了 3.2%,但训练初期损失下降很慢。后来我发现是注意力图的归一化方式有问题——用 L2 归一化比用 Softmax 更稳定。你想想看,Softmax 会让注意力图变成概率分布,丢失了空间上的相对强度信息。

常见陷阱:

  • 不要同时对齐太多层。我见过有人一次对齐 5 个中间层,结果学生模型直接崩了。建议从 1~2 层开始。
  • 教师模型要固定住。如果教师也在训练,特征分布会不断变化,学生根本追不上。
  • 注意特征图的尺寸。如果教师和学生下采样倍数不同,需要插值对齐,但插值会引入额外误差。

3.7 小结

特征迁移是知识蒸馏中非常重要的一环。Hint Regression 适合结构相似的师生模型,直接高效;Attention Transfer 更鲁棒,适合结构差异大的情况。我个人习惯在项目初期先用 AT Loss 快速验证效果,再根据情况切换到 Hint Regression 做精细调优。记住,特征迁移不是万能的,它需要和 logit 蒸馏配合使用,才能达到最佳效果。