MoE-KD: Your Teacher Model is Worth Mixture-of-Experts for Knowledge Distillation

Knowledge distillation (KD) aims to transfer useful informa-tion from a large-scale model (teacher) to a lightweight model (student).Classical KD focuses on leveraging the teacher’s predictions as soft labelsto regularize student training. However, the exact match of predictionsin Kullback-Leibler (KL) divergence could be somewhat in conflict withthe classification objective, given that the distribution discrepancies be-tween teacher-generated predictions and ground-truth annotations tendto be fairly severe. In this paper, we rethink the role of teacher predic-tions from a Mixture-of-Experts (MoE) perspective and transfer knowl-edge by introducing teacher predictions as latent variables to reformulatethe classification objective. This MoE strategy results in breaking downthe vanilla classification task into a mixture of easier subtasks with theteacher classifier as a gating function to weigh the importance of sub-tasks. Each subtask is efficiently conquered by distinct experts that areeffectively implemented by resorting to multi-level teacher outputs. Wefurther develop a theoretical framework to formulate our method, termedMoE-KD, as an Expectation-Maximization (EM) algorithm and provideproof of the convergence. Extensive experiments manifest that MoE-KDoutperforms advanced knowledge distillers on mainstream benchmarks.