基于特征的知识蒸馏算法：从理论到工程实践

1次阅读

共计 2146 个字符，预计需要花费 6 分钟才能阅读完成。

在移动端和 IoT 设备上部署深度学习模型时，计算资源受限是主要瓶颈。传统 logits 蒸馏（如 Hinton 提出的原始方法）只能利用模型最后一层的输出信息，忽视了深层网络中层特征包含的丰富语义信息。这导致学生模型难以复现教师模型的中间表示能力，在复杂任务上表现下降明显。

通过实验发现，在 ImageNet 数据集上，仅使用 logits 蒸馏时 ResNet50→MobileNetV2 的准确率差距达 12.3%，而引入特征蒸馏后差距缩小到 6.8%。这种提升源于特征蒸馏能够：

保留教师模型的多尺度特征表示
增强学生模型中间层的判别能力
改善梯度流动路径

特征蒸馏的核心是最小化教师模型与学生模型在特定层的特征分布差异。对于选定的第 $l$ 层，损失函数可表示为：

$$
\mathcal{L}{feat} = \frac{1}{2N}\sum^N | \phi(T_l(x_i)) – \phi(S_l(x_i)) |_2^2
$$

其中 $\phi(\cdot)$ 是特征适配函数，常用形式包括：

直接映射 （FitNets）：线性投影矩阵
注意力转移 （AT）：基于 Gram 矩阵的注意力图
激活边界 （AB）：带 margin 的余弦相似度

图：特征蒸馏中的 KL 散度计算过程（温度系数 τ =3）

# 多层级特征蒸馏实现
class FeatureDistiller(nn.Module):
    def __init__(self, teacher, student, feat_layers=[3,7,10]):
        super().__init__()
        self.teacher = teacher
        self.student = student
        # 特征适配器（每个蒸馏层独立）self.adapters = nn.ModuleList([nn.Conv2d(s_ch, t_ch, 1) 
            for t_ch, s_ch in zip(teacher.get_channel_dims(feat_layers),
                student.get_channel_dims(feat_layers)
            )
        ])

    def forward(self, x):
        # 获取多层级特征 [batch, channel, H, W]
        with torch.no_grad():
            t_feats = self.teacher.extract_features(x) 
        s_feats = self.student.extract_features(x)

        losses = []
        for idx, (t_f, s_f, adapter) in enumerate(zip(t_feats, s_feats, self.adapters)):
            # 特征对齐和损失计算
            adapted_s = adapter(s_f)
            loss = F.mse_loss(F.normalize(t_f, dim=1),
                F.normalize(adapted_s, dim=1)
            ) * (0.5 ** idx)  # 层级权重衰减
            losses.append(loss)

        return sum(losses)

关键参数说明：
– feat_layers: 选择进行蒸馏的层索引
– 0.5**idx: 深层特征权重衰减系数
– F.normalize: 特征 L2 归一化避免量纲差异

方法	学生模型准确率	FLOPs
Baseline	92.1%	0.56G
Logits 蒸馏	93.4%	0.56G
FitNets	94.2%	0.58G
AT+AB(本文)	95.7%	0.57G

测试环境：RTX 3090, PyTorch 1.12, batch_size=128

图：原始学生模型（左）与特征蒸馏后（右）的 t -SNE 特征分布对比

早期层 （conv1-3）：提升低级特征提取能力
计算成本低但收益有限
适合处理纹理丰富的输入
中间层 （conv4-7）：平衡语义与计算量
推荐作为主要蒸馏目标
需注意 feature map 尺寸对齐
深层（conv8+）：增强语义判别力
计算开销大但提升明显
需要梯度截断防止震荡

scaler = GradScaler()  # 初始化梯度缩放器

with autocast():
    feat_loss = distiller(inputs)
    total_loss = 0.3*feat_loss + 0.7*cls_loss

scaler.scale(total_loss).backward()
scaler.step(optimizer)
scaler.update()

注意事项：
– 特征蒸馏损失需要更大的缩放因子
– 确保适配器模块使用 FP32 精度
– 监控梯度幅值避免溢出