知识蒸馏实战：从attention transfer论文入门模型压缩技术

1次阅读

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

在移动端和边缘设备部署深度学习模型时，我们常常面临两个核心矛盾：

模型精度 与计算资源 的冲突：大模型（如 ResNet50）在 ImageNet 上能达到 78% 的 top- 1 准确率，但需要超过 4GFLOPs 的计算量
响应速度 与模型体积 的矛盾：自动驾驶场景要求模型在 100ms 内完成推理，但原始模型可能达到 200MB 以上

知识蒸馏（Knowledge Distillation）通过让轻量化的学生模型（Student）模仿教师模型（Teacher）的行为来解决这个问题。而 Attention Transfer（AT）作为 2016 年的经典方法，首次提出通过迁移注意力图（Attention Maps）来实现更高效的知识传递。

Attention Transfer 的核心是让学生网络模仿教师网络的注意力分布，其损失函数由两部分组成：

$$
\mathcal{L}{AT} = \mathcal{L} ||_p
$$}(y, \sigma(z_s)) + \beta \sum_{j=1}^4 || \frac{Q_s^j}{||Q_s^j||_2} – \frac{Q_t^j}{||Q_t^j||_2

其中：
– $Q_s^j$ 和 $Q_t^j$ 分别表示学生和教师在第 j 层的注意力图
– $\beta$ 是调节权重（建议初始值 1.0）
– $p$ 通常取 2（L2 范数）

方法	计算开销	CIFAR-100 准确率	参数量
Teacher(ResNet34)	1.16GFLOPs	73.2%	21.3M
Baseline KD	+15%	69.8%	2.3M
AT（本文方法）	+18%	71.4%	2.3M

# 教师模型（ResNet34）teacher = models.resnet34(pretrained=True)
for param in teacher.parameters():
    param.requires_grad = False  # 冻结参数

# 学生模型（MobileNetV2）student = models.mobilenet_v2(num_classes=100)

class AttentionLoss(nn.Module):
    def __init__(self, p=2):
        super().__init__()
        self.p = p

    def forward(self, fs, ft):
        # fs: 学生特征图 [B,C,H,W]
        # ft: 教师特征图 [B,C,H,W]
        loss = (fs.norm(p=2, dim=(2,3)) - ft.norm(p=2, dim=(2,3))).pow(2).mean()
        return loss

# 温度参数调节（建议初始 τ =4）def softmax_with_temperature(logits, tau=4):
    return torch.softmax(logits/tau, dim=1)

# 多任务损失组合
total_loss = ce_loss + 1.0 * at_loss  # β=1.0

当教师和学生网络结构差异较大时（如 ResNet 与 MobileNet），可以：

在损失计算前添加自适应池化层

adaptive_pool = nn.AdaptiveAvgPool2d((h, w))

使用 1 ×1 卷积统一通道数

温度参数 τ ：从高温（τ=4）开始，每 10 个 epoch 线性降至 τ =1
学习率：建议初始 lr=0.01，与 τ 同步衰减
β 系数 ：在[0.5, 2.0] 范围内网格搜索

# 梯度累积（batch_size=64 时）for i, (inputs, labels) in enumerate(dataloader):
    outputs = student(inputs)
    loss = criterion(outputs, labels)/4  # 除以累积步数
    loss.backward()

    if (i+1)%4 == 0:  # 每 4 步更新一次
        optimizer.step()
        optimizer.zero_grad()

与量化结合：能否在 AT 训练时模拟 8bit 量化噪声？参考论文《Quantization-Aware Knowledge Distillation》
NLP 适配：Transformer 中的注意力头是否可以直接迁移？可能需要处理序列长度动态变化的问题
跨模态应用：视频分类中如何利用时序注意力转移？

经过在 CIFAR-100 上的测试，发现几个有趣现象：