MoKD: Multi-Task Optimization for Knowledge Distillation

📄 arXiv: 2505.08170v2 📥 PDF

作者: Zeeshan Hayder, Ali Cheraghian, Lars Petersson, Mehrtash Harandi

分类: cs.CV

发布日期: 2025-05-13 (更新: 2025-08-02)

备注: Major changes to the paper and the authorship

💡 一句话要点

提出MoKD，通过多任务优化知识蒸馏解决梯度冲突和知识鸿沟问题。

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

关键词: 知识蒸馏 多任务学习 模型压缩 目标检测 图像分类 梯度平衡 子空间学习

📋 核心要点

  1. 现有知识蒸馏方法难以平衡教师模型的指导和任务本身的目标，容易出现梯度冲突和梯度主导问题。
  2. MoKD将知识蒸馏重新建模为多目标优化问题，通过平衡不同目标的梯度来解决梯度冲突和梯度主导问题。
  3. MoKD引入子空间学习框架，将特征投影到高维空间，弥合教师和学生模型之间的知识表示差距，提升知识迁移效果。

📝 摘要（中文）

知识蒸馏(KD)是一种有效的训练紧凑模型的方法，它将知识从大型、高性能的教师模型转移到学生模型。KD的两个关键挑战是：1)平衡从教师指导和任务目标中学习；2)处理教师和学生模型之间知识表示的差异。为了解决这些问题，我们提出了多任务优化知识蒸馏(MoKD)。MoKD解决了两个主要的梯度问题：a)梯度冲突，即特定任务梯度和蒸馏梯度未对齐；b)梯度主导，即一个目标的梯度占主导地位，导致不平衡。MoKD将KD重新定义为一个多目标优化问题，从而更好地平衡目标。此外，它还引入了一个子空间学习框架，将特征表示投影到高维空间，从而提高知识转移。通过在ImageNet-1K数据集上进行图像分类和在COCO数据集上进行目标检测的大量实验表明，我们的MoKD优于现有方法，以更高的效率实现了最先进的性能。据我们所知，与从头开始训练的模型相比，MoKD模型也实现了最先进的性能。

🔬 方法详解

问题定义：知识蒸馏旨在将大型教师模型的知识迁移到小型学生模型，但现有方法在平衡任务目标和蒸馏目标时存在困难。梯度冲突（task-specific梯度和蒸馏梯度方向不一致）和梯度主导（某个目标的梯度远大于其他目标）是主要痛点，导致学生模型性能受限。此外，教师和学生模型之间的知识表示差异也阻碍了有效的知识迁移。

核心思路：MoKD的核心思路是将知识蒸馏问题转化为一个多目标优化问题，通过优化多个目标函数（包括任务损失和蒸馏损失）来平衡任务学习和知识迁移。通过显式地考虑不同目标之间的关系，MoKD能够更好地解决梯度冲突和梯度主导问题，从而提高学生模型的性能。同时，利用子空间学习来对齐教师和学生模型的特征空间，从而更好地进行知识迁移。

技术框架：MoKD的整体框架包括以下几个主要模块：1) 特征提取模块：分别从教师和学生模型中提取特征表示。2) 子空间学习模块：将教师和学生模型的特征表示投影到高维子空间，以减小知识表示的差异。3) 多目标优化模块：通过优化多个目标函数（包括任务损失和蒸馏损失）来训练学生模型。4) 梯度平衡模块：用于平衡不同目标函数的梯度，避免梯度冲突和梯度主导。

关键创新：MoKD的关键创新在于将知识蒸馏问题重新建模为多目标优化问题，并引入了梯度平衡机制和子空间学习框架。与传统的知识蒸馏方法相比，MoKD能够更好地平衡任务学习和知识迁移，并有效地解决梯度冲突和梯度主导问题。子空间学习模块则进一步提升了知识迁移的效率。

关键设计：MoKD的关键设计包括：1) 多目标损失函数的设计，需要仔细权衡任务损失和蒸馏损失的权重。2) 子空间学习模块的结构和训练方式，例如使用线性投影或非线性映射。3) 梯度平衡模块的具体实现，例如使用梯度归一化或动态权重调整。4) 高维子空间的维度选择，需要根据具体任务和数据集进行调整。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

MoKD在ImageNet-1K图像分类和COCO目标检测任务上取得了显著的性能提升。在ImageNet-1K上，MoKD模型超越了现有的知识蒸馏方法，并达到了与从头训练的模型相媲美的性能。在COCO目标检测任务上，MoKD也取得了state-of-the-art的结果，证明了其在复杂任务上的有效性。

🎯 应用场景

MoKD可应用于各种需要模型压缩和加速的场景，例如移动设备上的图像识别、自动驾驶中的目标检测、以及资源受限环境下的模型部署。该方法能够有效提升小模型的性能，使其在计算资源有限的场景下也能达到较高的精度，具有重要的实际应用价值。

📄 摘要（原文）

Compact models can be effectively trained through Knowledge Distillation (KD), a technique that transfers knowledge from larger, high-performing teacher models. Two key challenges in Knowledge Distillation (KD) are: 1) balancing learning from the teacher's guidance and the task objective, and 2) handling the disparity in knowledge representation between teacher and student models. To address these, we propose Multi-Task Optimization for Knowledge Distillation (MoKD). MoKD tackles two main gradient issues: a) Gradient Conflicts, where task-specific and distillation gradients are misaligned, and b) Gradient Dominance, where one objective's gradient dominates, causing imbalance. MoKD reformulates KD as a multi-objective optimization problem, enabling better balance between objectives. Additionally, it introduces a subspace learning framework to project feature representations into a high-dimensional space, improving knowledge transfer. Our MoKD is demonstrated to outperform existing methods through extensive experiments on image classification using the ImageNet-1K dataset and object detection using the COCO dataset, achieving state-of-the-art performance with greater efficiency. To the best of our knowledge, MoKD models also achieve state-of-the-art performance compared to models trained from scratch.