Every Expert Matters: Towards Effective Knowledge Distillation for Mixture-of-Experts Language Models

2025-02-18

AI Summary

To address the high memory overhead of Mixture-of-Experts (MoE) large language models, which hinders their deployment in resource-constrained settings, this paper proposes the first knowledge distillation framework tailored for MoE teacher models—uniquely exploiting implicit knowledge embedded in non-activated experts. Our method introduces two key components: (1) a Knowledge Augmentation (KA) module that explicitly models latent representations from inactive experts; and (2) a Student-Aware Router (SAR), enabling differentiable joint optimization of the router and student model to facilitate holistic expert knowledge distillation. Extensive experiments across multiple benchmarks demonstrate that the distilled student model achieves substantial parameter reduction while maintaining—and in many cases improving—generation quality, thereby validating the efficacy and superiority of full-expert knowledge distillation in MoE architectures.

Technology Category

• Machine Learning: Mixture of Experts (MoE)
• Data Mining & Knowledge Management: Knowledge Acquisition from the Web
• Search and Optimization: Learning to Search

Application Category

• Semantics and Knowledge: Methods to enhance, augment, integrate or synergize semantic models such as knowledge graphs and LLMs
• Search and Retrieval-Augmented AI: Web learning to rank, online learning, and counterfactual learning for ranking
• User Modeling, Personalization and Recommendation: Explainable and interpretable methods for personalization

Abstract

With the emergence of Mixture-of-Experts (MoE), the efficient scaling of model size has accelerated the development of large language models in recent years. However, their high memory requirements prevent their use in resource-constrained environments. While knowledge distillation (KD) has been a proven method for model compression, its application to MoE teacher models remains underexplored. Through our investigation, we discover that non-activated experts in MoE models possess valuable knowledge that benefits student models. We further demonstrate that existing KD methods are not optimal for compressing MoE models, as they fail to leverage this knowledge effectively. To address this, we propose two intuitive MoE-specific KD methods for the first time: Knowledge Augmentation (KA) and Student-Aware Router (SAR), both designed to effectively extract knowledge from all experts. Specifically, KA augments knowledge by sampling experts multiple times, while SAR uses all experts and adjusts the expert weights through router training to provide optimal knowledge. Extensive experiments show that our methods outperform conventional KD methods, demonstrating their effectiveness for MoE teacher models.

Problem

• Optimizing knowledge distillation for MoE models
• Leveraging knowledge from non-activated MoE experts
• Developing MoE-specific KD methods for efficiency

Innovation

• Knowledge Augmentation technique
• Student-Aware Router training
• Expert knowledge extraction