AdaMoE: Token-Adaptive Routing with Null Experts for Mixture-of-Experts Language Models
Zihao Zeng, Yibo Miao, Hongcheng Gao, Hao Zhang, Zhijie Deng · 2024
Mixture of experts (MoE) has become the standard for constructing production-level large language models (LLMs) due to its promise to boost model capacity without causing significant overheads.Nevertheless, existing MoE methods usually enforce a constant top-k routing for all tokens, which is arguably restrictive because various tokens (e.g., " " vs. "apple") may require various numbers of experts for feature abstraction.Lifting such a constraint can help make the most of limited resources and unleash the potential of the model for downstream tasks.In this sense, we introduce AdaMOE to realize token-adaptive routing for MoE, where different tokens are permitted to select a various number of experts.AdaMOE makes minimal modifications to the vanilla MoE with top-k routing-it simply introduces a fixed number of null experts, which do not consume any FLOPs, to the expert set and increases the value of k.AdaMOE does not force each token to occupy a fixed number of null experts but ensures the average usage of the null experts with a load-balancing loss, leading to an adaptive number of null/true experts used by each token.AdaMOE exhibits a strong resemblance to MoEs with expert choice routing while allowing for trivial auto-regressive modeling.AdaMOE is easy to implement and can be effectively applied to pre-trained (MoE-)LLMs.Extensive studies show that AdaMOE can reduce average expert load (FLOPs) while achieving superior performance.For example, on the ARC-C dataset, applying our method to fine-tuning Mixtral-8x7B can reduce FLOPs by 14.5% while increasing accuracy by 1.69%.Code is available at this link.