m2mKD:面向模块化 Transformer 的模块到模块知识蒸馏
机器学习
2024-07-09 v3 计算机视觉与模式识别
摘要
模块化神经架构因其强大的泛化能力和对新领域的高效适应性而备受关注。然而,由于内在的稀疏连接性导致的优化困难,训练这些模型面临挑战。利用知识蒸馏等技术从单体模型中汲取知识,可以促进训练并实现多样化知识的整合。尽管如此,传统的知识蒸馏方法并非为模块化模型量身定制,难以应对独特的架构和巨大的参数量。受这些挑战的启发,我们提出了模块到模块知识蒸馏(m2mKD),用于在模块间传递知识。m2mKD 将预训练单体模型的教师模块与模块化模型的学生模块分别结合到一个共享的元模型中,以促使学生模块模仿教师模块的行为。我们在两种模块化神经架构上评估了 m2mKD:神经注意电路(NACs)和视觉混合专家模型(V-MoE)。将 m2mKD 应用于 NACs 在 Tiny-ImageNet 上的同分布(IID)准确率(最高提升 5.6%)和在 Tiny-ImageNet-R 上的分布外(OOD)鲁棒性(最高提升 4.2%)方面取得了显著改进。此外,使用 m2mKD 训练的 V-MoE-Base 模型在 ImageNet-1k 上的准确率比端到端训练高出 3.5%。代码地址:https://github.com/kamanphoebe/m2mKD。
引用
@article{arxiv.2402.16918,
title = {m2mKD: Module-to-Module Knowledge Distillation for Modular Transformers},
author = {Ka Man Lo and Yiming Liang and Wenyu Du and Yuantao Fan and Zili Wang and Wenhao Huang and Lei Ma and Jie Fu},
journal= {arXiv preprint arXiv:2402.16918},
year = {2024}
}