4. 特征层蒸馏：中间层特征对齐、Attention Transfer、FitNets与NST方法对比
前面我们聊了logits层面的蒸馏，说白了就是让学生的软输出模仿老师的软输出。但有个问题——你想想看，如果老师模型很深，学生模型很浅，光靠最后那层logits，学生很难学到老师中间那些抽象特征。这时候，特征层蒸馏就派上用场了。
我个人习惯把特征层蒸馏分成四类：FitNets、Attention Transfer、NST（Neural Selective Transfer），以及最朴素的中间层特征对齐。它们的目标都一样——让学生模型的中间层特征图，尽量接近老师模型的中间层特征图。但具体怎么对齐、对齐哪一层、用什么损失函数，差别还挺大的。
4.1 为什么需要特征层蒸馏？
先看一个我踩过的坑。几年前我做一个小型语音识别模型的蒸馏，直接用KD（知识蒸馏）训练，学生模型精度死活上不去。后来我把老师模型中间层的特征图可视化出来，发现学生模型根本学不到那些高频细节的激活模式。说白了，logits层面的信息太“浓缩”了，学生消化不了。
特征层蒸馏的核心思路很简单：
- 选层：从老师模型中挑出若干中间层
- 对齐：让学生模型对应层的特征图与老师接近
- 损失：用L2、L1或更复杂的度量计算差异
但这里有个关键问题——老师和学生的特征图尺寸往往不一样。怎么办？加一个适配层（adaptation layer），把学生的特征图映射到和老师一样的维度。
核心公式：特征蒸馏损失 = Σ || f_t(x) - g(f_s(x)) ||²
其中 f_t 是老师特征，f_s 是学生特征，g 是适配层（通常是1x1卷积或全连接层）
4.2 FitNets：最早的特征蒸馏方法
FitNets是2015年提出的，算是特征层蒸馏的开山之作。它的做法很直接——选老师模型的某几个中间层，让学生模型对应层的特征图去拟合。
我记得第一次实现FitNets时，遇到一个坑：适配层的初始化。如果适配层初始化不好，训练初期梯度就爆炸了。后来我改用Xavier初始化，才稳定下来。
FitNets的关键设计：
- 使用L2损失作为特征对齐的度量
- 引入适配层（adaptation layer）解决维度不匹配
- 通常只对齐1-2个中间层，避免训练不稳定
- 训练分两阶段：先特征蒸馏预训练，再logits蒸馏微调
我的经验：FitNets对适配层的设计很敏感。如果适配层太深（比如3层以上），学生模型容易过拟合到适配层上，反而学不到老师的特征。我一般只用1层1x1卷积或1层全连接。
4.3 Attention Transfer：让注意力图对齐
Attention Transfer（AT）是2017年提出的，思路比FitNets更巧妙——不直接对齐特征图，而是对齐特征图的注意力图。说白了，就是让学生的“注意力分布”和老师一样。
注意力图怎么算？很简单：对特征图在通道维度上做sum或mean，得到一个二维的注意力热力图。然后让学生的热力图和老师的热力图尽量接近。
# 注意力图计算示例
def attention_map(feature_map, method='sum'):
# feature_map shape: [B, C, H, W]
if method == 'sum':
# 通道维度求和
attn = torch.sum(feature_map ** 2, dim=1) # [B, H, W]
elif method == 'mean':
# 通道维度求平均
attn = torch.mean(feature_map ** 2, dim=1) # [B, H, W]
# 归一化
attn = attn / attn.view(attn.size(0), -1).sum(dim=1).view(-1, 1, 1)
return attn
# AT损失
def at_loss(student_feat, teacher_feat):
student_attn = attention_map(student_feat)
teacher_attn = attention_map(teacher_feat)
return F.mse_loss(student_attn, teacher_attn)
AT相比FitNets的优势：
- 注意力图是二维的，维度远小于特征图，训练更稳定
- 不需要适配层（因为注意力图尺寸天然一致）
- 对特征图的噪声更鲁棒
注意：AT虽然稳定，但信息量也少了。因为注意力图丢失了通道维度的信息。我做过实验，在图像分类任务上，AT单独使用效果不如FitNets+KD的组合。建议AT作为辅助损失，不要单独用。
4.4 NST：选择性传递关键神经元
NST（Neural Selective Transfer）是2019年提出的，思路更激进——不是所有神经元都值得传递，只传递那些“激活值高”的神经元。
为什么这么做？因为老师模型的特征图里，大部分神经元激活值很低（接近0），这些神经元对最终决策贡献很小。如果强制学生去拟合这些低激活值，反而会引入噪声。
NST的做法：
- 计算老师特征图每个位置的激活值
- 选出激活值最高的k%位置（比如top 10%）
- 只在这些位置上计算蒸馏损失
- 其他位置忽略（或者给很小的权重）
# NST选择性掩码计算
def nst_mask(teacher_feat, keep_ratio=0.1):
# teacher_feat: [B, C, H, W]
# 计算每个位置的激活强度
act_strength = torch.sum(teacher_feat ** 2, dim=1) # [B, H, W]
# 展平
B, H, W = act_strength.shape
flat = act_strength.view(B, -1) # [B, H*W]
# 计算阈值
k = int(H * W * keep_ratio)
threshold = torch.topk(flat, k, dim=1)[0][:, -1:] # [B, 1]
# 生成掩码
mask = (flat > threshold).float() # [B, H*W]
mask = mask.view(B, 1, H, W)
return mask
# NST损失
def nst_loss(student_feat, teacher_feat, mask):
diff = (student_feat - teacher_feat) ** 2
masked_diff = diff * mask
return masked_diff.sum() / mask.sum()
我的经验：keep_ratio这个超参数很关键。我试过0.05到0.5之间的值，发现0.1-0.2效果最好。如果keep_ratio太小（比如0.01），学生学不到足够信息；如果太大（比如0.5），又退化成普通特征对齐。
4.5 四种方法对比
下面这个表格是我在实际项目中总结的对比，供你参考：
| 方法 | 对齐目标 | 是否需要适配层 | 训练稳定性 | 信息量 | 推荐场景 |
|---|---|---|---|---|---|
| FitNets | 特征图（L2） | 是 | 中等 | 高 | 学生模型较深时 |
| Attention Transfer | 注意力图 | 否 | 高 | 中 | 轻量级模型、快速实验 |
| NST | 选择性特征图 | 是 | 中等 | 中高 | 特征图稀疏、噪声大时 |
| 朴素特征对齐 | 特征图（L1/L2） | 是 | 低 | 高 | 老师和学生结构相似时 |
4.6 特征层蒸馏的整体流程
下面这张图展示了特征层蒸馏的完整流程，从老师模型提取中间特征，经过适配层映射，再与学生特征计算损失：
4.7 实践中的选择建议
说了这么多，到底该用哪种？我根据实际项目经验，给你几个建议：
- 快速验证阶段：先用Attention Transfer，因为它不需要适配层，训练稳定，代码量最少
- 追求精度：用FitNets + KD的组合，但要注意适配层的设计和训练稳定性
- 特征图稀疏场景（比如检测模型、分割模型）：用NST，它能过滤掉大量无效区域
- 老师和学生结构相似：直接用朴素特征对齐，省掉适配层的麻烦
我曾经踩过的坑：有一次我同时用了FitNets、AT和KD三个损失，结果训练时损失函数互相打架，模型精度反而下降了。后来我用了梯度裁剪和损失权重调整（给特征蒸馏损失0.1的权重，KD损失0.9），才稳定下来。记住：不是损失越多越好，关键是平衡。
4.8 小结
特征层蒸馏的核心，就是让学生模型在中间层就学会老师模型的“思维方式”。FitNets直接对齐特征图，AT对齐注意力分布，NST选择性对齐关键区域——各有各的适用场景。
我个人建议：先从AT入手，因为它最简单、最稳定。如果精度不够，再换成FitNets或NST。记住，特征蒸馏通常要和logits蒸馏配合使用，单独用效果往往不够好。