m2mKD: Module-to-Module Knowledge Distillation for Modular Transformers
Ka Man Lo Yiming Liang Wenyu Du Yuantao Fan Zili Wang Wenhao Huang Lei Ma Jie Fu
Modular neural architectures are gaining increasing attention due to their
powerful capability for generalization and sample-efficient adaptation to new
domains. However, training modular models, particularly in the early stages,
poses challenges due to the optimization difficulties arising from their
intrinsic sparse connectivity. Leveraging the knowledge from monolithic models,
using techniques such as knowledge distillation, is likely to facilitate the
training of modular models and enable them to integrate knowledge from multiple
models pretrained on diverse sources. Nevertheless, conventional knowledge
distillation approaches are not tailored to modular models and can fail when
directly applied due to the unique architectures and the enormous number of
parameters involved. Motivated by these challenges, we propose a general
module-to-module knowledge distillation (m2mKD) method for transferring
knowledge between modules. Our approach involves teacher modules split from a
pretrained monolithic model, and student modules of a modular model. m2mKD
separately combines these modules with a shared meta model and encourages the
student module to mimic the behaviour of the teacher module. We evaluate the
effectiveness of m2mKD on two distinct modular neural architectures: Neural
Attentive Circuits (NACs) and Vision Mixture-of-Experts (V-MoE). By applying
m2mKD to NACs, we achieve significant improvements in IID accuracy on
Tiny-ImageNet (up to 5.6%) and OOD robustness on Tiny-ImageNet-R (up to 4.2%).
On average, we observe a 1% gain in both ImageNet and ImageNet-R. The
V-MoE-Base model trained using m2mKD also achieves 3.5% higher accuracy than
end-to-end training on ImageNet. The experimental results demonstrate that our
method offers a promising solution for connecting modular networks with
pretrained monolithic models. Code is available at
https://github.com/kamanphoebe/m2mKD.