跳转至

ACAM-KD: Adaptive and Cooperative Attention Masking for Knowledge Distillation

会议: ICCV 2025
arXiv: 2503.06307
代码: 无
领域: 模型压缩
关键词: 知识蒸馏, 注意力遮蔽, 学生-教师交互, 自适应特征选择, 目标检测

一句话总结

提出 ACAM-KD,通过学生-教师交叉注意力特征融合(STCA-FF)和自适应空间-通道遮蔽(ASCM)两个模块,使知识蒸馏中的特征选择能随学生学习状态动态演化,在 COCO 检测上 RetinaNet R50 从 R101 蒸馏时 mAP 达 41.2(+1.4 超越 SOTA),Cityscapes 分割上 DeepLabV3-MBV2 mIoU 提升 3.09。

研究背景与动机

  1. 领域现状:知识蒸馏(KD)是模型压缩的主流技术,其中特征蒸馏特别适合检测、分割等密集预测任务,通过传递中间特征表示保留空间信息。

  2. 现有痛点:现有特征蒸馏方法依赖固定的、教师驱动的特征选择策略。FKD 使用教师高注意力区域,FGD 结合真值框和全局上下文,MasKD 用教师特征离线学习 token 生成遮蔽。这些方法的共同问题是:(1) 遮蔽对同一图片在每个 epoch 都相同,忽略学生进度;(2) 教师认为重要的区域不一定是学生当前最需要的;(3) 只关注空间维度,忽略通道维度。

  3. 核心矛盾:作者可视化发现一个惊人现象——学生在 epoch 12 时注意力比教师更好地聚焦前景目标,但到 epoch 24 却被迫退化到与教师相同的固定注意力模式。这说明固定的教师驱动蒸馏可能反过来限制甚至损害学生的学习。

  4. 本文目标:设计一种自适应的、动态演化的蒸馏遮蔽机制,让蒸馏关注区域根据学生-教师的交互状态实时调整。

  5. 切入角度:蒸馏应该是师生"协作"而非教师"单方面传授"。通过交叉注意力让两者的特征交互,生成的遮蔽自然反映双方的联合状态。

  6. 核心 idea:教师特征做 query、学生特征做 key/value 的交叉注意力融合产生师生交互特征,在此基础上用可学习选择单元动态生成空间和通道双维度的蒸馏遮蔽。

方法详解

整体框架

ACAM-KD 作为插件模块嵌入标准特征蒸馏框架。教师和学生网络各自提取特征后,STCA-FF 模块对两者的特征进行交叉注意力融合,得到反映师生交互状态的融合特征。然后 ASCM 模块在融合特征上动态生成空间和通道遮蔽 mask,用于加权蒸馏损失。整个过程在每个训练迭代中实时执行,遮蔽随学生学习状态持续演化。

关键设计

  1. 学生-教师交叉注意力特征融合 (STCA-FF):

    • 功能:生成反映师生交互状态的融合特征,为后续动态遮蔽提供基础。
    • 核心思路:教师特征 \(F^T\) 生成 query \(Q = W_q F^T\),学生特征 \(F^S\) 生成 key \(K = W_k F^S\) 和 value \(V = W_v F^S\)。注意力矩阵 \(A = \text{softmax}(QK/\sqrt{C_q})\),融合特征 \(F_{fused} = AV\)\(W_q, W_k\) 将通道降维到 \(C_q = C/2\) 节省计算,\(W_v\) 保持原维度。
    • 设计动机:教师做 query 意味着教师指导"在哪里看",但学生做 key/value 意味着实际内容来自学生,两者交互产生的融合特征自然捕捉了"教师认为重要但学生还未掌握的区域"。实验证明教师做 query 优于学生做 query(41.2 vs 41.0 mAP)。
  2. 自适应空间-通道遮蔽 (ASCM):

    • 功能:在空间和通道双维度上动态生成特征选择遮蔽。
    • 核心思路:设 M 组可学习的选择单元——通道选择 \(m^c \in \mathbb{R}^{M \times 1}\)、空间选择 \(m^s \in \mathbb{R}^{M \times C}\)。通道 mask \(M^c = \sigma(m^c \cdot v)\)\(v\)\(F_{fused}\) 的空间平均池化向量;空间 mask \(M^s = \sigma(m^s \cdot z)\)\(z\)\(F_{fused}\) 沿空间展平的矩阵。两组 mask 分别生成通道维和空间维的蒸馏损失。关键在于 \(m^c, m^s\) 是持续训练更新的参数,而 \(F_{fused}\) 随学生特征变化,因此 mask 会在训练中动态演化。
    • 设计动机:空间选择关注"哪个位置重要",通道选择关注"哪个语义通道重要",两者互补。检测中 M=6,分割中 M=19(与类别数匹配)。
  3. 遮蔽多样性损失:

    • 功能:防止多组 mask 退化为相同模式。
    • 核心思路:用 Dice 系数度量 mask 之间的相似度并作为正则损失:\(L_{div} = \frac{2\sum_{i}\sum_{j \neq i} M_i \cdot M_j}{\sum_i M_i^2 + \sum_j M_j^2}\),最小化该损失使不同 mask 关注不同区域。
    • 设计动机:没有多样性约束时,所有 mask 可能收敛到同一模式或退化为零值。

损失函数 / 训练策略

  • 总损失:\(L = L_{task} + \alpha(L_{distill}^c + L_{distill}^s) + \lambda L_{div}\),其中 \(\alpha = 1, \lambda = 1\)
  • 通道蒸馏损失按通道维加权,空间蒸馏损失按空间维加权
  • 使用继承策略(inheritance)稳定早期训练
  • 检测:SGD,momentum=0.9,weight decay=1e-4,MMDetection 框架
  • 分割:SGD,weight decay=5e-4,40K 迭代,poly 学习率调度

实验关键数据

主实验

检测器 教师→学生 ACAM-KD mAP 前SOTA mAP 提升
RetinaNet R101→R50 41.2 39.9 (FreeKD) +1.3
Faster-RCNN R101→R50 41.4 40.8 (MasKD) +0.6
RepPoints R101→R50 42.5 41.1 (MasKD) +1.4
RetinaNet X101→R50 41.5 41.0 (FreeKD) +0.5
RepPoints X101→R50 42.8 42.4 (FreeKD) +0.4
分割框架 学生 ACAM-KD mIoU 基线 mIoU 提升
DeepLabV3 R18 77.53 72.96 +4.57
DeepLabV3 MBV2 76.21 73.12 +3.09
PSPNet R18 75.99 72.55 +3.44

消融实验

配置 mAP 说明
Spatial only 40.9 仅空间遮蔽
Channel only 40.4 仅通道遮蔽
Spatial + Channel 41.2 两者互补,最优
No masking 37.4 无遮蔽=基线学生
Fixed masking from teacher 39.8 固定教师遮蔽
Adaptive masking from teacher 39.9 自适应但仅教师
ACAM-KD 41.2 师生协作自适应

关键发现

  • 空间遮蔽对小目标检测贡献最大(\(AP_s\) 24.5→25.4),通道遮蔽对中大目标更有效
  • 固定遮蔽 vs 自适应遮蔽差异不大(39.8 vs 39.9),但加入学生交互后跃升至 41.2,证明关键在于师生协作
  • 推理时零额外开销——STCA-FF 和 ASCM 仅在训练时使用
  • 轻量学生 MBV2(3.2M 参数)蒸馏后达 76.21 mIoU,接近 84.7M 教师的 78.07

亮点与洞察

  • "学生可能比教师更聪明"的洞察:可视化证据表明学生在训练中期的注意力可能优于教师,固定教师驱动的蒸馏反而限制了学生的潜力。这一发现对 KD 领域有深远影响。
  • 师生协作范式:从"教师单方面传授"转变为"师生协商聚焦",交叉注意力是实现这一转变的优雅机制。
  • 空间+通道双维度遮蔽:两个维度的互补效果清晰可量化(40.9+40.4→41.2),且实现简洁。

局限与展望

  • 训练时交叉注意力引入额外计算和显存开销
  • M 的选择(检测=6,分割=19)依赖任务先验,通用性有限
  • 未探索跨层蒸馏场景
  • 可考虑将动态遮蔽扩展到时序任务(视频检测等)

相关工作与启发

  • vs MasKD:MasKD 用教师特征离线学习 token 生成 mask,训练过程中 mask 固定;ACAM-KD 的 mask 随训练持续演化
  • vs FreeKD:FreeKD 用频域语义提示引导空间选择,但同样是固定的;ACAM-KD 通过师生交互实现动态适应
  • vs CWD:CWD 引入通道维度对齐但无空间选择;ACAM-KD 统一了空间和通道选择

评分

  • 新颖性: ⭐⭐⭐⭐ 师生交叉注意力+动态遮蔽设计有新意,"学生超越教师"洞察有价值
  • 实验充分度: ⭐⭐⭐⭐⭐ 3种检测器+3种分割架构+2种教师骨干+详细消融
  • 写作质量: ⭐⭐⭐⭐ 动机分析直观有力,注意力可视化说服力强
  • 价值: ⭐⭐⭐⭐ 通用KD改进方法,推理无额外开销

相关论文