One-for-All: Bridge the Gap Between Heterogeneous Architectures in Knowledge Distillation

📄 arXiv: 2310.19444v1 📥 PDF

作者: Zhiwei Hao, Jianyuan Guo, Kai Han, Yehui Tang, Han Hu, Yunhe Wang, Chang Xu

分类: cs.CV

发布日期: 2023-10-30

🔗 代码/项目: GITHUB

💡 一句话要点

提出OFA-KD框架以解决异构架构间知识蒸馏问题

🎯 匹配领域: 支柱二：RL算法与架构 (RL & Architecture)

关键词: 知识蒸馏 异构架构 深度学习 模型压缩 特征对齐

📋 核心要点

  1. 现有的知识蒸馏方法大多假设教师和学生模型属于同一架构，导致在异构模型间蒸馏时效果不佳。
  2. 论文提出了一种OFA-KD框架，通过将中间特征投影到对齐的潜在空间来解决异构架构间的蒸馏问题。
  3. 实验结果显示，OFA-KD框架在CIFAR-100和ImageNet-1K数据集上分别实现了最高8.0%和0.7%的性能提升。

📝 摘要（中文）

知识蒸馏（KD）是一种通过教师-学生训练方案提升模型性能的有效方法。然而，现有的蒸馏方法大多假设教师和学生模型属于同一模型家族，尤其是基于提示的方法。通过使用中心核对齐（CKA）比较异构教师和学生模型之间的特征，我们观察到显著的特征差异，表明以往的基于提示的方法在跨架构蒸馏中效果不佳。为了解决异构模型的蒸馏挑战，我们提出了一种简单而有效的OFA-KD框架，显著提升了异构架构间的蒸馏性能。具体而言，我们将中间特征投影到对齐的潜在空间，如logits空间，从而丢弃架构特定信息。此外，我们引入了一种自适应目标增强方案，以防止学生受到无关信息的干扰。大量实验表明，OFA-KD在不同架构（包括CNN、Transformer和MLP）间的蒸馏效果优越，学生模型在CIFAR-100数据集上最大提升8.0%，在ImageNet-1K数据集上提升0.7%。

🔬 方法详解

问题定义：本论文旨在解决异构架构间知识蒸馏的有效性问题。现有方法主要基于教师和学生模型同构的假设，导致在不同架构间蒸馏时特征差异显著，效果不理想。

核心思路：论文提出的OFA-KD框架通过将中间特征投影到一个对齐的潜在空间，消除架构特定的信息，从而增强异构模型间的蒸馏效果。该设计旨在减少特征间的差异，提高学生模型的学习效率。

技术框架：OFA-KD框架主要包括两个模块：特征对齐模块和自适应目标增强模块。特征对齐模块负责将中间特征映射到对齐的潜在空间，而自适应目标增强模块则通过动态调整目标信息，确保学生模型不受无关信息的干扰。

关键创新：OFA-KD的核心创新在于其特征对齐机制和自适应目标增强方案，这与传统的基于提示的蒸馏方法有本质区别，后者通常假设教师和学生模型具有相似的特征分布。

关键设计：在OFA-KD中，特征对齐通过中心核对齐（CKA）实现，损失函数设计为结合对齐损失和目标增强损失，以确保学生模型在学习过程中能够有效聚焦于重要特征。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

OFA-KD框架在CIFAR-100数据集上实现了最高8.0%的性能提升，在ImageNet-1K数据集上提升0.7%。这些结果表明，OFA-KD在异构架构间的蒸馏效果显著优于传统方法，具有良好的实用性和推广价值。

🎯 应用场景

OFA-KD框架在多种深度学习模型架构间的知识蒸馏中展现出优越性能，具有广泛的应用潜力。该研究可用于模型压缩、迁移学习及多模态学习等领域，帮助提升模型在资源受限环境下的表现，推动智能系统的普及与应用。

📄 摘要（原文）

Knowledge distillation~(KD) has proven to be a highly effective approach for enhancing model performance through a teacher-student training scheme. However, most existing distillation methods are designed under the assumption that the teacher and student models belong to the same model family, particularly the hint-based approaches. By using centered kernel alignment (CKA) to compare the learned features between heterogeneous teacher and student models, we observe significant feature divergence. This divergence illustrates the ineffectiveness of previous hint-based methods in cross-architecture distillation. To tackle the challenge in distilling heterogeneous models, we propose a simple yet effective one-for-all KD framework called OFA-KD, which significantly improves the distillation performance between heterogeneous architectures. Specifically, we project intermediate features into an aligned latent space such as the logits space, where architecture-specific information is discarded. Additionally, we introduce an adaptive target enhancement scheme to prevent the student from being disturbed by irrelevant information. Extensive experiments with various architectures, including CNN, Transformer, and MLP, demonstrate the superiority of our OFA-KD framework in enabling distillation between heterogeneous architectures. Specifically, when equipped with our OFA-KD, the student models achieve notable performance improvements, with a maximum gain of 8.0% on the CIFAR-100 dataset and 0.7% on the ImageNet-1K dataset. PyTorch code and checkpoints can be found at https://github.com/Hao840/OFAKD.