AMD: Automatic Multi-step Distillation of Large-scale Vision Models

📄 arXiv: 2407.04208v1 📥 PDF

作者: Cheng Han, Qifan Wang, Sohail A. Dianat, Majid Rabbani, Raghuveer M. Rao, Yi Fang, Qiang Guan, Lifu Huang, Dongfang Liu

分类: cs.CV

发布日期: 2024-07-05

备注: 19 pages, 5 figures

💡 一句话要点

提出AMD：自动多步蒸馏方法，用于大规模视觉模型压缩

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

关键词: 知识蒸馏 模型压缩 多步蒸馏 教师助理模型 自动优化

📋 核心要点

  1. 现有知识蒸馏方法在大模型压缩时，当教师模型和学生模型容量差距大时，性能下降明显。
  2. 提出自动多步蒸馏（AMD）方法，通过引入教师助理模型，分阶段进行知识传递，缓解容量差距问题。
  3. 在CIFAR和ImageNet等数据集上实验表明，AMD方法优于现有基线方法，提升了学生模型的性能。

📝 摘要（中文）

Transformer架构因其卓越的性能已成为各种视觉任务的标准模型。随着模型规模的持续扩大，模型蒸馏在实际应用中变得至关重要，尤其是在计算资源受限的设备上。然而，当教师模型和学生模型之间存在较大容量差距时（例如，10倍压缩率），现有的知识蒸馏方法效果会降低。本文提出了一种名为自动多步蒸馏（AMD）的新方法，用于大规模视觉模型压缩。我们的蒸馏过程分多个步骤进行。首先，教师模型被蒸馏成一个中间的教师助理模型，然后进一步蒸馏到学生模型。引入了一种高效的优化框架，以自动识别能够最大化学生模型性能的最佳教师助理模型。我们在多个图像分类数据集（包括CIFAR-10、CIFAR-100和ImageNet）上进行了大量实验。结果表明，我们的方法优于几种已建立的基线方法，为未来大规模视觉模型的知识蒸馏方法铺平了道路。

🔬 方法详解

问题定义：论文旨在解决大规模视觉模型压缩中，教师模型和学生模型容量差距过大导致知识蒸馏效果不佳的问题。现有方法直接将大型教师模型的知识迁移到小型学生模型，容易造成信息损失和训练困难。

核心思路：论文的核心思路是引入一个中间的“教师助理”模型，将知识蒸馏过程分解为多个步骤。首先将教师模型的知识传递给教师助理模型，然后再由教师助理模型将知识传递给学生模型。通过逐步缩小模型容量差距，提高知识传递的效率和效果。

技术框架：AMD方法的整体框架包含三个主要阶段：1) 教师模型训练完成；2) 自动搜索最佳教师助理模型：通过优化框架自动确定教师助理模型的结构和参数，使其能够有效地从教师模型中学习知识；3) 学生模型蒸馏：使用训练好的教师助理模型作为新的教师，对学生模型进行蒸馏训练。

关键创新：AMD的关键创新在于自动搜索最佳教师助理模型。传统的知识蒸馏方法通常需要手动设计教师助理模型，这需要大量的经验和实验。AMD通过引入优化框架，自动搜索能够最大化学生模型性能的教师助理模型，从而提高了知识蒸馏的效率和效果。

关键设计：AMD的关键设计包括：1) 教师助理模型的结构选择：论文可能采用了多种不同的网络结构作为教师助理模型的候选，例如不同深度的ResNet或MobileNet等；2) 优化框架的设计：优化框架需要能够有效地搜索最佳教师助理模型，可能采用了进化算法、强化学习等方法；3) 损失函数的设计：损失函数需要能够引导教师助理模型学习教师模型的知识，并提高学生模型的性能，可能采用了KL散度、Hinton损失等。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

论文在CIFAR-10、CIFAR-100和ImageNet等图像分类数据集上进行了实验，结果表明AMD方法显著优于现有的知识蒸馏方法。具体性能提升数据未知，但摘要强调了优于多个已建立的基线方法，表明了AMD方法的有效性。

🎯 应用场景

该研究成果可广泛应用于移动设备、嵌入式系统等计算资源受限的场景，例如移动端图像识别、智能监控、自动驾驶等。通过模型压缩，可以在保证模型性能的同时，降低计算成本和功耗，加速AI技术的落地应用。未来，该方法可以进一步扩展到其他视觉任务，如目标检测、语义分割等。

📄 摘要（原文）

Transformer-based architectures have become the de-facto standard models for diverse vision tasks owing to their superior performance. As the size of the models continues to scale up, model distillation becomes extremely important in various real applications, particularly on devices limited by computational resources. However, prevailing knowledge distillation methods exhibit diminished efficacy when confronted with a large capacity gap between the teacher and the student, e.g, 10x compression rate. In this paper, we present a novel approach named Automatic Multi-step Distillation (AMD) for large-scale vision model compression. In particular, our distillation process unfolds across multiple steps. Initially, the teacher undergoes distillation to form an intermediate teacher-assistant model, which is subsequently distilled further to the student. An efficient and effective optimization framework is introduced to automatically identify the optimal teacher-assistant that leads to the maximal student performance. We conduct extensive experiments on multiple image classification datasets, including CIFAR-10, CIFAR-100, and ImageNet. The findings consistently reveal that our approach outperforms several established baselines, paving a path for future knowledge distillation methods on large-scale vision models.