A da M o E : Token-Adaptive Routing with Null Experts for Mixture-of-Experts Language Models (2024)

Abstract: 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., "<EOS>” 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](https://github.com/CengZihao/AdaMoE).

## introduction large language models (llms) have exhibited exceptional performance across diverse tasks and do-* equal contribution. moe routinely routes each token to a fixed amount of experts, typically the k ones with top routing probabilities. however, not all tokens require the same number of experts for feature abstraction. intuitively, the semantic tokens deserve a higher concentration of experts, while others with less significant meaning can be processed more swiftly. lifting the top-k routing constraint can help make the most of limited resources and unleash the potential of the model. to achieve this, moe with expert choice routing @xcite performs expert-level routing, where each expert chooses a fixed number of tokens for processing and different tokens could be processed by different numbers of experts. yet, an unacceptable drawback figure 1: the number of selected experts for various tokens in an adamoe variant of mixtral-8x7b. as shown, after applying adamoe, the model possesses the ability to perform token-adaptive routing. also note that some tokens only require 1 expert for feature abstraction, which offers the opportunity for inference acceleration. this work introduces adamoe, a novel method designed to achieve token-level adaptive routing in moe, allowing different tokens to select varying numbers of experts. an illustrative example is presented in figure 1 . adamoe requires minimal changes to the vanilla moe with top-k routing by incorporating a fixed number of null experts into the expert set. these null experts do not consume any computational resources. by increasing the value of k, more experts can be activated. to encourage the average usage rate of the null experts, adamoe minimizes a load balancing loss. this leads to an adaptive number of null experts and true experts being employed by each token. notably, adamoe shares similarities with existing moe with expert choice routing while also enabling straightforward causal language modeling. adamoe is easy to implement and can be applied to both pre-trained regular llms and moe-llms for supervised fine-tuning. for the former, we experiment on llama2-7b @xcite by introducing lora experts and corresponding routers to the model. for the latter, we experiment on mixtral-8x7b @xcite by augmenting the original router with extra weights for the null experts. the results underscore the effectiveness of adamoe's token-adaptive mechanism in enhancing both computational efficiency and model performance. for example, when finetuning llama2-7b, adamoe achieves much higher accuracy across almost all evaluated datasets. moreover, when fine-tuning mixtral-8x7b with adamoe on arc-challenge @xcite , we observed a 14.5% reduction in total flops, accompanied by a 1.69% increase in accuracy. ## method in this section, we introduce adamoe, which incorporates null experts to allow for more flexible and efficient expert selection for various tokens. an illustrative comparison between vanilla moe and adamoe is presented in figure 2 . ## experiments in this section, we demonstrate the superior performance of adamoe across various benchmarks, particularly in reducing expert load and enhancing task performance through its token-adaptive routing strategy. we first apply our method to regular llms with the mo-lora architecture (section 4.1). then apply it to traditional moe-llms (section 4.2). additionally, we also include extensive ablation studies to provide further insights into our approach's effectiveness. ## conclusion moe has been a promising method for training powerful models with fewer parameters. in this paper, we introduced adamoe, which uses null experts to enable token-level adaptive expert allocation and overcome the drawbacks of fixed expert allocation. extensive experiments validate its effectiveness. the adamoe approach significantly enhances efficiency and adaptability, paving the way for more capable large language models.