m2mKD: Module-to-Module Knowledge Distillation for Modular Transformers

📄 arXiv: 2402.16918v3 📥 PDF

作者: Ka Man Lo, Yiming Liang, Wenyu Du, Yuantao Fan, Zili Wang, Wenhao Huang, Lei Ma, Jie Fu

分类: cs.LG, cs.CV

发布日期: 2024-02-26 (更新: 2024-07-07)

🔗 代码/项目: GITHUB

💡 一句话要点

提出模块间知识蒸馏方法以提升模块化变换器性能

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

关键词: 模块化神经网络 知识蒸馏 计算机视觉 模型训练 深度学习

📋 核心要点

  1. 现有的知识蒸馏方法未针对模块化模型进行优化，导致在训练时面临架构复杂性和参数量庞大的挑战。
  2. 本文提出的模块间知识蒸馏（m2mKD）方法，通过教师模块与学生模块的互动，促进知识的有效转移。
  3. 实验结果表明，m2mKD在NACs和V-MoE上均显著提升了模型的准确率和鲁棒性，验证了其有效性。

📝 摘要（中文）

模块化神经架构因其强大的泛化能力和高效的领域适应性而受到关注。然而，由于内在稀疏连接带来的优化困难，训练这些模型面临挑战。通过知识蒸馏等技术从单体模型中获取知识可以促进训练并整合多样化知识。传统的知识蒸馏方法并未针对模块化模型进行优化，难以应对独特的架构和庞大的参数量。为此，本文提出了模块间知识蒸馏（m2mKD），用于在模块之间转移知识。m2mKD结合了预训练单体模型的教师模块和模块化模型的学生模块，通过共享的元模型鼓励学生模块模仿教师模块的行为。我们在两个模块化神经架构上评估了m2mKD，结果显示在Tiny-ImageNet上，NACs的IID准确率提升了5.6%，OOD鲁棒性提升了4.2%；而使用m2mKD训练的V-MoE-Base模型在ImageNet-1k上比端到端训练提高了3.5%的准确率。

🔬 方法详解

问题定义：本文旨在解决模块化神经网络训练中的优化困难，尤其是由于稀疏连接导致的知识转移效率低下的问题。现有的知识蒸馏方法未能有效适应模块化架构，限制了其性能提升。

核心思路：提出模块间知识蒸馏（m2mKD）方法，通过将预训练单体模型的教师模块与模块化模型的学生模块相结合，利用共享的元模型促进知识的有效转移，从而提升学生模块的性能。

技术框架：m2mKD的整体架构包括教师模块、学生模块和共享的元模型。教师模块提供指导，学生模块通过模仿教师模块的行为进行学习，整个过程通过元模型进行协调和优化。

关键创新：m2mKD的主要创新在于针对模块化模型的特性设计了知识蒸馏过程，解决了传统方法在模块间知识转移时的局限性，能够有效应对复杂的模块化结构和参数量。

关键设计：在m2mKD中，损失函数设计为鼓励学生模块与教师模块之间的行为一致性，同时在训练过程中采用了动态调整的学习率策略，以适应不同模块的学习进度。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

实验结果显示，应用m2mKD后，NACs在Tiny-ImageNet上的IID准确率提高了5.6%，OOD鲁棒性提升了4.2%；而V-MoE-Base模型在ImageNet-1k上的准确率比传统端到端训练提高了3.5%。这些结果表明m2mKD在提升模块化模型性能方面的显著效果。

🎯 应用场景

该研究的潜在应用领域包括计算机视觉、自然语言处理等需要模块化架构的任务。通过提升模块化模型的训练效率和性能，m2mKD有助于推动智能系统在新领域的快速适应与应用，具有重要的实际价值和未来影响。

📄 摘要（原文）

Modular neural architectures are gaining attention for their powerful generalization and efficient adaptation to new domains. However, training these models poses challenges due to optimization difficulties arising from intrinsic sparse connectivity. Leveraging knowledge from monolithic models through techniques like knowledge distillation can facilitate training and enable integration of diverse knowledge. Nevertheless, conventional knowledge distillation approaches are not tailored to modular models and struggle with unique architectures and enormous parameter counts. Motivated by these challenges, we propose module-to-module knowledge distillation (m2mKD) for transferring knowledge between modules. m2mKD combines teacher modules of a pretrained monolithic model and student modules of a modular model with a shared meta model respectively to encourage the student module to mimic the behaviour of the teacher module. We evaluate m2mKD on two modular neural architectures: Neural Attentive Circuits (NACs) and Vision Mixture-of-Experts (V-MoE). Applying m2mKD to NACs yields significant improvements in IID accuracy on Tiny-ImageNet (up to 5.6%) and OOD robustness on Tiny-ImageNet-R (up to 4.2%). Additionally, the V-MoE-Base model trained with m2mKD achieves 3.5% higher accuracy than end-to-end training on ImageNet-1k. Code is available at https://github.com/kamanphoebe/m2mKD.