中文

One-for-All:弥合知识蒸馏中异构架构间的鸿沟

计算机视觉与模式识别 2023-10-31 v1

摘要

知识蒸馏(KD)已被证明是一种通过教师-学生训练方案提升模型性能的高效方法。然而,大多数现有蒸馏方法是在教师与学生模型属于同一模型族的假设下设计的,尤其是基于提示(hint)的方法。通过使用中心化核对齐(CKA)比较异构教师与学生模型之间学到的特征,我们观察到显著的特征差异。该差异说明了先前基于提示的方法在跨架构蒸馏中的无效性。为应对异构模型蒸馏中的挑战,我们提出了一种简单而有效的通用 KD 框架,称为 OFA-KD,其显著提升了异构架构间的蒸馏性能。具体而言,我们将中间特征投影到对齐的潜空间(如 logits 空间),从而丢弃架构特定信息。此外,我们引入了一种自适应目标增强方案,以防止学生被无关信息干扰。基于包括 CNN、Transformer 和 MLP 在内的多种架构的大量实验,证明了我们的 OFA-KD 框架在实现异构架构间蒸馏方面的优越性。具体而言,当配备我们的 OFA-KD 时,学生模型取得了显著的性能提升,在 CIFAR-100 数据集上最大增益为 8.0%,在 ImageNet-1K 数据集上为 0.7%。PyTorch 代码与检查点可在 https://github.com/Hao840/OFAKD 获取。

关键词

引用

@article{arxiv.2310.19444,
  title  = {One-for-All: Bridge the Gap Between Heterogeneous Architectures in Knowledge Distillation},
  author = {Zhiwei Hao and Jianyuan Guo and Kai Han and Yehui Tang and Han Hu and Yunhe Wang and Chang Xu},
  journal= {arXiv preprint arXiv:2310.19444},
  year   = {2023}
}